Skip to content

fix: record streaming token usage in v2 profile stats#21

Open
himorishige wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
himorishige:fix/v2-streaming-usage-stats
Open

fix: record streaming token usage in v2 profile stats#21
himorishige wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
himorishige:fix/v2-streaming-usage-stats

Conversation

@himorishige

Copy link
Copy Markdown
Contributor

What

Adds a native v2 stream tap (switchyard-components-v2/src/stats/stream_tap.rs) and wires it into all five v2 profiles (passthrough, random-routing, llm-routing, cascade, latency-service), so /v1/routing/stats records token usage for streaming responses. Fixes the 108 calls / 0 tokens shape reported in #18.

Root cause

All v2 profiles read usage from the buffered body only:

let usage = response.body().map(usage_from_body).unwrap_or_default();

ChatResponse::body() returns None for the three stream variants, so streaming usage silently defaulted to zero while record_success still counted the call.

Approach

Per the guidance in #18, this does not wire in the v1 StatsResponseProcessor (it depends on ProxyContext state that v2 doesn't have). Instead:

  • UsageAttribution captures what v2 already holds locally before the stream is consumed: the stats accumulator, model, tier, request-start Instant, and backend latency.
  • record_usage_or_tap_stream(response, attribution) records buffered bodies immediately and wraps the three stream variants so usage is committed once when the stream completes.
  • Usage parsing reuses the existing switchyard-components::stats parsers (openai_chat_usage_from_stream_event, openai_responses_usage_from_stream_event, AnthropicStreamUsage).
  • Usage is recorded via record_usage_after_success_attribution, so the tier call counted by record_success is not double-counted.

Profile wiring stays local: passthrough / random-routing / latency-service replace the inline usage read; llm-routing / cascade change record_success to take the response by value and return the (possibly wrapped) response. The cascade enable_stats=false guard is preserved (no tap when stats are disabled).

Semantics notes (for review)

  • Last usage frame wins, committed at stream completion. Spec-compliant upstreams send usage only on the final chunk; cumulative emitters (e.g. vLLM stream_options.continuous_usage_stats) update every chunk, so committing the first frame would under-count.
  • Usage-less streams still record latency. Streams without a usage frame commit zero usage at completion, so total-latency samples aren't lost (matches the previous v2 behavior of recording zero usage).
  • Total latency spans until the stream completes, matching the v1 stream taps. This includes client read time for streams — happy to change if you'd prefer header-time latency.
  • Client-disconnected streams commit nothing — same limitation as the v1 taps.
  • Classifier-path usage is untouched (llm-routing / cascade classifiers are always non-streaming, so record_classifier_usage still reads the buffered body).

Testing

  • Unit tests for the tap: buffered, OpenAI Chat / Responses / Anthropic streams, cumulative usage frames, usage-less streams (latency-sample preservation).
  • Per-profile streaming tests for all five profiles, asserting usage is zero before the stream is consumed and recorded after.
  • cargo fmt --all --check, cargo clippy --workspace --all-targets -- -D warnings, and cargo test --workspace all pass (toolchain 1.94.0).
  • Live re-run of the /v1/routing/stats records zero token usage for streaming responses (v2 profile server) #18 reproduction (passthrough profile → OpenRouter, Linux aarch64):
Request before (v0.1.0) after this fix
non-streaming +49 +43
"stream": true +0 +50
"stream": true + stream_options.include_usage +0 +50

(deltas are total_tokens in /v1/routing/stats per request; absolute counts vary with model output length)

Closes #18

Signed-off-by: Hiroshi Morishige <hiroshi.morishige@gmail.com>
@ayushag-nv

Copy link
Copy Markdown
Contributor

@himorishige Thanks for putting up the PR. Can you please rebase and resolve the merge conflicts once

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

/v1/routing/stats records zero token usage for streaming responses (v2 profile server)

2 participants