diff --git a/desktop/src/electron/service-bridge/modular-state.ts b/desktop/src/electron/service-bridge/modular-state.ts index ab6c5fb5..88f8e6fc 100644 --- a/desktop/src/electron/service-bridge/modular-state.ts +++ b/desktop/src/electron/service-bridge/modular-state.ts @@ -14,7 +14,7 @@ import type { LogEntry, LogPage } from '@/shared/types/log' import type { NodeItemMetrics } from '@/shared/types/metrics' import type { NodeItem } from '@/shared/types/nodes' import type { ServiceError, ServiceErrorAction, ServiceErrorSeverity } from '@/shared/types/errors' -import type { Workload, WorkloadState } from '@/shared/types/workloads' +import type { Workload, WorkloadState, WorkloadStats } from '@/shared/types/workloads' import type { ClusterNode, Invite } from '@/shared/types/cluster' import type { WsInvokeResponse } from '@/shared/types/ws-channels' import { @@ -368,9 +368,32 @@ function parseWorkload(value: JsonValue | undefined): Workload | null { } const scheduledOn = nullableStringValue(obj.scheduledOn) if (scheduledOn) workload.scheduledOn = scheduledOn + const stats = parseWorkloadStats(obj.stats) + if (stats) workload.stats = stats return workload } +/** + * Parse a workload's optional `stats` object. A field is copied only when + * present with the right type, so a measurement the proxy omitted stays absent + * rather than becoming a zero the UI would render. + */ +function parseWorkloadStats(value: JsonValue | undefined): WorkloadStats | null { + const obj = objectValue(value) + if (!obj) return null + const stats: WorkloadStats = {} + const promptTokens = nullableNumberValue(obj.promptTokens) + if (promptTokens !== null) stats.promptTokens = promptTokens + const completionTokens = nullableNumberValue(obj.completionTokens) + if (completionTokens !== null) stats.completionTokens = completionTokens + const tokensPerSecond = nullableNumberValue(obj.tokensPerSecond) + if (tokensPerSecond !== null) stats.tokensPerSecond = tokensPerSecond + const ttftMs = nullableNumberValue(obj.ttftMs) + if (ttftMs !== null) stats.ttftMs = ttftMs + if (booleanValue(obj.estimated)) stats.estimated = true + return Object.keys(stats).length > 0 ? stats : null +} + /** * Parse a broker `workloads:get-initial` baseline (`{ workloads: workloadInfo[] }`, * ordered by createdAt). Each element is a bare `workloadInfo` — the same shape diff --git a/desktop/src/shared/types/workloads.ts b/desktop/src/shared/types/workloads.ts index 71340812..adf501b1 100644 --- a/desktop/src/shared/types/workloads.ts +++ b/desktop/src/shared/types/workloads.ts @@ -30,4 +30,35 @@ export interface Workload { completedAt: number | null error: string | null requesterId: string | null + /** + * Inference statistics for a finished job. Present only on the terminal + * (`completed` / `failed`) transition, and only when the origin proxy could + * measure something; see {@link WorkloadStats}. + */ + stats?: WorkloadStats +} + +/** + * Inference statistics the origin proxy measures from the response body as it + * streams to the client (services/shared/inferstats). Every field is optional: + * absent means "not measured", never zero. + */ +export interface WorkloadStats { + /** Prompt size in tokens, as reported by the engine. */ + promptTokens?: number + /** + * Tokens generated. Reported by the engine when its response carried a usage + * object; otherwise counted from stream chunks, in which case `estimated` is + * set. + */ + completionTokens?: number + /** Decode throughput: `completionTokens` over generation time, one decimal. */ + tokensPerSecond?: number + /** Time to first token in ms: request start to the first response body byte. */ + ttftMs?: number + /** + * `completionTokens` (and so `tokensPerSecond`) were counted from stream + * chunks rather than reported by the engine, so treat them as approximate. + */ + estimated?: boolean } diff --git a/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx b/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx index 2fa7b476..8a277851 100644 --- a/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx +++ b/desktop/src/ui/components/Workloads/WorkloadItemCard.tsx @@ -8,6 +8,7 @@ import { workloadExecutionNodeId } from '@/shared/utils/workloads' import { useNodesStore } from '@/ui/stores/nodes.store' import { formatModelDisplayName } from '@/ui/utils/format-model-display-name' import { getWorkloadColorBar } from '@/ui/utils/colors' +import { formatWorkloadStats } from '@/ui/utils/format-workload-stats' import EngineIcon from '@/ui/components/EngineIcon' const formatDate = (timestamp: number) => { @@ -56,6 +57,7 @@ function WorkloadItemCard({ workload }: { workload: Workload }) { }) const ranOnLabel = workload.state === 'running' ? 'Running on' : 'Ran on' const barColor = useMemo(() => getWorkloadColorBar(workload.state), [workload.state]) + const statsParts = useMemo(() => formatWorkloadStats(workload.stats), [workload.stats]) const subtext = useMemo(() => { const state = workload.state @@ -156,6 +158,15 @@ function WorkloadItemCard({ workload }: { workload: Workload }) { {subtext} + {statsParts.length > 0 && ( + + {statsParts.map((part, index) => ( + + {index > 0 ? `\u00b7 ${part}` : part} + + ))} + + )} {workload.error && workload.state === 'failed' && ( = 100 + ? Math.round(tokensPerSecond).toString() + : tokensPerSecond.toFixed(1) +} + +function formatDuration(ms: number): string { + return ms < 1000 ? `${Math.round(ms)} ms` : `${(ms / 1000).toFixed(1)} s` +} diff --git a/desktop/tests/modular/format-workload-stats.test.ts b/desktop/tests/modular/format-workload-stats.test.ts new file mode 100644 index 00000000..dad220a8 --- /dev/null +++ b/desktop/tests/modular/format-workload-stats.test.ts @@ -0,0 +1,41 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it } from 'vitest' +import { formatWorkloadStats } from '@/ui/utils/format-workload-stats' + +describe('formatWorkloadStats', () => { + it('shows engine-reported throughput, token count and time to first token', () => { + expect( + formatWorkloadStats({ + promptTokens: 12, + completionTokens: 1234, + tokensPerSecond: 42.34, + ttftMs: 812 + }) + ).toEqual(['42.3 tok/s', '1,234 tokens', '812 ms to first token']) + }) + + it('marks estimated counts and rounds fast rates to whole tokens', () => { + expect( + formatWorkloadStats({ + completionTokens: 40, + tokensPerSecond: 133.7, + ttftMs: 1480, + estimated: true + }) + ).toEqual(['~134 tok/s', '~40 tokens', '1.5 s to first token']) + }) + + it('falls back to prompt tokens for a job that generated nothing', () => { + expect(formatWorkloadStats({ promptTokens: 8, ttftMs: 30 })).toEqual([ + '8 prompt tokens', + '30 ms to first token' + ]) + }) + + it('is empty when nothing was measured', () => { + expect(formatWorkloadStats(undefined)).toEqual([]) + expect(formatWorkloadStats({})).toEqual([]) + }) +}) diff --git a/services/lmstudio-proxy/README.md b/services/lmstudio-proxy/README.md index 71a8b70d..1562084f 100644 --- a/services/lmstudio-proxy/README.md +++ b/services/lmstudio-proxy/README.md @@ -179,6 +179,12 @@ One lifecycle transition per forwarded inference request, carrying a single `wor {"jsonrpc":"2.0","method":"workload:started","params":{"workloadInfo":{"id":"17","model":"lmstudio-community/Qwen3-8B-GGUF","engine":"lmstudio","runId":"3ce8a1740b62df95","state":"running","originatedFrom":"","scheduledOn":"22222222-2222-2222-2222-222222222222","createdAt":1716998400000,"startedAt":1716998400000,"completedAt":null,"error":null,"requesterId":null}}} ``` +The terminal `workload:completed` / `workload:errored` carries the same object plus an additive `stats` block measured from the response body as it streamed — token counts from the engine's usage report (OpenAI `usage`, Ollama `eval_count`) or, failing that, counted from stream chunks (`estimated`), decode throughput and time to first token. The body is never buffered: only its last few kilobytes are inspected (`nvpair-shared/inferstats`). + +```json +{"jsonrpc":"2.0","method":"workload:completed","params":{"workloadInfo":{"id":"17","model":"lmstudio-community/Qwen3-8B-GGUF","engine":"lmstudio","runId":"3ce8a1740b62df95","state":"completed","originatedFrom":"","scheduledOn":"22222222-2222-2222-2222-222222222222","createdAt":1716998400000,"startedAt":1716998400000,"completedAt":1716998412000,"error":null,"requesterId":null,"stats":{"promptTokens":12,"completionTokens":340,"tokensPerSecond":42.3,"ttftMs":812}}}} +``` + #### `node/activity` Raised while a node's engine is streaming a response back through the proxy: every successful write of upstream body bytes reports the node that produced them. The broker relays it to `nvpair-node-scanner`, which treats it as proof of life and cancels that node's eviction — a node saturated by inference cannot answer a liveness probe, but it is demonstrably alive precisely because it is streaming. Coalesced to one report per node per 2s (`nvpair-shared/nodeactivity`), since a generation writes hundreds of chunks and the scanner treats one report as good for a minute. `msSince` is the age of the observation; the broker adds its own relay delay before passing it on. diff --git a/services/lmstudio-proxy/proxy.go b/services/lmstudio-proxy/proxy.go index 6e619e77..5675980b 100644 --- a/services/lmstudio-proxy/proxy.go +++ b/services/lmstudio-proxy/proxy.go @@ -29,6 +29,7 @@ import ( "nvpair-shared/clustertrust" "nvpair-shared/cors" "nvpair-shared/errors" + "nvpair-shared/inferstats" "nvpair-shared/netmon" "nvpair-shared/netpick" "nvpair-shared/nodeactivity" @@ -182,6 +183,11 @@ type Workload struct { CompletedAt *int64 `json:"completedAt"` Error *string `json:"error"` RequesterID *string `json:"requesterId"` + // Stats carries the inference statistics measured for the terminal + // (completed/errored) transition: token counts, decode throughput and time + // to first token. Additive: omitted while the workload is running and when + // nothing could be measured (see nvpair-shared/inferstats). + Stats *inferstats.Stats `json:"stats,omitempty"` } // workloadParams is the params envelope for a workload:* notification @@ -238,6 +244,12 @@ type statusCapture struct { // discovery cannot obtain for itself while the node is too busy to answer a // probe. Called on the reverse proxy's copy goroutine, so it must be cheap. upstreamAlive func() + + // tap, when set, observes every body byte written to the client so the + // workload's terminal event can carry token counts and throughput. Like + // upstreamAlive it is set at the commit point, and only for a successful + // inference response, so the proxy's own error bodies are never measured. + tap *inferstats.Tap } // Unwrap exposes the underlying ResponseWriter so http.ResponseController can @@ -277,6 +289,9 @@ func (sc *statusCapture) Write(b []byte) (int, error) { if err == nil && sc.upstreamAlive != nil { sc.upstreamAlive() } + if err == nil && sc.tap != nil { + sc.tap.Observe(b) + } return n, err } @@ -1026,6 +1041,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { finalStatus int started bool wl *Workload + tap *inferstats.Tap ) // Emit workload:started up front, the moment we begin forwarding, naming @@ -1050,6 +1066,10 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { StartedAt: &createdMs, } p.emitWorkload(workloadStartedMethod, *wl) + // The tap measures the response body for this workload's stats. It is + // allocated here, before the disconnect watcher goroutine starts, and + // armed only once a candidate commits (ModifyResponse below). + tap = inferstats.NewTap(start) } // The terminal workload transition (completed/errored) can be reached from @@ -1076,6 +1096,11 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { if errMsg != "" { wl.Error = &errMsg } + // Stats ride the terminal event only; a failed stream still reports + // what was generated before it broke. + if tap != nil { + wl.Stats = tap.Finish() + } snapshot := *wl wlMu.Unlock() method := workloadCompletedMethod @@ -1165,6 +1190,13 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // came from the node. Same goroutine as the body copy, so no // synchronization is needed. sc.upstreamAlive = func() { p.reportActivity(cand.id) } + // Measure the body only for a successful inference response: an + // error body carries no tokens, and a control request has no + // workload to attach stats to. + if tap != nil && resp.StatusCode < http.StatusBadRequest { + tap.Arm(resp.Header) + sc.tap = tap + } // The engine may enforce its own origin policy. Honor it: // overwriting a declared Access-Control-Allow-Origin would // silently widen the user's policy, and a wildcard is invalid diff --git a/services/lmstudio-proxy/stats_test.go b/services/lmstudio-proxy/stats_test.go new file mode 100644 index 00000000..fd2ce61f --- /dev/null +++ b/services/lmstudio-proxy/stats_test.go @@ -0,0 +1,152 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// recordedWorkload decodes the workloadInfo of the last frame rec captured for +// method, failing the test when none was emitted. +func recordedWorkload(t *testing.T, rec *prRec, method string) Workload { + t.Helper() + rec.mu.Lock() + raw := append([]byte(nil), rec.b...) + rec.mu.Unlock() + var found []byte + for _, line := range bytes.Split(raw, []byte{'\n'}) { + if bytes.Contains(line, []byte(`"method":"`+method+`"`)) { + found = line + } + } + if found == nil { + t.Fatalf("no %s frame recorded; frames:\n%s", method, raw) + } + var frame struct { + Params struct { + WorkloadInfo Workload `json:"workloadInfo"` + } `json:"params"` + } + if err := json.Unmarshal(found, &frame); err != nil { + t.Fatalf("decode %s frame: %v", method, err) + } + return frame.Params.WorkloadInfo +} + +// streamChunks serves body chunks as a flushed stream with a small gap between +// them, the way an engine streams tokens. +func streamChunks(contentType string, chunks ...string) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", contentType) + w.WriteHeader(http.StatusOK) + for _, c := range chunks { + io.WriteString(w, c) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + time.Sleep(2 * time.Millisecond) + } + }) +} + +// TestHandleHTTP_CompletedCarriesStats: a streamed inference response whose +// final chunk reports usage must surface engine-exact token counts and a +// throughput figure on workload:completed — and nothing on workload:started, +// which fires before a byte has flowed. +func TestHandleHTTP_CompletedCarriesStats(t *testing.T) { + upstream := httptest.NewServer(streamChunks("text/event-stream", + "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n", + "data: [DONE]\n\n", + )) + defer upstream.Close() + + rec := &prRec{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "node-a", upstream.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 1235) + + rr := httptest.NewRecorder() + p.handleHTTP(rr, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"llama","stream":true}`))) + if rr.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", rr.Code) + } + if !strings.Contains(rr.Body.String(), "[DONE]") { + t.Fatalf("client did not receive the full stream:\n%s", rr.Body.String()) + } + + if started := recordedWorkload(t, rec, "workload:started"); started.Stats != nil { + t.Fatalf("workload:started carried stats %+v before any byte flowed", started.Stats) + } + done := recordedWorkload(t, rec, "workload:completed") + if done.Stats == nil { + t.Fatal("workload:completed carried no stats") + } + s := done.Stats + if s.PromptTokens != 5 || s.CompletionTokens != 2 || s.Estimated { + t.Fatalf("stats = %+v, want engine-reported prompt=5 completion=2", s) + } + if s.TokensPerSecond <= 0 { + t.Fatalf("stats = %+v, want a positive tokensPerSecond", s) + } +} + +// TestHandleHTTP_StreamWithoutUsageEstimates: an engine that streams without a +// usage report still yields a token count, counted from chunks and flagged as +// an estimate. +func TestHandleHTTP_StreamWithoutUsageEstimates(t *testing.T) { + upstream := httptest.NewServer(streamChunks("text/event-stream", + "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n", + "data: [DONE]\n\n", + )) + defer upstream.Close() + + rec := &prRec{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "node-a", upstream.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 1235) + + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"llama","stream":true}`))) + + s := recordedWorkload(t, rec, "workload:completed").Stats + if s == nil || !s.Estimated || s.CompletionTokens != 3 || s.PromptTokens != 0 { + t.Fatalf("stats = %+v, want estimated completion=3 and no prompt count", s) + } +} + +// TestHandleHTTP_ErrorBodyNotMeasured: an upstream error body is never mined for +// tokens, even when it happens to contain usage-looking fields. +func TestHandleHTTP_ErrorBodyNotMeasured(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + io.WriteString(w, `{"error":"boom","usage":{"prompt_tokens":1,"completion_tokens":99}}`) + })) + defer upstream.Close() + + rec := &prRec{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "node-a", upstream.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 1235) + + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"llama"}`))) + + if failed := recordedWorkload(t, rec, "workload:errored"); failed.Stats != nil { + t.Fatalf("workload:errored carried stats %+v from an error body", failed.Stats) + } +} diff --git a/services/nvpair-tui/ui/workloads.go b/services/nvpair-tui/ui/workloads.go index ea7931df..1df9950a 100644 --- a/services/nvpair-tui/ui/workloads.go +++ b/services/nvpair-tui/ui/workloads.go @@ -4,6 +4,8 @@ package ui import ( + "strconv" + "nvpair-tui/rpc" "github.com/charmbracelet/bubbles/key" @@ -19,6 +21,14 @@ type workload struct { State string `json:"state"` OriginatedFrom string `json:"originatedFrom"` CreatedAt int64 `json:"createdAt"` // Unix millis + // Stats arrives with the terminal event only (see nvpair-shared/inferstats). + Stats *workloadStats `json:"stats"` +} + +// workloadStats is the subset of a workload's stats block the view shows. +type workloadStats struct { + TokensPerSecond float64 `json:"tokensPerSecond"` + Estimated bool `json:"estimated"` } // workloadsView shows cluster-wide inference workloads. The table is built @@ -54,15 +64,16 @@ func (v *workloadsView) Init() tea.Cmd { func (v *workloadsView) SetSize(w, h int) { v.width, v.height = w, h - const engine, state, age = 10, 10, 6 - id := clampWidth((w-engine-state-age-2)/3, 8) - model := clampWidth(w-engine-state-age-id-2, 10) + const engine, state, age, speed = 10, 10, 6, 7 + id := clampWidth((w-engine-state-age-speed-2)/3, 8) + model := clampWidth(w-engine-state-age-speed-id-2, 10) v.table.SetColumns([]table.Column{ {Title: "ID", Width: id}, {Title: "MODEL", Width: model}, {Title: "ENGINE", Width: engine}, {Title: "STATE", Width: state}, {Title: "AGE", Width: age}, + {Title: "TOK/S", Width: speed}, }) v.table.SetWidth(w) v.table.SetHeight(clampWidth(h-1, 1)) @@ -135,6 +146,7 @@ func (v *workloadsView) refreshRows() { w.Engine, w.State, ageLabel(w.CreatedAt), + speedLabel(w.Stats), }) } v.table.SetRows(rows) @@ -153,3 +165,22 @@ func (v *workloadsView) View() string { func (v *workloadsView) Help() []key.Binding { return nil } func workloadKey(origin, id string) string { return origin + "/" + id } + +// speedLabel renders decode throughput for the TOK/S column: one decimal below +// 100 tok/s, whole tokens above, prefixed "~" when the proxy only estimated the +// token count from stream chunks. Empty until the terminal event brings the +// measurement. +func speedLabel(s *workloadStats) string { + if s == nil || s.TokensPerSecond <= 0 { + return "" + } + prec := 1 + if s.TokensPerSecond >= 100 { + prec = 0 + } + label := strconv.FormatFloat(s.TokensPerSecond, 'f', prec, 64) + if s.Estimated { + label = "~" + label + } + return label +} diff --git a/services/nvpair-ui-broker/README.md b/services/nvpair-ui-broker/README.md index d2b0e174..84f25cbc 100644 --- a/services/nvpair-ui-broker/README.md +++ b/services/nvpair-ui-broker/README.md @@ -198,7 +198,7 @@ Two classes of proxy notification are **not** re-emitted under the `proxy:` name - **Inferred workloads** — a `workloads:upsert` transitioning a workload to `failed` that **no origin ever sent**. The broker synthesizes one in two situations: when a node leaves discovery while workloads are pinned to it, and when a remote origin that is still present stops re-asserting a workload this node believes is running (the origin's re-sync heartbeat asserts each of its active workloads indefinitely, so prolonged silence about one means it is finished or the origin is gone). Both are recorded as *inferred*, so the origin's next authoritative event overrides them; a client should treat a `failed` as the broker's best current answer rather than proof the origin reported a failure, and its `error` text names the reason. Workloads this node originated or is itself executing are never inferred about. -`workloads:upsert` carries `params.workloadInfo` (a full `Workload`); `workloads:remove` carries `params.workloadId` and the origin `params.originatedFrom`. See [`nvpair-workload-manager`](../nvpair-workload-manager/README.md) for the `Workload` shape. +`workloads:upsert` carries `params.workloadInfo` (a full `Workload`); `workloads:remove` carries `params.workloadId` and the origin `params.originatedFrom`. See [`nvpair-workload-manager`](../nvpair-workload-manager/README.md) for the `Workload` shape. A terminal upsert may carry `stats` (token counts, decode throughput, time to first token) measured by the origin proxy; the broker passes it through unchanged. ```json {"jsonrpc":"2.0","method":"workloads:upsert","params":{"workloadInfo":{"id":"42","model":"llama-3-70b","engine":"ollama","state":"running","originatedFrom":"MY-PC","scheduledOn":"GPU-RIG","createdAt":1716998400000,"startedAt":1716998400000,"completedAt":null,"error":null,"requesterId":null}}} diff --git a/services/nvpair-workload-manager/README.md b/services/nvpair-workload-manager/README.md index 474d9383..2034c09a 100644 --- a/services/nvpair-workload-manager/README.md +++ b/services/nvpair-workload-manager/README.md @@ -88,6 +88,7 @@ Defined in [`workload.go`](workload.go): | `createdAt`, `startedAt`, `completedAt` | Epoch milliseconds; the last two are nullable | | `error` | Normalized failure text, nullable | | `requesterId` | Optional client attribution, nullable | +| `stats` | Optional inference statistics the origin proxy measured from the response body, present only on the terminal transition: `promptTokens`, `completionTokens`, `tokensPerSecond`, `ttftMs` (time to first token) and `estimated` (counts came from stream chunks, not the engine). Each field is omitted when not measured | Optional and nullable fields use pointers so a peer's payload round-trips without inventing zero values. diff --git a/services/ollama-proxy/README.md b/services/ollama-proxy/README.md index 35f959b5..f9558b15 100644 --- a/services/ollama-proxy/README.md +++ b/services/ollama-proxy/README.md @@ -184,6 +184,12 @@ One lifecycle transition per forwarded inference request, carrying a single `wor {"jsonrpc":"2.0","method":"workload:started","params":{"workloadInfo":{"id":"17","model":"llama3:latest","engine":"ollama","runId":"3ce8a1740b62df95","state":"running","originatedFrom":"","scheduledOn":"22222222-2222-2222-2222-222222222222","createdAt":1716998400000,"startedAt":1716998400000,"completedAt":null,"error":null,"requesterId":null}}} ``` +The terminal `workload:completed` / `workload:errored` carries the same object plus an additive `stats` block measured from the response body as it streamed — token counts from the engine's usage report (OpenAI `usage`, Ollama `eval_count`) or, failing that, counted from stream chunks (`estimated`), decode throughput and time to first token. The body is never buffered: only its last few kilobytes are inspected (`nvpair-shared/inferstats`). + +```json +{"jsonrpc":"2.0","method":"workload:completed","params":{"workloadInfo":{"id":"17","model":"llama3:latest","engine":"ollama","runId":"3ce8a1740b62df95","state":"completed","originatedFrom":"","scheduledOn":"22222222-2222-2222-2222-222222222222","createdAt":1716998400000,"startedAt":1716998400000,"completedAt":1716998412000,"error":null,"requesterId":null,"stats":{"promptTokens":12,"completionTokens":340,"tokensPerSecond":42.3,"ttftMs":812}}}} +``` + #### `node/activity` Raised while a node's engine is streaming a response back through the proxy: every successful write of upstream body bytes reports the node that produced them. The broker relays it to `nvpair-node-scanner`, which treats it as proof of life and cancels that node's eviction — a node saturated by inference cannot answer a liveness probe, but it is demonstrably alive precisely because it is streaming. Coalesced to one report per node per 2s (`nvpair-shared/nodeactivity`), since a generation writes hundreds of chunks and the scanner treats one report as good for a minute. `msSince` is the age of the observation; the broker adds its own relay delay before passing it on. diff --git a/services/ollama-proxy/proxy.go b/services/ollama-proxy/proxy.go index ad4ac7a2..c491be84 100644 --- a/services/ollama-proxy/proxy.go +++ b/services/ollama-proxy/proxy.go @@ -30,6 +30,7 @@ import ( "nvpair-shared/clustertrust" "nvpair-shared/cors" "nvpair-shared/errors" + "nvpair-shared/inferstats" "nvpair-shared/netmon" "nvpair-shared/netpick" "nvpair-shared/nodeactivity" @@ -187,6 +188,11 @@ type Workload struct { CompletedAt *int64 `json:"completedAt"` Error *string `json:"error"` RequesterID *string `json:"requesterId"` + // Stats carries the inference statistics measured for the terminal + // (completed/errored) transition: token counts, decode throughput and time + // to first token. Additive: omitted while the workload is running and when + // nothing could be measured (see nvpair-shared/inferstats). + Stats *inferstats.Stats `json:"stats,omitempty"` } // workloadParams is the params envelope for a workload:* notification @@ -243,6 +249,12 @@ type statusCapture struct { // discovery cannot obtain for itself while the node is too busy to answer a // probe. Called on the reverse proxy's copy goroutine, so it must be cheap. upstreamAlive func() + + // tap, when set, observes every body byte written to the client so the + // workload's terminal event can carry token counts and throughput. Like + // upstreamAlive it is set at the commit point, and only for a successful + // inference response, so the proxy's own error bodies are never measured. + tap *inferstats.Tap } // Unwrap exposes the underlying ResponseWriter so http.ResponseController can @@ -282,6 +294,9 @@ func (sc *statusCapture) Write(b []byte) (int, error) { if err == nil && sc.upstreamAlive != nil { sc.upstreamAlive() } + if err == nil && sc.tap != nil { + sc.tap.Observe(b) + } return n, err } @@ -1221,6 +1236,7 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { finalStatus int started bool wl *Workload + tap *inferstats.Tap ) // Emit workload:started up front, the moment we begin forwarding, naming @@ -1245,6 +1261,10 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { StartedAt: &createdMs, } p.emitWorkload(workloadStartedMethod, *wl) + // The tap measures the response body for this workload's stats. It is + // allocated here, before the disconnect watcher goroutine starts, and + // armed only once a candidate commits (ModifyResponse below). + tap = inferstats.NewTap(start) } // The terminal workload transition (completed/errored) can be reached from @@ -1271,6 +1291,11 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { if errMsg != "" { wl.Error = &errMsg } + // Stats ride the terminal event only; a failed stream still reports + // what was generated before it broke. + if tap != nil { + wl.Stats = tap.Finish() + } snapshot := *wl wlMu.Unlock() method := workloadCompletedMethod @@ -1360,6 +1385,13 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // came from the node. Same goroutine as the body copy, so no // synchronization is needed. sc.upstreamAlive = func() { p.reportActivity(cand.id) } + // Measure the body only for a successful inference response: an + // error body carries no tokens, and a control request has no + // workload to attach stats to. + if tap != nil && resp.StatusCode < http.StatusBadRequest { + tap.Arm(resp.Header) + sc.tap = tap + } // The engine may enforce its own origin policy (Ollama's // OLLAMA_ORIGINS). Honor it: overwriting a declared // Access-Control-Allow-Origin would silently widen the user's diff --git a/services/ollama-proxy/stats_test.go b/services/ollama-proxy/stats_test.go new file mode 100644 index 00000000..7fde99ff --- /dev/null +++ b/services/ollama-proxy/stats_test.go @@ -0,0 +1,151 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// recordedWorkload decodes the workloadInfo of the last frame rec captured for +// method, failing the test when none was emitted. +func recordedWorkload(t *testing.T, rec *recRW, method string) Workload { + t.Helper() + rec.mu.Lock() + raw := append([]byte(nil), rec.b...) + rec.mu.Unlock() + var found []byte + for _, line := range bytes.Split(raw, []byte{'\n'}) { + if bytes.Contains(line, []byte(`"method":"`+method+`"`)) { + found = line + } + } + if found == nil { + t.Fatalf("no %s frame recorded; frames:\n%s", method, raw) + } + var frame struct { + Params struct { + WorkloadInfo Workload `json:"workloadInfo"` + } `json:"params"` + } + if err := json.Unmarshal(found, &frame); err != nil { + t.Fatalf("decode %s frame: %v", method, err) + } + return frame.Params.WorkloadInfo +} + +// streamChunks serves body chunks as a flushed stream with a small gap between +// them, the way an engine streams tokens. +func streamChunks(contentType string, chunks ...string) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", contentType) + w.WriteHeader(http.StatusOK) + for _, c := range chunks { + io.WriteString(w, c) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + time.Sleep(2 * time.Millisecond) + } + }) +} + +// TestHandleHTTP_CompletedCarriesStats: a streamed inference response whose +// final chunk reports usage must surface engine-exact token counts and a +// throughput figure on workload:completed — and nothing on workload:started, +// which fires before a byte has flowed. +func TestHandleHTTP_CompletedCarriesStats(t *testing.T) { + upstream := httptest.NewServer(streamChunks("text/event-stream", + "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n", + "data: [DONE]\n\n", + )) + defer upstream.Close() + + rec := &recRW{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "node-a", upstream.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 11434) + + rr := httptest.NewRecorder() + p.handleHTTP(rr, httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"llama","stream":true}`))) + if rr.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", rr.Code) + } + if !strings.Contains(rr.Body.String(), "[DONE]") { + t.Fatalf("client did not receive the full stream:\n%s", rr.Body.String()) + } + + if started := recordedWorkload(t, rec, "workload:started"); started.Stats != nil { + t.Fatalf("workload:started carried stats %+v before any byte flowed", started.Stats) + } + done := recordedWorkload(t, rec, "workload:completed") + if done.Stats == nil { + t.Fatal("workload:completed carried no stats") + } + s := done.Stats + if s.PromptTokens != 5 || s.CompletionTokens != 2 || s.Estimated { + t.Fatalf("stats = %+v, want engine-reported prompt=5 completion=2", s) + } + if s.TokensPerSecond <= 0 { + t.Fatalf("stats = %+v, want a positive tokensPerSecond", s) + } +} + +// TestHandleHTTP_OllamaNativeStreamStats: Ollama's own API reports counts on +// the terminal "done":true line, so an /api/chat stream yields engine-exact +// counts with throughput taken from the engine's eval_duration. +func TestHandleHTTP_OllamaNativeStreamStats(t *testing.T) { + upstream := httptest.NewServer(streamChunks("application/x-ndjson", + `{"model":"llama","message":{"role":"assistant","content":"Hi"},"done":false}`+"\n", + `{"model":"llama","message":{"role":"assistant","content":" there"},"done":false}`+"\n", + `{"model":"llama","message":{"role":"assistant","content":""},"done_reason":"stop","done":true,"total_duration":900000000,"prompt_eval_count":9,"eval_count":40,"eval_duration":500000000}`+"\n", + )) + defer upstream.Close() + + rec := &recRW{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "node-a", upstream.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 11434) + + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/api/chat", + strings.NewReader(`{"model":"llama"}`))) + + s := recordedWorkload(t, rec, "workload:completed").Stats + if s == nil || s.Estimated || s.PromptTokens != 9 || s.CompletionTokens != 40 || s.TokensPerSecond != 80 { + t.Fatalf("stats = %+v, want prompt=9 completion=40 tokensPerSecond=80 from the done line", s) + } +} + +// TestHandleHTTP_ErrorBodyNotMeasured: an upstream error body is never mined for +// tokens, even when it happens to contain usage-looking fields. +func TestHandleHTTP_ErrorBodyNotMeasured(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + io.WriteString(w, `{"error":"boom","usage":{"prompt_tokens":1,"completion_tokens":99}}`) + })) + defer upstream.Close() + + rec := &recRW{} + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "node-a", upstream.URL, "llama")) + p := NewProxy(NewCodec(rec), disc, 11434) + + p.handleHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"llama"}`))) + + if failed := recordedWorkload(t, rec, "workload:errored"); failed.Stats != nil { + t.Fatalf("workload:errored carried stats %+v from an error body", failed.Stats) + } +} diff --git a/services/shared/inferstats/inferstats.go b/services/shared/inferstats/inferstats.go new file mode 100644 index 00000000..0d82c7c7 --- /dev/null +++ b/services/shared/inferstats/inferstats.go @@ -0,0 +1,257 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Package inferstats derives per-request inference statistics — token counts, +// decode throughput, and time to first token — from an inference response as +// the proxy streams it to the client. Both inference proxies front engines that +// speak the OpenAI-compatible API (and Ollama's native one), so keeping one +// implementation is what stops the two from drifting apart. +// +// The design goal is negligible overhead on the hot path. A Tap never buffers +// or parses the whole body: per body write it records a timestamp, counts +// stream events, and keeps only the last few kilobytes. Every engine puts its +// usage report at the END of the response — the OpenAI `usage` object trails +// `choices` in a non-streaming reply and rides the final chunk of a stream +// (`stream_options.include_usage`), and Ollama's native `eval_count` / +// `eval_duration` sit on the terminal `"done":true` line — so that tail is all +// Finish needs. Token counts are pulled from the tail with a byte search, not a +// JSON decode, so a truncated leading object (a long non-streaming completion) +// is not a problem. +// +// Prompts, messages, and response text are never retained beyond the rolling +// tail, which is discarded at Finish; only numbers leave this package. +package inferstats + +import ( + "bytes" + "math" + "net/http" + "strconv" + "strings" + "sync" + "time" +) + +// Stats is the additive `stats` object carried on a workload's terminal event. +// Every field is optional: a reader must treat an absent field as "not +// measured", never as zero. A field the engine did not report is omitted rather +// than invented. +type Stats struct { + // PromptTokens is the prompt size as reported by the engine. + PromptTokens int64 `json:"promptTokens,omitempty"` + // CompletionTokens is the number of tokens generated. Reported by the + // engine when its response carried a usage object; otherwise counted from + // stream chunks, in which case Estimated is set. + CompletionTokens int64 `json:"completionTokens,omitempty"` + // TokensPerSecond is decode throughput: CompletionTokens over generation + // time. Generation time is the engine's own measurement when it reports one + // (Ollama's eval_duration), else first-to-last body byte for a stream, else + // the whole request for a non-streaming reply. Rounded to one decimal. + TokensPerSecond float64 `json:"tokensPerSecond,omitempty"` + // TTFTMs is time to first token: request start to the first response body + // byte. Headers are deliberately not the boundary — a streaming engine sends + // them before prefill finishes. + TTFTMs int64 `json:"ttftMs,omitempty"` + // Estimated marks CompletionTokens (and so TokensPerSecond) as counted from + // stream chunks rather than reported by the engine. A chunk usually carries + // one token, but an engine may batch several, so treat the count as + // approximate. + Estimated bool `json:"estimated,omitempty"` +} + +// tailSize bounds the bytes retained from the end of the body. An OpenAI usage +// chunk with its details objects plus the `[DONE]` sentinel is a few hundred +// bytes; Ollama's terminal line likewise. 4 KiB leaves ample slack for trailing +// engine-specific fields (vLLM's `prompt_logprobs`, LM Studio's `stats`). +const tailSize = 4096 + +// Tap observes one response body. It is safe for concurrent use: the reverse +// proxy's copy goroutine calls Observe while a disconnect watcher may call +// Finish. The mutex is uncontended in practice, so it costs a few nanoseconds +// per body write. +type Tap struct { + mu sync.Mutex + start time.Time + armed bool + // opaque is set when the body is content-encoded (a client that asked for + // gzip and an engine that obliged). Only timing is usable then; counting + // or searching compressed bytes would be noise. + opaque bool + // sse is set for a text/event-stream body, where a stream event is a + // `data:` line; otherwise (Ollama's NDJSON) an event is a line. + sse bool + firstAt time.Time + lastAt time.Time + events int64 + atLineStart bool + tail []byte +} + +// NewTap returns a Tap for a request that began at start. It observes nothing +// until Arm is called. +func NewTap(start time.Time) *Tap { + return &Tap{start: start, atLineStart: true} +} + +// Arm starts observation once the proxy has committed to an upstream response, +// using the response headers to decide how to read the body. Call it only for a +// successful inference response — an error body carries no tokens. +func (t *Tap) Arm(h http.Header) { + t.mu.Lock() + defer t.mu.Unlock() + t.armed = true + enc := strings.ToLower(strings.TrimSpace(h.Get("Content-Encoding"))) + t.opaque = enc != "" && enc != "identity" + t.sse = strings.HasPrefix(strings.ToLower(h.Get("Content-Type")), "text/event-stream") + if !t.opaque { + t.tail = make([]byte, 0, tailSize) + } +} + +// Observe records one body write. It must be cheap: it runs on the reverse +// proxy's copy goroutine between the upstream read and the client write. +func (t *Tap) Observe(p []byte) { + if len(p) == 0 { + return + } + t.mu.Lock() + defer t.mu.Unlock() + if !t.armed { + return + } + now := time.Now() + if t.firstAt.IsZero() { + t.firstAt = now + } + t.lastAt = now + if t.opaque { + return + } + if t.sse { + if t.atLineStart && bytes.HasPrefix(p, []byte("data:")) { + t.events++ + } + t.events += int64(bytes.Count(p, []byte("\ndata:"))) + } else { + t.events += int64(bytes.Count(p, []byte{'\n'})) + } + t.atLineStart = p[len(p)-1] == '\n' + t.retain(p) +} + +// retain appends p to the rolling tail, keeping only the last tailSize bytes. +// The buffer is fixed-capacity, so this never allocates after Arm. +func (t *Tap) retain(p []byte) { + if len(p) >= tailSize { + t.tail = t.tail[:tailSize] + copy(t.tail, p[len(p)-tailSize:]) + return + } + if overflow := len(t.tail) + len(p) - tailSize; overflow > 0 { + copy(t.tail, t.tail[overflow:]) + t.tail = t.tail[:len(t.tail)-overflow] + } + t.tail = append(t.tail, p...) +} + +// Finish derives the statistics and releases the tail. It returns nil when the +// Tap was never armed or no body byte was observed, so a caller can attach the +// result directly to an omitempty field. +func (t *Tap) Finish() *Stats { + t.mu.Lock() + defer t.mu.Unlock() + if !t.armed || t.firstAt.IsZero() { + return nil + } + s := &Stats{TTFTMs: t.firstAt.Sub(t.start).Milliseconds()} + tail := t.tail + t.tail = nil + + var genDur time.Duration + switch { + case t.opaque: + // Compressed body: nothing to count. Timing alone is still worth + // reporting. + default: + if n, ok := lastNumber(tail, `"completion_tokens":`); ok { + // OpenAI-compatible usage object (non-streaming reply, or the + // final chunk when the client asked for stream_options.include_usage). + s.CompletionTokens = n + if p, ok := lastNumber(tail, `"prompt_tokens":`); ok { + s.PromptTokens = p + } + } else if n, ok := lastNumber(tail, `"eval_count":`); ok { + // Ollama native terminal line. eval_duration is the engine's own + // decode-time measurement in nanoseconds — more precise than our + // wall clock, so prefer it. + s.CompletionTokens = n + if p, ok := lastNumber(tail, `"prompt_eval_count":`); ok { + s.PromptTokens = p + } + if d, ok := lastNumber(tail, `"eval_duration":`); ok && d > 0 { + genDur = time.Duration(d) + } + } else if t.events > 0 { + // No usage report: approximate from stream events, discounting the + // SSE `[DONE]` sentinel, which carries no token. + n := t.events + if t.sse && bytes.Contains(tail, []byte("[DONE]")) { + n-- + } + if n > 0 { + s.CompletionTokens = n + s.Estimated = true + } + } + } + + if s.CompletionTokens > 0 { + if genDur == 0 { + if t.events > 1 { + // Streamed: the first byte marks the end of prefill, so + // first-to-last is decode time. + genDur = t.lastAt.Sub(t.firstAt) + } else { + // Non-streaming: the whole reply arrived at once, so the only + // honest denominator is the full request. + genDur = t.lastAt.Sub(t.start) + } + } + if genDur > 0 { + tps := float64(s.CompletionTokens) / genDur.Seconds() + s.TokensPerSecond = math.Round(tps*10) / 10 + } + } + return s +} + +// lastNumber finds the LAST occurrence of key (a quoted JSON key including its +// trailing colon, e.g. `"completion_tokens":`) and parses the non-negative +// integer that follows it. The last occurrence is the right one: every engine +// emits its usage report after the generated content, so a model that happens +// to write the same key in its output cannot shadow it. Including the closing +// quote and colon in the key keeps `"completion_tokens":` from matching inside +// `"completion_tokens_details":`, and the opening quote keeps `"eval_count":` +// from matching inside `"prompt_eval_count":`. +func lastNumber(b []byte, key string) (int64, bool) { + i := bytes.LastIndex(b, []byte(key)) + if i < 0 { + return 0, false + } + j := i + len(key) + for j < len(b) && (b[j] == ' ' || b[j] == '\t') { + j++ + } + k := j + for k < len(b) && b[k] >= '0' && b[k] <= '9' { + k++ + } + if k == j || k-j > 18 { + return 0, false + } + n, err := strconv.ParseInt(string(b[j:k]), 10, 64) + if err != nil { + return 0, false + } + return n, true +} diff --git a/services/shared/inferstats/inferstats_test.go b/services/shared/inferstats/inferstats_test.go new file mode 100644 index 00000000..bd8b8f69 --- /dev/null +++ b/services/shared/inferstats/inferstats_test.go @@ -0,0 +1,230 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package inferstats + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + "time" +) + +func headers(contentType, encoding string) http.Header { + h := http.Header{} + h.Set("Content-Type", contentType) + if encoding != "" { + h.Set("Content-Encoding", encoding) + } + return h +} + +// feed arms a tap and observes each chunk in order, sleeping gap between +// chunks so first-to-last timing is measurable. +func feed(t *testing.T, h http.Header, gap time.Duration, chunks ...string) *Stats { + t.Helper() + tap := NewTap(time.Now()) + tap.Arm(h) + for i, c := range chunks { + if i > 0 && gap > 0 { + time.Sleep(gap) + } + tap.Observe([]byte(c)) + } + return tap.Finish() +} + +func TestFinish_NilBeforeArmOrBody(t *testing.T) { + tap := NewTap(time.Now()) + tap.Observe([]byte("ignored: not armed")) + if got := tap.Finish(); got != nil { + t.Fatalf("unarmed tap produced stats %+v", got) + } + tap = NewTap(time.Now()) + tap.Arm(headers("application/json", "")) + if got := tap.Finish(); got != nil { + t.Fatalf("armed tap with no body produced stats %+v", got) + } +} + +func TestFinish_OpenAINonStreaming(t *testing.T) { + body := `{"id":"chatcmpl-1","object":"chat.completion","model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],` + + `"usage":{"prompt_tokens":12,"completion_tokens":34,"total_tokens":46,"completion_tokens_details":{"reasoning_tokens":5}},"prompt_logprobs":null}` + s := feed(t, headers("application/json", ""), 0, body) + if s == nil { + t.Fatal("no stats") + } + if s.PromptTokens != 12 || s.CompletionTokens != 34 || s.Estimated { + t.Fatalf("got %+v, want prompt=12 completion=34 estimated=false", s) + } + if s.TokensPerSecond <= 0 { + t.Fatalf("tokensPerSecond not derived for a non-streaming reply: %+v", s) + } +} + +func TestFinish_OpenAIStreamWithUsageChunk(t *testing.T) { + chunks := []string{ + "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":2,\"total_tokens\":9}}\n\n", + "data: [DONE]\n\n", + } + s := feed(t, headers("text/event-stream", ""), 2*time.Millisecond, chunks...) + if s == nil { + t.Fatal("no stats") + } + if s.PromptTokens != 7 || s.CompletionTokens != 2 || s.Estimated { + t.Fatalf("got %+v, want engine-reported prompt=7 completion=2", s) + } + if s.TokensPerSecond <= 0 { + t.Fatalf("tokensPerSecond not derived for a stream: %+v", s) + } +} + +func TestFinish_OpenAIStreamWithoutUsage_Estimates(t *testing.T) { + var chunks []string + for i := 0; i < 10; i++ { + chunks = append(chunks, fmt.Sprintf("data: {\"choices\":[{\"delta\":{\"content\":\"t%d\"}}]}\n\n", i)) + } + chunks = append(chunks, "data: [DONE]\n\n") + s := feed(t, headers("text/event-stream", ""), time.Millisecond, chunks...) + if s == nil { + t.Fatal("no stats") + } + if !s.Estimated || s.CompletionTokens != 10 { + t.Fatalf("got %+v, want estimated completion=10 ([DONE] discounted)", s) + } + if s.PromptTokens != 0 { + t.Fatalf("prompt tokens invented without a usage object: %+v", s) + } +} + +// A `data:` line split across two body writes at the newline boundary must +// still count once — the common coalescing pattern is a write ending in "\n" +// followed by one starting with "data:". +func TestObserve_SSEEventSplitAtLineStart(t *testing.T) { + s := feed(t, headers("text/event-stream", ""), 0, + "data: {\"a\":1}\n\n", "data: {\"b\":2}\n", "\ndata: {\"c\":3}\n\n", "data: [DONE]\n\n") + if s == nil || s.CompletionTokens != 3 || !s.Estimated { + t.Fatalf("got %+v, want estimated completion=3", s) + } +} + +func TestFinish_OllamaNativeStream(t *testing.T) { + chunks := []string{ + `{"model":"llama3","message":{"role":"assistant","content":"Hi"},"done":false}` + "\n", + `{"model":"llama3","message":{"role":"assistant","content":" there"},"done":false}` + "\n", + `{"model":"llama3","message":{"role":"assistant","content":""},"done_reason":"stop","done":true,` + + `"total_duration":1500000000,"load_duration":100000000,"prompt_eval_count":9,"prompt_eval_duration":200000000,` + + `"eval_count":40,"eval_duration":500000000}` + "\n", + } + s := feed(t, headers("application/x-ndjson", ""), 0, chunks...) + if s == nil { + t.Fatal("no stats") + } + if s.PromptTokens != 9 || s.CompletionTokens != 40 || s.Estimated { + t.Fatalf("got %+v, want prompt=9 completion=40 from the terminal line", s) + } + // 40 tokens over the engine's own 0.5s eval_duration. + if s.TokensPerSecond != 80 { + t.Fatalf("tokensPerSecond = %v, want 80 (eval_count/eval_duration)", s.TokensPerSecond) + } +} + +// The usage report is always at the end, so a long non-streaming completion +// whose head has scrolled out of the tail still yields exact counts. +func TestFinish_LongBodyBeyondTail(t *testing.T) { + content := strings.Repeat("lorem ipsum ", 2000) // ~24 KiB, well past tailSize + body := `{"choices":[{"message":{"content":"` + content + `"}}],"usage":{"prompt_tokens":3,"completion_tokens":5000,"total_tokens":5003}}` + s := feed(t, headers("application/json", ""), 0, body) + if s == nil || s.CompletionTokens != 5000 || s.PromptTokens != 3 { + t.Fatalf("got %+v, want completion=5000 prompt=3", s) + } +} + +// Generated text that happens to contain a usage-looking key must not shadow +// the engine's real report, which always comes later in the body. +func TestFinish_ContentMentioningUsageKeyIsIgnored(t *testing.T) { + body := `{"choices":[{"message":{"content":"the API returns \"completion_tokens\":999 in usage"}}],"usage":{"prompt_tokens":1,"completion_tokens":20,"total_tokens":21}}` + s := feed(t, headers("application/json", ""), 0, body) + if s == nil || s.CompletionTokens != 20 { + t.Fatalf("got %+v, want completion=20 from the trailing usage object", s) + } +} + +func TestFinish_CompressedBodyReportsTimingOnly(t *testing.T) { + s := feed(t, headers("text/event-stream", "gzip"), 0, "\x1f\x8b\x08 not really gzip data: data: data:") + if s == nil { + t.Fatal("no stats for a compressed body; timing should still be reported") + } + if s.CompletionTokens != 0 || s.PromptTokens != 0 || s.TokensPerSecond != 0 || s.Estimated { + t.Fatalf("token fields derived from compressed bytes: %+v", s) + } +} + +func TestFinish_TTFTMeasuredFromRequestStart(t *testing.T) { + tap := NewTap(time.Now().Add(-250 * time.Millisecond)) + tap.Arm(headers("application/json", "")) + tap.Observe([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":1}}`)) + s := tap.Finish() + if s == nil || s.TTFTMs < 250 { + t.Fatalf("ttftMs = %v, want >= 250 (first body byte relative to request start)", s) + } +} + +func TestStats_JSONOmitsUnmeasuredFields(t *testing.T) { + b, err := json.Marshal(&Stats{TTFTMs: 12}) + if err != nil { + t.Fatal(err) + } + if string(b) != `{"ttftMs":12}` { + t.Fatalf("got %s, want only the measured field", b) + } + b, _ = json.Marshal(&Stats{CompletionTokens: 3, TokensPerSecond: 1.5, Estimated: true}) + if !bytes.Contains(b, []byte(`"estimated":true`)) || bytes.Contains(b, []byte(`"promptTokens"`)) { + t.Fatalf("got %s", b) + } +} + +func TestRetain_KeepsExactlyTheTail(t *testing.T) { + tap := NewTap(time.Now()) + tap.Arm(headers("application/json", "")) + // Many small writes, then one oversized write; the tail must end with the + // most recent bytes in both regimes and never exceed tailSize. + for i := 0; i < 100; i++ { + tap.Observe(bytes.Repeat([]byte{'a'}, 100)) + } + if len(tap.tail) != tailSize { + t.Fatalf("tail len %d after overflow, want %d", len(tap.tail), tailSize) + } + big := append(bytes.Repeat([]byte{'b'}, tailSize*2), []byte("END")...) + tap.Observe(big) + if len(tap.tail) != tailSize || !bytes.HasSuffix(tap.tail, []byte("END")) || tap.tail[0] != 'b' { + t.Fatalf("oversized write not retained as its last %d bytes", tailSize) + } +} + +func TestLastNumber(t *testing.T) { + cases := []struct { + in, key string + want int64 + ok bool + }{ + {`{"completion_tokens":34}`, `"completion_tokens":`, 34, true}, + {`{"completion_tokens": 34}`, `"completion_tokens":`, 34, true}, + {`{"completion_tokens_details":{"x":1}}`, `"completion_tokens":`, 0, false}, + {`{"prompt_eval_count":9}`, `"eval_count":`, 0, false}, + {`{"eval_count":}`, `"eval_count":`, 0, false}, + {`{"a":1}{"a":2}`, `"a":`, 2, true}, + } + for _, c := range cases { + got, ok := lastNumber([]byte(c.in), c.key) + if got != c.want || ok != c.ok { + t.Errorf("lastNumber(%q, %q) = %d,%v want %d,%v", c.in, c.key, got, ok, c.want, c.ok) + } + } +} diff --git a/services/versions.json b/services/versions.json index 29d8c230..b5e456f1 100644 --- a/services/versions.json +++ b/services/versions.json @@ -1,10 +1,10 @@ { "$comment": "Single source of truth for all version numbers. See VERSIONING.md for bump rules.", - "product": "0.91.7", - "installer": "0.91.7", + "product": "0.92.0", + "installer": "0.92.0", "components": { - "ollama-proxy": "0.26.2", - "lmstudio-proxy": "0.16.2", + "ollama-proxy": "0.27.0", + "lmstudio-proxy": "0.17.0", "nvpair-node-info": "0.13.3", "nvpair-node-scanner": "0.20.3", "nvpair-manual-nodes": "0.11.1", @@ -15,6 +15,6 @@ "nvpair-engine-manager": "0.17.4", "nvpair-cluster-manager": "1.1.4", "nvpair-job-scheduler": "0.4.1", - "nvpair-tui": "0.7.2" + "nvpair-tui": "0.8.0" } }