From c8aaaf9722be26786d96d9f3f12579abc3ac9612 Mon Sep 17 00:00:00 2001 From: Can GULDOGAN Date: Wed, 9 Sep 2026 01:36:49 +0100 Subject: [PATCH] Show per-job inference stats (tok/s, tokens, TTFT) in the Jobs list The proxies now measure each inference response as it streams to the client and attach an additive `stats` object to the workload's terminal event: prompt/completion tokens, decode throughput and time to first token. Token counts come from the engine's own usage report (OpenAI `usage`, Ollama `eval_count`/`eval_duration`); a stream without one is counted from its chunks and flagged `estimated`. Overhead is a timestamp, a byte count and a copy into a fixed 4 KiB tail per body write; the body is never buffered or JSON-decoded. Only the trailing bytes are inspected, since every engine puts its usage report last. Error bodies and compressed bodies are never mined for tokens, and the stream content is discarded at the terminal event. - services/shared/inferstats: the tap and extraction, with unit tests - ollama-proxy, lmstudio-proxy: arm the tap at the commit point, attach stats in emitTerminal; proxy-level tests for SSE with and without usage, Ollama NDJSON, and error bodies - desktop: WorkloadStats type, Electron parsing, a stats row on the job card with a formatter and unit test - nvpair-tui: TOK/S column - READMEs and versions.json (ollama-proxy 0.27.0, lmstudio-proxy 0.17.0, nvpair-tui 0.8.0, product 0.92.0) Co-Authored-By: Claude Fable 5.1 Signed-off-by: Can GULDOGAN --- .../electron/service-bridge/modular-state.ts | 25 +- desktop/src/shared/types/workloads.ts | 31 +++ .../components/Workloads/WorkloadItemCard.tsx | 11 + desktop/src/ui/utils/format-workload-stats.ts | 39 +++ .../modular/format-workload-stats.test.ts | 41 +++ services/lmstudio-proxy/README.md | 6 + services/lmstudio-proxy/proxy.go | 32 +++ services/lmstudio-proxy/stats_test.go | 152 +++++++++++ services/nvpair-tui/ui/workloads.go | 37 ++- services/nvpair-ui-broker/README.md | 2 +- services/nvpair-workload-manager/README.md | 1 + services/ollama-proxy/README.md | 6 + services/ollama-proxy/proxy.go | 32 +++ services/ollama-proxy/stats_test.go | 151 ++++++++++ services/shared/inferstats/inferstats.go | 257 ++++++++++++++++++ services/shared/inferstats/inferstats_test.go | 230 ++++++++++++++++ services/versions.json | 10 +- 17 files changed, 1053 insertions(+), 10 deletions(-) create mode 100644 desktop/src/ui/utils/format-workload-stats.ts create mode 100644 desktop/tests/modular/format-workload-stats.test.ts create mode 100644 services/lmstudio-proxy/stats_test.go create mode 100644 services/ollama-proxy/stats_test.go create mode 100644 services/shared/inferstats/inferstats.go create mode 100644 services/shared/inferstats/inferstats_test.go 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" } }