From f8afa8843a7d62f5bce657bd2e7f328733e7866e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Mon, 5 Oct 2026 23:08:35 +0300 Subject: [PATCH 01/11] fix(pylon): exclude cached prompt tokens from fallback TPS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- .../stargate/crates/mock-dynamo/src/openai.rs | 20 +++++- .../stargate/crates/mock-dynamo/src/tests.rs | 34 ++++++---- .../pylon-lib/src/quic_http_tunnel/core.rs | 2 +- .../pylon-lib/src/quic_http_tunnel/tests.rs | 30 +++++++++ .../crates/pylon-lib/src/request_observer.rs | 41 ++++++++---- .../crates/pylon-lib/src/runtime_state.rs | 8 +++ .../pylon-lib/src/sse_message_stream.rs | 66 ++++++++++++++++++- .../crates/pylon-lib/src/stats/collector.rs | 40 +++++++++++ .../crates/pylon-lib/src/stats/projection.rs | 4 +- .../stargate/docs/runtime-stats-interface.md | 4 +- 10 files changed, 215 insertions(+), 34 deletions(-) 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/pylon-lib/src/quic_http_tunnel/core.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs index b2cae508b7..845d6c0a1e 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); 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..06e4acc01e 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,7 @@ pub(crate) struct RequestObserver { request_input_tokens: u64, input_tokens: u64, input_tokens_explicit: bool, + uncached_input_tokens: Option, generation: Option, embedding_items: Option, chat_calibration: Option, @@ -206,6 +207,7 @@ impl RequestObserver { request_input_tokens: input_tokens, input_tokens, input_tokens_explicit: false, + uncached_input_tokens: None, generation, embedding_items: None, chat_calibration: (output_token_calibration_enabled @@ -291,24 +293,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 +557,7 @@ impl RequestObserver { changed_generations: Vec::new(), input_interval, input_tokens_explicit: self.input_tokens_explicit, + uncached_input_tokens: self.uncached_input_tokens, output_calibration: backend .map_or_else(OutputCalibrationFacts::default, |backend| { backend.output_calibration 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..fe39e5078e 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 @@ -123,6 +123,7 @@ 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) output_calibration: OutputCalibrationFacts, pub(crate) upstream_duration: Option, } @@ -520,6 +521,7 @@ impl PylonRuntimeState { changed_generations: Vec::new(), input_interval: None, input_tokens_explicit: false, + uncached_input_tokens: None, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, @@ -579,6 +581,7 @@ impl PylonRuntimeState { changed_generations: Vec::new(), input_interval: None, input_tokens_explicit: false, + uncached_input_tokens: None, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, @@ -728,6 +731,10 @@ impl RequestObservationEvent { self.input_tokens_explicit } + pub(crate) fn uncached_input_tokens(&self) -> Option { + self.uncached_input_tokens + } + pub(crate) fn output_calibration(&self) -> OutputCalibrationFacts { self.output_calibration } @@ -795,6 +802,7 @@ mod tests { changed_generations: Vec::new(), input_interval: None, input_tokens_explicit: false, + uncached_input_tokens: None, 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..0ddca37fd7 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,18 @@ 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 = match ( + invalid_cached_input_tokens, + input_tokens, + cached_input_tokens, + ) { + (false, Some(input), Some(cached)) => Some(input.saturating_sub(cached)), + _ => None, + }; let (output_tokens, invalid_output_tokens) = merge_u64_fields([ &usage["completion_tokens"], &response_usage["output_tokens"], @@ -780,6 +793,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 +817,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 +842,7 @@ fn exact_usage( let calibration_ineligible = usage_protocol_mismatch || invalid_container || invalid_input_tokens + || invalid_cached_input_tokens || invalid_output_tokens || invalid_reasoning_tokens || invalid_audio_tokens @@ -832,6 +856,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 +1458,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 +1468,47 @@ 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":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 +1518,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 +1605,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/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 2f8d42f79e..1e457a7d45 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,7 @@ mod tests { changed_generations: Vec::new(), input_interval, input_tokens_explicit: false, + uncached_input_tokens: None, output_calibration, upstream_duration: None, }, @@ -1352,6 +1353,45 @@ 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_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); + } + #[test] fn chat_observation_without_backend_interval_does_not_record_throughput() { let mut aggregator = test_aggregator(StatsCollectorConfig::default()); 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..01dcc62362 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 @@ -146,7 +146,9 @@ impl StatsAggregator { && let Some(input_tps) = model_state.request_input_intervals.observe( &observation.request_id, interval, - observation.input_tokens, + event + .uncached_input_tokens() + .unwrap_or(observation.input_tokens), event.input_tokens_explicit(), config, ) diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index a55bc8adc9..61d3d44b1a 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -118,7 +118,9 @@ 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`. When usage reports cached prompt tokens, input +throughput uses only the uncached prompt tokens; without that detail it keeps +using the total prompt tokens. Text peeking is last resort. ## Optional KV Stats From 2d6c3ec592dc991cfbb766ff6b4556a69cefadb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Mon, 5 Oct 2026 23:54:15 +0300 Subject: [PATCH 02/11] fix(pylon): reject inconsistent cached token counts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- .../pylon-lib/src/sse_message_stream.rs | 30 +++++++++++++++++-- 1 file changed, 27 insertions(+), 3 deletions(-) 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 0ddca37fd7..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 @@ -776,13 +776,16 @@ fn exact_usage( &usage["prompt_tokens_details"]["cached_tokens"], &response_usage["input_tokens_details"]["cached_tokens"], ]); - let uncached_input_tokens = match ( + let (uncached_input_tokens, cached_input_exceeds_total) = match ( invalid_cached_input_tokens, input_tokens, cached_input_tokens, ) { - (false, Some(input), Some(cached)) => Some(input.saturating_sub(cached)), - _ => None, + (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"], @@ -843,6 +846,7 @@ fn exact_usage( || invalid_container || invalid_input_tokens || invalid_cached_input_tokens + || cached_input_exceeds_total || invalid_output_tokens || invalid_reasoning_tokens || invalid_audio_tokens @@ -1494,6 +1498,26 @@ mod tests { }, 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 { From e944a5b6052b85b61c52497b48b3f12d3c495346 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Wed, 7 Oct 2026 21:20:43 +0300 Subject: [PATCH 03/11] fix(pylon): defer fallback rates until usage arrives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- .../pylon-lib/src/quic_http_tunnel/core.rs | 54 ++++++++++++++ .../crates/pylon-lib/src/request_observer.rs | 10 +++ .../pylon-lib/src/request_observer/tunnel.rs | 4 ++ .../crates/pylon-lib/src/runtime_state.rs | 8 +++ .../crates/pylon-lib/src/stats/collector.rs | 72 +++++++++++++++++++ .../crates/pylon-lib/src/stats/projection.rs | 3 +- .../stargate/docs/runtime-stats-interface.md | 11 ++- 7 files changed, 158 insertions(+), 4 deletions(-) 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 845d6c0a1e..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 @@ -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/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index 06e4acc01e..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 @@ -173,6 +173,7 @@ pub(crate) struct RequestObserver { input_tokens: u64, input_tokens_explicit: bool, uncached_input_tokens: Option, + input_usage_expected: bool, generation: Option, embedding_items: Option, chat_calibration: Option, @@ -208,6 +209,7 @@ impl RequestObserver { input_tokens, input_tokens_explicit: false, uncached_input_tokens: None, + input_usage_expected: false, generation, embedding_items: None, chat_calibration: (output_token_calibration_enabled @@ -224,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 { @@ -558,6 +567,7 @@ impl RequestObserver { 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 fe39e5078e..8948d35e5e 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 @@ -124,6 +124,7 @@ pub struct RequestObservationEvent { 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, } @@ -522,6 +523,7 @@ impl PylonRuntimeState { input_interval: None, input_tokens_explicit: false, uncached_input_tokens: None, + input_usage_expected: false, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, @@ -582,6 +584,7 @@ impl PylonRuntimeState { input_interval: None, input_tokens_explicit: false, uncached_input_tokens: None, + input_usage_expected: false, output_calibration: OutputCalibrationFacts::default(), upstream_duration: None, }, @@ -735,6 +738,10 @@ impl RequestObservationEvent { 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 } @@ -803,6 +810,7 @@ mod tests { 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/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 1e457a7d45..94b5a6f690 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 @@ -808,6 +808,7 @@ mod tests { input_interval, input_tokens_explicit: false, uncached_input_tokens: None, + input_usage_expected: false, output_calibration, upstream_duration: None, }, @@ -1373,6 +1374,77 @@ mod tests { 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(), + 0, + "provisional total-token rate must not enter the window" + ); + + 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()); 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 01dcc62362..bab08d965e 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 @@ -142,7 +142,8 @@ impl StatsAggregator { if matches!( observation.endpoint, RequestObservationEndpoint::ChatCompletions | RequestObservationEndpoint::Responses - ) && let Some(interval) = event.input_interval() + ) && (!event.input_usage_expected() || observation.is_terminal()) + && let Some(interval) = event.input_interval() && let Some(input_tps) = model_state.request_input_intervals.observe( &observation.request_id, interval, diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index 61d3d44b1a..1d001e18c2 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -118,9 +118,14 @@ 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`. When usage reports cached prompt tokens, input -throughput uses only the uncached prompt tokens; without that detail it keeps -using the total prompt tokens. 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). For requests expected +to report usage, input throughput is recorded at stream completion rather than +published provisionally with total prompt tokens; cached-token details use only +the uncached prompt count, and absent details fall back to the total count. +Requests that do not request usage retain the existing total-token timing +behavior. Text peeking is last resort. ## Optional KV Stats From 6ad412dc91ab849cf61eba82fd55d1ab009e6286 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Fri, 9 Oct 2026 21:39:35 +0300 Subject: [PATCH 04/11] fix(pylon): guard fallback input TPS peaks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- .../crates/pylon-lib/src/stats/aggregator.rs | 80 ++++++---- .../crates/pylon-lib/src/stats/collector.rs | 148 ++++++++++++++++-- .../crates/pylon-lib/src/stats/projection.rs | 31 ++-- .../stargate/docs/runtime-stats-interface.md | 12 +- 4 files changed, 209 insertions(+), 62 deletions(-) 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..dd48cd8441 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 @@ -145,6 +145,16 @@ struct RetainedInputInterval { interval: RequestInputInterval, input_tokens: u64, input_tokens_explicit: bool, + cache_aware: 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, + pub(super) cache_aware: bool, + pub(super) deferred_until_usage: bool, } #[derive(Debug, Default)] @@ -164,14 +174,23 @@ impl RequestInputIntervalWindow { self.has_observed_rate } + pub(super) fn all_intervals_cache_aware(&self) -> bool { + self.intervals.iter().all(|entry| entry.cache_aware) + } + pub(super) fn observe( &mut self, - request_id: &str, - interval: RequestInputInterval, - input_tokens: u64, - input_tokens_explicit: bool, + sample: InputIntervalSample<'_>, config: &StatsCollectorConfig, ) -> Option { + let InputIntervalSample { + request_id, + interval, + input_tokens, + input_tokens_explicit, + cache_aware, + deferred_until_usage, + } = sample; if config.smoothing_window_size == 0 || interval.first_generated_output_at <= interval.submitted_at { @@ -200,36 +219,31 @@ impl RequestInputIntervalWindow { if entry.input_tokens_explicit && !input_tokens_explicit { return None; } - let rate_changed = entry.interval != interval || entry.input_tokens != input_tokens; + let rate_changed = entry.interval != interval + || entry.input_tokens != input_tokens + || entry.cache_aware != cache_aware; entry.interval = interval; entry.input_tokens = input_tokens; entry.input_tokens_explicit |= input_tokens_explicit; + entry.cache_aware = cache_aware; if !rate_changed { return None; } } else { - if self - .evicted_through - .is_some_and(|evicted| interval.first_generated_output_at <= evicted) + if !deferred_until_usage + && self + .evicted_through + .is_some_and(|evicted| interval.first_generated_output_at <= evicted) { return None; } - let insertion_index = self - .intervals - .iter() - .position(|entry| { - entry.interval.first_generated_output_at > interval.first_generated_output_at - }) - .unwrap_or(self.intervals.len()); - self.intervals.insert( - insertion_index, - RetainedInputInterval { - request_id: request_id.to_string(), - interval, - input_tokens, - input_tokens_explicit, - }, - ); + self.intervals.push_back(RetainedInputInterval { + request_id: request_id.to_string(), + interval, + input_tokens, + input_tokens_explicit, + cache_aware, + }); while self.intervals.len() > config.smoothing_window_size { let evicted = self .intervals @@ -1062,7 +1076,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 +1105,18 @@ 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) { + 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 { 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 94b5a6f690..0fd4891b94 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 @@ -2482,28 +2482,136 @@ 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_is_retained_in_completion_order() { + 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, 500.0); + + 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, 400.0); + assert_eq!( + aggregator.per_model["model-a"] + .metrics + .request_input_intervals + .len(), + 4 + ); } #[test] @@ -2575,10 +2683,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, + cache_aware: false, + deferred_until_usage: false, + }, &config, ); @@ -2820,7 +2932,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 @@ -2886,7 +2998,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 @@ -3947,7 +4059,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 bab08d965e..3c6541cae2 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; @@ -145,18 +145,23 @@ impl StatsAggregator { ) && (!event.input_usage_expected() || observation.is_terminal()) && let Some(interval) = event.input_interval() && let Some(input_tps) = model_state.request_input_intervals.observe( - &observation.request_id, - interval, - event - .uncached_input_tokens() - .unwrap_or(observation.input_tokens), - event.input_tokens_explicit(), + 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(), + cache_aware: event.uncached_input_tokens().is_some(), + deferred_until_usage: 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; + let update_max_input_tps = model_state + .request_input_intervals + .all_intervals_cache_aware(); + input_tps_changed = model_state.publish_mean_input_tps(input_tps, update_max_input_tps); } 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 1d001e18c2..8391f2c8eb 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -163,10 +163,14 @@ 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. +model generation, including a configured initial TPS and authoritative engine +samples. For OpenAI fallback, a mean raises the maximum only when every retained +request interval has a valid cached-token breakdown. Total-token-only fallback +observations still update the mean but do not raise the maximum, because cached +input may be included. 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 +mean is published. 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. From a954ec75afc2b0801b54b2014ff7b83190b85dea Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 13:58:27 -0700 Subject: [PATCH 05/11] fix(pylon): keep deferred fallback input TPS windows contiguous Requests that expect usage now reserve a pending slot in first-output order instead of joining the window at completion. The published rate covers the latest resolved entries before the oldest pending one, so long decodes no longer leave gaps that inflate the interval union and under-read the rate. A retained-entry cap provisionally resolves the oldest pending entry with its total tokens so one stuck request cannot freeze the window forever. Calibration observations remain max-eligible, so a calibrated backend still publishes max_input_tps without cached-usage traffic. Eligibility is not downgraded by a later observation without cached-token data, and evicted requests cannot re-enter the window. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/proto/proto/stargate.proto | 5 +- .../crates/pylon-lib/src/runtime_state.rs | 4 +- .../crates/pylon-lib/src/stats/aggregator.rs | 439 +++++++++++++++--- .../crates/pylon-lib/src/stats/collector.rs | 18 +- .../crates/pylon-lib/src/stats/projection.rs | 17 +- 5 files changed, 410 insertions(+), 73 deletions(-) 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/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index 8948d35e5e..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. 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 dd48cd8441..42ced01992 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,13 +139,19 @@ fn calibration_factor(ratios: &VecDeque) -> Option { } } +// Entries newer than a pending request wait for its usage. Past this many +// windows of retained entries, the oldest pending request resolves with the +// tokens it has so the published rate cannot stall indefinitely. +const MAX_RETAINED_INPUT_INTERVAL_WINDOWS: usize = 8; + #[derive(Debug)] struct RetainedInputInterval { request_id: String, interval: RequestInputInterval, input_tokens: u64, input_tokens_explicit: bool, - cache_aware: bool, + max_input_tps_eligible: bool, + pending: bool, } pub(super) struct InputIntervalSample<'a> { @@ -153,14 +159,28 @@ pub(super) struct InputIntervalSample<'a> { pub(super) interval: RequestInputInterval, pub(super) input_tokens: u64, pub(super) input_tokens_explicit: bool, - pub(super) cache_aware: bool, - pub(super) deferred_until_usage: 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, + published: Option, has_observed_rate: bool, } @@ -174,22 +194,18 @@ impl RequestInputIntervalWindow { self.has_observed_rate } - pub(super) fn all_intervals_cache_aware(&self) -> bool { - self.intervals.iter().all(|entry| entry.cache_aware) - } - pub(super) fn observe( &mut self, sample: InputIntervalSample<'_>, config: &StatsCollectorConfig, - ) -> Option { + ) -> Option { let InputIntervalSample { request_id, interval, input_tokens, input_tokens_explicit, - cache_aware, - deferred_until_usage, + max_input_tps_eligible, + pending, } = sample; if config.smoothing_window_size == 0 || interval.first_generated_output_at <= interval.submitted_at @@ -200,82 +216,128 @@ impl RequestInputIntervalWindow { 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 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 rate_changed = entry.interval != interval - || entry.input_tokens != input_tokens - || entry.cache_aware != cache_aware; - entry.interval = interval; - entry.input_tokens = input_tokens; - entry.input_tokens_explicit |= input_tokens_explicit; - entry.cache_aware = cache_aware; - if !rate_changed { + 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 !deferred_until_usage - && self - .evicted_through - .is_some_and(|evicted| interval.first_generated_output_at <= evicted) + if self + .evicted_through + .is_some_and(|evicted| interval.first_generated_output_at <= evicted) { return None; } - self.intervals.push_back(RetainedInputInterval { - request_id: request_id.to_string(), - interval, - input_tokens, - input_tokens_explicit, - cache_aware, - }); - while self.intervals.len() > config.smoothing_window_size { + let insertion_index = self + .intervals + .iter() + .position(|entry| { + entry.interval.first_generated_output_at > interval.first_generated_output_at + }) + .unwrap_or(self.intervals.len()); + self.intervals.insert( + insertion_index, + RetainedInputInterval { + request_id: request_id.to_string(), + interval, + input_tokens, + input_tokens_explicit, + max_input_tps_eligible, + pending, + }, + ); + } + 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 = window_size.saturating_mul(MAX_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"); + .expect("resolved input interval prefix should not be empty"); + let first_output = evicted.interval.first_generated_output_at; self.evicted_through = Some( self.evicted_through - .map_or(evicted.interval.first_generated_output_at, |prior| { - prior.max(evicted.interval.first_generated_output_at) - }), + .map_or(first_output, |prior| prior.max(first_output)), ); } + if self.intervals.len() <= retained_limit { + return; + } + // Provisionally resolve with the tokens observed so far. Usage that + // arrives while the entry is still retained upgrades it in place. + let Some(oldest_pending) = self.intervals.iter_mut().find(|entry| entry.pending) else { + return; + }; + oldest_pending.pending = false; } - 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 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); @@ -297,7 +359,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)] @@ -1105,6 +1170,8 @@ 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. + // 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; @@ -1312,3 +1379,267 @@ 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_resolves_the_oldest_pending_entry_provisionally() { + let config = config(2); + let start = Instant::now(); + let mut window = RequestInputIntervalWindow::default(); + window.observe(usage("z", interval(start, 0)), &config); + window.observe(usage("y", interval(start, 100)), &config); + window.observe(placeholder("a", interval(start, 200)), &config); + window.observe(placeholder("b", interval(start, 300)), &config); + let limit = config.smoothing_window_size * MAX_RETAINED_INPUT_INTERVAL_WINDOWS; + let ids = (0..limit - 3) + .map(|index| format!("n-{index}")) + .collect::>(); + let mut provisional = None; + for (index, request_id) in ids.iter().enumerate() { + let request_interval = interval(start, 400 + index as u64 * 100); + if let Some(rate) = window.observe(usage(request_id, request_interval), &config) { + provisional = Some(rate); + } + } + + assert_eq!(window.len(), limit); + assert_eq!( + &retained(&window)[..3], + [("y", false), ("a", false), ("b", true)] + ); + assert_tps(provisional, (UNCACHED + TOTAL) as f64 / 0.6, false); + + assert_tps( + window.observe(usage("a", interval(start, 200)), &config), + 200.0 / 0.6, + true, + ); + } + + #[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 0fd4891b94..39310bea29 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 @@ -1406,8 +1406,8 @@ mod tests { .metrics .request_input_intervals .len(), - 0, - "provisional total-token rate must not enter the window" + 1, + "the request holds its window position while its usage is pending" ); let mut terminal = live; @@ -2562,7 +2562,7 @@ mod tests { } #[test] - fn delayed_usage_observation_is_retained_in_completion_order() { + 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(); @@ -2596,7 +2596,11 @@ mod tests { terminal.uncached_input_tokens = Some(100); aggregator.apply_fallback_observation(&terminal); } - assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 500.0); + 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; @@ -2604,7 +2608,7 @@ mod tests { delayed.uncached_input_tokens = Some(100); aggregator.apply_fallback_observation(&delayed); - assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 400.0); + assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 500.0); assert_eq!( aggregator.per_model["model-a"] .metrics @@ -2688,8 +2692,8 @@ mod tests { interval, input_tokens: later_cumulative.input_tokens, input_tokens_explicit: false, - cache_aware: false, - deferred_until_usage: false, + max_input_tps_eligible: false, + pending: false, }, &config, ); 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 3c6541cae2..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 @@ -142,9 +142,8 @@ impl StatsAggregator { if matches!( observation.endpoint, RequestObservationEndpoint::ChatCompletions | RequestObservationEndpoint::Responses - ) && (!event.input_usage_expected() || observation.is_terminal()) - && let Some(interval) = event.input_interval() - && let Some(input_tps) = model_state.request_input_intervals.observe( + ) && let Some(interval) = event.input_interval() + && let Some(window_rate) = model_state.request_input_intervals.observe( InputIntervalSample { request_id: &observation.request_id, interval, @@ -152,16 +151,16 @@ impl StatsAggregator { .uncached_input_tokens() .unwrap_or(observation.input_tokens), input_tokens_explicit: event.input_tokens_explicit(), - cache_aware: event.uncached_input_tokens().is_some(), - deferred_until_usage: event.input_usage_expected() && observation.is_terminal(), + 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, ) { - let update_max_input_tps = model_state - .request_input_intervals - .all_intervals_cache_aware(); - input_tps_changed = model_state.publish_mean_input_tps(input_tps, update_max_input_tps); + 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 { From 250aa3ef65a2ce8fd12e51ec80a3d24b42c2041c Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:01:53 -0700 Subject: [PATCH 06/11] test(pylon): cover cache-aware max and contiguous deferred input windows Add collector tests for total-token observations that move the mean without raising max_input_tps, windows that wait for total-token entries to leave before raising the maximum, deferred requests whose usage never arrives, calibration windows that still raise the maximum, a pending request holding newer completions out of the mean, and the concurrent long-decode schedule where the deferred rate must match the undeferred 500 TPS. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/pylon-lib/src/stats/collector.rs | 514 ++++++++++++++++++ 1 file changed, 514 insertions(+) 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 39310bea29..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 @@ -1464,6 +1464,520 @@ mod tests { 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()); From c3a1f5709f676958b391b105a844905332125501 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:01:54 -0700 Subject: [PATCH 07/11] docs(pylon): describe pending fallback slots and max-eligible input means Document that requests expecting usage hold a pending slot in first-output order, that the published rate covers resolved requests before the oldest pending one with a retained-entry cap, and which means may raise max_input_tps, including calibration. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../stargate/docs/runtime-stats-interface.md | 50 ++++++++++++------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index 8391f2c8eb..156c3a9b60 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -120,12 +120,22 @@ 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`. 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). For requests expected -to report usage, input throughput is recorded at stream completion rather than -published provisionally with total prompt tokens; cached-token details use only -the uncached prompt count, and absent details fall back to the total count. -Requests that do not request usage retain the existing total-token timing -behavior. Text peeking is last resort. +`--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 is the mean over the latest smoothing-window requests +(8 by default), 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. +Retained requests are capped at 8 times the window size. Past the cap, the +oldest pending request resolves provisionally with total prompt tokens and is +not max-eligible. Usage that arrives while it is still retained replaces the +provisional values. Requests that do not request usage resolve at first output +with total prompt tokens. ## Optional KV Stats @@ -162,15 +172,19 @@ 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 and authoritative engine -samples. For OpenAI fallback, a mean raises the maximum only when every retained -request interval has a valid cached-token breakdown. Total-token-only fallback -observations still update the mean but do not raise the maximum, because cached -input may be included. 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 -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`, no calibration, and Chat +traffic without `include_usage` never reports one. 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. @@ -193,8 +207,10 @@ 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. 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: From 02238b6774025b6ae9c9f4e77f4c366503f530d5 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:28:04 -0700 Subject: [PATCH 08/11] fix(pylon): drop stale pending fallback input slots At serving concurrency the retained-entry cap fired on every arrival and resolved pending entries with total prompt tokens, so the published mean counted cached tokens again and max_input_tps never rose. A lost terminal event could also freeze the mean until the cap filled. Bounds now drop the oldest pending entry instead of resolving it, and its late usage is ignored. The count cap holds 1024 entries (or 8 windows if larger), and a pending entry is also dropped once newer first outputs are more than 120 seconds past its own. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/pylon-lib/src/stats/aggregator.rs | 203 ++++++++++++++---- 1 file changed, 164 insertions(+), 39 deletions(-) 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 42ced01992..39b8decd35 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,10 +139,17 @@ fn calibration_factor(ratios: &VecDeque) -> Option { } } -// Entries newer than a pending request wait for its usage. Past this many -// windows of retained entries, the oldest pending request resolves with the -// tokens it has so the published rate cannot stall indefinitely. -const MAX_RETAINED_INPUT_INTERVAL_WINDOWS: usize = 8; +// Entries newer than a pending request wait for its usage. Past either bound +// the oldest pending request is dropped from the rate rather than resolved with +// total tokens that may include cached prompt tokens. +// 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 are dropped from the rate. +const MAX_PENDING_INPUT_INTERVAL_LAG: Duration = Duration::from_secs(120); #[derive(Debug)] struct RetainedInputInterval { @@ -180,6 +187,7 @@ pub(super) struct InputIntervalRate { pub(super) struct RequestInputIntervalWindow { intervals: VecDeque, evicted_through: Option, + newest_first_output: Option, published: Option, has_observed_rate: bool, } @@ -212,6 +220,12 @@ impl RequestInputIntervalWindow { { 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 }); @@ -298,7 +312,8 @@ impl RequestInputIntervalWindow { } fn enforce_bounds(&mut self, window_size: usize) { - let retained_limit = window_size.saturating_mul(MAX_RETAINED_INPUT_INTERVAL_WINDOWS); + let retained_limit = MAX_RETAINED_INPUT_INTERVALS + .max(window_size.saturating_mul(MIN_RETAINED_INPUT_INTERVAL_WINDOWS)); loop { let resolved_prefix = self .intervals @@ -310,24 +325,33 @@ impl RequestInputIntervalWindow { .intervals .pop_front() .expect("resolved input interval prefix should not be empty"); - let first_output = evicted.interval.first_generated_output_at; - self.evicted_through = Some( - self.evicted_through - .map_or(first_output, |prior| prior.max(first_output)), - ); + self.advance_evicted_through(evicted.interval.first_generated_output_at); } - if self.intervals.len() <= retained_limit { - return; - } - // Provisionally resolve with the tokens observed so far. Usage that - // arrives while the entry is still retained upgrades it in place. - let Some(oldest_pending) = self.intervals.iter_mut().find(|entry| entry.pending) else { + let oldest_pending_index = resolved_prefix.min(window_size); + let Some(oldest_pending) = self.intervals.get(oldest_pending_index) else { return; }; - oldest_pending.pending = false; + 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; + } + // Accept a one-request gap instead of counting cached prompt tokens. + // Advancing the guard makes the dropped request's late usage ignored. + self.intervals.remove(oldest_pending_index); + self.advance_evicted_through(first_output); } } + 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 rate_window = || self.intervals.iter().take_while(|entry| !entry.pending); let input_tokens = rate_window().fold(0_u64, |total, entry| { @@ -1542,38 +1566,139 @@ mod request_input_interval_window_tests { } #[test] - fn retained_limit_resolves_the_oldest_pending_entry_provisionally() { - let config = config(2); + fn retained_limit_drops_the_oldest_pending_entry() { + let config = config(1); let start = Instant::now(); let mut window = RequestInputIntervalWindow::default(); - window.observe(usage("z", interval(start, 0)), &config); - window.observe(usage("y", interval(start, 100)), &config); - window.observe(placeholder("a", interval(start, 200)), &config); - window.observe(placeholder("b", interval(start, 300)), &config); - let limit = config.smoothing_window_size * MAX_RETAINED_INPUT_INTERVAL_WINDOWS; - let ids = (0..limit - 3) + 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::>(); - let mut provisional = None; for (index, request_id) in ids.iter().enumerate() { - let request_interval = interval(start, 400 + index as u64 * 100); - if let Some(rate) = window.observe(usage(request_id, request_interval), &config) { - provisional = Some(rate); - } + 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); - assert_eq!(window.len(), limit); - assert_eq!( - &retained(&window)[..3], - [("y", false), ("a", false), ("b", true)] - ); - assert_tps(provisional, (UNCACHED + TOTAL) as f64 / 0.6, false); - + let last = format!("n-{}", ids.len()); + let last_interval = interval(start, 200 + ids.len() as u64 * 100); assert_tps( - window.observe(usage("a", interval(start, 200)), &config), - 200.0 / 0.6, + 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")); + } + + #[test] + fn high_concurrency_long_decodes_publish_only_uncached_rates() { + let config = config(8); + let start = Instant::now(); + let ids = (0..600) + .map(|index| format!("req-{index}")) + .collect::>(); + // 10 requests per second, 500 ms prefill, 10 s decode: about 100 in flight. + let mut events = Vec::new(); + for index in 0..600_u64 { + let first_output_ms = index * 100 + 500; + events.push((first_output_ms, 1, index as usize)); + events.push((first_output_ms + 10_000, 0, index as usize)); + } + events.sort_unstable(); + + let mut window = RequestInputIntervalWindow::default(); + let steady = 800.0 / 1.2; + let mut steady_state = false; + for (_, kind, index) in events { + let request_interval = interval(start, index as u64 * 100); + let sample = if kind == 1 { + placeholder(&ids[index], request_interval) + } else { + usage(&ids[index], request_interval) + }; + if let Some(rate) = window.observe(sample, &config) { + assert!( + rate.max_input_tps_eligible, + "published a rate that counts cached tokens" + ); + assert!((rate.input_tps - 1_000.0).abs() > 1e-3); + } + if steady_state { + assert_tps(current(&window), steady, true); + } else { + steady_state = + current(&window).is_some_and(|rate| (rate.input_tps - steady).abs() < 1e-6); + } + } + assert!(steady_state); + assert_tps(window.published, steady, true); } #[test] From 5d8e5c84e6b40f45346c8cf4438a792c9e2b3732 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:31:50 -0700 Subject: [PATCH 09/11] docs(pylon): document pending slot bounds for fallback input TPS Describe the 120 second lag bound and the retained-entry cap that drop the oldest pending request from the fallback input rate, and note that engines without cached-token details do not raise max_input_tps. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../stargate/docs/runtime-stats-interface.md | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index 156c3a9b60..f6e8d6e6d2 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -126,16 +126,17 @@ client sets `stream_options.include_usage=true` or the deployment opts in to Without cached details, it counts total prompt tokens. Text peeking is last resort. -Fallback input throughput is the mean over the latest smoothing-window requests -(8 by default), 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. -Retained requests are capped at 8 times the window size. Past the cap, the -oldest pending request resolves provisionally with total prompt tokens and is -not max-eligible. Usage that arrives while it is still retained replaces the -provisional values. Requests that do not request usage resolve at first output -with total prompt tokens. +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 is then +ignored. Requests that do not request usage resolve at first output with total +prompt tokens. ## Optional KV Stats @@ -182,9 +183,11 @@ 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`, no calibration, and Chat -traffic without `include_usage` never reports one. Its cluster then has no -maximum capacity, which default Pulsar ranking requires (see below). +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. From e9d0b86a062bbc964616787155f2f89b2ae72d3b Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:42:15 -0700 Subject: [PATCH 10/11] fix(pylon): calibrate input TPS with unique uncached prompts Calibration sent the same "1"-repeated prompt to every request in a batch, and each ramp step repeated or extended the previous prompt. With prefix caching enabled (vLLM automatic prefix caching, SGLang radix cache, KV-aware routing), nearly every calibration prefill after the first was a cache hit, so the calibrated input rate overstated uncached prefill throughput. Calibration windows are max-eligible, so the inflated rate could also set max_input_tps. Each calibration request now starts with a random per-request prefix followed by filler, keeping the prompt exactly prompt_units characters so the input header estimate and ramp semantics are unchanged. The non-streaming calibration path also parsed only completion_tokens, so calibration windows used the character-count header estimate. It now parses usage.prompt_tokens and prompt_tokens_details.cached_tokens and feeds them to the request observer, matching the streaming parser: malformed cached counts or cached counts above the prompt total leave the uncached count unavailable. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../stargate/crates/pylon-lib/src/bringup.rs | 262 +++++++++++++++++- .../pylon-lib/src/bringup/calibration.rs | 53 ++-- .../crates/pylon-lib/src/bringup/upstream.rs | 34 ++- 3 files changed, 318 insertions(+), 31 deletions(-) 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)] From 777b4aefca4a5e96aea43bd54fce01ea9e8ce561 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 14:55:08 -0700 Subject: [PATCH 11/11] fix(pylon): let dropped fallback input slots re-enter with usage When every usage-reporting request outlived the lag bound or the retained cap, each pending slot was dropped before its usage arrived and the late usage was rejected, so no runtime mean was published while traffic flowed. Bound drops now record the request instead of advancing the eviction guard. Later live events for it cannot reserve a new slot, and its usage re-enters at its first-output position unless newer resolved entries have already left the window. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/pylon-lib/src/stats/aggregator.rs | 142 ++++++++++++++---- .../stargate/docs/runtime-stats-interface.md | 13 +- 2 files changed, 124 insertions(+), 31 deletions(-) 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 39b8decd35..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 @@ -140,15 +140,15 @@ 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 from the rate rather than resolved with -// total tokens that may include cached prompt tokens. +// 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 are dropped from the rate. +// first output stop holding back newer requests. const MAX_PENDING_INPUT_INTERVAL_LAG: Duration = Duration::from_secs(120); #[derive(Debug)] @@ -187,6 +187,9 @@ pub(super) struct InputIntervalRate { 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, @@ -272,6 +275,14 @@ impl RequestInputIntervalWindow { 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) @@ -338,10 +349,20 @@ impl RequestInputIntervalWindow { if self.intervals.len() <= retained_limit && !lagging { return; } - // Accept a one-request gap instead of counting cached prompt tokens. - // Advancing the guard makes the dropped request's late usage ignored. - self.intervals.remove(oldest_pending_index); - self.advance_evicted_through(first_output); + // 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(); + } } } @@ -1657,50 +1678,117 @@ mod request_input_interval_window_tests { assert!(retained(&window).iter().all(|(id, _)| *id != "stuck")); } - #[test] - fn high_concurrency_long_decodes_publish_only_uncached_rates() { + // 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..600) + let ids = (0..requests) .map(|index| format!("req-{index}")) .collect::>(); - // 10 requests per second, 500 ms prefill, 10 s decode: about 100 in flight. let mut events = Vec::new(); - for index in 0..600_u64 { - let first_output_ms = index * 100 + 500; - events.push((first_output_ms, 1, index as usize)); - events.push((first_output_ms + 10_000, 0, index as usize)); + 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 steady = 800.0 / 1.2; + let mut first_published_ms = None; let mut steady_state = false; - for (_, kind, index) in events { - let request_interval = interval(start, index as u64 * 100); + 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(&ids[index], request_interval) + placeholder(request_id, request_interval) } else { - usage(&ids[index], request_interval) + 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" + "published a rate that counts cached tokens at {at_ms} ms" ); - assert!((rate.input_tps - 1_000.0).abs() > 1e-3); + first_published_ms.get_or_insert(at_ms); } if steady_state { - assert_tps(current(&window), steady, true); - } else { + assert_tps(current(&window), expected, true); + } else if at_ms <= arrivals_end_ms { steady_state = - current(&window).is_some_and(|rate| (rate.input_tps - steady).abs() < 1e-6); + current(&window).is_some_and(|rate| (rate.input_tps - expected).abs() < 1e-6); } } - assert!(steady_state); - assert_tps(window.published, steady, true); + 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); diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index f6e8d6e6d2..149e3d856f 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -134,9 +134,10 @@ 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 is then -ignored. Requests that do not request usage resolve at first output with total -prompt tokens. +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 @@ -211,7 +212,11 @@ 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. Calibration request windows are max-eligible, so the ramp can -raise the maximum without cached-token details. Duplicate engine events for +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.