diff --git a/src/libraries/rust/stargate/benches/backend-degradation.yaml b/src/libraries/rust/stargate/benches/backend-degradation.yaml index c17ff479a..33af7876f 100644 --- a/src/libraries/rust/stargate/benches/backend-degradation.yaml +++ b/src/libraries/rust/stargate/benches/backend-degradation.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7000 decode_tokens_per_s: 140 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 degradation: actions: diff --git a/src/libraries/rust/stargate/benches/bursty-8-backends.yaml b/src/libraries/rust/stargate/benches/bursty-8-backends.yaml index 21d0ce1e9..5a36bc5d5 100644 --- a/src/libraries/rust/stargate/benches/bursty-8-backends.yaml +++ b/src/libraries/rust/stargate/benches/bursty-8-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7200 decode_tokens_per_s: 150 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 2 profile: name: constrained @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 10 prefill_tokens_per_s: 2200 decode_tokens_per_s: 50 - registration: - last_mean_input_tps: 60 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml b/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml index 482895b8f..95e5d6c1b 100644 --- a/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml +++ b/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 3500 decode_tokens_per_s: 85 - registration: - last_mean_input_tps: 100 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml b/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml index dc588cfb8..759be16fb 100644 --- a/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml +++ b/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7000 decode_tokens_per_s: 140 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 2 profile: name: slow @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 2200 decode_tokens_per_s: 55 - registration: - last_mean_input_tps: 70 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/hotset-8-backends.yaml b/src/libraries/rust/stargate/benches/hotset-8-backends.yaml index 3c1d09041..02af1ec6b 100644 --- a/src/libraries/rust/stargate/benches/hotset-8-backends.yaml +++ b/src/libraries/rust/stargate/benches/hotset-8-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7000 decode_tokens_per_s: 140 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 2 profile: name: slow @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 2200 decode_tokens_per_s: 55 - registration: - last_mean_input_tps: 70 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml b/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml index a7cbc5df5..9a666d2d9 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 7000 decode_tokens_per_s: 170 - registration: - last_mean_input_tps: 160 profiles: - count: 2 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 12000 decode_tokens_per_s: 260 - registration: - last_mean_input_tps: 280 - count: 4 profile: name: balanced-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 7000 decode_tokens_per_s: 170 - registration: - last_mean_input_tps: 160 - count: 2 profile: name: constrained-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 7 prefill_tokens_per_s: 3400 decode_tokens_per_s: 85 - registration: - last_mean_input_tps: 80 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml b/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml index 7c28f4a43..0e986f847 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 8200 decode_tokens_per_s: 190 - registration: - last_mean_input_tps: 185 profiles: - count: 4 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 14000 decode_tokens_per_s: 285 - registration: - last_mean_input_tps: 310 - count: 8 profile: name: balanced-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 8200 decode_tokens_per_s: 190 - registration: - last_mean_input_tps: 185 - count: 4 profile: name: constrained-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 6 prefill_tokens_per_s: 4000 decode_tokens_per_s: 100 - registration: - last_mean_input_tps: 90 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml index b4fdb7d1a..1d85af82a 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 profiles: - count: 2 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 1 prefill_tokens_per_s: 65536 decode_tokens_per_s: 320 - registration: - last_mean_input_tps: 65536 - count: 4 profile: name: balanced-prefix-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 - count: 2 profile: name: constrained-prefix-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 16384 decode_tokens_per_s: 120 - registration: - last_mean_input_tps: 16384 traffic_pattern: kind: prefix_reuse diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml index 447e4db95..1800d1277 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 profiles: - count: 2 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 1 prefill_tokens_per_s: 65536 decode_tokens_per_s: 320 - registration: - last_mean_input_tps: 65536 - count: 4 profile: name: balanced-prefix-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 - count: 2 profile: name: constrained-prefix-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 16384 decode_tokens_per_s: 120 - registration: - last_mean_input_tps: 16384 traffic_pattern: kind: prefix_reuse diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml index 2019da9fa..75d6b10b8 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 1 prefill_tokens_per_s: 32768 decode_tokens_per_s: 512 - registration: - last_mean_input_tps: 32768 traffic_pattern: kind: prefix_reuse diff --git a/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml index 0e31c2394..b02085a27 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 9000 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 200 traffic_pattern: kind: uniform diff --git a/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml b/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml index 06c36fbbc..9cb7d8d55 100644 --- a/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml +++ b/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 8000 decode_tokens_per_s: 160 - registration: - last_mean_input_tps: 260 - count: 3 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 1 profile: name: small-cache @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 6 prefill_tokens_per_s: 2600 decode_tokens_per_s: 70 - registration: - last_mean_input_tps: 80 traffic_pattern: kind: mixed_size diff --git a/src/libraries/rust/stargate/benches/overload-6-backends.yaml b/src/libraries/rust/stargate/benches/overload-6-backends.yaml index 6bc9f8206..c61d28485 100644 --- a/src/libraries/rust/stargate/benches/overload-6-backends.yaml +++ b/src/libraries/rust/stargate/benches/overload-6-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 1800 decode_tokens_per_s: 45 - registration: - last_mean_input_tps: 60 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 130 - count: 4 profile: name: constrained @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 1800 decode_tokens_per_s: 45 - registration: - last_mean_input_tps: 60 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml b/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml index 713027202..01560fa0e 100644 --- a/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml +++ b/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml @@ -15,10 +15,10 @@ name: queue-mismatch-retry-ab metadata: - description: A/B diagnostic for pylon rejection of stale Stargate queue estimates under burst load. The scenario runs 2048 requests per arm for benchmark evidence. Mock backends enforce real concurrent-service, prefill, and decode delays; benchmark pylons pin prompt throughput to the matching 2200 tokens/s profile value. Queue admission intentionally models prompt work only because output residency is unknowable at admission time. + description: A/B diagnostic for pylon rejection of stale Stargate queue estimates under burst load. The scenario runs 2048 requests per arm for benchmark evidence. Mock backends enforce real concurrent-service, prefill, and decode delays, and pylons measure the resulting prompt throughput. Queue admission intentionally models prompt work only because output residency is unknowable at admission time. tags: [queue-admission, queue-mismatch, ab, wait-and-widen, diagnostic] expected_runtime: medium - expected_signal: With the same fixed prompt-throughput estimate and real mock service delays in both arms, the enabled row should record pylon rejections and Stargate queue-mismatch retries; the disabled row should record disabled decisions without those retries. + expected_signal: With the same mock service delays in both arms, the enabled row should record pylon rejections and Stargate queue-mismatch retries; the disabled row should record disabled decisions without those retries. model: dummy-model seed: 154 request_count: 2048 @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 2200 decode_tokens_per_s: 100 - registration: - last_mean_input_tps: 2200 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml b/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml index c280d752b..6f8c7de71 100644 --- a/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml +++ b/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4500 decode_tokens_per_s: 95 - registration: - last_mean_input_tps: 130 traffic_pattern: kind: stair_step diff --git a/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml b/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml index 7b3072764..cd264c0e9 100644 --- a/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml +++ b/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 4500 decode_tokens_per_s: 95 - registration: - last_mean_input_tps: 130 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/uniform-4-backends.yaml b/src/libraries/rust/stargate/benches/uniform-4-backends.yaml index 419fb13dc..833b72bdc 100644 --- a/src/libraries/rust/stargate/benches/uniform-4-backends.yaml +++ b/src/libraries/rust/stargate/benches/uniform-4-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 5000 decode_tokens_per_s: 100 - registration: - last_mean_input_tps: 140 traffic_pattern: kind: uniform 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 7604e2307..9eb3c80d0 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs @@ -202,12 +202,11 @@ pub(crate) async fn chat_completions( info!(id = %id, model = %model, stream = stream, "received chat/completions request"); let request_id = optional_header(&headers, "x-request-id").unwrap_or_else(|| id.clone()); let cache_affinity_key = optional_header(&headers, "x-cache-affinity-key"); - if stream { - state.emit_counters(&request_id, &model, 0, 0, false); - } + state.emit_counters(&request_id, &model, 0, 0, false); let kv_cache_access = state .process_input_with_cache(cache_affinity_key.as_deref(), input_tokens) .await; + state.emit_counters(&request_id, &model, input_tokens, 0, false); let first_token_delay = state.ttft + Duration::from_millis(jitter_ms(&request_id, "ttft", state.ttft_jitter_ms)); info!( @@ -303,6 +302,7 @@ pub(crate) async fn responses( let kv_cache_access = state .process_input_with_cache(cache_affinity_key.as_deref(), input_tokens) .await; + state.emit_counters(&request_id, &model, input_tokens, 0, false); let first_token_delay = state.ttft + Duration::from_millis(jitter_ms(&request_id, "ttft", state.ttft_jitter_ms)); @@ -514,8 +514,6 @@ fn stream_response(config: StreamResponseConfig) -> Response { } tokio::time::sleep(first_token_delay).await; - state.emit_counters(&request_id, &model, input_tokens, 0, false); - if kind == StreamKind::Chat { yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Role)); } 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 d4ba751a0..2382cd705 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs @@ -631,6 +631,116 @@ async fn chat_completion_retains_prefix_only_after_modeled_prefill_completes() { ); } +#[tokio::test(start_paused = true)] +async fn non_streaming_chat_reports_input_when_prefill_completes() { + let state = AppState { + prefill_tokens_per_s: 100.0, + ttft: Duration::from_secs(10), + ..test_state() + }; + let mut stats_events = state.stats_events.subscribe(); + let mut headers = HeaderMap::new(); + headers.insert( + "x-request-id", + HeaderValue::from_static("req-prefill-stats"), + ); + headers.insert("x-input-tokens", HeaderValue::from_static("100")); + let mut chat_request = request(Some(1)); + chat_request.stream = Some(false); + + let request = tokio::spawn(chat_completions(State(state), headers, Json(chat_request))); + let StatsStreamEvent::Stats { + tokens_processed, + tokens_generated, + finished, + .. + } = stats_events + .recv() + .await + .expect("baseline event should arrive") + else { + panic!("expected a stats event"); + }; + assert_eq!(tokens_processed, Some(0)); + assert_eq!(tokens_generated, Some(0)); + assert!(!finished); + + tokio::time::advance(Duration::from_millis(999)).await; + assert!(matches!( + stats_events.try_recv(), + Err(broadcast::error::TryRecvError::Empty) + )); + tokio::time::advance(Duration::from_millis(1)).await; + + let StatsStreamEvent::Stats { + tokens_processed, + tokens_generated, + finished, + .. + } = stats_events + .recv() + .await + .expect("prefill event should arrive") + else { + panic!("expected a stats event"); + }; + assert_eq!(tokens_processed, Some(100)); + assert_eq!(tokens_generated, Some(0)); + assert!(!finished); + assert!( + !request.is_finished(), + "TTFT should still delay the response" + ); +} + +#[tokio::test(start_paused = true)] +async fn responses_reports_input_without_waiting_for_stream_polling() { + let state = AppState { + prefill_tokens_per_s: 100.0, + ttft: Duration::from_secs(10), + ..test_state() + }; + let mut stats_events = state.stats_events.subscribe(); + let mut headers = HeaderMap::new(); + headers.insert( + "x-request-id", + HeaderValue::from_static("req-response-prefill"), + ); + headers.insert("x-input-tokens", HeaderValue::from_static("100")); + let response_request = ResponsesRequest { + stream: Some(true), + model: Some("dummy-model".to_string()), + max_output_tokens: Some(1), + input: None, + }; + + let request = tokio::spawn(responses(State(state), headers, Json(response_request))); + stats_events + .recv() + .await + .expect("baseline event should arrive"); + tokio::time::advance(Duration::from_secs(1)).await; + let StatsStreamEvent::Stats { + tokens_processed, + tokens_generated, + finished, + .. + } = stats_events + .recv() + .await + .expect("prefill event should arrive") + else { + panic!("expected a stats event"); + }; + + assert_eq!(tokens_processed, Some(100)); + assert_eq!(tokens_generated, Some(0)); + assert!(!finished); + request + .await + .expect("responses request should return a stream"); +} + #[test] fn kv_cache_evicts_least_recently_used_entry() { let mut cache = KvCacheState::new(300); diff --git a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto index ef4fe5d1f..0df216f55 100644 --- a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto +++ b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto @@ -126,6 +126,11 @@ message ModelStats { // Sticky source labels for stat observations seen by this backend since the // model metrics state was initialized. repeated string stats_sources = 17; + // Maximum input tokens per second recorded for the current backend model + // generation. Absent until a valid value is configured or observed. + // Stargate sums this field only when every active backend in a shared + // cluster reports it. + optional double max_input_tps = 18; } enum InferenceServerStatus { diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs index 2af6f5c30..58c43cf62 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs @@ -47,7 +47,7 @@ pub enum ModelSource { #[derive(Clone, Debug)] pub enum ModelInitialization { Calibration(CalibrationConfig), - ConfiguredInputTps { input_tps: f64, pin: bool }, + ConfiguredInputTps { input_tps: f64 }, Uncalibrated, } @@ -348,10 +348,9 @@ impl ModelLifecycleSupervisor { ModelInitialization::Calibration(_) | ModelInitialization::Uncalibrated => { ModelStatsInitialization::Empty } - ModelInitialization::ConfiguredInputTps { input_tps, pin } => { + ModelInitialization::ConfiguredInputTps { input_tps } => { ModelStatsInitialization::ConfiguredInputTps { input_tps: *input_tps, - pin: *pin, } } }; @@ -1014,10 +1013,7 @@ mod tests { poll_interval: Duration::from_millis(10), request_timeout: Duration::from_secs(1), }), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1135,10 +1131,7 @@ mod tests { "model-a".to_string(), "model-b".to_string(), ])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1201,10 +1194,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: upstream.to_string(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { active_canary_interval: Duration::ZERO, ..BringupConfig::default() @@ -1242,10 +1232,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: "http://127.0.0.1:1".to_string(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { active_canary_interval: Duration::ZERO, ..BringupConfig::default() @@ -1317,10 +1304,7 @@ mod tests { config: ModelLifecycleConfig { upstream_http_base_url: "http://127.0.0.1:1".to_string(), source: ModelSource::Static(BTreeSet::new()), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1512,10 +1496,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: upstream.base_url.clone(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs index cba367b0f..3800ba81f 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs @@ -450,6 +450,7 @@ fn runtime_snapshot_forwards_bootstrap_and_collected_stats_exactly() { "model-a", CurrentModelStats { last_mean_input_tps: 3.5, + max_input_tps: Some(4.5), output_tps: 2.5, queue_size: 4, queued_input_size: 5, @@ -475,6 +476,7 @@ fn runtime_snapshot_forwards_bootstrap_and_collected_stats_exactly() { assert_eq!(model.status, InferenceServerStatus::Active as i32); let stats = model.stats.as_ref().expect("stats should be present"); assert_eq!(stats.last_mean_input_tps, 3.5); + assert_eq!(stats.max_input_tps, Some(4.5)); assert_eq!(stats.output_tps, 2.5); assert_eq!( stats.queue_time_estimate_ms_by_priority, 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 0cc849340..a0b8b3049 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 @@ -31,6 +31,8 @@ use reqwest::header::HeaderMap; pub struct CurrentModelStats { // Sticky runtime-observed mean input TPS for this backend. pub last_mean_input_tps: f64, + // Max input TPS for this model generation. None until configured or observed. + 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. pub output_tps: f64, @@ -366,6 +368,7 @@ impl PylonRuntimeState { let registration = InferenceServerModelRegistration { stats: Some(ModelStats { last_mean_input_tps: stats.last_mean_input_tps, + max_input_tps: stats.max_input_tps, output_tps: stats.output_tps, max_output_tps: stats.max_output_tps, queue_size: stats.queue_size, 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 c6c680883..ea19712b8 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 @@ -55,7 +55,6 @@ pub(super) struct ModelMetricsState { pub(super) struct GenerationMetricsState { pub(super) generation: ModelGeneration, pub(super) metrics: ModelMetricsState, - pub(super) pinned_input_tps: Option, } #[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq)] pub(super) struct KvCacheStatsSnapshot { @@ -237,22 +236,17 @@ impl StatsAggregator { if self.per_model.contains_key(generation.model_id()) { return None; } - let (metrics, pinned_input_tps) = match initialization { - super::collector::ModelStatsInitialization::Empty => { - (ModelMetricsState::default(), None) - } - super::collector::ModelStatsInitialization::ConfiguredInputTps { input_tps, pin } => { + let metrics = match initialization { + super::collector::ModelStatsInitialization::Empty => ModelMetricsState::default(), + super::collector::ModelStatsInitialization::ConfiguredInputTps { input_tps } => { let input_tps_distribution = TpsDistribution::bootstrap(input_tps) .expect("configured input TPS must be positive and finite"); - ( - ModelMetricsState { - last_mean_input_tps: input_tps, - input_tps_distribution, - aggregate_state_counted: true, - ..ModelMetricsState::default() - }, - pin.then_some(input_tps), - ) + ModelMetricsState { + last_mean_input_tps: input_tps, + input_tps_distribution, + aggregate_state_counted: true, + ..ModelMetricsState::default() + } } }; self.aggregate_model_state_count += usize::from(metrics.aggregate_state_counted); @@ -261,7 +255,6 @@ impl StatsAggregator { GenerationMetricsState { generation: generation.clone(), metrics, - pinned_input_tps, }, ); let stats = self.snapshot(generation.model_id()); @@ -649,7 +642,6 @@ impl StatsAggregator { dirty |= apply_input_throughput_sample( config, model_state, - generation_state.pinned_input_tps, InputThroughputSample { units, duration, @@ -805,7 +797,6 @@ fn adjust_live_count(count: &mut usize, delta: isize) { pub(super) fn apply_input_throughput_sample( config: &StatsCollectorConfig, model_state: &mut ModelMetricsState, - pinned_input_tps: Option, sample: InputThroughputSample, ) -> bool { if sample.units < config.min_input_tokens { @@ -819,18 +810,19 @@ pub(super) fn apply_input_throughput_sample( let Some(input_tps) = tps_for_units(sample.units, duration, config.duration_floor) else { return false; }; + let previous_max = model_state.input_tps_distribution.max; model_state.input_tps_distribution.update(input_tps); + let max_changed = model_state.input_tps_distribution.max != previous_max; let mean_input_tps = model_state.input_tps_distribution.mean; if !model_state.input_tps_distribution.has_sufficient_data() || !valid_last_mean_input_tps(mean_input_tps) { - return false; + return max_changed; } - let last_mean_input_tps = pinned_input_tps.unwrap_or(mean_input_tps); - if model_state.last_mean_input_tps == last_mean_input_tps { - return false; + if model_state.last_mean_input_tps == mean_input_tps { + return max_changed; } - model_state.last_mean_input_tps = last_mean_input_tps; + model_state.last_mean_input_tps = mean_input_tps; true } @@ -877,6 +869,8 @@ impl ModelMetricsState { }; CurrentModelStats { last_mean_input_tps: self.last_mean_input_tps, + max_input_tps: (self.input_tps_distribution.count > 0) + .then_some(self.input_tps_distribution.max), output_tps: active_chat_output_tps.max(average_with_sum( &self.chat_output_tps_samples, self.chat_output_tps_sum, 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 1c77c6958..2b924af0d 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 @@ -25,7 +25,6 @@ use crate::{CurrentModelStats, PylonRuntimeState, RequestObservationEvent}; use stargate_runtime::OwnedTask; use super::aggregator::{ENGINE_STATS_SOURCE, KvCacheStatsSnapshot, StatsAggregator}; - const DEFAULT_OBSERVATION_CHANNEL_CAPACITY: usize = 1024; const DEFAULT_SMOOTHING_WINDOW_SIZE: usize = 8; const DEFAULT_MIN_INPUT_TOKENS: u64 = 1; @@ -110,7 +109,7 @@ impl StatsCollectorHandle { #[derive(Debug, Clone, Copy)] pub(crate) enum ModelStatsInitialization { Empty, - ConfiguredInputTps { input_tps: f64, pin: bool }, + ConfiguredInputTps { input_tps: f64 }, } enum StatsCollectorCommand { @@ -637,7 +636,7 @@ mod tests { } } - async fn begin_configured_model(&self, model_id: &str, input_tps: f64, pin: bool) { + async fn begin_configured_model(&self, model_id: &str, input_tps: f64) { let generation = ModelGeneration::new(model_id, 0); assert!(self.runtime_state.begin_generation(generation.clone())); assert!( @@ -645,7 +644,7 @@ mod tests { .control() .begin_generation( generation, - ModelStatsInitialization::ConfiguredInputTps { input_tps, pin } + ModelStatsInitialization::ConfiguredInputTps { input_tps } ) .await .expect("collector should acknowledge configured generation") @@ -1239,7 +1238,7 @@ mod tests { aggregator.stream("req-a", (0, 0), false, Duration::ZERO); let updates = aggregator.stream("req-a", (10, 4), false, milliseconds(100)); let stats = published_stats(updates); - assert_stats!(stats; output_tps: 40.0, max_output_tps: 40.0, stats_sources: ["engine_stats_stream"]); + assert_stats!(stats; max_input_tps: Some(100.0), output_tps: 40.0, max_output_tps: 40.0, stats_sources: ["engine_stats_stream"]); for tick in 2..=5 { let updates = aggregator.stream("req-a", (tick * 10, 4), false, milliseconds(tick * 100)); @@ -1252,37 +1251,32 @@ mod tests { } #[test] - fn pinned_configured_input_tps_is_preserved_across_engine_stats_updates() { + fn configured_input_tps_moves_with_the_first_real_sample() { let mut aggregator = test_aggregator_with_initialization( StatsCollectorConfig::default(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 2_200.0, - pin: true, - }, + ModelStatsInitialization::ConfiguredInputTps { input_tps: 100.0 }, ); - let stats = aggregator.stream_stats("req-a", (0, 0), false, Duration::ZERO); - assert_eq!(stats.last_mean_input_tps, 2_200.0); - for tick in 1..=5 { - aggregator.stream("req-a", (tick * 10, 0), false, milliseconds(tick * 100)); - } - assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 2_200.0); - } - - #[test] - fn unpinned_configured_input_tps_moves_with_the_first_real_sample() { - let mut aggregator = test_aggregator_with_initialization( - StatsCollectorConfig::default(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 100.0, - pin: false, - }, - ); - assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 100.0); + let bootstrapped = aggregator.snapshot("model-a"); + assert_eq!(bootstrapped.last_mean_input_tps, 100.0); + assert_eq!(bootstrapped.max_input_tps, Some(100.0)); aggregator.stream("req-a", (0, 0), false, Duration::ZERO); let stats = published_stats(aggregator.stream("req-a", (20, 0), false, milliseconds(100))); assert!((stats.last_mean_input_tps - (700.0 / 6.0)).abs() < f64::EPSILON); + assert_eq!(stats.max_input_tps, Some(200.0)); + } + + #[test] + fn max_input_tps_persists_for_the_generation() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + aggregator.stream("req-a", (0, 0), false, Duration::ZERO); + aggregator.stream("req-a", (20, 0), false, milliseconds(100)); + aggregator.stream("req-a", (30, 0), false, milliseconds(200)); + + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(200.0)); + aggregator.sweep(seconds(24 * 60 * 60)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(200.0)); } #[test] @@ -1336,12 +1330,13 @@ mod tests { .is_empty(), "second output counter is still below the duration floor" ); - let input_only_updates = - aggregator.partial_stream("req-partial", (Some(1), None), false, milliseconds(11)); - assert!( - input_only_updates.is_empty(), - "input-only updates must not publish a stale output TPS sample" - ); + let input_only_stats = aggregator + .partial_stream("req-partial", (Some(1), None), false, milliseconds(11)) + .pop() + .expect("first input counter should publish max input TPS") + .1; + assert_eq!(input_only_stats.output_tps, 100.0); + assert!(input_only_stats.max_input_tps.is_some()); } #[test] @@ -1898,6 +1893,28 @@ mod tests { assert_stats!(stats; last_mean_input_tps: 100.0, output_tps: 0.0, queue_size: 0, queued_input_size: 0, num_running_queries: 0, input_processing_queries: 0, output_generation_queries: 0, stats_sources: ["engine_stats_stream"]); } + #[test] + fn stats_aggregator_publishes_generation_max_input_tps() { + let mut aggregator = test_aggregator(config!( + engine_stats_request_ttl: Duration::ZERO, + engine_stats_model_ttl: Duration::ZERO, + )); + aggregator.stream("req-generation-max", (0, 0), false, Duration::ZERO); + + let stats = published_stats(aggregator.stream( + "req-generation-max", + (10, 0), + false, + milliseconds(100), + )); + assert_eq!(stats.last_mean_input_tps, 0.0); + assert_eq!(stats.max_input_tps, Some(100.0)); + + let updates = aggregator.stream("req-generation-max", (15, 0), false, milliseconds(200)); + assert!(updates.is_empty()); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(100.0)); + } + #[test] fn stats_aggregator_tombstones_stale_request_before_late_finish() { let config = config!( @@ -2465,9 +2482,7 @@ mod tests { #[tokio::test] async fn stats_collector_bootstraps_input_tps_for_queue_admission() { let collector = RunningCollector::spawn_empty(StatsCollectorConfig::default(), None, false); - collector - .begin_configured_model("model-a", 2_200.0, false) - .await; + collector.begin_configured_model("model-a", 2_200.0).await; let stats = collector .wait_for_stats("bootstrap TPS stats should be published", |stats| { stats.last_mean_input_tps == 2_200.0 @@ -2531,6 +2546,7 @@ mod tests { r#"pylon_requests_total{model="model-a",routing_key="rk-1",status="complete"} 5"# )); assert!(body.contains(r#"pylon_model_last_mean_input_tps{model="model-a"} 10"#)); + assert!(body.contains(r#"pylon_model_max_input_tps{model="model-a"} 10"#)); assert!(body.contains(r#"pylon_model_output_tps{model="model-a"} 5"#)); } @@ -2566,10 +2582,7 @@ mod tests { control .begin_generation( first.clone(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 100.0, - pin: false, - }, + ModelStatsInitialization::ConfiguredInputTps { input_tps: 100.0 }, ) .await .expect("stats collector should acknowledge initialization") diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs index 3ac3727b1..cfe392c79 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs @@ -151,6 +151,7 @@ metrics! { float_gauge output_tps("model_output_tps", "Current output TPS by model", ["model"]); float_gauge embedding_item_tps("model_embedding_item_tps", "Current embeddings item throughput by model", ["model"]); float_gauge last_mean_input_tps("model_last_mean_input_tps", "Last valid mean input TPS by model", ["model"]); + float_gauge max_input_tps("model_max_input_tps", "Generation max input TPS by model", ["model"]); float_gauge max_output_tps("model_max_output_tps", "Observed max output TPS by model", ["model"]); float_gauge max_embedding_item_tps("model_max_embedding_item_tps", "Observed max embeddings item throughput by model", ["model"]); float_gauge queue_size("model_queue_size", "Current queued request count by model", ["model"]); @@ -311,6 +312,12 @@ impl PylonMetrics { ] { gauge.with_label_values(&[model_id]).set(value); } + match stats.max_input_tps { + Some(value) => self.max_input_tps.with_label_values(&[model_id]).set(value), + None => { + let _ = self.max_input_tps.remove_label_values(&[model_id]); + } + } for (gauge, values) in [ (&self.stats_capability, &stats.stats_capabilities), (&self.stats_source, &stats.stats_sources), @@ -352,6 +359,7 @@ impl PylonMetrics { &self.output_tps, &self.embedding_item_tps, &self.last_mean_input_tps, + &self.max_input_tps, &self.max_output_tps, &self.max_embedding_item_tps, &self.queue_size, @@ -748,6 +756,7 @@ mod tests { let metrics = PylonMetrics::new().expect("metrics should initialize"); let stats = CurrentModelStats { last_mean_input_tps: 12.0, + max_input_tps: Some(13.0), stats_capabilities: vec!["input_tps".to_string()], stats_sources: vec!["engine_stats_stream".to_string()], ..CurrentModelStats::default() @@ -915,6 +924,7 @@ mod tests { output_tps: 20.0, embedding_item_tps: 25.0, last_mean_input_tps: 30.0, + max_input_tps: Some(35.0), max_output_tps: 40.0, max_embedding_item_tps: 45.0, queue_size: 2, @@ -934,6 +944,7 @@ mod tests { r#"pylon_model_output_tps{model="model-a"} 20"#, r#"pylon_model_embedding_item_tps{model="model-a"} 25"#, r#"pylon_model_last_mean_input_tps{model="model-a"} 30"#, + r#"pylon_model_max_input_tps{model="model-a"} 35"#, r#"pylon_model_max_embedding_item_tps{model="model-a"} 45"#, r#"pylon_model_queue_size{model="model-a"} 2"#, r#"pylon_model_queued_input_tokens{model="model-a"} 17"#, @@ -942,6 +953,10 @@ mod tests { r#"pylon_model_stats_source{model="model-a",source="engine_stats_stream"} 1"#, ], ); + + metrics.observe_model_stats("model-a", &CurrentModelStats::default()); + let body = metrics.gather_text().expect("metrics should encode"); + assert!(!body.contains(r#"pylon_model_max_input_tps{model="model-a"}"#)); } #[test] 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 d416c86e5..8ced5f997 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 @@ -157,7 +157,6 @@ impl StatsAggregator { let Some(generation_state) = self.per_model.get_mut(&observation.model_id) else { return changed_models; }; - let pinned_input_tps = generation_state.pinned_input_tps; let model_state = &mut generation_state.metrics; model_state.chunk_usage_stats_observed |= observation.output_tokens_from_chunk_usage; let record_sample = |samples, sum: &mut f64, max: &mut f64, sample| { @@ -200,9 +199,9 @@ impl StatsAggregator { } } } - if input_sample.is_some_and(|sample| { - apply_input_throughput_sample(&self.config, model_state, pinned_input_tps, sample) - }) { + if input_sample + .is_some_and(|sample| apply_input_throughput_sample(&self.config, model_state, sample)) + { push_changed_model(&mut changed_models, observation.model_id.clone()); } changed_models diff --git a/src/libraries/rust/stargate/crates/pylon/src/main.rs b/src/libraries/rust/stargate/crates/pylon/src/main.rs index fd08dc930..c92dba818 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/main.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/main.rs @@ -116,9 +116,6 @@ struct Args { /// Upstream HTTP path for the engine stats stream #[arg(long, default_value = "/pylon/v1/stats/stream", value_name = "PATH")] engine_stats_stream_path: String, - /// Keep --initial-input-tps fixed for deterministic benchmark/test experiments - #[arg(long, default_value_t = false, hide = true)] - benchmark_pin_input_tps: bool, /// Minimum interval between registration/stat updates to stargate #[arg(long, default_value_t = 1000, value_name = "MS")] min_update_interval_ms: u64, @@ -535,17 +532,6 @@ mod tests { assert!(startup::PylonStartupPlan::from_args(&args).is_err()); } - #[test] - fn benchmark_pin_requires_initial_input_tps() { - let uncalibrated = parse_args("--benchmark-pin-input-tps"); - let calibration = parse_args("--do-calibration --benchmark-pin-input-tps"); - let initial = parse_args("--initial-input-tps 2200 --benchmark-pin-input-tps"); - - assert!(startup::PylonStartupPlan::from_args(&uncalibrated).is_err()); - assert!(startup::PylonStartupPlan::from_args(&calibration).is_err()); - assert!(startup::PylonStartupPlan::from_args(&initial).is_ok()); - } - #[test] fn engine_stats_stream_defaults_to_auto_mode_and_v1_path() { let args = parse_args(""); diff --git a/src/libraries/rust/stargate/crates/pylon/src/startup.rs b/src/libraries/rust/stargate/crates/pylon/src/startup.rs index 9a5196501..edbd9d97e 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/startup.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/startup.rs @@ -598,10 +598,6 @@ fn model_initialization_from_args(args: &Args) -> Result { .is_none_or(|input_tps| input_tps.is_finite() && input_tps > 0.0), "initial input TPS must be finite and positive" ); - ensure!( - !args.benchmark_pin_input_tps || args.initial_input_tps.is_some(), - "--benchmark-pin-input-tps requires --initial-input-tps" - ); if args.do_calibration { ensure!( args.calibration_requests > 0, @@ -617,10 +613,7 @@ fn model_initialization_from_args(args: &Args) -> Result { } Ok(match args.initial_input_tps { - Some(input_tps) => ModelInitialization::ConfiguredInputTps { - input_tps, - pin: args.benchmark_pin_input_tps, - }, + Some(input_tps) => ModelInitialization::ConfiguredInputTps { input_tps }, None => ModelInitialization::Uncalibrated, }) } @@ -1016,10 +1009,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: "http://127.0.0.1:1".to_string(), source: ModelSource::Static(BTreeSet::new()), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1373,10 +1363,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: plan.upstream.clone(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1_000.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1_000.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs index 683498349..f628c9f82 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs @@ -337,16 +337,10 @@ fn validate_profile(profile: &BackendProfile) -> anyhow::Result<()> { profile.service_time_ms.decode_tokens_per_s > 0, "backend decode_tokens_per_s must be > 0" ); - if let Some(prefill_tokens_per_s) = profile.service_time_ms.prefill_tokens_per_s { - ensure!( - prefill_tokens_per_s > 0.0 && prefill_tokens_per_s.is_finite(), - "backend prefill_tokens_per_s must be finite and > 0 when set" - ); - } ensure!( - profile.registration.last_mean_input_tps > 0.0 - && profile.registration.last_mean_input_tps.is_finite(), - "backend registration.last_mean_input_tps must be finite and > 0" + profile.service_time_ms.prefill_tokens_per_s > 0.0 + && profile.service_time_ms.prefill_tokens_per_s.is_finite(), + "backend prefill_tokens_per_s must be finite and > 0" ); Ok(()) } @@ -387,7 +381,6 @@ config_struct!(BackendProfile { #[serde(default)] pub kv_cache_capacity_tokens: u64, pub service_time_ms: ServiceTimeConfig, - pub registration: RegistrationConfig, }); fn default_backend_name() -> String { @@ -405,11 +398,7 @@ config_struct!(ServiceTimeConfig { pub decode_tokens_per_s: u64, #[serde(default)] pub decode_jitter_ms: u64, - pub prefill_tokens_per_s: Option, -}); - -config_struct!(RegistrationConfig { - pub last_mean_input_tps: f64, + pub prefill_tokens_per_s: f64, }); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -569,10 +558,7 @@ mod tests { ttft_jitter_ms: 0, decode_tokens_per_s: 100, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 100.0, + prefill_tokens_per_s: 100.0, }, } } @@ -683,7 +669,7 @@ mod tests { ] { remove_yaml_field(&mut yaml["backends"]["profile"], key); } - for key in ["ttft_jitter_ms", "decode_jitter_ms", "prefill_tokens_per_s"] { + for key in ["ttft_jitter_ms", "decode_jitter_ms"] { remove_yaml_field(&mut yaml["backends"]["profile"]["service_time_ms"], key); } @@ -704,10 +690,6 @@ mod tests { assert_eq!(config.backends.profile.kv_cache_capacity_tokens, 0); assert_eq!(config.backends.profile.service_time_ms.ttft_jitter_ms, 0); assert_eq!(config.backends.profile.service_time_ms.decode_jitter_ms, 0); - assert_eq!( - config.backends.profile.service_time_ms.prefill_tokens_per_s, - None - ); } #[test] @@ -865,13 +847,10 @@ action: pause_backend } #[test] - fn rejects_invalid_registered_input_throughput() { + fn rejects_invalid_prefill_throughput() { let mut profile = profile("invalid-throughput"); - profile.registration.last_mean_input_tps = 0.0; - assert_validation_error( - validate_profile(&profile), - "registration.last_mean_input_tps", - ); + profile.service_time_ms.prefill_tokens_per_s = 0.0; + assert_validation_error(validate_profile(&profile), "prefill_tokens_per_s"); } #[test] diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs index 4630a044e..81194f407 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs @@ -126,14 +126,14 @@ pub(super) fn render_manifest(render: RenderManifestConfig<'_>) -> RenderedManif )); } backends.push_str(&format!( - "apiVersion: apps/v1\nkind: Deployment\nmetadata:\n name: {inference_server_id}-pylon\n namespace: {backends_ns}\nspec:\n replicas: 1\n selector:\n matchLabels:\n app: {inference_server_id}-pylon\n template:\n metadata:\n labels:\n app: {inference_server_id}-pylon\n benchmark.stargate/profile: {profile_name}\n spec:\n containers:\n - name: pylon\n image: {pylon_image}\n imagePullPolicy: IfNotPresent\n args:\n - --upstream-http-base-url=http://{upstream_backend_name}-http.{backends_ns}.svc.cluster.local:8090\n - --model-name={model}\n - --stargate-address=stargate.{stargate_ns}.svc.cluster.local:50071\n - --inference-server-id={inference_server_id}\n{cluster_id_arg} - --backend-connectivity=reverse\n - --quic-insecure\n - --tunnel-protocol={tunnel_protocol}\n - --kv-cache-stats-path=/kv-cache/stats\n - --min-update-interval-ms=100\n - --disable-bringup\n - --active-canary-interval-ms=0\n - --initial-input-tps={last_mean_input_tps}\n - --benchmark-pin-input-tps\n", + "apiVersion: apps/v1\nkind: Deployment\nmetadata:\n name: {inference_server_id}-pylon\n namespace: {backends_ns}\nspec:\n replicas: 1\n selector:\n matchLabels:\n app: {inference_server_id}-pylon\n template:\n metadata:\n labels:\n app: {inference_server_id}-pylon\n benchmark.stargate/profile: {profile_name}\n spec:\n containers:\n - name: pylon\n image: {pylon_image}\n imagePullPolicy: IfNotPresent\n args:\n - --upstream-http-base-url=http://{upstream_backend_name}-http.{backends_ns}.svc.cluster.local:8090\n - --model-name={model}\n - --stargate-address=stargate.{stargate_ns}.svc.cluster.local:50071\n - --inference-server-id={inference_server_id}\n{cluster_id_arg} - --backend-connectivity=reverse\n - --quic-insecure\n - --tunnel-protocol={tunnel_protocol}\n - --kv-cache-stats-path=/kv-cache/stats\n - --min-update-interval-ms=100\n - --disable-bringup\n - --active-canary-interval-ms=0\n - --initial-input-tps={initial_input_tps}\n", upstream_backend_name = pylon.upstream_backend_name, inference_server_id = pylon.inference_server_id, profile_name = pylon.profile_slug, pylon_image = image_refs.pylon, model = config.model, tunnel_protocol = config.tunnel_protocol, - last_mean_input_tps = pylon.last_mean_input_tps, + initial_input_tps = pylon.initial_input_tps, )); backends.push_str(&pylon_queue_admission_args); backends.push_str("---\n"); diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs index e48495859..ff0c6bad6 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs @@ -26,7 +26,7 @@ use super::render::{ use super::run::{BenchmarkK8sRun, prepare_benchmark_k8s_run_with_resolved_dependencies}; use crate::config::{ AlgorithmConfig, ArrivalPatternConfig, BackendConfig, BackendProfile, BenchmarkConfig, - DegradationConfig, RegistrationConfig, ScenarioMetadata, ServiceTimeConfig, StargateConfig, + DegradationConfig, ScenarioMetadata, ServiceTimeConfig, StargateConfig, TokenDistributionConfig, TrafficPatternConfig, UniformTrafficConfig, }; use serde::Deserialize; @@ -57,10 +57,7 @@ fn config() -> BenchmarkConfig { ttft_jitter_ms: 10, decode_tokens_per_s: 50, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 100.0, + prefill_tokens_per_s: 100.0, }, }, }, @@ -725,9 +722,9 @@ fn rendered_pylons_include_per_algorithm_queue_admission_args() { "- --disable-bringup", "- --active-canary-interval-ms=0", "- --initial-input-tps=100", - "- --benchmark-pin-input-tps", "- --pylon-queue-mismatch-tolerance-factor=1", "- --pylon-queue-mismatch-retry-after-ms=5", ], ); + assert!(!rendered.backends.contains("--benchmark-pin-input-tps")); } diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs index a54340e91..71d276968 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs @@ -1000,8 +1000,7 @@ backends: count: 1 profile: name: balanced - service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50 } - registration: { last_mean_input_tps: 100.0 } + service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50, prefill_tokens_per_s: 100.0 } traffic_pattern: kind: uniform routing_keys: 1 diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs index 258746d18..d75020bdf 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs @@ -329,8 +329,8 @@ mod tests { use super::*; use crate::config::{ AlgorithmConfig, ArrivalPatternConfig, BackendConfig, BackendProfile, HotsetTrafficConfig, - MixedSizeClassConfig, MixedSizeTrafficConfig, PrefixReuseTrafficConfig, RegistrationConfig, - ServiceTimeConfig, StargateConfig, TokenDistributionConfig, UniformTrafficConfig, + MixedSizeClassConfig, MixedSizeTrafficConfig, PrefixReuseTrafficConfig, ServiceTimeConfig, + StargateConfig, TokenDistributionConfig, UniformTrafficConfig, }; fn base_config() -> BenchmarkConfig { @@ -358,10 +358,7 @@ mod tests { ttft_jitter_ms: 20, decode_tokens_per_s: 80, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 100.0, + prefill_tokens_per_s: 100.0, }, }, }, diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs index 185127b72..ec94be266 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs @@ -295,8 +295,7 @@ fn build_compose_spec( "--min-update-interval-ms" => "100", "--disable-bringup", "--active-canary-interval-ms=0", - "--initial-input-tps" => pylon.last_mean_input_tps.to_string(), - "--benchmark-pin-input-tps", + "--initial-input-tps" => pylon.initial_input_tps.to_string(), ]); if let Some(pylon_queue_admission) = &algorithm.pylon_queue_admission { client_command.extend(pylon_queue_admission.pylon_args()); @@ -351,8 +350,7 @@ backends: count: 2 profile: name: balanced - service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50 } - registration: { last_mean_input_tps: 100.0 } + service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50, prefill_tokens_per_s: 100.0 } traffic_pattern: kind: uniform routing_keys: 2 @@ -566,7 +564,7 @@ algorithms: Some("100") ); assert!( - client + !client .command .iter() .any(|candidate| candidate == "--benchmark-pin-input-tps") diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs index 0f2c9dd96..cf7151aeb 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs @@ -497,8 +497,8 @@ fn render_warnings(out: &mut String, context: &ReportContext, entries: &[ReportE mod tests { use super::*; use crate::config::{ - ArrivalPatternConfig, BackendConfig, BackendProfile, RegistrationConfig, ServiceTimeConfig, - StargateConfig, TokenDistributionConfig, TrafficPatternConfig, UniformTrafficConfig, + ArrivalPatternConfig, BackendConfig, BackendProfile, ServiceTimeConfig, StargateConfig, + TokenDistributionConfig, TrafficPatternConfig, UniformTrafficConfig, }; use crate::score::{ BackendSummary, CacheSummary, FailureSummary, QueueAdmissionSummary, @@ -531,10 +531,7 @@ mod tests { ttft_jitter_ms: 0, decode_tokens_per_s: 1, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 1.0, + prefill_tokens_per_s: 1.0, }, }, }, diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs index 04392009a..cc277bd7e 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs @@ -53,7 +53,7 @@ impl BackendRuntimeSpec { decode_jitter_ms: profile.service_time_ms.decode_jitter_ms, ttft_ms: profile.service_time_ms.ttft_mean, ttft_jitter_ms: profile.service_time_ms.ttft_jitter_ms, - prefill_tokens_per_s: profile.service_time_ms.prefill_tokens_per_s.unwrap_or(0.0), + prefill_tokens_per_s: profile.service_time_ms.prefill_tokens_per_s, max_concurrent_requests, kv_cache_capacity_tokens, } @@ -68,7 +68,7 @@ pub(crate) struct PylonRuntimeSpec { pub(crate) inference_server_id: String, pub(crate) cluster_id: Option, pub(crate) profile_slug: String, - pub(crate) last_mean_input_tps: f64, + pub(crate) initial_input_tps: f64, } impl PylonRuntimeSpec { @@ -82,7 +82,7 @@ impl PylonRuntimeSpec { inference_server_id: backend_name(backend_index), cluster_id: config.backends.cluster_id_for_index(backend_index), profile_slug: slugify(&profile.name), - last_mean_input_tps: profile.registration.last_mean_input_tps, + initial_input_tps: profile.service_time_ms.prefill_tokens_per_s, } } @@ -135,8 +135,7 @@ mod tests { "decode_tokens_per_s": 50, "decode_jitter_ms": 2, "prefill_tokens_per_s": 123.0 - }, - "registration": { "last_mean_input_tps": 100.0 } + } } }, "traffic_pattern": { @@ -165,7 +164,7 @@ mod tests { } #[test] - fn pylon_runtime_spec_targets_shared_upstream_and_keeps_registration_identity() { + fn pylon_runtime_spec_targets_shared_upstream() { let config = config(); let spec = PylonRuntimeSpec::for_backend(&config, 1); @@ -175,6 +174,7 @@ mod tests { assert_eq!(spec.upstream_backend_name, "backend-0"); assert_eq!(spec.inference_server_id, "backend-1"); assert_eq!(spec.cluster_id.as_deref(), Some("cluster-0")); + assert_eq!(spec.initial_input_tps, 123.0); assert!(!spec.owns_upstream_backend()); } diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs index 82e54f5e0..c68109123 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs @@ -563,8 +563,8 @@ pub fn topology_for(backends: &BackendConfig) -> RoutingTopology { for index in 0..backends.count { let capacity = backends .profile_for_index(index) - .registration - .last_mean_input_tps; + .service_time_ms + .prefill_tokens_per_s; if capacity > 0.0 && capacity.is_finite() { let backend_id = format!("backend-{index}"); let cluster_id = backends.effective_cluster_id_for_index(index); @@ -660,9 +660,7 @@ fn percentile(values: &[u64], q: f64) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::config::{ - BackendProfile, BackendProfileGroup, RegistrationConfig, ServiceTimeConfig, - }; + use crate::config::{BackendProfile, BackendProfileGroup, ServiceTimeConfig}; fn result(backend: &str, ttft: u64, ttlt: u64) -> RequestResult { RequestResult { @@ -705,7 +703,7 @@ mod tests { default_tps: f64, groups: &[(usize, f64)], ) -> BackendConfig { - let profile = |last_mean_input_tps| BackendProfile { + let profile = |prefill_tokens_per_s| BackendProfile { name: "test".to_string(), weight: 1.0, max_concurrent_requests: None, @@ -715,10 +713,7 @@ mod tests { ttft_jitter_ms: 0, decode_tokens_per_s: 1, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps, + prefill_tokens_per_s, }, }; BackendConfig { diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs index 512ef322d..71fd649f1 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs @@ -68,6 +68,14 @@ pub enum LoadBalancerAlgorithm { PulsarWaitAndWiden, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub(crate) enum PulsarRendezvousWeight { + #[default] + LastMeanInputTps, + MaxInputTps, +} + impl LoadBalancerAlgorithm { pub const ALL: [Self; 6] = [ Self::PowerOfN, @@ -280,6 +288,7 @@ pub struct LoadBalancerAlgorithmConfig { pub max_input_work_seconds: Option, pub request_algorithms: HashMap, pub settings: LoadBalancerAlgorithmSettings, + pub(crate) rendezvous_weight: PulsarRendezvousWeight, } impl LoadBalancerAlgorithmConfig { @@ -436,6 +445,7 @@ impl RawCommonAlgorithmConfig { self, settings: LoadBalancerAlgorithmSettings, consider_kv_free_tokens: Option, + rendezvous_weight: PulsarRendezvousWeight, ) -> Result { let algorithm = settings.algorithm(); if !self.unsupported_fields.is_empty() { @@ -458,6 +468,7 @@ impl RawCommonAlgorithmConfig { max_input_work_seconds: self.max_input_work_seconds, request_algorithms: self.request_algorithms, settings, + rendezvous_weight, }) } } @@ -493,6 +504,8 @@ enum RawLoadBalancerAlgorithmConfig { Pulsar { seed: Option, consider_kv_free_tokens: Option, + #[serde(default)] + rendezvous_weight: PulsarRendezvousWeight, #[serde(flatten)] common: RawCommonAlgorithmConfig, }, @@ -501,6 +514,8 @@ enum RawLoadBalancerAlgorithmConfig { #[serde(flatten)] settings: WaitAndWidenAlgorithmConfig, consider_kv_free_tokens: Option, + #[serde(default)] + rendezvous_weight: PulsarRendezvousWeight, #[serde(flatten)] common: RawCommonAlgorithmConfig, }, @@ -513,47 +528,64 @@ impl RawLoadBalancerAlgorithmConfig { RawCommonAlgorithmConfig, LoadBalancerAlgorithmSettings, Option, + PulsarRendezvousWeight, ) { match self { Self::PowerOfN { settings, common } => ( common, LoadBalancerAlgorithmSettings::PowerOfN(settings), None, + PulsarRendezvousWeight::default(), ), Self::WaitAndWiden { settings, common } => ( common, LoadBalancerAlgorithmSettings::WaitAndWiden(settings), None, + PulsarRendezvousWeight::default(), + ), + Self::RoundRobin(common) => ( + common, + LoadBalancerAlgorithmSettings::RoundRobin, + None, + PulsarRendezvousWeight::default(), + ), + Self::Random(common) => ( + common, + LoadBalancerAlgorithmSettings::Random, + None, + PulsarRendezvousWeight::default(), ), - Self::RoundRobin(common) => (common, LoadBalancerAlgorithmSettings::RoundRobin, None), - Self::Random(common) => (common, LoadBalancerAlgorithmSettings::Random, None), Self::Pulsar { seed, consider_kv_free_tokens, + rendezvous_weight, common, } => ( common, LoadBalancerAlgorithmSettings::Pulsar(seed), consider_kv_free_tokens, + rendezvous_weight, ), Self::PulsarWaitAndWiden { settings, consider_kv_free_tokens, + rendezvous_weight, common, } => ( common, LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(settings), consider_kv_free_tokens, + rendezvous_weight, ), } } fn into_config(self) -> Result { - let (common, settings, consider_kv_free_tokens) = self.normalized(); + let (common, settings, consider_kv_free_tokens, rendezvous_weight) = self.normalized(); if let LoadBalancerAlgorithmSettings::PowerOfN(config) = &settings { config.validated_sample_count()?; } - common.into_config(settings, consider_kv_free_tokens) + common.into_config(settings, consider_kv_free_tokens, rendezvous_weight) } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs index 7e6110f94..f986cb8d2 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs @@ -44,6 +44,7 @@ pub(crate) use algorithm::input_work_seconds_for_request; pub(super) use algorithm::{HashInputBuilder, cache_affinity_key_is_cacheable, input_work_units}; pub use cluster_comparator::ClusterComparator; pub(super) use cluster_comparator::{Ttft, ttft}; +pub(crate) use config::PulsarRendezvousWeight; pub use config::{ LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerAlgorithmOverride, LoadBalancerAlgorithmSettings, LoadBalancerConfig, LoadBalancerModelConfig, diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs index d0799661d..4e1f3d826 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs @@ -17,13 +17,15 @@ mod ranking; use std::sync::Arc; +use stargate_protocol::common::valid_last_mean_input_tps; + use super::{ LoadBalancer, LoadBalancerAlgorithmConfig, LoadBalancerCandidateChoice, LoadBalancerRequest, }; use crate::routing_state::RoutedClusterSnapshot; use ranking::{ PulsarRankingLookup, PulsarRankingStore, PulsarScorer, ScoredCandidate, - compare_ranked_candidate, has_valid_input_capacity, + compare_ranked_candidate, }; #[cfg(test)] @@ -56,11 +58,12 @@ impl LoadBalancer for PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Option { - let (lookup, cached_choice) = - self.rankings - .lookup_choice(request, candidates, |ranking| { - self.choose_from_ranked_indices(request, candidates, ranking) - })?; + let (lookup, cached_choice) = self.rankings.lookup_choice( + self.config.rendezvous_weight, + request, + candidates, + |ranking| self.choose_from_ranked_indices(request, candidates, ranking), + )?; match lookup { PulsarRankingLookup::Hit => cached_choice, PulsarRankingLookup::MissCacheable => { @@ -70,9 +73,13 @@ impl LoadBalancer for PulsarLoadBalancer { } let choice = self.choose_from_ranked_indices(request, candidates, &ranking); self.rankings - .insert_or_choose_existing(request, candidates, ranking, |cached| { - self.choose_from_ranked_indices(request, candidates, cached) - }) + .insert_or_choose_existing( + self.config.rendezvous_weight, + request, + candidates, + ranking, + |cached| self.choose_from_ranked_indices(request, candidates, cached), + ) .or(choice) } PulsarRankingLookup::MissBypass => self.choose_by_score_scan(request, candidates), @@ -116,7 +123,12 @@ impl PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Vec { - ranking::pulsar_ranked_indices(self.config.seed(), request, candidates) + ranking::pulsar_ranked_indices( + self.config.seed(), + self.config.rendezvous_weight, + request, + candidates, + ) } fn choose_by_score_scan( @@ -124,7 +136,8 @@ impl PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Option { - let mut scorer = PulsarScorer::new(self.config.seed(), request); + let mut scorer = + PulsarScorer::new(self.config.seed(), request, self.config.rendezvous_weight); let mut best_overall = None; let mut best_feasible = None; for (candidate_index, candidate) in candidates.iter().enumerate() { @@ -182,7 +195,8 @@ impl PulsarLoadBalancer { chosen: &RoutedClusterSnapshot, chosen_score: f64, ) -> (usize, bool) { - let mut scorer = PulsarScorer::new(self.config.seed(), request); + let mut scorer = + PulsarScorer::new(self.config.seed(), request, self.config.rendezvous_weight); let mut rank_depth = 1; let mut skipped_for_kv_free_tokens = false; for candidate in candidates { @@ -201,7 +215,7 @@ impl PulsarLoadBalancer { #[cfg(test)] pub(super) fn weight(&self, candidate: &RoutedClusterSnapshot) -> Option { - ranking::pulsar_weight(candidate) + ranking::pulsar_weight(self.config.rendezvous_weight, candidate) } pub(super) fn feasibility( @@ -218,7 +232,7 @@ pub(super) fn input_work_admission_candidate( request: &LoadBalancerRequest<'_>, candidate: &RoutedClusterSnapshot, ) -> bool { - has_valid_input_capacity(candidate) + valid_last_mean_input_tps(candidate.stats.last_mean_input_tps) && candidate_feasibility(config, request, candidate).is_eligible() } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs index 65e11df47..a68bb11ee 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs @@ -18,11 +18,11 @@ use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::Arc; use parking_lot::RwLock; -use stargate_protocol::common::valid_last_mean_input_tps; use xxhash_rust::xxh3::xxh3_64; use crate::load_balancer::{ - LoadBalancerCandidateChoice, LoadBalancerRequest, cache_affinity_key_is_cacheable, + LoadBalancerCandidateChoice, LoadBalancerRequest, PulsarRendezvousWeight, + cache_affinity_key_is_cacheable, }; use crate::routing_state::{RoutedClusterSnapshot, RoutingTargetKey}; @@ -37,6 +37,7 @@ pub(super) struct PulsarRankingStore { impl PulsarRankingStore { pub(super) fn lookup_choice( &self, + selector: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], choose: impl Fn(&[usize]) -> Option, @@ -52,7 +53,7 @@ impl PulsarRankingStore { { let cache = self.cache.read(); - if cache.matches(request.routing_target, candidates) + if cache.matches(request.routing_target, candidates, selector) && let Some(ranking) = cache.get_ref(cache_affinity_key) { // Choose under the guard to avoid cloning the cached Arc. @@ -61,7 +62,7 @@ impl PulsarRankingStore { } let mut cache = self.cache.write(); - cache.refresh_if_needed(request.routing_target, candidates); + cache.refresh_if_needed(request.routing_target, candidates, selector); if let Some(ranking) = cache.get_ref(cache_affinity_key) { // Another thread populated the ranking after the read miss. return Some((PulsarRankingLookup::Hit, choose(ranking))); @@ -71,6 +72,7 @@ impl PulsarRankingStore { pub(super) fn insert_or_choose_existing( &self, + selector: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ranking: Arc>, @@ -78,7 +80,7 @@ impl PulsarRankingStore { ) -> Option { let cache_affinity_key = request.cache_affinity_key.unwrap_or(""); let mut cache = self.cache.write(); - cache.refresh_if_needed(request.routing_target, candidates); + cache.refresh_if_needed(request.routing_target, candidates, selector); if let Some(cached) = cache.get_ref(cache_affinity_key) { return choose(cached); } @@ -103,7 +105,7 @@ pub(super) enum PulsarRankingLookup { #[derive(Clone, Debug, PartialEq)] struct PulsarCandidateSignature { cluster_id: String, - last_mean_input_tps_bits: u64, + weight_bits: Option, } #[derive(Debug, Default)] @@ -121,8 +123,9 @@ impl PulsarRankingCache { &mut self, target: &RoutingTargetKey, candidates: &[RoutedClusterSnapshot], + selector: PulsarRendezvousWeight, ) { - if self.matches(target, candidates) { + if self.matches(target, candidates, selector) { return; } @@ -131,7 +134,7 @@ impl PulsarRankingCache { .iter() .map(|candidate| PulsarCandidateSignature { cluster_id: candidate.cluster_id.clone(), - last_mean_input_tps_bits: candidate.stats.last_mean_input_tps.to_bits(), + weight_bits: pulsar_weight(selector, candidate).map(f64::to_bits), }) .collect(); self.rankings.clear(); @@ -140,7 +143,12 @@ impl PulsarRankingCache { self.probation_order.clear(); } - fn matches(&self, target: &RoutingTargetKey, candidates: &[RoutedClusterSnapshot]) -> bool { + fn matches( + &self, + target: &RoutingTargetKey, + candidates: &[RoutedClusterSnapshot], + selector: PulsarRendezvousWeight, + ) -> bool { self.target.as_ref() == Some(target) && self.candidate_signature.len() == candidates.len() && self @@ -149,8 +157,8 @@ impl PulsarRankingCache { .zip(candidates) .all(|(cached, candidate)| { cached.cluster_id == candidate.cluster_id - && cached.last_mean_input_tps_bits - == candidate.stats.last_mean_input_tps.to_bits() + && cached.weight_bits + == pulsar_weight(selector, candidate).map(f64::to_bits) }) } @@ -232,10 +240,15 @@ pub(super) fn compare_ranked_candidate( pub(super) struct PulsarScorer { hash_bytes: Vec, prefix_len: usize, + selector: PulsarRendezvousWeight, } impl PulsarScorer { - pub(super) fn new(seed: Option<&str>, request: &LoadBalancerRequest<'_>) -> Self { + pub(super) fn new( + seed: Option<&str>, + request: &LoadBalancerRequest<'_>, + selector: PulsarRendezvousWeight, + ) -> Self { let hash_bytes = pulsar_hash_prefix( seed, &request.routing_target.routing_key, @@ -245,11 +258,12 @@ impl PulsarScorer { Self { prefix_len: hash_bytes.len(), hash_bytes, + selector, } } pub(super) fn score(&mut self, candidate: &RoutedClusterSnapshot) -> Option { - let weight = pulsar_weight(candidate)?; + let weight = pulsar_weight(self.selector, candidate)?; let u = self.hash_to_unit_interval(candidate); let e = -u.ln(); (e.is_finite() && e > 0.0).then(|| weight / e) @@ -269,10 +283,11 @@ impl PulsarScorer { pub(in crate::load_balancer) fn pulsar_ranked_indices( seed: Option<&str>, + selector: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Vec { - let mut scorer = PulsarScorer::new(seed, request); + let mut scorer = PulsarScorer::new(seed, request, selector); let mut scored = Vec::with_capacity(candidates.len()); for (candidate_index, candidate) in candidates.iter().enumerate() { if let Some(score) = scorer.score(candidate) { @@ -297,18 +312,15 @@ pub(in crate::load_balancer) fn pulsar_ranked_indices( .collect() } -pub(super) fn pulsar_weight(candidate: &RoutedClusterSnapshot) -> Option { - // Default to a stable capacity signal rather than live load. PULSAR needs a - // deterministic per-key ranking for cache affinity; if ranking follows transient - // load, hot prefixes flap between backends and destroy locality. Relative load - // belongs in feasibility gates, not in the base rendezvous weight. `last_mean_input_tps` - // is the built-in stable capacity proxy we already have, and for PULSAR it is - // required: a backend without valid capacity metadata does not participate. - has_valid_input_capacity(candidate).then_some(candidate.stats.last_mean_input_tps) -} - -pub(super) fn has_valid_input_capacity(candidate: &RoutedClusterSnapshot) -> bool { - valid_last_mean_input_tps(candidate.stats.last_mean_input_tps) +pub(super) fn pulsar_weight( + selector: PulsarRendezvousWeight, + candidate: &RoutedClusterSnapshot, +) -> Option { + let value = match selector { + PulsarRendezvousWeight::LastMeanInputTps => Some(candidate.stats.last_mean_input_tps), + PulsarRendezvousWeight::MaxInputTps => candidate.stats.max_input_tps, + }?; + (value > 0.0 && value.is_finite()).then_some(value) } const PULSAR_HASH_VERSION: u8 = 1; diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs index c55cb9446..3164f357c 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs @@ -141,6 +141,12 @@ fn seeded_pulsar_algorithm_config(seed: &str) -> LoadBalancerAlgorithmConfig { config } +fn max_weight_pulsar_algorithm_config(seed: &str) -> LoadBalancerAlgorithmConfig { + parse_json(&format!( + r#"{{"algorithm":"pulsar","seed":"{seed}","rendezvous_weight":"max-input-tps"}}"# + )) +} + #[test] fn set_seed_reports_unsupported_algorithms_without_panicking() { for algorithm in [ @@ -669,6 +675,20 @@ fn input_work_seconds_for_pulsar_excludes_low_free_kv_when_considered() { assert!((seconds - 1.5).abs() < f64::EPSILON); } +#[test] +fn max_rendezvous_weight_keeps_mean_input_work_capacity() { + let config = max_weight_pulsar_algorithm_config("seed-1"); + let target = target(); + let request = request(&target, Some("prefix-a"), Some(100)); + let mut backend = work_candidate("backend", 5, 50.0, 50); + backend.stats.max_input_tps = Some(500.0); + + assert_eq!( + input_work_seconds_for_request(&config, &request, &[backend]), + Some(3.0) + ); +} + #[test] fn invalid_algorithm_name_fails_during_parse() { assert_json_rejected::(r#"{"default":"not-a-real-lb"}"#, "not-a-real-lb"); @@ -714,6 +734,10 @@ fn algorithm_specific_load_balancer_fields_are_rejected_for_other_algorithms() { r#"{"algorithm":"wait-and-widen","consider_kv_free_tokens":true}"#, "consider_kv_free_tokens", ), + ( + r#"{"algorithm":"round-robin","rendezvous_weight":"max-input-tps"}"#, + "rendezvous_weight", + ), (r#"{"algorithm":"random","sample_count":4}"#, "sample_count"), ( r#"{"algorithm":"random","comparator":"ttft"}"#, @@ -937,6 +961,29 @@ fn detailed_algorithm_configs_preserve_all_variant_identities() { } } +#[test] +fn pulsar_rendezvous_weight_parses_for_both_pulsar_variants() { + for algorithm in ["pulsar", "pulsar-wait-and-widen"] { + let config: LoadBalancerAlgorithmConfig = parse_json(&format!( + r#"{{"algorithm":"{algorithm}","rendezvous_weight":"max-input-tps"}}"# + )); + assert_eq!( + config.rendezvous_weight, + PulsarRendezvousWeight::MaxInputTps + ); + } + + for algorithm in [ + LoadBalancerAlgorithm::Pulsar, + LoadBalancerAlgorithm::PulsarWaitAndWiden, + ] { + assert_eq!( + LoadBalancerAlgorithmConfig::from(algorithm).rendezvous_weight, + PulsarRendezvousWeight::LastMeanInputTps + ); + } +} + #[test] fn unknown_load_balancer_config_fields_are_rejected() { assert_json_rejected::( @@ -2405,7 +2452,12 @@ fn pulsar_ranking_returns_candidate_slice_indices() { work_candidate("fast", 5, 10.0, 0), ]; - let ranking = pulsar_ranked_indices(config.seed(), &request, &candidates); + let ranking = pulsar_ranked_indices( + config.seed(), + PulsarRendezvousWeight::default(), + &request, + &candidates, + ); assert_eq!(ranking.len(), 2); assert_eq!( @@ -2560,6 +2612,51 @@ fn pulsar_ranking_cache_invalidates_when_capacity_weight_changes() { panic!("expected to find an affinity key whose ranking changes after capacity changes"); } +#[test] +fn pulsar_max_ranking_cache_tracks_only_the_selected_weight() { + let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); + let target = target(); + let first_key = "max-cache-a"; + let second_key = "max-cache-b"; + let mut snapshots = candidates(&["inst-a", "inst-b"]); + snapshots[0].stats.max_input_tps = Some(100.0); + snapshots[1].stats.max_input_tps = Some(200.0); + choose( + &pulsar, + &request(&target, Some(first_key), Some(128)), + &snapshots, + ); + choose( + &pulsar, + &request(&target, Some(second_key), Some(128)), + &snapshots, + ); + assert_eq!( + pulsar.cached_affinity_key_bytes(), + first_key.len() + second_key.len() + ); + + snapshots[0].stats.last_mean_input_tps = 1.0; + snapshots[1].stats.last_mean_input_tps = 10_000.0; + choose( + &pulsar, + &request(&target, Some(first_key), Some(128)), + &snapshots, + ); + assert_eq!( + pulsar.cached_affinity_key_bytes(), + first_key.len() + second_key.len() + ); + + snapshots[0].stats.max_input_tps = Some(300.0); + choose( + &pulsar, + &request(&target, Some(first_key), Some(128)), + &snapshots, + ); + assert_eq!(pulsar.cached_affinity_key_bytes(), first_key.len()); +} + #[test] fn pulsar_does_not_cache_oversized_affinity_key() { let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); @@ -2596,6 +2693,26 @@ fn pulsar_uses_last_mean_input_tps_as_weight() { assert_eq!(pulsar.weight(&candidate), Some(123.0)); } +#[test] +fn pulsar_can_use_generation_max_input_tps_as_weight() { + let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); + let mut candidate = work_candidate("inst-a", 5, 123.0, 0); + candidate.stats.max_input_tps = Some(456.0); + + assert_eq!(pulsar.weight(&candidate), Some(456.0)); + + candidate.stats.last_mean_input_tps = 0.0; + assert_eq!(pulsar.weight(&candidate), Some(456.0)); +} + +#[test] +fn pulsar_max_weight_does_not_fall_back_to_mean() { + let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); + let candidate = work_candidate("inst-a", 5, 123.0, 0); + + assert_eq!(pulsar.weight(&candidate), None); +} + #[test] fn pulsar_excludes_candidate_with_invalid_last_mean_input_tps() { let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs index ed781cbe0..2a15fa564 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs @@ -103,6 +103,7 @@ impl ClusterRoutingGeneration { } let backend_count = active_backend_count as u128; let mut backend_stats = ModelStats::default(); + let mut max_input_tps = Some(0.0); let mut rtt_mean_nanos = 0_u128; let mut rtt_remainder_nanos = 0_u128; @@ -111,6 +112,13 @@ impl ClusterRoutingGeneration { if valid_last_mean_input_tps(backend.stats.last_mean_input_tps) { backend_stats.last_mean_input_tps += backend.stats.last_mean_input_tps; } + max_input_tps = match (max_input_tps, backend.stats.max_input_tps) { + (Some(sum), Some(value)) if value > 0.0 && value.is_finite() => { + let total = sum + value; + total.is_finite().then_some(total) + } + _ => None, + }; backend_stats.queue_size += backend.stats.queue_size; backend_stats.queued_input_size += backend.stats.queued_input_size; backend_stats.input_processing_queries += backend.stats.input_processing_queries; @@ -143,6 +151,7 @@ impl ClusterRoutingGeneration { (rtt_mean_nanos / 1_000_000_000) as u64, (rtt_mean_nanos % 1_000_000_000) as u32, ); + backend_stats.max_input_tps = max_input_tps; Some((backend_stats, rtt, active_backend_count)) } diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs index 522f417a1..7735d46cf 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs @@ -307,6 +307,7 @@ fn shared_backend_a_stats() -> ModelStats { ModelStats { output_tps: 2.0, last_mean_input_tps: 100.0, + max_input_tps: Some(110.0), max_output_tps: 50.0, queue_size: 1, queued_input_size: 100, @@ -329,6 +330,7 @@ fn shared_backend_b_stats() -> ModelStats { ModelStats { output_tps: 5.0, last_mean_input_tps: 120.0, + max_input_tps: Some(130.0), max_output_tps: 60.0, queue_size: 2, queued_input_size: 200, @@ -1720,6 +1722,7 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { backend_a.stats = ModelStats { output_tps: 1.5, last_mean_input_tps: 10.0, + max_input_tps: Some(15.0), queue_size: 2, queued_input_size: 20, input_processing_queries: 1, @@ -1739,6 +1742,7 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { backend_b.stats = ModelStats { output_tps: 2.5, last_mean_input_tps: 30.0, + max_input_tps: Some(35.0), queue_size: 4, queued_input_size: 40, input_processing_queries: 5, @@ -1760,6 +1764,7 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { assert_stats!(stats, output_tps: 4.0, last_mean_input_tps: 40.0, + max_input_tps: Some(50.0), queue_size: 6, queued_input_size: 60, input_processing_queries: 6, @@ -1777,6 +1782,73 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { assert_eq!(stats.stats_sources.len(), 3); } +#[test] +fn cluster_backend_aggregate_requires_every_backend_max() { + let mut backend_a = backend!("cluster-max", "backend-a", 1.0, 10.0, 5); + backend_a.stats.max_input_tps = Some(15.0); + let cluster_state = RoutedClusterState::new(backend_a.registration.cluster_generation.clone()); + cluster_state.upsert_backend(Arc::new(backend_a.clone())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("one backend max should aggregate") + .0 + .max_input_tps, + Some(15.0) + ); + + let mut backend_b = backend!( + backend_a.registration.cluster_generation.clone() => + "cluster-max", "backend-b", 1.0, 10.0, 5 + ); + backend_b.stats.max_input_tps = Some(25.0); + cluster_state.upsert_backend(Arc::new(backend_b.clone())); + let complete_bits = cluster_state + .backend_aggregate() + .expect("complete backend max should aggregate") + .0 + .max_input_tps + .map(f64::to_bits); + assert_eq!(complete_bits, Some(40.0_f64.to_bits())); + + let reversed = RoutedClusterState::new(backend_a.registration.cluster_generation.clone()); + reversed.upsert_backend(Arc::new(backend_b.clone())); + reversed.upsert_backend(Arc::new(backend_a)); + assert_eq!( + reversed + .backend_aggregate() + .expect("backend order should not change aggregate") + .0 + .max_input_tps + .map(f64::to_bits), + complete_bits + ); + + backend_b.stats.max_input_tps = None; + cluster_state.upsert_backend(Arc::new(backend_b.clone())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("missing max should not remove the cluster") + .0 + .max_input_tps, + None + ); + + for invalid in [0.0, -1.0, f64::NAN, f64::INFINITY] { + backend_b.stats.max_input_tps = Some(invalid); + cluster_state.upsert_backend(Arc::new(backend_b.clone())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("invalid max should not remove the cluster") + .0 + .max_input_tps, + None + ); + } +} + #[tokio::test] async fn list_active_models_filters_by_routing_key() { let scenario = RegistrationScenario::new(None); @@ -2153,6 +2225,7 @@ async fn shared_cluster_registration_exposes_one_aggregated_cluster_candidate() assert_eq!(cluster.active_backend_count, 2); assert_stats!(cluster.stats, last_mean_input_tps: 220.0, + max_input_tps: Some(240.0), output_tps: 7.0, queue_size: 3, queued_input_size: 300, diff --git a/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs b/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs index 6a72031a8..06793745d 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs @@ -1269,10 +1269,7 @@ pub async fn start_and_register_backend_with_bringup( ModelLifecycleConfig { upstream_http_base_url: upstream_http_base_url.clone(), source: ModelSource::Static(BTreeSet::from([model.to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1.0, - pin: true, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1.0 }, bringup, health_paths: pylon_lib::UpstreamHealthPaths::default(), startup_health_wait: std::time::Duration::ZERO, diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs index d685908a2..082a7bff0 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs @@ -1264,10 +1264,7 @@ impl LifecycleBackendOptions<'_> { ModelLifecycleConfig { upstream_http_base_url: upstream_http_base_url.clone(), source: ModelSource::Static(model_ids.into_iter().collect::>()), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1.0, - pin: true, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1.0 }, bringup, health_paths: pylon_lib::UpstreamHealthPaths::default(), startup_health_wait: std::time::Duration::ZERO,