diff --git a/rust/src/bench/src/backends/mod.rs b/rust/src/bench/src/backends/mod.rs index b5400c5d0f4..da23eb46a0a 100644 --- a/rust/src/bench/src/backends/mod.rs +++ b/rust/src/bench/src/backends/mod.rs @@ -52,6 +52,21 @@ pub struct ChatDelta { #[derive(Deserialize)] pub struct ChunkUsage { pub completion_tokens: Option, + pub prompt_tokens: Option, +} + +impl ChunkUsage { + /// Record the server-reported token counts on `output`. A zero + /// `prompt_tokens` keeps the client estimate: some OpenAI-compatible + /// servers send placeholder zeros in usage chunks. + pub fn apply(&self, output: &mut RequestFuncOutput) { + if let Some(ct) = self.completion_tokens { + output.output_tokens = ct as usize; + } + if let Some(pt) = self.prompt_tokens.filter(|&pt| pt > 0) { + output.prompt_len = pt as usize; + } + } } use crate::cli::BackendKind; @@ -246,6 +261,21 @@ pub fn build_headers( mod tests { use super::*; + /// A placeholder zero `prompt_tokens` must not replace the client estimate. + #[test] + fn chunk_usage_ignores_zero_prompt_tokens() { + let mut output = RequestFuncOutput { + prompt_len: 100, + ..Default::default() + }; + ChunkUsage { + completion_tokens: Some(3), + prompt_tokens: Some(0), + } + .apply(&mut output); + assert_eq!((output.prompt_len, output.output_tokens), (100, 3)); + } + #[test] fn test_update_server_metrics_converts_ms_to_seconds() { let chunk: CompletionChunk = serde_json::from_value(serde_json::json!({ diff --git a/rust/src/bench/src/backends/openai_chat.rs b/rust/src/bench/src/backends/openai_chat.rs index 30c6093296d..36ce961abf8 100644 --- a/rust/src/bench/src/backends/openai_chat.rs +++ b/rust/src/bench/src/backends/openai_chat.rs @@ -130,10 +130,8 @@ impl OpenAIChatBackend { } // Separate `if` (not `else if`) — Dynamo may send // both choices and usage in the same chunk. - if let Some(ref usage) = data.usage - && let Some(ct) = usage.completion_tokens - { - output.output_tokens = ct as usize; + if let Some(usage) = &data.usage { + usage.apply(&mut output); } } } @@ -363,7 +361,11 @@ mod tests { /// Serve one streaming chat response over a raw socket: `num_tokens` /// content chunks back-to-back, then `usage_delay` of silence, then the /// choice-less usage chunk and `[DONE]`. Returns the endpoint URL. - fn spawn_sse_server(num_tokens: usize, usage_delay: std::time::Duration) -> String { + fn spawn_sse_server( + num_tokens: usize, + usage_delay: std::time::Duration, + prompt_tokens: Option, + ) -> String { use std::io::{BufRead, BufReader, Read, Write}; let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); @@ -393,8 +395,10 @@ mod tests { send("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n"); } std::thread::sleep(usage_delay); + let prompt_tokens = + prompt_tokens.map_or(String::new(), |pt| format!(",\"prompt_tokens\":{pt}")); send(&format!( - "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}}}}}\n\n" + "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}{prompt_tokens}}}}}\n\n" )); send("data: [DONE]\n\n"); }); @@ -411,7 +415,7 @@ mod tests { // real delay is needed here; there is no observable event to wait on. let usage_delay = std::time::Duration::from_millis(150); let input = RequestFuncInput { - api_url: spawn_sse_server(3, usage_delay), + api_url: spawn_sse_server(3, usage_delay, None), model: "test-model".to_string(), output_len: 3, ..Default::default() @@ -442,4 +446,42 @@ mod tests { // keys already set above must not be overridden by extra_body assert_eq!(v["stream"], true); } + + /// The server's `usage.prompt_tokens` replaces the client-side estimate, + /// so input-token totals count what the server actually processed + /// (chat template, multimodal tokens). Mirrors Python #38654. + #[tokio::test] + async fn test_usage_prompt_tokens_override_prompt_len() { + let input = RequestFuncInput { + api_url: spawn_sse_server(3, std::time::Duration::ZERO, Some(128)), + model: "test-model".to_string(), + prompt_len: 100, + output_len: 3, + ..Default::default() + }; + + let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap(); + + assert!(output.success, "{}", output.error); + assert_eq!(output.prompt_len, 128); + } + + /// Servers that omit `usage.prompt_tokens` keep the client-side estimate, + /// and `completion_tokens` is still recorded. + #[tokio::test] + async fn test_missing_prompt_tokens_keeps_client_prompt_len() { + let input = RequestFuncInput { + api_url: spawn_sse_server(3, std::time::Duration::ZERO, None), + model: "test-model".to_string(), + prompt_len: 100, + output_len: 3, + ..Default::default() + }; + + let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap(); + + assert!(output.success, "{}", output.error); + assert_eq!(output.prompt_len, 100); + assert_eq!(output.output_tokens, 3); + } } diff --git a/rust/src/bench/src/backends/openai_completions.rs b/rust/src/bench/src/backends/openai_completions.rs index 0b62d342bcd..04998819131 100644 --- a/rust/src/bench/src/backends/openai_completions.rs +++ b/rust/src/bench/src/backends/openai_completions.rs @@ -156,10 +156,8 @@ impl OpenAICompletionsBackend { } // Separate `if` (not `else if`) — Dynamo may send // both choices and usage in the same chunk. - if let Some(ref usage) = data.usage - && let Some(ct) = usage.completion_tokens - { - output.output_tokens = ct as usize; + if let Some(usage) = &data.usage { + usage.apply(&mut output); } } } @@ -201,3 +199,69 @@ impl OpenAICompletionsBackend { Ok(output) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Serve `body` as one streaming completions response over a raw socket. + /// Returns the endpoint URL. + fn spawn_sse_server(body: &'static str) -> String { + use std::io::{BufRead, BufReader, Read, Write}; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut content_length = 0; + loop { + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") { + content_length = v.trim().parse().unwrap(); + } + if line == "\r\n" { + break; + } + } + reader.read_exact(&mut vec![0; content_length]).unwrap(); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n\ + connection: close\r\n\r\n{body}" + ) + .as_bytes(), + ) + .unwrap(); + }); + format!("http://{addr}/v1/completions") + } + + /// The server's `usage.prompt_tokens` replaces the client-side estimate, + /// as in the chat backend. + #[tokio::test] + async fn test_usage_prompt_tokens_override_prompt_len() { + let input = RequestFuncInput { + api_url: spawn_sse_server(concat!( + "data: {\"choices\":[{\"text\":\"a\"}]}\n\n", + "data: {\"choices\":[{\"text\":\"b\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"completion_tokens\":2,\"prompt_tokens\":128}}\n\n", + "data: [DONE]\n\n", + )), + model: "test-model".to_string(), + prompt_len: 100, + output_len: 2, + ..Default::default() + }; + + let output = OpenAICompletionsBackend + .send_request(&input, &reqwest::Client::new()) + .await + .unwrap(); + + assert!(output.success, "{}", output.error); + assert_eq!((output.prompt_len, output.output_tokens), (128, 2)); + } +} diff --git a/rust/src/bench/src/backends/openai_responses.rs b/rust/src/bench/src/backends/openai_responses.rs index 2a998263613..3178b5f4b3b 100644 --- a/rust/src/bench/src/backends/openai_responses.rs +++ b/rust/src/bench/src/backends/openai_responses.rs @@ -27,6 +27,7 @@ struct ResponsesEventResponse { #[derive(serde::Deserialize)] struct ResponsesUsage { + input_tokens: Option, output_tokens: Option, } @@ -180,8 +181,14 @@ impl ResponsesStream { // with `status: "incomplete"`; the spec also has `response.incomplete`. Some("response.completed" | "response.incomplete") => { self.saw_terminal_event = true; - if let Some(n) = event.response.and_then(|r| r.usage?.output_tokens) { - self.output.output_tokens = n as usize; + if let Some(usage) = event.response.and_then(|r| r.usage) { + if let Some(n) = usage.output_tokens { + self.output.output_tokens = n as usize; + } + // Same rule as `ChunkUsage::apply`: a zero keeps the client estimate. + if let Some(n) = usage.input_tokens.filter(|&n| n > 0) { + self.output.prompt_len = n as usize; + } } } kind if matches!(kind, Some("error" | "response.failed")) || event.error.is_some() => { @@ -253,6 +260,21 @@ mod tests { stream.finish() } + /// The terminal event's `usage.input_tokens` replaces the client estimate, + /// as in Python's `async_request_openai_responses`. + #[test] + fn test_terminal_usage_input_tokens_override_prompt_len() { + let output = run(&[ + text_delta("Hi"), + event( + "response.completed", + ",\"response\":{\"usage\":{\"input_tokens\":128,\"output_tokens\":1}}", + ), + ]); + assert!(output.success, "{}", output.error); + assert_eq!((output.prompt_len, output.output_tokens), (128, 1)); + } + #[test] fn test_reasoning_starts_ttft_and_latency_ends_at_last_token() { let output = run(&[ diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 5833a3bdcea..117e9fe72d9 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -645,7 +645,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let total_output_tokens: usize = input_requests.iter().map(|r| r.expected_output_len).sum(); println!("Dry run stats:"); println!(" Total prompts: {}", input_requests.len()); - println!(" Total input tokens: {total_input_tokens}"); + println!(" Estimated total input tokens: {total_input_tokens}"); println!(" Total expected output tokens: {total_output_tokens}"); println!( " Avg input tokens: {:.1}", @@ -1118,7 +1118,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { (m, Vec::new()) } else { calculate_metrics( - &input_requests, &outputs, benchmark_duration, &config.selected_percentiles, @@ -1148,7 +1147,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { ) { let ss = steady_state::compute( &outputs, - &input_requests, &window, &config.selected_percentiles, config.backend.is_pooling(), diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs index 7e3fc77565c..b55ed1cbde6 100644 --- a/rust/src/bench/src/metrics/calculator.rs +++ b/rust/src/bench/src/metrics/calculator.rs @@ -3,7 +3,6 @@ use crate::backends::RequestFuncOutput; use crate::config::GoodputConfig; -use crate::datasets::SampleRequest; use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; use crate::multi_turn::ConversationOutput; @@ -27,7 +26,6 @@ fn log_failed_requests(outputs: &[RequestFuncOutput]) { /// /// Mirrors Python's `calculate_metrics()` from serve.py:392-599. pub fn calculate_metrics( - input_requests: &[SampleRequest], outputs: &[RequestFuncOutput], dur_s: f64, selected_percentiles: &[f64], @@ -44,7 +42,7 @@ pub fn calculate_metrics( let mut ttfts: Vec = Vec::new(); let mut e2els: Vec = Vec::new(); - for (i, output) in outputs.iter().enumerate() { + for output in outputs.iter() { if output.success { let output_len = if output.output_tokens > 0 { output.output_tokens @@ -57,7 +55,7 @@ pub fn calculate_metrics( }; actual_output_lens.push(output_len); - total_input += input_requests[i].prompt_len; + total_input += output.prompt_len; if output_len > 1 { let latency_minus_ttft = output.latency - output.ttft; @@ -283,7 +281,6 @@ fn percentiles_from_sorted(sorted: &[f64], percentiles: &[f64]) -> Vec<(f64, f64 /// /// Mirrors Python's `calculate_metrics_for_embeddings()` from serve.py. /// Key differences from `calculate_metrics()`: -/// - Uses `output.prompt_len` (server-reported) instead of `input_requests[i].prompt_len` /// - Only computes E2EL (no TTFT/TPOT/ITL since pooling is non-streaming) /// - `total_output` is 0 (pooling produces no output tokens) /// - `total_token_throughput` = `total_input / dur_s` (input tokens only) @@ -380,7 +377,7 @@ pub fn calculate_embedding_metrics( /// Calculate multi-turn benchmark metrics with per-turn breakdown. /// -/// 1. Flatten all turns into SampleRequest/RequestFuncOutput pairs → overall metrics +/// 1. Flatten all turns' outputs → overall metrics /// 2. Group by turn_index → per-turn metrics /// 3. Conversation-level stats pub fn calculate_multi_turn_metrics( @@ -391,24 +388,15 @@ pub fn calculate_multi_turn_metrics( max_turn_count: usize, ) -> MultiTurnMetrics { // Flatten all turns for overall metrics - let mut all_requests: Vec = Vec::new(); let mut all_outputs: Vec = Vec::new(); for conv in conversation_outputs { for turn in &conv.turns { - all_requests.push(SampleRequest { - prompt: std::sync::Arc::from(""), - prompt_len: turn.request_output.prompt_len, - expected_output_len: turn.request_output.output_tokens, - request_id: None, - ..Default::default() - }); all_outputs.push(turn.request_output.clone()); } } let (overall, _) = calculate_metrics( - &all_requests, &all_outputs, dur_s, selected_percentiles, @@ -419,25 +407,16 @@ pub fn calculate_multi_turn_metrics( // Per-turn breakdown let mut per_turn: Vec = Vec::new(); for turn_idx in 0..max_turn_count { - let mut turn_requests: Vec = Vec::new(); let mut turn_outputs: Vec = Vec::new(); for conv in conversation_outputs { if let Some(turn) = conv.turns.get(turn_idx) { - turn_requests.push(SampleRequest { - prompt: std::sync::Arc::from(""), - prompt_len: turn.request_output.prompt_len, - expected_output_len: turn.request_output.output_tokens, - request_id: None, - ..Default::default() - }); turn_outputs.push(turn.request_output.clone()); } } if !turn_outputs.is_empty() { let (turn_metrics, _) = calculate_metrics( - &turn_requests, &turn_outputs, dur_s, selected_percentiles, @@ -509,6 +488,26 @@ mod tests { assert!((p99 - 98.01).abs() < 0.1); } + /// Input-token totals come from each output's `prompt_len` (server-reported + /// via `usage.prompt_tokens`), not the client-side request estimate. + /// Mirrors Python's `serve.py`. + #[test] + fn test_total_input_uses_output_prompt_len() { + let outputs = vec![ + RequestFuncOutput { + success: true, + prompt_len: 128, + output_tokens: 3, + ..Default::default() + }; + 2 + ]; + + let (metrics, _) = calculate_metrics(&outputs, 1.0, &[], false, &GoodputConfig::default()); + + assert_eq!(metrics.total_input, 256); + } + #[test] fn test_multi_turn_metrics_use_request_prompt_lengths() { let turn = |turn_index, prompt_len, output_tokens| crate::multi_turn::TurnOutput { diff --git a/rust/src/bench/src/metrics/steady_state.rs b/rust/src/bench/src/metrics/steady_state.rs index a403275b7ac..f108bb7c850 100644 --- a/rust/src/bench/src/metrics/steady_state.rs +++ b/rust/src/bench/src/metrics/steady_state.rs @@ -170,16 +170,12 @@ pub fn detect_window( }) } -use crate::datasets::SampleRequest; - -/// Compute steady-state metrics from per-request outputs, the detected window, -/// and the original request list (for input-token attribution). +/// Compute steady-state metrics from per-request outputs and the detected window. /// /// `is_pooling` gates off TTFT because pooling backends write `ttft = latency` /// as a placeholder; reporting TTFT over that would silently be E2EL. pub fn compute( outputs: &[RequestFuncOutput], - input_requests: &[SampleRequest], window: &SteadyStateWindow, percentiles: &[f64], is_pooling: bool, @@ -200,7 +196,7 @@ pub fn compute( let mut ttfts_in_window: Vec = Vec::new(); let mut tpots_in_window: Vec = Vec::new(); - for (idx, o) in outputs.iter().enumerate() { + for o in outputs.iter() { if !o.success { continue; } @@ -213,7 +209,7 @@ pub fn compute( } if started_in { - input_tokens_in_window += input_requests[idx].prompt_len; + input_tokens_in_window += o.prompt_len; if !is_pooling { ttfts_in_window.push(o.ttft); // Per-request TPOT: (latency - ttft) / (output_tokens - 1). @@ -417,10 +413,6 @@ mod tests { assert!((w.end_s - 10.0).abs() < 1e-9); } - use std::sync::Arc; - - use crate::datasets::SampleRequest; - fn mk_full( start: f64, latency: f64, @@ -440,16 +432,6 @@ mod tests { } } - fn mk_sample(prompt_len: usize) -> SampleRequest { - SampleRequest { - prompt: Arc::from(""), - prompt_len, - expected_output_len: 0, - request_id: None, - ..Default::default() - } - } - #[test] fn compute_constant_concurrency_matches_full_run() { // 3 identical requests, each 10s, each emitting 5 tokens uniformly via itl=2s. @@ -459,13 +441,10 @@ mod tests { mk_full(0.0, 10.0, 2.0, itl.clone(), 5), mk_full(0.0, 10.0, 2.0, itl.clone(), 5), ]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); let m = compute( &outputs, - &requests, &w, &[99.0], false, // is_pooling @@ -502,12 +481,10 @@ mod tests { let late_itl = vec![0.02; 999]; // 1000 tokens; emission starts after end_s let late = mk_full(9.9, 30.0, 0.05, late_itl, 1000); outputs.push(late); - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); // Target = 2 so threshold_abs = 2; window opens when both are in flight. let w = detect_window(&outputs, Some(2), 1.0, 0.1, 40.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], false); + let m = compute(&outputs, &w, &[99.0], false); // Late request started in window but most of its tokens are emitted AFTER end_s. // Only a handful (≤3) should count. @@ -534,10 +511,8 @@ mod tests { mk_full(0.0, 5.0, 5.0, Vec::new(), 0), mk_full(0.0, 5.0, 5.0, Vec::new(), 0), ]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], /* is_pooling */ true); + let m = compute(&outputs, &w, &[99.0], /* is_pooling */ true); assert_eq!(m.mean_ttft_ms, 0.0); assert_eq!(m.median_ttft_ms, 0.0); assert!(m.percentiles_ttft_ms.is_empty()); @@ -560,10 +535,8 @@ mod tests { mk_full(0.0, 10.0, 2.0, itl.clone(), 5), mk_full(0.0, 10.0, 2.0, itl.clone(), 5), ]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], false); + let m = compute(&outputs, &w, &[99.0], false); assert!((m.mean_tpot_ms - 2000.0).abs() < 1e-6); assert!((m.median_tpot_ms - 2000.0).abs() < 1e-6); @@ -574,10 +547,8 @@ mod tests { #[test] fn json_roundtrip_some() { let outputs = vec![mk_full(0.0, 10.0, 2.0, vec![2.0; 4], 5); 3]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], false); + let m = compute(&outputs, &w, &[99.0], false); let s = serde_json::to_string(&m).unwrap(); let parsed: SteadyStateMetrics = serde_json::from_str(&s).unwrap(); assert!((parsed.output_throughput - m.output_throughput).abs() < 1e-12); @@ -605,14 +576,13 @@ mod tests { let itl = vec![1.0; 9]; let outputs: Vec = (0..5).map(|i| mk_full(i as f64, 20.0, 1.0, itl.clone(), 10)).collect(); - let requests: Vec = (0..5).map(|_| mk_sample(100)).collect(); let w = detect_window(&outputs, Some(5), 1.0, 1.0, 25.0).unwrap(); assert_eq!(w.threshold_abs, 5); assert!((w.start_s - 4.0).abs() < 1e-9); assert!((w.end_s - 20.0).abs() < 1e-9); - let m = compute(&outputs, &requests, &w, &[50.0, 99.0], false); + let m = compute(&outputs, &w, &[50.0, 99.0], false); // Round-trip through JSON. let s = serde_json::to_string(&m).unwrap();