diff --git a/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs b/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs index 5e44d3d727..8c4a87df49 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs @@ -159,10 +159,16 @@ struct AssistantMessage<'a> { #[derive(Serialize)] struct ChatUsage { prompt_tokens: usize, + prompt_tokens_details: ChatPromptTokensDetails, completion_tokens: usize, total_tokens: usize, } +#[derive(Serialize)] +struct ChatPromptTokensDetails { + cached_tokens: u64, +} + #[derive(Serialize)] #[serde(untagged)] pub(crate) enum EmbeddingValue { @@ -351,6 +357,9 @@ pub(crate) async fn chat_completions( }], usage: ChatUsage { prompt_tokens: input_tokens, + prompt_tokens_details: ChatPromptTokensDetails { + cached_tokens: kv_cache_access.reused_input_tokens, + }, completion_tokens: output_tokens, total_tokens: input_tokens.saturating_add(output_tokens), }, @@ -608,6 +617,7 @@ pub(crate) enum ChatStreamChunk<'a> { Stop, Usage { input_tokens: usize, + cached_input_tokens: u64, output_tokens: usize, }, } @@ -625,10 +635,14 @@ pub(crate) fn chat_chunk_json( ChatStreamChunk::Stop => Some((None, None, Some("stop"))), ChatStreamChunk::Usage { input_tokens, + cached_input_tokens, output_tokens, } => { usage = Some(Some(ChatUsage { prompt_tokens: input_tokens, + prompt_tokens_details: ChatPromptTokensDetails { + cached_tokens: cached_input_tokens, + }, completion_tokens: output_tokens, total_tokens: input_tokens.saturating_add(output_tokens), })); @@ -666,6 +680,7 @@ fn stream_response(config: StreamResponseConfig) -> Response { mut pacing, kind, } = config; + let cached_input_tokens = kv_cache_access.reused_input_tokens; let stream = async_stream::stream! { let mut output_text = String::new(); if let StreamKind::Responses { created_at } = kind { @@ -747,6 +762,9 @@ fn stream_response(config: StreamResponseConfig) -> Response { }], "usage": { "input_tokens": input_tokens, + "input_tokens_details": { + "cached_tokens": cached_input_tokens, + }, "output_tokens": output_tokens, "total_tokens": input_tokens.saturating_add(output_tokens), }, @@ -760,7 +778,7 @@ fn stream_response(config: StreamResponseConfig) -> Response { if let StreamKind::Chat { include_usage, .. } = kind { if include_usage { - yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Usage { input_tokens, output_tokens }, true)); + yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Usage { input_tokens, cached_input_tokens, output_tokens }, true)); } yield Ok(Event::default().data("[DONE]")); } diff --git a/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs b/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs index b512fa9c20..216103684f 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs @@ -414,7 +414,7 @@ async fn test_control_http_api_updates_one_model_and_reports_request_counters() assert_eq!( body, serde_json::from_str::( - r#"{"id":"chatcmpl-mock-id","object":"chat.completion","model":"model-a","choices":[{"index":0,"message":{"role":"assistant","content":"Hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}"#, + r#"{"id":"chatcmpl-mock-id","object":"chat.completion","model":"model-a","choices":[{"index":0,"message":{"role":"assistant","content":"Hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"prompt_tokens_details":{"cached_tokens":0},"completion_tokens":1,"total_tokens":2}}"#, ) .expect("expected chat response fixture should be valid JSON") ); @@ -597,6 +597,7 @@ async fn streaming_chat_usage_reports_actual_output_only_when_requested() { distribution: OutputTokenDistribution::Uniform, }, context_length_tokens: 5, + kv_cache: Arc::new(Mutex::new(KvCacheState::new(10))), ..test_state() }; let app = Router::new() @@ -614,7 +615,7 @@ async fn streaming_chat_usage_reports_actual_output_only_when_requested() { address, "POST", "/v1/chat/completions", - "connection: close\r\nx-input-tokens: 2\r\nx-output-tokens: 100", + "connection: close\r\nx-input-tokens: 2\r\nx-output-tokens: 100\r\nx-cache-affinity-key: chat-usage", &body.to_string(), ) .await; @@ -640,7 +641,12 @@ async fn streaming_chat_usage_reports_actual_output_only_when_requested() { assert_eq!(usage["choices"], serde_json::json!([])); assert_eq!( usage["usage"], - serde_json::json!({"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}) + serde_json::json!({ + "prompt_tokens": 2, + "prompt_tokens_details": {"cached_tokens": 2}, + "completion_tokens": 3, + "total_tokens": 5 + }) ); assert!( output_events @@ -1137,13 +1143,14 @@ async fn responses_endpoint_streams_response_events_without_private_stats_header let (addr, server) = spawn_test_app(app).await; let cases = [ - (None, 1), - (Some(serde_json::json!("hello")), 5), - (Some(serde_json::json!(["hi", {"x": 1}])), 9), - (Some(serde_json::json!({"x": 1})), 7), + (None, 1, 1), + (Some(serde_json::json!("hello")), 5, 4), + (Some(serde_json::json!(["hi", {"x": 1}])), 9, 4), + (Some(serde_json::json!({"x": 1})), 7, 0), ]; let mut response = String::new(); - for (index, (input, expected_tokens)) in cases.into_iter().enumerate() { + for (index, (input, expected_tokens, expected_uncached_tokens)) in cases.into_iter().enumerate() + { let mut request_body = serde_json::json!({ "model": "request-model", "max_output_tokens": 2, @@ -1156,19 +1163,19 @@ async fn responses_endpoint_streams_response_events_without_private_stats_header addr, "POST", "/v1/responses", - &format!("connection: close\r\nx-request-id: req-responses-{index}\r\nx-input-tokens: {expected_tokens}\r\nx-cache-affinity-key: cache-{index}"), + &format!("connection: close\r\nx-request-id: req-responses-{index}\r\nx-input-tokens: {expected_tokens}\r\nx-cache-affinity-key: cache-shared"), &request_body.to_string(), ) .await; assert!(response.contains(&format!( - "x-kv-cache-uncached-input-tokens: {expected_tokens}" + "x-kv-cache-uncached-input-tokens: {expected_uncached_tokens}" ))); assert!(response.contains(&format!(r#""input_tokens":{expected_tokens}"#))); } assert!(response.starts_with("HTTP/1.1 200 OK")); - assert!(response.contains("x-kv-cache-hit: false")); - assert!(response.contains("x-kv-cache-reused-input-tokens: 0")); - assert!(response.contains("x-kv-cache-uncached-input-tokens: 7")); + assert!(response.contains("x-kv-cache-hit: true")); + assert!(response.contains("x-kv-cache-reused-input-tokens: 7")); + assert!(response.contains("x-kv-cache-uncached-input-tokens: 0")); assert!(response.contains("content-type: text/event-stream")); let (_, body_text) = response @@ -1182,6 +1189,7 @@ async fn responses_endpoint_streams_response_events_without_private_stats_header assert!(body_text.contains(r#""status":"completed""#)); assert!(body_text.contains(r#""model":"request-model""#)); assert!(body_text.contains(r#""input_tokens":7"#)); + assert!(body_text.contains(r#""input_tokens_details":{"cached_tokens":7}"#)); assert!(body_text.contains(r#""output_tokens":2"#)); assert!(body_text.contains(r#""total_tokens":9"#)); diff --git a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto index e33537de7c..abf40b207f 100644 --- a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto +++ b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto @@ -123,8 +123,9 @@ message ModelStats { // Sticky source labels for stat observations seen by this backend since the // model metrics state was initialized. repeated string stats_sources = 17; - // Greatest last_mean_input_tps published for the current backend model - // generation. Absent until the first mean is configured or observed. + // Greatest max-eligible last_mean_input_tps published for the current backend + // model generation. Means that may count cached prompt tokens as prefill do + // not raise it. Absent until the first eligible mean is configured or observed. optional double max_input_tps = 18; } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup.rs index be8330e459..0e4d26661b 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup.rs @@ -145,11 +145,25 @@ mod tests { ) -> ( Result, Vec, + ) { + let (result, events) = send_observed_completion_events(upstream_http_base_url).await; + let states = events + .into_iter() + .map(|event| event.into_observation().state) + .collect(); + (result, states) + } + + async fn send_observed_completion_events( + upstream_http_base_url: &str, + ) -> ( + Result, + Vec, ) { let (runtime_state, observations) = PylonRuntimeState::observed( InferenceServerStatus::Active, &["test-model".to_string()], - 8, + 16, None, ); let result = send_completion_request( @@ -167,11 +181,7 @@ mod tests { Some(&runtime_state), ) .await; - let states = observations - .try_iter() - .map(|event| event.into_observation().state) - .collect(); - (result, states) + (result, observations.try_iter().collect()) } fn test_stats_collector() -> (PylonRuntimeState, StatsCollectorHandle) { @@ -778,6 +788,231 @@ mod tests { assert_eq!(max_in_flight.load(Ordering::SeqCst), 3); } + #[test] + fn calibration_prompts_are_unique_and_keep_the_requested_length() { + for prompt_units in [8, 33, CALIBRATION_PROMPT_UNITS_FLOOR, 4096] { + let prompts = (0..64) + .map(|_| calibration_prompt(prompt_units)) + .collect::>(); + assert!(prompts.iter().all(|prompt| prompt.len() == prompt_units)); + assert!(prompts.iter().all(|prompt| prompt.is_ascii())); + let prefix_len = prompt_units.min(32); + let prefixes = prompts + .iter() + .map(|prompt| &prompt[..prefix_len]) + .collect::>(); + assert_eq!( + prefixes.len(), + prompts.len(), + "every calibration prompt must start with a distinct prefix" + ); + } + } + + #[tokio::test] + async fn calibration_prompts_do_not_share_prefixes_across_requests_or_steps() { + let prompts = Arc::new(Mutex::new(Vec::new())); + let base_url = spawn_test_server(TestServerState { + prompts: Some(prompts.clone()), + ..TestServerState::default() + }) + .await; + let client = reqwest::Client::new(); + let runtime_state = test_runtime_state(); + let steps = calibration_ramp(&CalibrationConfig { + calibration_requests: 2, + calibration_prompt_units: 1024, + calibration_max_concurrency: 2, + ..CalibrationConfig::default() + }) + .take(4) + .collect::>(); + for step in &steps { + let outcome = send_calibration_batch( + &client, + &base_url, + &test_generation(), + Duration::from_secs(1), + *step, + &runtime_state, + ) + .await + .expect("calibration batch should succeed"); + assert_eq!(outcome, CalibrationStepOutcome::Completed); + } + + let prompts = prompts.lock().await.clone(); + let mut expected_lengths = steps + .iter() + .flat_map(|step| std::iter::repeat_n(step.prompt_units, step.request_count)) + .collect::>(); + let mut lengths = prompts.iter().map(String::len).collect::>(); + expected_lengths.sort_unstable(); + lengths.sort_unstable(); + assert_eq!( + lengths, expected_lengths, + "unique prompts must keep each step's prompt size" + ); + for (index, prompt) in prompts.iter().enumerate() { + for other in &prompts[index + 1..] { + assert_ne!( + prompt[..32], + other[..32], + "calibration prompts must not share a cacheable prefix" + ); + } + } + } + + fn parse_input_usage(usage: Value) -> Option { + serde_json::from_value::(serde_json::json!({"usage": usage})) + .expect("usage should deserialize") + .usage + .input_usage() + } + + #[test] + fn calibration_usage_reports_uncached_prompt_tokens() { + assert_eq!( + parse_input_usage(serde_json::json!({ + "prompt_tokens": 100, + "completion_tokens": 1, + "prompt_tokens_details": {"cached_tokens": 64}, + })), + Some(InputUsage { + input_tokens: 100, + uncached_input_tokens: Some(36), + }) + ); + assert_eq!( + parse_input_usage(serde_json::json!({ + "prompt_tokens": 100, + "completion_tokens": 1, + "prompt_tokens_details": {"cached_tokens": 0}, + })), + Some(InputUsage { + input_tokens: 100, + uncached_input_tokens: Some(100), + }) + ); + } + + #[test] + fn calibration_usage_without_cached_details_reports_only_the_total() { + for usage in [ + serde_json::json!({"prompt_tokens": 100, "completion_tokens": 1}), + serde_json::json!({ + "prompt_tokens": 100, + "completion_tokens": 1, + "prompt_tokens_details": null, + }), + serde_json::json!({ + "prompt_tokens": 100, + "completion_tokens": 1, + "prompt_tokens_details": {"audio_tokens": 0}, + }), + ] { + assert_eq!( + parse_input_usage(usage), + Some(InputUsage { + input_tokens: 100, + uncached_input_tokens: None, + }) + ); + } + } + + #[test] + fn calibration_usage_ignores_invalid_cached_tokens() { + for cached_tokens in [serde_json::json!(128), serde_json::json!("64")] { + assert_eq!( + parse_input_usage(serde_json::json!({ + "prompt_tokens": 100, + "completion_tokens": 1, + "prompt_tokens_details": {"cached_tokens": cached_tokens}, + })), + Some(InputUsage { + input_tokens: 100, + uncached_input_tokens: None, + }), + "cached tokens that are malformed or exceed the prompt total are unavailable" + ); + } + } + + #[test] + fn calibration_usage_without_prompt_tokens_is_unavailable() { + assert_eq!( + parse_input_usage(serde_json::json!({"completion_tokens": 1})), + None + ); + assert_eq!( + parse_input_usage(serde_json::json!({ + "prompt_tokens": "100", + "completion_tokens": 1, + "prompt_tokens_details": {"cached_tokens": 64}, + })), + None + ); + } + + async fn observed_calibration_input(usage: Value) -> (u64, bool, Option) { + let server = spawn_test_server(TestServerState { + usage: Some(usage), + ..TestServerState::default() + }) + .await; + let (result, events) = send_observed_completion_events(server.as_str()).await; + result.expect("calibration request should succeed"); + let terminal = events + .into_iter() + .last() + .expect("calibration request should be observed"); + let explicit = terminal.input_tokens_explicit; + let uncached = terminal.uncached_input_tokens; + let observation = terminal.into_observation(); + assert_eq!(observation.state, crate::RequestObservationState::Complete); + (observation.input_tokens, explicit, uncached) + } + + #[tokio::test] + async fn calibration_response_usage_feeds_exact_uncached_input_tokens() { + assert_eq!( + observed_calibration_input(serde_json::json!({ + "prompt_tokens": 13, + "completion_tokens": 1, + "prompt_tokens_details": {"cached_tokens": 8}, + })) + .await, + (13, true, Some(5)) + ); + assert_eq!( + observed_calibration_input(serde_json::json!({ + "prompt_tokens": 13, + "completion_tokens": 1, + })) + .await, + (13, true, None) + ); + assert_eq!( + observed_calibration_input(serde_json::json!({ + "prompt_tokens": 13, + "completion_tokens": 1, + "prompt_tokens_details": {"cached_tokens": 21}, + })) + .await, + (13, true, None) + ); + } + + #[tokio::test] + async fn calibration_response_without_prompt_usage_keeps_the_header_estimate() { + assert_eq!( + observed_calibration_input(serde_json::json!({"completion_tokens": 1})).await, + (5, false, None) + ); + } + #[tokio::test] async fn calibration_top_level_usage_does_not_claim_chunk_usage_stats() { let base_url = spawn_test_server(TestServerState::default()).await; @@ -1816,6 +2051,8 @@ mod tests { request_ids: Option>>>, prompt_lengths: Option>>>, prompt_rejections: Option>>>, + prompts: Option>>>, + usage: Option, } impl Default for TestServerState { @@ -1838,6 +2075,8 @@ mod tests { request_ids: None, prompt_lengths: None, prompt_rejections: None, + prompts: None, + usage: None, } } } @@ -1922,6 +2161,9 @@ mod tests { if let Some(prompt_lengths) = &state.prompt_lengths { prompt_lengths.lock().await.push(prompt_len); } + if let Some(prompts) = &state.prompts { + prompts.lock().await.push(prompt.to_string()); + } let reject_prompt = if let Some(prompt_rejections) = &state.prompt_rejections { let mut prompt_rejections = prompt_rejections.lock().await; if prompt_rejections.first() == Some(&prompt_len) { @@ -2046,9 +2288,9 @@ mod tests { .into_response(); } - Json(serde_json::json!({ - "usage": {"completion_tokens": completion_tokens} - })) - .into_response() + let usage = state + .usage + .unwrap_or_else(|| serde_json::json!({"completion_tokens": completion_tokens})); + Json(serde_json::json!({ "usage": usage })).into_response() } } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/calibration.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/calibration.rs index 57223c51f8..0b958618bc 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/calibration.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/calibration.rs @@ -17,6 +17,7 @@ use std::time::{Duration, Instant}; use futures::stream::{self, StreamExt, TryStreamExt}; use stargate_protocol::common::valid_last_mean_input_tps; +use uuid::Uuid; use crate::generated_request_id::GeneratedRequestKind; use crate::runtime_state::{ModelGeneration, PylonRuntimeState}; @@ -191,26 +192,28 @@ pub(super) async fn send_calibration_batch( request_count, concurrency, } = batch; - let request = serde_json::json!({ - "model": generation.model_id(), - "messages": [{"role": "user", "content": "1".repeat(prompt_units)}], - "max_tokens": 1, - "seed": 33, - "temperature": 0.7, - "top_p": 1.0, - "stream": false, - }); - let step = stream::iter((0..request_count).map(|_| { - send_completion_request( - http_client, - upstream_http_base_url, - None, - &request, - GeneratedRequestKind::Calibration, - generation, - Some(runtime_state), - ) + let request = serde_json::json!({ + "model": generation.model_id(), + "messages": [{"role": "user", "content": calibration_prompt(prompt_units)}], + "max_tokens": 1, + "seed": 33, + "temperature": 0.7, + "top_p": 1.0, + "stream": false, + }); + async move { + send_completion_request( + http_client, + upstream_http_base_url, + None, + &request, + GeneratedRequestKind::Calibration, + generation, + Some(runtime_state), + ) + .await + } })) .buffer_unordered(concurrency.min(request_count)) .try_for_each(|_| async { Ok(()) }); @@ -219,3 +222,15 @@ pub(super) async fn send_calibration_batch( Err(_) => Ok(CalibrationStepOutcome::Saturated), } } + +/// Builds a `prompt_units`-character calibration prompt that starts with a +/// random per-request prefix. Engines with prefix caching would otherwise +/// serve repeated or extended prompts from cache, so calibration would report +/// cache-hit throughput instead of uncached prefill throughput. +pub(super) fn calibration_prompt(prompt_units: usize) -> String { + let mut prompt = format!("{} ", Uuid::new_v4().simple()); + prompt.truncate(prompt_units); + let filler = prompt_units - prompt.len(); + prompt.extend(std::iter::repeat_n('1', filler)); + prompt +} diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs index 87a0bdda56..b5e07d552d 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs @@ -302,6 +302,10 @@ fn finish_observation( let generation = observer .generation_mut() .expect("chat completion observer should expose generation progress"); + if let Some(input_usage) = completion.usage.input_usage() { + generation + .observe_input_usage(input_usage.input_tokens, input_usage.uncached_input_tokens); + } generation.observe_generated_output( Instant::now(), completion.usage.completion_tokens > 0, @@ -386,12 +390,38 @@ impl BringupError { #[derive(Debug, Deserialize)] pub(super) struct ChatCompletionResponse { - usage: Usage, + pub(super) usage: Usage, } #[derive(Debug, Deserialize)] -struct Usage { +pub(super) struct Usage { completion_tokens: u32, + #[serde(default)] + prompt_tokens: serde_json::Value, + #[serde(default)] + prompt_tokens_details: serde_json::Value, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct InputUsage { + pub(super) input_tokens: u64, + pub(super) uncached_input_tokens: Option, +} + +impl Usage { + /// Returns the exact prompt usage, matching the streaming usage parser: + /// cached tokens that are malformed or exceed the prompt total leave the + /// uncached count unavailable. + pub(super) fn input_usage(&self) -> Option { + let input_tokens = self.prompt_tokens.as_u64()?; + let uncached_input_tokens = self.prompt_tokens_details["cached_tokens"] + .as_u64() + .and_then(|cached| input_tokens.checked_sub(cached)); + Some(InputUsage { + input_tokens, + uncached_input_tokens, + }) + } } #[derive(Debug, Deserialize)] diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs index b2cae508b7..b6d418a30b 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs @@ -497,7 +497,7 @@ impl TunnelRequestLifecycle { obs.observe_output_calibration_details(exact_usage.reasoning_tokens, false); } if let Some(input_tokens) = exact_usage.input_tokens { - obs.observe_input_tokens_total(input_tokens); + obs.observe_input_usage(input_tokens, exact_usage.uncached_input_tokens); } if let Some(output_tokens) = exact_usage.output_tokens { let update = parser.observe_exact_output_tokens(output_tokens); @@ -780,6 +780,13 @@ pub(super) async fn forward_tunnel_request( } else { None }; + if let Some(observer) = lifecycle.observer.as_mut() { + observer.set_input_usage_expected(input_usage_expected( + observation_endpoint, + &body_bytes, + app.force_chat_completions_include_usage, + )); + } if let Some(decision) = lifecycle.admit_queue(app, &request_headers) { let QueueAdmissionDecision::Rejected { expected_ms, @@ -1080,6 +1087,24 @@ fn validate_request_body( .ok_or(stream_error) } +fn input_usage_expected( + endpoint: Option, + body_bytes: &[u8], + force_chat_completions_include_usage: bool, +) -> bool { + match endpoint { + Some(RequestObservationEndpoint::Responses) => true, + Some(RequestObservationEndpoint::ChatCompletions) => { + force_chat_completions_include_usage + || serde_json::from_slice::(body_bytes) + .ok() + .and_then(|body| body.get("stream_options")?.get("include_usage")?.as_bool()) + == Some(true) + } + Some(RequestObservationEndpoint::Embeddings) | None => false, + } +} + fn prepare_chat_completions_usage_rewrite( body_bytes: &[u8], ) -> Result, &'static str> { @@ -1743,6 +1768,35 @@ mod tests { assert_eq!(inserted["stream_options"]["include_usage"], true); } + #[test] + fn input_usage_expectation_matches_endpoint_and_request_options() { + assert!(input_usage_expected( + Some(RequestObservationEndpoint::ChatCompletions), + br#"{"stream_options":{"include_usage":true}}"#, + false, + )); + assert!(input_usage_expected( + Some(RequestObservationEndpoint::ChatCompletions), + br#"{"stream":true}"#, + true, + )); + assert!(!input_usage_expected( + Some(RequestObservationEndpoint::ChatCompletions), + br#"{"stream_options":{"include_usage":false}}"#, + false, + )); + assert!(input_usage_expected( + Some(RequestObservationEndpoint::Responses), + br#"{"stream":true}"#, + false, + )); + assert!(!input_usage_expected( + Some(RequestObservationEndpoint::Embeddings), + b"{}", + true, + )); + } + #[test] fn forced_chat_usage_rejects_invalid_stream_options() { assert_eq!( diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs index 88e5ed66ed..f2f9f459a1 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs @@ -2332,6 +2332,36 @@ async fn quic_tunnel_counts_terminal_only_usage_tokens() { tunnel.shutdown().await; } +#[tokio::test] +async fn quic_tunnel_preserves_cached_prompt_usage_for_fallback_stats() { + let app = sse_app( + "/v1/chat/completions", + &[ + r#"{"object":"chat.completion.chunk","choices":[{"delta":{"content":"Hello"}}],"usage":null}"#, + r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":13,"completion_tokens":1,"prompt_tokens_details":{"cached_tokens":8}}}"#, + "[DONE]", + ], + ); + let (mut tunnel, rx) = observed_tunnel_for(app).await; + let headers = tunnel_request_headers( + "/v1/chat/completions", + "model-terminal-usage", + "req-cached-usage", + "13", + ); + tunnel + .send(headers, br#"{"messages":[],"stream":true}"#) + .await; + tunnel.response_head(StatusCode::OK).await; + tunnel.drain().await; + + let event = recv_terminal_event(&rx).await; + assert_eq!(event.observation().input_tokens, 13); + assert_eq!(event.uncached_input_tokens(), Some(5)); + + tunnel.shutdown().await; +} + #[tokio::test] async fn quic_tunnel_retains_split_usage_safety_details_for_calibration() { let app = sse_app( diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index 5cca12e87e..24ce113275 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs @@ -172,6 +172,8 @@ pub(crate) struct RequestObserver { request_input_tokens: u64, input_tokens: u64, input_tokens_explicit: bool, + uncached_input_tokens: Option, + input_usage_expected: bool, generation: Option, embedding_items: Option, chat_calibration: Option, @@ -206,6 +208,8 @@ impl RequestObserver { request_input_tokens: input_tokens, input_tokens, input_tokens_explicit: false, + uncached_input_tokens: None, + input_usage_expected: false, generation, embedding_items: None, chat_calibration: (output_token_calibration_enabled @@ -222,6 +226,13 @@ impl RequestObserver { self.runtime_state.output_token_calibration_enabled() } + pub(crate) fn set_input_usage_expected(&mut self, expected: bool) { + if self.input_usage_expected != expected { + self.input_usage_expected = expected; + self.emit(); + } + } + pub(super) fn update_embedding_items(&mut self, embedding_items: Option) { debug_assert_eq!(self.endpoint, RequestObservationEndpoint::Embeddings); if self.embedding_items != embedding_items { @@ -291,24 +302,34 @@ impl RequestObserver { ); } + #[cfg(test)] pub(crate) fn observe_input_tokens_total(&mut self, input_tokens: u64) { - if self.input_tokens_explicit { - if input_tokens < self.input_tokens { - tracing::warn!( - request_id = self.request_id, - prior_input_tokens = self.input_tokens, - input_tokens, - "ignoring regressing explicit input token counter" - ); - return; - } - if input_tokens == self.input_tokens { - return; - } + self.observe_input_usage(input_tokens, None); + } + + pub(crate) fn observe_input_usage( + &mut self, + input_tokens: u64, + uncached_input_tokens: Option, + ) { + if self.input_tokens_explicit && input_tokens < self.input_tokens { + tracing::warn!( + request_id = self.request_id, + prior_input_tokens = self.input_tokens, + input_tokens, + "ignoring regressing explicit input token counter" + ); + return; } + let changed = input_tokens != self.input_tokens + || !self.input_tokens_explicit + || self.uncached_input_tokens != uncached_input_tokens; self.input_tokens = input_tokens; self.input_tokens_explicit = true; - self.emit(); + self.uncached_input_tokens = uncached_input_tokens; + if changed { + self.emit(); + } } pub(crate) fn observe_estimated_output_tokens_total(&mut self, output_tokens: u64) { @@ -545,6 +566,8 @@ impl RequestObserver { changed_generations: Vec::new(), input_interval, input_tokens_explicit: self.input_tokens_explicit, + uncached_input_tokens: self.uncached_input_tokens, + input_usage_expected: self.input_usage_expected, output_calibration: backend .map_or_else(OutputCalibrationFacts::default, |backend| { backend.output_calibration diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/tunnel.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/tunnel.rs index 6c16a4e738..850798de38 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/tunnel.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/tunnel.rs @@ -51,6 +51,10 @@ impl TunnelRequestObserver { } } + pub(crate) fn set_input_usage_expected(&mut self, expected: bool) { + self.observer.set_input_usage_expected(expected); + } + pub(crate) fn on_backend_submission(&mut self, submitted_at: Instant) { self.observer.on_backend_submission(submitted_at); } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index 2c691057d0..415fd07cf9 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -33,7 +33,9 @@ use reqwest::header::HeaderMap; pub struct CurrentModelStats { // Sticky runtime-observed mean input TPS for this backend. pub last_mean_input_tps: f64, - // Greatest last_mean_input_tps of this model generation. None until the first mean. + // Greatest max-eligible last_mean_input_tps of this model generation: a + // configured seed, engine counter or embedding samples, calibration, or + // fallback windows whose usage excluded cached prompt tokens. None until then. pub max_input_tps: Option, // Token/sec output rate for streaming generation endpoints. Embeddings item // cardinality is observed separately and is not exported through this field. @@ -123,6 +125,8 @@ pub struct RequestObservationEvent { pub(crate) changed_generations: Vec, pub(crate) input_interval: Option, pub(crate) input_tokens_explicit: bool, + pub(crate) uncached_input_tokens: Option, + pub(crate) input_usage_expected: bool, pub(crate) output_calibration: OutputCalibrationFacts, pub(crate) upstream_duration: Option, } @@ -520,6 +524,8 @@ impl PylonRuntimeState { changed_generations: Vec::new(), input_interval: None, input_tokens_explicit: false, + uncached_input_tokens: None, + input_usage_expected: false, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, @@ -579,6 +585,8 @@ impl PylonRuntimeState { changed_generations: Vec::new(), input_interval: None, input_tokens_explicit: false, + uncached_input_tokens: None, + input_usage_expected: false, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, @@ -728,6 +736,14 @@ impl RequestObservationEvent { self.input_tokens_explicit } + pub(crate) fn uncached_input_tokens(&self) -> Option { + self.uncached_input_tokens + } + + pub(crate) fn input_usage_expected(&self) -> bool { + self.input_usage_expected + } + pub(crate) fn output_calibration(&self) -> OutputCalibrationFacts { self.output_calibration } @@ -795,6 +811,8 @@ mod tests { changed_generations: Vec::new(), input_interval: None, input_tokens_explicit: false, + uncached_input_tokens: None, + input_usage_expected: false, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/sse_message_stream.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/sse_message_stream.rs index 8cea901169..7ed15c14db 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/sse_message_stream.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/sse_message_stream.rs @@ -36,6 +36,7 @@ pub(crate) struct GeneratedOutput { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct ExactUsage { pub(crate) input_tokens: Option, + pub(crate) uncached_input_tokens: Option, pub(crate) output_tokens: Option, pub(crate) reasoning_tokens: Option, } @@ -771,6 +772,21 @@ fn exact_usage( let response_usage = &value["response"]["usage"]; let (input_tokens, invalid_input_tokens) = merge_u64_fields([&usage["prompt_tokens"], &response_usage["input_tokens"]]); + let (cached_input_tokens, invalid_cached_input_tokens) = merge_u64_fields([ + &usage["prompt_tokens_details"]["cached_tokens"], + &response_usage["input_tokens_details"]["cached_tokens"], + ]); + let (uncached_input_tokens, cached_input_exceeds_total) = match ( + invalid_cached_input_tokens, + input_tokens, + cached_input_tokens, + ) { + (false, Some(input), Some(cached)) => match input.checked_sub(cached) { + Some(uncached) => (Some(uncached), false), + None => (None, true), + }, + _ => (None, false), + }; let (output_tokens, invalid_output_tokens) = merge_u64_fields([ &usage["completion_tokens"], &response_usage["output_tokens"], @@ -780,6 +796,7 @@ fn exact_usage( return ( (input_tokens.is_some() || output_tokens.is_some()).then_some(ExactUsage { input_tokens, + uncached_input_tokens, output_tokens, reasoning_tokens: None, }), @@ -803,9 +820,18 @@ fn exact_usage( ]); let (accepted_prediction_tokens, invalid_accepted_prediction_tokens) = merge_u64_fields([&completion_details["accepted_prediction_tokens"]]); - let invalid_container = [usage, response_usage, completion_details, output_details] - .into_iter() - .any(|value| !value.is_null() && value.as_object().is_none()); + let chat_input_details = &usage["prompt_tokens_details"]; + let responses_input_details = &response_usage["input_tokens_details"]; + let invalid_container = [ + usage, + response_usage, + completion_details, + output_details, + chat_input_details, + responses_input_details, + ] + .into_iter() + .any(|value| !value.is_null() && value.as_object().is_none()); let unknown_output_details = output_details_have_unknown_fields(completion_details, CHAT_OUTPUT_DETAIL_FIELDS) || output_details_have_unknown_fields(output_details, RESPONSES_OUTPUT_DETAIL_FIELDS); @@ -819,6 +845,8 @@ fn exact_usage( let calibration_ineligible = usage_protocol_mismatch || invalid_container || invalid_input_tokens + || invalid_cached_input_tokens + || cached_input_exceeds_total || invalid_output_tokens || invalid_reasoning_tokens || invalid_audio_tokens @@ -832,6 +860,7 @@ fn exact_usage( (input_tokens.is_some() || output_tokens.is_some() || reasoning_tokens.is_some()) .then_some(ExactUsage { input_tokens, + uncached_input_tokens, output_tokens, reasoning_tokens, }), @@ -1433,6 +1462,7 @@ mod tests { r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":8,"completion_tokens":3,"completion_tokens_details":{"reasoning_tokens":2}}}"#, ExactUsage { input_tokens: Some(8), + uncached_input_tokens: None, output_tokens: Some(3), reasoning_tokens: Some(2), }, @@ -1442,15 +1472,67 @@ mod tests { r#"{"type":"response.completed","response":{"output":[{"type":"message"}],"usage":{"input_tokens":5,"output_tokens":2,"output_tokens_details":{"reasoning_tokens":1}}}}"#, ExactUsage { input_tokens: Some(5), + uncached_input_tokens: None, output_tokens: Some(2), reasoning_tokens: Some(1), }, false, ), + ( + r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":100,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":64}}}"#, + ExactUsage { + input_tokens: Some(100), + uncached_input_tokens: Some(36), + output_tokens: Some(3), + reasoning_tokens: None, + }, + false, + ), + ( + r#"{"type":"response.completed","response":{"output":[{"type":"message"}],"usage":{"input_tokens":80,"output_tokens":2,"input_tokens_details":{"cached_tokens":30}}}}"#, + ExactUsage { + input_tokens: Some(80), + uncached_input_tokens: Some(50), + output_tokens: Some(2), + reasoning_tokens: None, + }, + false, + ), + ( + r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":8}}}"#, + ExactUsage { + input_tokens: Some(5), + uncached_input_tokens: None, + output_tokens: Some(3), + reasoning_tokens: None, + }, + true, + ), + ( + r#"{"type":"response.completed","response":{"output":[{"type":"message"}],"usage":{"input_tokens":5,"output_tokens":2,"input_tokens_details":{"cached_tokens":8}}}}"#, + ExactUsage { + input_tokens: Some(5), + uncached_input_tokens: None, + output_tokens: Some(2), + reasoning_tokens: None, + }, + true, + ), + ( + r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":100,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":"64"}}}"#, + ExactUsage { + input_tokens: Some(100), + uncached_input_tokens: None, + output_tokens: Some(3), + reasoning_tokens: None, + }, + true, + ), ( r#"{"object":"chat.completion.chunk","output_tokens_so_far":7,"choices":[]}"#, ExactUsage { input_tokens: None, + uncached_input_tokens: None, output_tokens: Some(7), reasoning_tokens: None, }, @@ -1460,6 +1542,7 @@ mod tests { r#"{"object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":8,"completion_tokens_details":{"reasoning_tokens":2,"rejected_prediction_tokens":1}}}"#, ExactUsage { input_tokens: Some(8), + uncached_input_tokens: None, output_tokens: None, reasoning_tokens: Some(2), }, @@ -1546,6 +1629,7 @@ mod tests { parsed.facts.exact_usage, Some(ExactUsage { input_tokens: None, + uncached_input_tokens: None, output_tokens: Some(5), reasoning_tokens: None, }) diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs index a75bce1b46..7c187da734 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs @@ -139,18 +139,59 @@ fn calibration_factor(ratios: &VecDeque) -> Option { } } +// Entries newer than a pending request wait for its usage. Past either bound +// the oldest pending request is dropped rather than resolved with total tokens +// that may include cached prompt tokens. Its usage can still re-enter later. +// At roughly 100 bytes per entry, 1024 entries cost about 100 KiB per model. +const MAX_RETAINED_INPUT_INTERVALS: usize = 1024; +// Large smoothing windows still keep room for held-back entries. +const MIN_RETAINED_INPUT_INTERVAL_WINDOWS: usize = 8; +// Measured against the newest observed first output, not a wall clock, so it +// bounds staleness relative to newer requests. Decodes longer than this after +// first output stop holding back newer requests. +const MAX_PENDING_INPUT_INTERVAL_LAG: Duration = Duration::from_secs(120); + #[derive(Debug)] struct RetainedInputInterval { request_id: String, interval: RequestInputInterval, input_tokens: u64, input_tokens_explicit: bool, + max_input_tps_eligible: bool, + pending: bool, +} + +pub(super) struct InputIntervalSample<'a> { + pub(super) request_id: &'a str, + pub(super) interval: RequestInputInterval, + pub(super) input_tokens: u64, + pub(super) input_tokens_explicit: bool, + // True when `input_tokens` is known to count prefill work: usage split out + // cached prompt tokens, or the request is a Pylon calibration probe. + pub(super) max_input_tps_eligible: bool, + // True while the request still expects usage that has not arrived. The entry + // keeps its first-output position but is excluded from the rate until resolved. + pub(super) pending: bool, } +#[derive(Clone, Copy, Debug, PartialEq)] +pub(super) struct InputIntervalRate { + pub(super) input_tps: f64, + pub(super) max_input_tps_eligible: bool, +} + +// Entries are ordered by first output. The rate covers the resolved entries +// before the oldest pending entry, so it always describes a contiguous run of +// requests; newer resolved entries are held back until that request resolves. #[derive(Debug, Default)] pub(super) struct RequestInputIntervalWindow { intervals: VecDeque, evicted_through: Option, + // Pending requests dropped at a bound. Their later events must not reserve a + // new slot, but their usage may re-enter as a resolved entry. + dropped_pending: VecDeque, + newest_first_output: Option, + published: Option, has_observed_rate: bool, } @@ -166,48 +207,82 @@ impl RequestInputIntervalWindow { pub(super) fn observe( &mut self, - request_id: &str, - interval: RequestInputInterval, - input_tokens: u64, - input_tokens_explicit: bool, + sample: InputIntervalSample<'_>, config: &StatsCollectorConfig, - ) -> Option { + ) -> Option { + let InputIntervalSample { + request_id, + interval, + input_tokens, + input_tokens_explicit, + max_input_tps_eligible, + pending, + } = sample; if config.smoothing_window_size == 0 || interval.first_generated_output_at <= interval.submitted_at { return None; } + self.newest_first_output = Some( + self.newest_first_output + .map_or(interval.first_generated_output_at, |newest| { + newest.max(interval.first_generated_output_at) + }), + ); let existing = self.intervals.iter().position(|entry| { entry.request_id == request_id && entry.interval.submitted_at == interval.submitted_at }); let min_input_tokens = config.min_input_tokens.max(1); - if input_tokens_explicit && input_tokens < min_input_tokens { - if let Some(index) = existing { + if input_tokens < min_input_tokens { + let resolves_pending = + !pending && existing.is_some_and(|index| self.intervals[index].pending); + if let Some(index) = existing + && (input_tokens_explicit || resolves_pending) + { self.intervals.remove(index); - return self.record_current_rate(config.duration_floor); + return self.refresh(config); } return None; } - if input_tokens < min_input_tokens { - return None; - } if let Some(index) = existing { let entry = self .intervals .get_mut(index) .expect("located input interval should remain retained"); - if entry.input_tokens_explicit && !input_tokens_explicit { + if pending && !entry.pending { return None; } - let rate_changed = entry.interval != interval || entry.input_tokens != input_tokens; - entry.interval = interval; - entry.input_tokens = input_tokens; - entry.input_tokens_explicit |= input_tokens_explicit; - if !rate_changed { + let keep_values = (entry.input_tokens_explicit && !input_tokens_explicit) + || (entry.max_input_tps_eligible && !max_input_tps_eligible); + let resolves = entry.pending && !pending; + if keep_values && !resolves { + return None; + } + let changed = resolves + || (!keep_values + && (entry.interval != interval + || entry.input_tokens != input_tokens + || entry.max_input_tps_eligible != max_input_tps_eligible)); + entry.pending = pending; + if !keep_values { + entry.interval = interval; + entry.input_tokens = input_tokens; + entry.input_tokens_explicit |= input_tokens_explicit; + entry.max_input_tps_eligible = max_input_tps_eligible; + } + if !changed { return None; } } else { + if let Some(index) = self.dropped_pending.iter().position(|key| { + key.request_id == request_id && key.submitted_at == interval.submitted_at + }) { + if pending { + return None; + } + self.dropped_pending.remove(index); + } if self .evicted_through .is_some_and(|evicted| interval.first_generated_output_at <= evicted) @@ -228,40 +303,86 @@ impl RequestInputIntervalWindow { interval, input_tokens, input_tokens_explicit, + max_input_tps_eligible, + pending, }, ); - while self.intervals.len() > config.smoothing_window_size { + } + self.refresh(config) + } + + fn refresh(&mut self, config: &StatsCollectorConfig) -> Option { + self.enforce_bounds(config.smoothing_window_size); + let rate = self.rate(config.duration_floor); + self.has_observed_rate |= rate.is_some(); + if rate.is_none() || rate == self.published { + return None; + } + self.published = rate; + rate + } + + fn enforce_bounds(&mut self, window_size: usize) { + let retained_limit = MAX_RETAINED_INPUT_INTERVALS + .max(window_size.saturating_mul(MIN_RETAINED_INPUT_INTERVAL_WINDOWS)); + loop { + let resolved_prefix = self + .intervals + .iter() + .position(|entry| entry.pending) + .unwrap_or(self.intervals.len()); + for _ in window_size..resolved_prefix { let evicted = self .intervals .pop_front() - .expect("oversized input interval window should not be empty"); - self.evicted_through = Some( - self.evicted_through - .map_or(evicted.interval.first_generated_output_at, |prior| { - prior.max(evicted.interval.first_generated_output_at) - }), - ); + .expect("resolved input interval prefix should not be empty"); + self.advance_evicted_through(evicted.interval.first_generated_output_at); + } + let oldest_pending_index = resolved_prefix.min(window_size); + let Some(oldest_pending) = self.intervals.get(oldest_pending_index) else { + return; + }; + let first_output = oldest_pending.interval.first_generated_output_at; + let lagging = self.newest_first_output.is_some_and(|newest| { + newest.saturating_duration_since(first_output) > MAX_PENDING_INPUT_INTERVAL_LAG + }); + if self.intervals.len() <= retained_limit && !lagging { + return; + } + // Release newer entries without counting cached prompt tokens. The + // dropped request's usage re-enters at its first-output position + // unless newer resolved entries have already been evicted. + let dropped = self + .intervals + .remove(oldest_pending_index) + .expect("oldest pending input interval should be retained"); + self.dropped_pending.push_back(RequestIntervalKey { + request_id: dropped.request_id, + submitted_at: dropped.interval.submitted_at, + }); + while self.dropped_pending.len() > retained_limit { + self.dropped_pending.pop_front(); } } - self.record_current_rate(config.duration_floor) } - fn record_current_rate(&mut self, duration_floor: Duration) -> Option { - let rate = self.rate(duration_floor); - self.has_observed_rate |= rate.is_some(); - rate + fn advance_evicted_through(&mut self, first_output: Instant) { + self.evicted_through = Some( + self.evicted_through + .map_or(first_output, |prior| prior.max(first_output)), + ); } - fn rate(&self, duration_floor: Duration) -> Option { - let input_tokens = self.intervals.iter().fold(0_u64, |total, entry| { + fn rate(&self, duration_floor: Duration) -> Option { + let rate_window = || self.intervals.iter().take_while(|entry| !entry.pending); + let input_tokens = rate_window().fold(0_u64, |total, entry| { total.saturating_add(entry.input_tokens) }); if input_tokens == 0 { return None; } - let mut intervals = self - .intervals - .iter() + let max_input_tps_eligible = rate_window().all(|entry| entry.max_input_tps_eligible); + let mut intervals = rate_window() .map(|entry| entry.interval) .collect::>(); intervals.sort_unstable_by_key(|interval| interval.submitted_at); @@ -283,7 +404,10 @@ impl RequestInputIntervalWindow { union_duration = union_duration.saturating_add(end.saturating_duration_since(start)); let duration = union_duration.max(duration_floor); let input_tps = input_tokens as f64 / duration.as_secs_f64(); - valid_last_mean_input_tps(input_tps).then_some(input_tps) + valid_last_mean_input_tps(input_tps).then_some(InputIntervalRate { + input_tps, + max_input_tps_eligible, + }) } } #[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq)] @@ -1062,7 +1186,7 @@ pub(super) fn apply_input_throughput_sample( if model_state.last_mean_input_tps == mean_input_tps { return false; } - model_state.publish_mean_input_tps(mean_input_tps); + model_state.publish_mean_input_tps(mean_input_tps, true); true } @@ -1091,12 +1215,20 @@ pub(super) struct ModelStatsSnapshotInputs { impl ModelMetricsState { // The maximum follows the published smoothed mean rather than raw samples, // so the mean dilutes a single outlier sample before it can raise the weight. - pub(super) fn publish_mean_input_tps(&mut self, input_tps: f64) { + // Callers pass update_max = false for a fallback window that may count cached + // prompt tokens as prefill; the mean still moves, the maximum does not. + pub(super) fn publish_mean_input_tps(&mut self, input_tps: f64, update_max: bool) -> bool { + let mean_changed = self.last_mean_input_tps != input_tps; self.last_mean_input_tps = input_tps; - self.max_input_tps = Some( - self.max_input_tps - .map_or(input_tps, |max| max.max(input_tps)), - ); + if !update_max { + return mean_changed; + } + let next_max = self + .max_input_tps + .map_or(input_tps, |max| max.max(input_tps)); + let max_changed = self.max_input_tps != Some(next_max); + self.max_input_tps = Some(next_max); + mean_changed || max_changed } pub(super) fn clear_live_output_tps(&mut self) -> bool { @@ -1292,3 +1424,435 @@ mod output_token_calibration_tests { assert_eq!(calibration.scale(1), 1); } } + +#[cfg(test)] +mod request_input_interval_window_tests { + use super::*; + + const UNCACHED: u64 = 100; + const TOTAL: u64 = 150; + + fn config(smoothing_window_size: usize) -> StatsCollectorConfig { + StatsCollectorConfig { + smoothing_window_size, + ..Default::default() + } + } + + fn interval(start: Instant, submitted_ms: u64) -> RequestInputInterval { + let submitted_at = start + Duration::from_millis(submitted_ms); + RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + Duration::from_millis(500), + } + } + + // A deferred request before usage: header token estimate, held as pending. + fn placeholder(request_id: &str, interval: RequestInputInterval) -> InputIntervalSample<'_> { + InputIntervalSample { + request_id, + interval, + input_tokens: TOTAL, + input_tokens_explicit: false, + max_input_tps_eligible: false, + pending: true, + } + } + + // A terminal observation whose usage reported cached prompt tokens. + fn usage(request_id: &str, interval: RequestInputInterval) -> InputIntervalSample<'_> { + InputIntervalSample { + request_id, + interval, + input_tokens: UNCACHED, + input_tokens_explicit: true, + max_input_tps_eligible: true, + pending: false, + } + } + + // A request without expected usage: resolved immediately with total tokens. + fn live(request_id: &str, interval: RequestInputInterval) -> InputIntervalSample<'_> { + InputIntervalSample { + pending: false, + ..placeholder(request_id, interval) + } + } + + fn retained(window: &RequestInputIntervalWindow) -> Vec<(&str, bool)> { + window + .intervals + .iter() + .map(|entry| (entry.request_id.as_str(), entry.pending)) + .collect() + } + + fn current(window: &RequestInputIntervalWindow) -> Option { + window.rate(config(1).duration_floor) + } + + fn assert_tps(rate: Option, expected: f64, eligible: bool) { + let rate = rate.expect("input rate should be available"); + assert!( + (rate.input_tps - expected).abs() < 1e-6, + "expected {expected} TPS, got {}", + rate.input_tps + ); + assert_eq!(rate.max_input_tps_eligible, eligible); + } + + #[test] + fn long_decodes_keep_the_rate_window_contiguous() { + let config = config(4); + let start = Instant::now(); + let ids = (0..40) + .map(|index| format!("req-{index}")) + .collect::>(); + // Requests arrive every 100 ms with 500 ms prefill. Even requests finish + // 50 ms after first output; odd requests decode for about 2 s. + let mut events = Vec::new(); + for index in 0..40_u64 { + let first_output_ms = index * 100 + 500; + let decode_ms = if index % 2 == 0 { 50 } else { 2_050 }; + events.push((first_output_ms, index as usize, true)); + events.push((first_output_ms + decode_ms, index as usize, false)); + } + events.sort_by_key(|(at_ms, _, _)| *at_ms); + + let mut window = RequestInputIntervalWindow::default(); + let mut published = Vec::new(); + let mut contiguous = false; + for (_, index, first_output) in events { + let request_interval = interval(start, index as u64 * 100); + let sample = if first_output { + placeholder(&ids[index], request_interval) + } else { + usage(&ids[index], request_interval) + }; + if let Some(rate) = window.observe(sample, &config) { + published.push(rate.input_tps); + } + if contiguous { + assert_tps(current(&window), 500.0, true); + } else { + contiguous = + current(&window).is_some_and(|rate| (rate.input_tps - 500.0).abs() < 1e-6); + } + } + + assert!(contiguous); + assert_tps(window.published, 500.0, true); + for rate in published { + assert!( + (rate - 400.0 / 1.1).abs() > 1e-3 && (rate - 400.0 / 1.4).abs() > 1e-3, + "published a non-contiguous window rate {rate}" + ); + } + } + + #[test] + fn pending_entry_holds_back_newer_entries_until_resolved() { + let config = config(2); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + + assert_tps( + window.observe(usage("z", interval(start, 0)), &config), + 200.0, + true, + ); + assert!( + window + .observe(placeholder("a", interval(start, 100)), &config) + .is_none() + ); + assert!( + window + .observe(usage("b", interval(start, 200)), &config) + .is_none() + ); + assert!( + window + .observe(placeholder("a", interval(start, 100)), &config) + .is_none() + ); + assert_eq!(retained(&window), [("z", false), ("a", true), ("b", false)]); + + assert_tps( + window.observe(usage("a", interval(start, 100)), &config), + 200.0 / 0.6, + true, + ); + assert_eq!(retained(&window), [("a", false), ("b", false)]); + } + + #[test] + fn retained_limit_drops_the_oldest_pending_entry() { + let config = config(1); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + let earlier = InputIntervalSample { + input_tokens: 400, + ..usage("z", interval(start, 0)) + }; + assert_tps(window.observe(earlier, &config), 800.0, true); + window.observe(placeholder("a", interval(start, 100)), &config); + let ids = (0..MAX_RETAINED_INPUT_INTERVALS - 2) + .map(|index| format!("n-{index}")) + .collect::>(); + for (index, request_id) in ids.iter().enumerate() { + let request_interval = interval(start, 200 + index as u64 * 100); + assert!( + window + .observe(usage(request_id, request_interval), &config) + .is_none(), + "held-back entries must not publish while under the limit" + ); + } + assert_eq!(window.len(), MAX_RETAINED_INPUT_INTERVALS); + + let last = format!("n-{}", ids.len()); + let last_interval = interval(start, 200 + ids.len() as u64 * 100); + assert_tps( + window.observe(usage(&last, last_interval), &config), + 200.0, + true, + ); + assert_eq!(retained(&window), [(last.as_str(), false)]); + + assert!( + window + .observe(usage("a", interval(start, 100)), &config) + .is_none(), + "usage for a dropped request must be ignored" + ); + assert_eq!(window.len(), 1); + } + + #[test] + fn lagging_pending_entry_is_dropped_by_newer_first_outputs() { + let config = config(8); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + let ids = (0..8).map(|index| format!("r-{index}")).collect::>(); + for (index, request_id) in ids.iter().enumerate() { + window.observe( + usage(request_id, interval(start, index as u64 * 1_000)), + &config, + ); + } + assert_tps(window.published, 200.0, true); + window.observe(placeholder("stuck", interval(start, 8_000)), &config); + + // Faster requests arrive every second after the stuck one. + let lag_ms = MAX_PENDING_INPUT_INTERVAL_LAG.as_millis() as u64; + let ids = (0..=lag_ms / 1_000) + .map(|index| format!("fast-{index}")) + .collect::>(); + let mut released = None; + for (index, request_id) in ids.iter().enumerate() { + let submitted_ms = 9_000 + index as u64 * 1_000; + let sample = InputIntervalSample { + input_tokens: 400, + ..usage(request_id, interval(start, submitted_ms)) + }; + let rate = window.observe(sample, &config); + if submitted_ms - 8_000 <= lag_ms { + assert!( + rate.is_none(), + "published before the lag bound at {submitted_ms} ms" + ); + } else { + released = rate; + } + } + assert_tps(released, 800.0, true); + assert!(retained(&window).iter().all(|(id, _)| *id != "stuck")); + + assert!( + window + .observe(usage("stuck", interval(start, 8_000)), &config) + .is_none(), + "usage for a dropped request must be ignored" + ); + assert!(retained(&window).iter().all(|(id, _)| *id != "stuck")); + } + + // Requests arrive every `spacing_ms` with 500 ms prefill and report 100 + // uncached of 150 total prompt tokens `decode_ms` after first output. The + // window must publish only eligible rates and settle on `expected` while + // requests are still arriving. + fn assert_steady_schedule(requests: u64, spacing_ms: u64, decode_ms: u64, expected: f64) { + let config = config(8); + let start = Instant::now(); + let ids = (0..requests) + .map(|index| format!("req-{index}")) + .collect::>(); + let mut events = Vec::new(); + for index in 0..requests { + let first_output_ms = index * spacing_ms + 500; + events.push((first_output_ms, 1, index)); + events.push((first_output_ms + decode_ms, 0, index)); + } + events.sort_unstable(); + let arrivals_end_ms = (requests - 1) * spacing_ms + 500; + + let mut window = RequestInputIntervalWindow::default(); + let mut first_published_ms = None; + let mut steady_state = false; + for (at_ms, kind, index) in events { + let request_id = &ids[index as usize]; + let request_interval = interval(start, index * spacing_ms); + let sample = if kind == 1 { + placeholder(request_id, request_interval) + } else { + usage(request_id, request_interval) + }; + if let Some(rate) = window.observe(sample, &config) { + assert!( + rate.max_input_tps_eligible, + "published a rate that counts cached tokens at {at_ms} ms" + ); + first_published_ms.get_or_insert(at_ms); + } + if steady_state { + assert_tps(current(&window), expected, true); + } else if at_ms <= arrivals_end_ms { + steady_state = + current(&window).is_some_and(|rate| (rate.input_tps - expected).abs() < 1e-6); + } + } + assert!( + first_published_ms.is_some_and(|at_ms| at_ms < arrivals_end_ms), + "first publication at {first_published_ms:?} ms, arrivals end at {arrivals_end_ms} ms" + ); + assert!( + steady_state, + "steady state must be reached while requests arrive" + ); + assert_tps(window.published, expected, true); + } + + #[test] + fn high_concurrency_long_decodes_publish_only_uncached_rates() { + // 10 requests per second with 10 s decodes: about 100 in flight. + assert_steady_schedule(600, 100, 10_000, 800.0 / 1.2); + } + + #[test] + fn decodes_past_the_lag_bound_still_publish_while_requests_arrive() { + // 1 request per second with 150 s decodes: every slot hits the lag bound. + assert_steady_schedule(600, 1_000, 150_000, 800.0 / 4.0); + } + + #[test] + fn concurrency_past_the_retained_limit_still_publishes_while_requests_arrive() { + // 20 requests per second with 60 s decodes: about 1200 in flight. + assert_steady_schedule(4_000, 50, 60_000, 800.0 / 0.85); + } + + #[test] + fn dropped_pending_entry_ignores_live_events_and_reenters_with_usage() { + let config = config(2); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + assert_tps( + window.observe(usage("z", interval(start, 0)), &config), + 200.0, + true, + ); + window.observe(placeholder("a", interval(start, 1_000)), &config); + window.observe(placeholder("b", interval(start, 2_000)), &config); + let lag_ms = MAX_PENDING_INPUT_INTERVAL_LAG.as_millis() as u64; + window.observe( + placeholder("c", interval(start, 1_000 + lag_ms + 1_000)), + &config, + ); + assert_eq!(retained(&window), [("z", false), ("b", true), ("c", true)]); + + assert!( + window + .observe(placeholder("a", interval(start, 1_000)), &config) + .is_none(), + "a dropped request must not reserve a new slot" + ); + assert_eq!(window.len(), 3); + + let late_usage = InputIntervalSample { + input_tokens: 300, + ..usage("a", interval(start, 1_000)) + }; + assert_tps(window.observe(late_usage, &config), 400.0, true); + assert_eq!( + retained(&window), + [("z", false), ("a", false), ("b", true), ("c", true)] + ); + assert!(window.dropped_pending.is_empty()); + } + #[test] + fn terminal_after_eviction_is_ignored() { + let config = config(1); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + window.observe(placeholder("a", interval(start, 0)), &config); + window.observe(usage("a", interval(start, 0)), &config); + window.observe(live("b", interval(start, 100)), &config); + assert_eq!(retained(&window), [("b", false)]); + + assert!( + window + .observe(usage("a", interval(start, 0)), &config) + .is_none() + ); + assert!( + window + .observe(placeholder("a", interval(start, 0)), &config) + .is_none() + ); + assert_eq!(retained(&window), [("b", false)]); + } + + #[test] + fn mixed_live_and_deferred_entries_share_one_ordered_window() { + let config = config(2); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + + assert_tps( + window.observe(live("l0", interval(start, 0)), &config), + 300.0, + false, + ); + window.observe(placeholder("d1", interval(start, 100)), &config); + assert!( + window + .observe(live("l2", interval(start, 200)), &config) + .is_none(), + "a live entry newer than a pending entry is held back" + ); + + assert_tps( + window.observe(usage("d1", interval(start, 100)), &config), + (UNCACHED + TOTAL) as f64 / 0.6, + false, + ); + assert_eq!(retained(&window), [("d1", false), ("l2", false)]); + + assert!( + window + .observe(live("l0", interval(start, 0)), &config) + .is_none(), + "an evicted live request must not re-enter" + ); + let downgrade = InputIntervalSample { + input_tokens: TOTAL, + max_input_tps_eligible: false, + ..usage("d1", interval(start, 100)) + }; + assert!( + window.observe(downgrade, &config).is_none(), + "usage without cached-token data must not downgrade a resolved entry" + ); + assert!(window.intervals[0].max_input_tps_eligible); + } +} diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 2f8d42f79e..2d109e7f42 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -807,6 +807,8 @@ mod tests { changed_generations: Vec::new(), input_interval, input_tokens_explicit: false, + uncached_input_tokens: None, + input_usage_expected: false, output_calibration, upstream_duration: None, }, @@ -1352,6 +1354,630 @@ mod tests { assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 2_000.0); } + #[test] + fn fallback_input_tps_uses_uncached_tokens_when_usage_reports_them() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let submitted_at = std::time::Instant::now(); + let observation = completed_observation(100, 1, 10, seconds(1), seconds(2)); + let mut event = aggregator + .runtime_state + .transition_request_observation(observation); + event.input_interval = Some(crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }); + event.input_tokens_explicit = true; + event.uncached_input_tokens = Some(40); + + aggregator.apply_fallback_observation(&event); + + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 40.0); + } + + #[test] + fn fallback_input_tps_waits_for_expected_terminal_usage() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let submitted_at = std::time::Instant::now(); + let interval = crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }; + let mut live = + aggregator + .runtime_state + .transition_request_observation(completed_observation( + 100, + 1, + 10, + seconds(1), + seconds(2), + )); + live.observation.state = RequestObservationState::OutputGeneration; + live.input_interval = Some(interval); + live.input_usage_expected = true; + aggregator.apply_fallback_observation(&live); + + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 0.0); + assert_eq!( + aggregator + .per_model + .get("model-a") + .expect("model state should exist") + .metrics + .request_input_intervals + .len(), + 1, + "the request holds its window position while its usage is pending" + ); + + let mut terminal = live; + terminal.observation.state = RequestObservationState::Complete; + terminal.input_tokens_explicit = true; + terminal.uncached_input_tokens = Some(40); + aggregator.apply_fallback_observation(&terminal); + + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 40.0); + } + + #[test] + fn expected_usage_without_cached_breakdown_falls_back_at_completion() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let submitted_at = std::time::Instant::now(); + let mut event = + aggregator + .runtime_state + .transition_request_observation(completed_observation( + 100, + 1, + 10, + seconds(1), + seconds(2), + )); + event.input_interval = Some(crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }); + event.input_tokens_explicit = true; + event.input_usage_expected = true; + + aggregator.apply_fallback_observation(&event); + + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 100.0); + } + + #[test] + fn fallback_input_tps_keeps_total_tokens_when_cache_usage_is_absent() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let submitted_at = std::time::Instant::now(); + let observation = completed_observation(100, 1, 10, seconds(1), seconds(2)); + let mut event = aggregator + .runtime_state + .transition_request_observation(observation); + event.input_interval = Some(crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }); + event.input_tokens_explicit = true; + + aggregator.apply_fallback_observation(&event); + + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 100.0); + } + + /// Prompt usage carried by one fallback chat observation. + #[derive(Clone, Copy, Debug)] + enum TestInputUsage { + /// No usage payload; the request-side prompt estimate is used. + Absent, + /// Usage reports prompt tokens without a cached-token breakdown. + TotalOnly, + /// Usage reports a cached-token breakdown. + Cached { uncached: u64 }, + } + + fn input_usage_observation( + request_id: &str, + input_tokens: u64, + state: RequestObservationState, + ) -> RequestObservation { + RequestObservation { + state, + ..identified( + completed_observation(input_tokens, 1, 10, seconds(1), seconds(2)), + request_id, + ) + } + } + + fn prefill_interval( + start: std::time::Instant, + submitted_after: Duration, + prefill: Duration, + ) -> crate::runtime_state::RequestInputInterval { + let submitted_at = start + submitted_after; + crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + prefill, + } + } + + fn apply_input_usage_observation( + aggregator: &mut StatsAggregator, + observation: &RequestObservation, + interval: crate::runtime_state::RequestInputInterval, + usage: TestInputUsage, + input_usage_expected: bool, + ) -> CurrentModelStats { + let mut event = aggregator + .runtime_state + .transition_request_observation(observation.clone()); + event.input_interval = Some(interval); + event.input_usage_expected = input_usage_expected; + (event.input_tokens_explicit, event.uncached_input_tokens) = match usage { + TestInputUsage::Absent => (false, None), + TestInputUsage::TotalOnly => (true, None), + TestInputUsage::Cached { uncached } => (true, Some(uncached)), + }; + aggregator.apply_fallback_observation(&event); + aggregator.snapshot("model-a") + } + + #[track_caller] + fn assert_tps_near(actual: f64, expected: f64, context: &str) { + assert!( + (actual - expected).abs() <= expected.abs().max(1.0) * 1e-9, + "{context}: expected {expected} TPS, got {actual} TPS" + ); + } + + #[test] + fn fallback_total_token_input_tps_moves_mean_without_raising_max() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let start = std::time::Instant::now(); + + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation("total-absent", 100, RequestObservationState::Complete), + prefill_interval(start, Duration::ZERO, seconds(1)), + TestInputUsage::Absent, + false, + ); + assert_tps_near(stats.last_mean_input_tps, 100.0, "usage-absent mean"); + assert_eq!( + stats.max_input_tps, None, + "a request without usage must not raise max_input_tps" + ); + + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation("total-only", 300, RequestObservationState::Complete), + prefill_interval(start, seconds(2), seconds(1)), + TestInputUsage::TotalOnly, + false, + ); + assert_tps_near(stats.last_mean_input_tps, 200.0, "total-only mean"); + assert_eq!( + stats.max_input_tps, None, + "usage without cached_tokens must not raise max_input_tps" + ); + + // A one-entry window isolates the total-token rule from mixed windows: + // the prior cache-aware maximum stays while total-token means exceed it. + let mut aggregator = test_aggregator(config!(smoothing_window_size: 1)); + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation("cached", 100, RequestObservationState::Complete), + prefill_interval(start, Duration::ZERO, seconds(1)), + TestInputUsage::Cached { uncached: 40 }, + false, + ); + assert_tps_near(stats.last_mean_input_tps, 40.0, "cache-aware mean"); + assert_eq!(stats.max_input_tps, Some(40.0)); + + for (index, (input_tokens, usage)) in [ + (500, TestInputUsage::TotalOnly), + (900, TestInputUsage::Absent), + ] + .into_iter() + .enumerate() + { + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation( + &format!("total-after-cached-{index}"), + input_tokens, + RequestObservationState::Complete, + ), + prefill_interval(start, seconds(2 + 2 * index as u64), seconds(1)), + usage, + false, + ); + assert_tps_near( + stats.last_mean_input_tps, + input_tokens as f64, + &format!("{usage:?} mean"), + ); + assert_eq!( + stats.max_input_tps, + Some(40.0), + "{usage:?} observation must keep the prior cache-aware maximum" + ); + } + } + + #[test] + fn fallback_max_input_tps_waits_for_total_token_entries_to_leave_the_window() { + let mut aggregator = test_aggregator(config!(smoothing_window_size: 2)); + let start = std::time::Instant::now(); + // Disjoint one-second prefills make each window mean the average of + // the retained entries' token counts. + for (index, (input_tokens, usage, mean, maximum, context)) in [ + ( + 400, + TestInputUsage::Cached { uncached: 400 }, + 400.0, + 400.0, + "cache-aware entry sets the maximum", + ), + ( + 100, + TestInputUsage::TotalOnly, + 250.0, + 400.0, + "total-token entry lowers the mean", + ), + ( + 900, + TestInputUsage::Cached { uncached: 900 }, + 500.0, + 400.0, + "mean above the maximum still includes the total-token entry", + ), + ( + 700, + TestInputUsage::Cached { uncached: 700 }, + 800.0, + 800.0, + "maximum rises once the total-token entry leaves the window", + ), + ] + .into_iter() + .enumerate() + { + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation( + &format!("mixed-window-{index}"), + input_tokens, + RequestObservationState::Complete, + ), + prefill_interval(start, seconds(3 * index as u64), seconds(1)), + usage, + false, + ); + assert_tps_near(stats.last_mean_input_tps, mean, context); + assert_eq!(stats.max_input_tps, Some(maximum), "{context}"); + } + } + + /// Total-token prompt counts without a cached-token breakdown, as reported + /// by calibration requests: a usage-absent estimate, then prompt-only usage. + /// Each entry is (input tokens, usage, expected window mean). + const TOTAL_TOKEN_WINDOW: [(u64, TestInputUsage, f64); 2] = [ + (100, TestInputUsage::Absent, 100.0), + (300, TestInputUsage::TotalOnly, 200.0), + ]; + + fn apply_total_token_window(calibration: bool) -> Vec { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let generation = aggregator + .current_generation("model-a") + .cloned() + .expect("test model generation should exist"); + let start = std::time::Instant::now(); + TOTAL_TOKEN_WINDOW + .into_iter() + .enumerate() + .map(|(index, (input_tokens, usage, _))| { + let request_id = if calibration { + next_generated_request_id(GeneratedRequestKind::Calibration, &generation) + } else { + format!("client-total-token-{index}") + }; + apply_input_usage_observation( + &mut aggregator, + &input_usage_observation( + &request_id, + input_tokens, + RequestObservationState::Complete, + ), + prefill_interval(start, seconds(2 * index as u64), seconds(1)), + usage, + false, + ) + }) + .collect() + } + + #[test] + fn calibration_total_token_window_publishes_mean_and_max_input_tps() { + for (stats, (_, usage, mean)) in apply_total_token_window(true) + .into_iter() + .zip(TOTAL_TOKEN_WINDOW) + { + assert_tps_near( + stats.last_mean_input_tps, + mean, + &format!("calibration {usage:?} mean"), + ); + assert_eq!( + stats.max_input_tps, + Some(mean), + "calibration {usage:?} window must raise max_input_tps" + ); + } + + // Stream mode routes calibration observations through the same + // fallback input window. + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let generation = aggregator + .current_generation("model-a") + .cloned() + .expect("test model generation should exist"); + let mut event = + aggregator + .runtime_state + .transition_request_observation(input_usage_observation( + &next_generated_request_id(GeneratedRequestKind::Calibration, &generation), + 100, + RequestObservationState::Complete, + )); + event.input_interval = Some(prefill_interval( + std::time::Instant::now(), + Duration::ZERO, + seconds(1), + )); + aggregator.apply_stream_observation(&event); + let stats = aggregator.snapshot("model-a"); + assert_tps_near(stats.last_mean_input_tps, 100.0, "stream calibration mean"); + assert_eq!(stats.max_input_tps, Some(100.0)); + } + + #[test] + fn client_total_token_window_matching_calibration_raises_only_the_mean() { + for (stats, (_, usage, mean)) in apply_total_token_window(false) + .into_iter() + .zip(TOTAL_TOKEN_WINDOW) + { + assert_tps_near( + stats.last_mean_input_tps, + mean, + &format!("client {usage:?} mean"), + ); + assert_eq!( + stats.max_input_tps, None, + "client {usage:?} window must not raise max_input_tps" + ); + } + } + + #[test] + fn deferred_input_usage_without_usage_keeps_total_tokens_without_raising_max() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + let interval = prefill_interval(std::time::Instant::now(), Duration::ZERO, seconds(1)); + + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation( + "deferred-no-usage", + 100, + RequestObservationState::OutputGeneration, + ), + interval, + TestInputUsage::Absent, + true, + ); + assert_eq!( + stats.last_mean_input_tps, 0.0, + "expected usage defers the rate until the terminal event" + ); + assert_eq!(stats.max_input_tps, None); + + // The terminal event never carried usage, so the request-side prompt + // estimate is the only token count available. + let stats = apply_input_usage_observation( + &mut aggregator, + &input_usage_observation("deferred-no-usage", 100, RequestObservationState::Complete), + interval, + TestInputUsage::Absent, + true, + ); + assert_tps_near( + stats.last_mean_input_tps, + 100.0, + "missing usage keeps total-token behavior at completion", + ); + assert_eq!( + stats.max_input_tps, None, + "missing usage must not raise max_input_tps" + ); + } + + #[test] + fn pending_deferred_request_holds_newer_completions_out_of_the_mean() { + let mut aggregator = test_aggregator(config!(smoothing_window_size: 4)); + let start = std::time::Instant::now(); + // Disjoint one-second prefills, each two seconds apart, in first-output order. + let interval = |slot: u64| prefill_interval(start, seconds(2 * slot), seconds(1)); + let deferred = |request_id: &str, + slot: u64, + uncached: Option, + aggregator: &mut TestAggregator| { + let (state, usage) = match uncached { + Some(uncached) => ( + RequestObservationState::Complete, + TestInputUsage::Cached { uncached }, + ), + None => ( + RequestObservationState::OutputGeneration, + TestInputUsage::Absent, + ), + }; + apply_input_usage_observation( + aggregator, + &input_usage_observation(request_id, uncached.unwrap_or(100), state), + interval(slot), + usage, + true, + ) + }; + + deferred("resolved-before", 0, None, &mut aggregator); + let stats = deferred("resolved-before", 0, Some(100), &mut aggregator); + assert_tps_near(stats.last_mean_input_tps, 100.0, "first resolved request"); + assert_eq!(stats.max_input_tps, Some(100.0)); + + deferred("pending", 1, None, &mut aggregator); + for (slot, request_id) in [(2, "newer-1"), (3, "newer-2")] { + deferred(request_id, slot, None, &mut aggregator); + let stats = deferred(request_id, slot, Some(300), &mut aggregator); + assert_tps_near( + stats.last_mean_input_tps, + 100.0, + &format!("{request_id} resolved behind a pending older request is held back"), + ); + assert_eq!( + stats.max_input_tps, + Some(100.0), + "{request_id} held back must not raise max_input_tps" + ); + } + + let stats = deferred("pending", 1, Some(200), &mut aggregator); + let released_tps = (100 + 200 + 300 + 300) as f64 / 4.0; + assert_tps_near( + stats.last_mean_input_tps, + released_tps, + "resolving the pending request releases the newer completions", + ); + assert_eq!( + stats.max_input_tps, + Some(released_tps), + "only the released all cache-aware window raises the maximum" + ); + } + + #[test] + fn deferred_input_usage_matches_undeferred_rate_when_long_requests_finish_late() { + // Requests start every 100 ms with a 500 ms prefill and 100 uncached + // prompt tokens each. Even requests finish 100 ms after first output; + // odd requests keep decoding for 2 s, so their usage arrives after + // more than the window size of newer requests completed. + const WINDOW_SIZE: usize = 4; + const REQUESTS: usize = 16; + let config = || config!(smoothing_window_size: WINDOW_SIZE); + let start = std::time::Instant::now(); + let prefill = milliseconds(500); + let schedule = (0..REQUESTS) + .map(|index| { + let interval = prefill_interval(start, milliseconds(100 * index as u64), prefill); + let decode = if index % 2 == 0 { + milliseconds(100) + } else { + seconds(2) + }; + (index, interval, interval.first_generated_output_at + decode) + }) + .collect::>(); + // Consecutive requests span 300 ms of staggered starts plus one prefill. + let consecutive_window_tps = + (100 * WINDOW_SIZE) as f64 / (milliseconds(300) + prefill).as_secs_f64(); + + let mut undeferred = test_aggregator(config()); + let mut stats = undeferred.snapshot("model-a"); + for (index, interval, _) in &schedule { + stats = apply_input_usage_observation( + &mut undeferred, + &input_usage_observation( + &format!("undeferred-{index}"), + 100, + RequestObservationState::OutputGeneration, + ), + *interval, + TestInputUsage::Absent, + false, + ); + } + assert_tps_near( + stats.last_mean_input_tps, + consecutive_window_tps, + "without deferral the window samples consecutive first outputs", + ); + + // Live events arrive at first output and terminal usage at completion. + let mut events = schedule + .iter() + .flat_map(|(index, interval, completed_at)| { + [ + (interval.first_generated_output_at, false, *index, *interval), + (*completed_at, true, *index, *interval), + ] + }) + .collect::>(); + events.sort_by_key(|(at, terminal, index, _)| (*at, *terminal, *index)); + + let mut deferred = test_aggregator(config()); + let mut resolved = [false; REQUESTS]; + let mut published = Vec::new(); + let mut first_contiguous_window = None; + for (_, terminal, index, interval) in events { + let (state, usage) = if terminal { + ( + RequestObservationState::Complete, + TestInputUsage::Cached { uncached: 100 }, + ) + } else { + ( + RequestObservationState::OutputGeneration, + TestInputUsage::Absent, + ) + }; + let stats = apply_input_usage_observation( + &mut deferred, + &input_usage_observation(&format!("deferred-{index}"), 100, state), + interval, + usage, + true, + ); + resolved[index] |= terminal; + published.push(stats.last_mean_input_tps); + if first_contiguous_window.is_none() && resolved[..WINDOW_SIZE].iter().all(|r| *r) { + first_contiguous_window = Some(published.len() - 1); + } + } + + let mut sequence = published + .iter() + .map(|mean| format!("{mean:.1}")) + .collect::>(); + sequence.dedup(); + let first_contiguous_window = first_contiguous_window + .expect("the first window-size requests should all resolve during the schedule"); + assert!( + published[first_contiguous_window..] + .iter() + .all(|mean| (mean - consecutive_window_tps).abs() <= 1e-6), + "once the first {WINDOW_SIZE} requests resolved every published mean must be \ + {consecutive_window_tps} TPS; published means were [{}]", + sequence.join(", ") + ); + assert_tps_near( + *published.last().expect("schedule should publish"), + consecutive_window_tps, + "final deferred mean", + ); + } + #[test] fn chat_observation_without_backend_interval_does_not_record_throughput() { let mut aggregator = test_aggregator(StatsCollectorConfig::default()); @@ -2370,28 +2996,140 @@ mod tests { .enumerate() { let submitted_at = start + seconds(index as u64 * 3); + let observation = RequestObservation { + input_tokens, + ..observation( + RequestObservationEndpoint::ChatCompletions, + &format!("max-input-{index}"), + RequestObservationState::OutputGeneration, + ) + }; + let mut event = aggregator + .runtime_state + .transition_request_observation(observation); + event.input_interval = Some(crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }); + event.input_tokens_explicit = true; + event.uncached_input_tokens = Some(input_tokens); + aggregator.apply_fallback_observation(&event); + let stats = aggregator.snapshot("model-a"); + assert_eq!(stats.last_mean_input_tps, input_tokens as f64); + assert_eq!(stats.max_input_tps, Some(maximum)); + } + aggregator.sweep(seconds(600)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(300.0)); + } + + #[test] + fn fallback_total_token_observations_do_not_raise_max_input_tps() { + let mut aggregator = test_aggregator(config!(smoothing_window_size: 1)); + let start = std::time::Instant::now(); + for input_tokens in [100, 300] { apply_fallback_observation_with_interval( &mut aggregator, &RequestObservation { input_tokens, ..observation( RequestObservationEndpoint::ChatCompletions, - &format!("max-input-{index}"), + &format!("total-input-{input_tokens}"), RequestObservationState::OutputGeneration, ) }, crate::runtime_state::RequestInputInterval { - submitted_at, - first_generated_output_at: submitted_at + seconds(1), + submitted_at: start + milliseconds(input_tokens), + first_generated_output_at: start + milliseconds(input_tokens + 1_000), }, true, ); - let stats = aggregator.snapshot("model-a"); - assert_eq!(stats.last_mean_input_tps, input_tokens as f64); - assert_eq!(stats.max_input_tps, Some(maximum)); } - aggregator.sweep(seconds(600)); - assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(300.0)); + + let stats = aggregator.snapshot("model-a"); + assert_eq!(stats.last_mean_input_tps, 300.0); + assert_eq!(stats.max_input_tps, None); + + let input_tokens = 250; + let submitted_at = start + seconds(3); + let mut cache_aware = + aggregator + .runtime_state + .transition_request_observation(RequestObservation { + input_tokens, + ..observation( + RequestObservationEndpoint::ChatCompletions, + "cache-aware-input", + RequestObservationState::Complete, + ) + }); + cache_aware.input_interval = Some(crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }); + cache_aware.input_tokens_explicit = true; + cache_aware.uncached_input_tokens = Some(input_tokens); + aggregator.apply_fallback_observation(&cache_aware); + + let stats = aggregator.snapshot("model-a"); + assert_eq!(stats.last_mean_input_tps, 250.0); + assert_eq!(stats.max_input_tps, Some(250.0)); + } + + #[test] + fn delayed_usage_observation_holds_back_newer_requests() { + let mut aggregator = test_aggregator(config!(smoothing_window_size: 4)); + let start = std::time::Instant::now(); + let mut pending = Vec::new(); + for index in 0..6 { + let submitted_at = start + milliseconds(index * 100); + let mut event = + aggregator + .runtime_state + .transition_request_observation(RequestObservation { + input_tokens: 100, + time_to_first_output: Some(milliseconds(500)), + ..observation( + RequestObservationEndpoint::ChatCompletions, + &format!("deferred-{index}"), + RequestObservationState::OutputGeneration, + ) + }); + event.input_interval = Some(crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + milliseconds(500), + }); + event.input_usage_expected = true; + aggregator.apply_fallback_observation(&event); + pending.push(event); + } + + for pending_event in pending.iter().skip(1) { + let mut terminal = pending_event.clone(); + terminal.observation.state = RequestObservationState::Complete; + terminal.input_tokens_explicit = true; + terminal.uncached_input_tokens = Some(100); + aggregator.apply_fallback_observation(&terminal); + } + assert_eq!( + aggregator.snapshot("model-a").last_mean_input_tps, + 0.0, + "newer requests wait for the oldest pending request" + ); + + let mut delayed = pending[0].clone(); + delayed.observation.state = RequestObservationState::Complete; + delayed.input_tokens_explicit = true; + delayed.uncached_input_tokens = Some(100); + aggregator.apply_fallback_observation(&delayed); + + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 500.0); + assert_eq!( + aggregator.per_model["model-a"] + .metrics + .request_input_intervals + .len(), + 4 + ); } #[test] @@ -2463,10 +3201,14 @@ mod tests { .metrics .request_input_intervals .observe( - &later_cumulative.request_id, - interval, - later_cumulative.input_tokens, - false, + super::super::aggregator::InputIntervalSample { + request_id: &later_cumulative.request_id, + interval, + input_tokens: later_cumulative.input_tokens, + input_tokens_explicit: false, + max_input_tps_eligible: false, + pending: false, + }, &config, ); @@ -2708,7 +3450,7 @@ mod tests { ) }; apply_fallback_observation_with_interval(&mut aggregator, &estimated, old_interval, false); - assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(100.0)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, None); let mut late_exact = aggregator .runtime_state @@ -2774,7 +3516,7 @@ mod tests { .is_empty() ); assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 50.0); - assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(50.0)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, None); assert_eq!( aggregator.per_model["model-a"] .metrics @@ -3835,7 +4577,13 @@ mod tests { ) .expect("test model stats should initialize"); let observation = completed_observation(20, 2, 10, seconds(2), seconds(4)); - let updated_stats = apply_fallback_observation(&mut aggregator, &observation); + let mut event = aggregator + .runtime_state + .transition_request_observation(observation); + event = event_with_test_metadata(event); + event.input_tokens_explicit = true; + event.uncached_input_tokens = Some(20); + let updated_stats = aggregator.apply_fallback_observation(&event); for (model_id, stats) in updated_stats { publish_model_stats_update(&runtime_state, model_id, stats); } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs index 385708bb4d..18ea73a62b 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs @@ -19,10 +19,10 @@ use crate::runtime_state::OutputCalibrationFacts; use crate::{CurrentModelStats, RequestObservation, RequestObservationEvent}; use super::aggregator::{ - EmbeddingThroughputSample, InputThroughputSample, KvCacheStatsSnapshot, ModelMetricsState, - ModelStatsSnapshotInputs, RequestIntervalKey, StatsAggregator, aggregate_model_state, - apply_input_throughput_sample, current_unix_millis, output_decode_duration, push_sample, - tps_for_units, + EmbeddingThroughputSample, InputIntervalSample, InputThroughputSample, KvCacheStatsSnapshot, + ModelMetricsState, ModelStatsSnapshotInputs, RequestIntervalKey, StatsAggregator, + aggregate_model_state, apply_input_throughput_sample, current_unix_millis, + output_decode_duration, push_sample, tps_for_units, }; use super::collector::StatsCollectorConfig; @@ -143,17 +143,24 @@ impl StatsAggregator { observation.endpoint, RequestObservationEndpoint::ChatCompletions | RequestObservationEndpoint::Responses ) && let Some(interval) = event.input_interval() - && let Some(input_tps) = model_state.request_input_intervals.observe( - &observation.request_id, - interval, - observation.input_tokens, - event.input_tokens_explicit(), + && let Some(window_rate) = model_state.request_input_intervals.observe( + InputIntervalSample { + request_id: &observation.request_id, + interval, + input_tokens: event + .uncached_input_tokens() + .unwrap_or(observation.input_tokens), + input_tokens_explicit: event.input_tokens_explicit(), + max_input_tps_eligible: event.uncached_input_tokens().is_some() + || generated_request_kind(&observation.request_id) + == Some(GeneratedRequestKind::Calibration), + pending: event.input_usage_expected() && !observation.is_terminal(), + }, config, ) - && model_state.last_mean_input_tps != input_tps { - model_state.publish_mean_input_tps(input_tps); - input_tps_changed = true; + input_tps_changed = model_state + .publish_mean_input_tps(window_rate.input_tps, window_rate.max_input_tps_eligible); } let mut completed_sample_recorded = false; if observation.state == RequestObservationState::Complete { diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index a55bc8adc9..149e3d856f 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -118,7 +118,26 @@ only controls calibration traffic and does not supply this fallback. Transient errors, malformed events, and EOF do not switch `auto` to fallback. Fallback reads streamed OpenAI usage fields such as `usage.completion_tokens` or -`output_tokens_so_far`. Text peeking is last resort. +`output_tokens_so_far`. For Chat Completions, usage is available only when the +client sets `stream_options.include_usage=true` or the deployment opts in to +`--force-chat-completions-include-usage` (off by default). When usage reports +`cached_tokens` (Chat `prompt_tokens_details` or Responses +`input_tokens_details`), input throughput counts only uncached prompt tokens. +Without cached details, it counts total prompt tokens. Text peeking is last +resort. + +Fallback input throughput divides the input tokens of the latest +smoothing-window requests (8 by default) by the union of their +submit-to-first-output intervals. Requests are ordered by first output. A +request expected to report usage reserves a pending slot at first output. The +rate covers the latest resolved requests before the oldest pending one, and +newer requests wait. A long decode therefore delays the published mean until +its usage arrives. Pylon drops the oldest pending request from the rate when a +newer first output is more than 120 seconds past its own, or when more than +1024 requests (or 8 windows, if larger) are retained. Its late usage can still +re-enter the rate unless newer resolved requests have already left the window. +Requests that do not request usage resolve at first output with total prompt +tokens. ## Optional KV Stats @@ -155,11 +174,21 @@ counter-derived output window. If request stats go stale, volatile output TPS is cleared. Sticky input TPS stays until a later valid sample replaces it. -The maximum is the greatest `last_mean_input_tps` published for the current -model generation, including a configured initial TPS. It follows the smoothed -mean rather than individual samples, so the mean dilutes one outlier sample. -Lower means do not reduce it, and generation replacement clears it. It is -absent until the first mean is published. +The maximum is the greatest max-eligible `last_mean_input_tps` published for +the current model generation. A configured `--initial-input-tps` seeds it. +Engine stats stream counters and embeddings samples raise it. An OpenAI +fallback mean raises it only when every request in the rate window is +max-eligible: usage reported cached-token details, or the request is a Pylon +calibration request. Other fallback means still update `last_mean_input_tps` +but do not raise the maximum, because cached input may be counted as prefill. +The maximum follows the smoothed mean rather than individual samples, so the +mean dilutes one outlier sample. Lower means do not reduce it, and generation +replacement clears it. It is absent until the first max-eligible mean. A +fallback-only backend with no `--initial-input-tps` and no calibration never +reports one unless its usage includes cached-token details. Chat traffic +without `include_usage`, or an engine that omits those details, is not enough. +Its cluster then has no maximum capacity, which default Pulsar ranking requires +(see below). Pylon publishes `pylon_model_max_input_tps` when the maximum is known and removes the series when the model is removed or replaced with unknown maximum state. @@ -182,8 +211,14 @@ Before registration, Pylon initializes each model generation from exactly one source. `--initial-input-tps` installs the configured value. Local calibration installs nothing: it runs an increasing request ramp until a load-step timeout, and those requests' exact-generation observer events build the same distribution -used at runtime. Duplicate engine events for calibration IDs are ignored. Pylon -logs the current stats at timeout without reinjecting them. +used at runtime. Calibration request windows are max-eligible, so the ramp can +raise the maximum without cached-token details. Each calibration request starts +with a unique prompt prefix, so it cannot reuse an earlier calibration prompt +from the prefix cache. Calibrated throughput counts the prompt tokens the +engine reports, uncached when reported, and falls back to the character +estimate when usage lacks `prompt_tokens`. Duplicate engine events for +calibration IDs are ignored. Pylon logs the current stats at timeout without +reinjecting them. Later valid runtime samples continue updating either unpinned distribution. Labels: