diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 3cd57acada..db2c5797e7 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1331,7 +1331,7 @@ "FILE:@@//src/invocation-plane-services/http-invocation/crates/server/Cargo.toml fd58955d67ba6d2f03a63f582f937a783a9ffa9590c52e5969ac7b993d95b97a", "FILE:@@//src/libraries/rust/stargate/Cargo.lock 4e02eb8522889835f86fd866b968c1ab7f9f36b73a541cba0037ee26c8ef2bff", "FILE:@@//src/libraries/rust/stargate/Cargo.toml b2314ab83630edb9b57bdd6a396a08a12f2d10966ed6a6fec3e5df700ffe0cfc", - "FILE:@@//src/libraries/rust/stargate/crates/mock-dynamo/Cargo.toml 775c65c930aae4b90cd0673fcae8f8004c69dab454175dba6a712a3dc331c4cf", + "FILE:@@//src/libraries/rust/stargate/crates/mock-dynamo/Cargo.toml 6dc7e969f92ca031f252a2d1afe9d7a37c9c06a872bf28198467fc98830ec94c", "FILE:@@//src/libraries/rust/stargate/crates/proto/Cargo.toml 14c9ed935d6daf6223c449bb026778c4395c2e8d6dc84a67fb372a90b24ffae0", "FILE:@@//src/libraries/rust/stargate/crates/protocol/Cargo.toml c08e2034ca5c02374a6bb44a3f4ec1aee43af82a097880624ed1fa77811ee241", "FILE:@@//src/libraries/rust/stargate/crates/pylon/Cargo.toml 681f4d77101aacb261932dc14e1ebaf0c3be5dc71307fed5500565ce4c0c8192", diff --git a/src/libraries/rust/stargate/README.md b/src/libraries/rust/stargate/README.md index d80d799491..52c2634c07 100644 --- a/src/libraries/rust/stargate/README.md +++ b/src/libraries/rust/stargate/README.md @@ -30,6 +30,11 @@ set. Every Pylon must also select exactly one local stats initialization source: [Pylon onboarding](docs/operations/pylon-onboarding.md) for the complete lifecycle and calibration contract. +Pulsar weights clusters by generation maximum input TPS by default. Upgrade +Pylons and registration relays before Stargates; see +[load-balancer configuration](docs/load-balancer-configuration.md#pulsar) for +the rollout order and the mean-weighting option. + Pylon gates startup on an upstream health probe. It tries `/health` and then `/v1/health/ready`, reuses whichever path answers first, and forwards Stargate's `/health` RTT probe to that same path, so engines that serve only the diff --git a/src/libraries/rust/stargate/benches/backend-degradation.yaml b/src/libraries/rust/stargate/benches/backend-degradation.yaml index c17ff479a3..33af7876f7 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 21d0ce1e9c..5a36bc5d51 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 482895b8ff..95e5d6c1b7 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 dc588cfb80..759be16fb0 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 3c1d090410..02af1ec6b0 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 a7cbc5df5a..9a666d2d90 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 7c28f4a438..0e986f847a 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 b4fdb7d1ac..1d85af82aa 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 447e4db955..1800d12773 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 2019da9fa2..75d6b10b8e 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 0e31c2394f..b02085a278 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 06c36fbbc9..9cb7d8d55f 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 6bc9f82065..c61d28485c 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 7130272021..01560fa0e9 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 c280d752b4..6f8c7de71f 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 7b30727645..cd264c0e95 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 419fb13dce..833b72bdc4 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/Cargo.toml b/src/libraries/rust/stargate/crates/mock-dynamo/Cargo.toml index 5534b63b58..f89d621f05 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/Cargo.toml +++ b/src/libraries/rust/stargate/crates/mock-dynamo/Cargo.toml @@ -40,3 +40,6 @@ serde_json = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } 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 2cb0e71982..2d7127a16a 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs @@ -231,12 +231,11 @@ pub(crate) async fn chat_completions( let stream = req.stream == Some(true); info!(id = %id, model = %model, stream = stream, "received chat/completions request"); 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!( @@ -360,6 +359,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)); @@ -594,8 +594,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 let StreamKind::Chat { include_usage, .. } = kind { yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Role, include_usage)); } 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 8df33f0bc8..c431241428 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs @@ -889,6 +889,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(); + 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 323ee393d7..e33537de7c 100644 --- a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto +++ b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto @@ -123,6 +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. + optional double max_input_tps = 18; } enum InferenceServerStatus { diff --git a/src/libraries/rust/stargate/crates/proto/src/lib.rs b/src/libraries/rust/stargate/crates/proto/src/lib.rs index 12ac63e65c..96baa498cd 100644 --- a/src/libraries/rust/stargate/crates/proto/src/lib.rs +++ b/src/libraries/rust/stargate/crates/proto/src/lib.rs @@ -92,6 +92,30 @@ pub mod pb { #[cfg(test)] mod tests { use crate::pb::{StargateInfo, WatchStargatesResponse}; + use prost::Message; + + #[test] + fn model_stats_maximum_is_additive_and_preserves_absence() { + #[derive(Clone, PartialEq, Message)] + struct LegacyStats { + #[prost(double, tag = "1")] + last_mean_input_tps: f64, + } + let current = crate::pb::ModelStats { + last_mean_input_tps: 50.0, + max_input_tps: Some(200.0), + ..Default::default() + }; + let bytes = current.encode_to_vec(); + let decoded = crate::pb::ModelStats::decode(bytes.as_slice()).unwrap(); + assert_eq!(decoded.max_input_tps, Some(200.0)); + let legacy = LegacyStats::decode(bytes.as_slice()).unwrap(); + assert_eq!(legacy.last_mean_input_tps, 50.0); + let old_bytes = legacy.encode_to_vec(); + let old = crate::pb::ModelStats::decode(old_bytes.as_slice()).unwrap(); + assert_eq!(old.last_mean_input_tps, 50.0); + assert_eq!(old.max_input_tps, None); + } #[test] fn watch_stargates_response_json_serde_normalizes_set_fields() { 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 2183ddbc01..3eeb51e32d 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 @@ -1477,6 +1477,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, @@ -1502,6 +1503,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 4719389a63..2c691057d0 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,6 +33,8 @@ 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. + 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, @@ -467,6 +469,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 7fc319ded0..a75bce1b46 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 @@ -34,6 +34,7 @@ pub(super) const ENGINE_STATS_SOURCE: &str = "engine_stats_stream"; #[derive(Debug, Default)] pub(super) struct ModelMetricsState { pub(super) last_mean_input_tps: f64, + pub(super) max_input_tps: Option, pub(super) max_engine_concurrency: Option, pub(super) chat_output_tps_samples: VecDeque, pub(super) chat_output_tps_sum: f64, @@ -472,6 +473,7 @@ impl StatsAggregator { .expect("configured input TPS must be positive and finite"); ModelMetricsState { last_mean_input_tps: input_tps, + max_input_tps: Some(input_tps), input_tps_distribution, aggregate_state_counted: true, ..ModelMetricsState::default() @@ -1060,7 +1062,7 @@ pub(super) fn apply_input_throughput_sample( if model_state.last_mean_input_tps == mean_input_tps { return false; } - model_state.last_mean_input_tps = mean_input_tps; + model_state.publish_mean_input_tps(mean_input_tps); true } @@ -1087,6 +1089,16 @@ 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) { + 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)), + ); + } + pub(super) fn clear_live_output_tps(&mut self) -> bool { self.last_stats_event_at = None; if self.chat_output_tps_samples.is_empty() { @@ -1107,6 +1119,7 @@ impl ModelMetricsState { }; CurrentModelStats { last_mean_input_tps: self.last_mean_input_tps, + max_input_tps: self.max_input_tps, 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 35abda75fd..2f8d42f79e 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 @@ -1510,12 +1510,38 @@ mod tests { StatsCollectorConfig::default(), ModelStatsInitialization::ConfiguredInputTps { input_tps: 100.0 }, ); - 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(stats.last_mean_input_tps)); + } + + #[test] + fn engine_max_input_tps_follows_the_mean_and_persists_for_the_generation() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + aggregator.stream("req-a", (0, 0), false, Duration::ZERO); + // Four 100 TPS ticks, one 1,100 TPS outlier, then three 100 TPS ticks. + let mut tokens = 0; + for (tick, delta) in [10, 10, 10, 10, 110, 10, 10, 10].into_iter().enumerate() { + tokens += delta; + aggregator.stream( + "req-a", + (tokens, 0), + false, + milliseconds(100 * (tick as u64 + 1)), + ); + } + + let stats = aggregator.snapshot("model-a"); + assert_eq!(stats.last_mean_input_tps, 225.0); + assert_eq!(stats.max_input_tps, Some(300.0)); + aggregator.sweep(seconds(24 * 60 * 60)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(300.0)); } #[test] @@ -2335,6 +2361,39 @@ mod tests { assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 1_000.0); } + #[test] + fn fallback_max_input_tps_retains_the_peak_when_the_window_mean_falls() { + let mut aggregator = test_aggregator(config!(smoothing_window_size: 1)); + let start = std::time::Instant::now(); + for (index, (input_tokens, maximum)) in [(200, 200.0), (50, 200.0), (300, 300.0)] + .into_iter() + .enumerate() + { + let submitted_at = start + seconds(index as u64 * 3); + apply_fallback_observation_with_interval( + &mut aggregator, + &RequestObservation { + input_tokens, + ..observation( + RequestObservationEndpoint::ChatCompletions, + &format!("max-input-{index}"), + RequestObservationState::OutputGeneration, + ) + }, + crate::runtime_state::RequestInputInterval { + submitted_at, + first_generated_output_at: submitted_at + seconds(1), + }, + 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)); + } + #[test] fn reused_request_id_retains_distinct_submission_intervals() { let mut aggregator = test_aggregator(StatsCollectorConfig::default()); @@ -2649,6 +2708,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)); let mut late_exact = aggregator .runtime_state @@ -2682,6 +2742,7 @@ mod tests { .begin_generation(replacement, ModelStatsInitialization::Empty) .is_some() ); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, None); assert_eq!( aggregator.per_model["model-a"] .metrics @@ -2713,6 +2774,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.per_model["model-a"] .metrics @@ -3782,6 +3844,7 @@ mod tests { r#"pylon_requests_total{model="model-a",routing_key="rk-1",status="complete"} 1"# )); 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"#)); } 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 f948125741..bbf5694ef9 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 @@ -165,6 +165,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"]); @@ -351,6 +352,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), @@ -392,6 +399,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, @@ -914,6 +922,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() @@ -1081,6 +1090,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, @@ -1100,6 +1110,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"#, @@ -1108,6 +1119,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 c62023dd51..385708bb4d 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 @@ -152,7 +152,7 @@ impl StatsAggregator { ) && model_state.last_mean_input_tps != input_tps { - model_state.last_mean_input_tps = input_tps; + model_state.publish_mean_input_tps(input_tps); input_tps_changed = true; } let mut completed_sample_recorded = false; 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 683498349d..f628c9f82f 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 2245e44fd8..81194f4078 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", + "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 5ec0ca8a1b..33e9edab0e 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, }, }, }, 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 786365d286..6a947614e7 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 258746d188..d75020bdf7 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/microbench/lb.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs index 78a83abdc5..b14f5fff99 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/microbench/lb.rs @@ -536,31 +536,35 @@ fn excluded_cluster_ids_for_scenario(scenario: LbMicrobenchScenario) -> Option Vec { (0..count) - .map(|index| RoutedClusterSnapshot { - cluster_id: format!("cluster-{index:04}"), - stats: ModelStats { - output_tps: 120.0 + (index % 11) as f64, - last_mean_input_tps: 1_200.0 + (index % 13) as f64 * 31.0, - max_output_tps: 500.0, - queue_size: (index % 7) as u64, - queued_input_size: (index % 5) as u64 * 256, - kv_cache_capacity_tokens: 131_072, - kv_cache_used_tokens: 8_192 + (index % 23) as u64 * 64, - kv_cache_free_tokens: 122_880 - (index % 23) as u64 * 64, - num_running_queries: (index % 9) as u64, - max_engine_concurrency: 16, - total_query_input_size: (index % 6) as u64 * 384, - queue_time_estimate_ms_by_priority: HashMap::from([ - (0, (index % 5) as u64), - (2, 2 + (index % 11) as u64), - (4, 4 + (index % 17) as u64), - ]), - ..ModelStats::default() - }, - rtt: Duration::from_micros(500 + (index % 19) as u64 * 75), - snapshot_updated_at: Instant::now(), - status: InferenceServerStatus::Active, - active_backend_count: 1, + .map(|index| { + let input_tps = 1_200.0 + (index % 13) as f64 * 31.0; + RoutedClusterSnapshot { + cluster_id: format!("cluster-{index:04}"), + stats: ModelStats { + output_tps: 120.0 + (index % 11) as f64, + last_mean_input_tps: input_tps, + max_input_tps: Some(input_tps), + max_output_tps: 500.0, + queue_size: (index % 7) as u64, + queued_input_size: (index % 5) as u64 * 256, + kv_cache_capacity_tokens: 131_072, + kv_cache_used_tokens: 8_192 + (index % 23) as u64 * 64, + kv_cache_free_tokens: 122_880 - (index % 23) as u64 * 64, + num_running_queries: (index % 9) as u64, + max_engine_concurrency: 16, + total_query_input_size: (index % 6) as u64 * 384, + queue_time_estimate_ms_by_priority: HashMap::from([ + (0, (index % 5) as u64), + (2, 2 + (index % 11) as u64), + (4, 4 + (index % 17) as u64), + ]), + ..ModelStats::default() + }, + rtt: Duration::from_micros(500 + (index % 19) as u64 * 75), + snapshot_updated_at: Instant::now(), + status: InferenceServerStatus::Active, + active_backend_count: 1, + } }) .collect() } 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 4c589178bb..79486237fd 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs @@ -295,7 +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(), + "--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()); @@ -350,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 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 9a500b12d4..119cf5afe1 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs @@ -498,8 +498,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, @@ -532,10 +532,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 04392009af..cc277bd7ef 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 c7d72ddc0f..38c040f488 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs @@ -594,8 +594,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); @@ -691,9 +691,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 { @@ -738,7 +736,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, @@ -748,10 +746,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-k8s-router/src/grpc.rs b/src/libraries/rust/stargate/crates/stargate-k8s-router/src/grpc.rs index 87cdc810a2..9af4b30717 100644 --- a/src/libraries/rust/stargate/crates/stargate-k8s-router/src/grpc.rs +++ b/src/libraries/rust/stargate/crates/stargate-k8s-router/src/grpc.rs @@ -388,6 +388,7 @@ mod tests { watch_hits: Arc, register_hits: Arc, metadata: MetadataRecords, + registrations: Arc>>, registration_errors: Arc>>, } @@ -499,7 +500,8 @@ mod tests { let stream = async_stream::stream! { while let Some(message) = inbound.next().await { match message { - Ok(_registration) => { + Ok(registration) => { + recorder.registrations.lock().expect("registration lock poisoned").push(registration); yield Ok(InferenceServerAck { reverse_tunnel_target: stargate_id.clone(), reverse_tunnel_pylon_dial_addr: String::new(), @@ -814,7 +816,19 @@ mod tests { let router = start_router(snapshot(&[("stargate-1", fake_b.addr)])).await; let mut client = router.client("stargate-1.stargate.external"); - let mut request = Request::new(tokio_stream::iter([registration()])); + let mut registration = registration(); + registration.models.insert( + "model-a".into(), + stargate_proto::pb::InferenceServerModelRegistration { + stats: Some(stargate_proto::pb::ModelStats { + last_mean_input_tps: 50.0, + max_input_tps: Some(200.0), + ..Default::default() + }), + status: stargate_proto::pb::InferenceServerStatus::Active as i32, + }, + ); + let mut request = Request::new(tokio_stream::iter([registration])); request.metadata_mut().insert( "authorization", "Bearer token".parse().expect("valid metadata"), @@ -836,6 +850,18 @@ mod tests { assert_eq!(ack.reverse_tunnel_target, "stargate-1"); assert_eq!(recorder_b.register_hits.load(Ordering::Relaxed), 1); + assert_eq!( + recorder_b + .registrations + .lock() + .expect("registration lock poisoned")[0] + .models["model-a"] + .stats + .as_ref() + .unwrap() + .max_input_tps, + Some(200.0) + ); assert_eq!( recorder_b .metadata diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs index ef65254e61..12fdc656f7 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/trace.rs @@ -45,6 +45,7 @@ pub(super) fn proxy_openai_request_span(headers: &HeaderMap) -> Span { selected_inst.id = field::Empty, selected_inst.output_tps = field::Empty, selected_inst.last_mean_input_tps = field::Empty, + selected_inst.max_input_tps = field::Empty, selected_inst.max_output_tps = field::Empty, selected_inst.queue_size = field::Empty, selected_inst.queued_input_size = field::Empty, @@ -126,6 +127,7 @@ pub(super) fn record_routing_to_span(span: &Span, routing: RoutingTraceFields<'_ "selected_inst.id" = routing.chosen.inference_server_id.as_str(), "selected_inst.output_tps" = stats.output_tps, "selected_inst.last_mean_input_tps" = stats.last_mean_input_tps, + "selected_inst.max_input_tps" = stats.max_input_tps, "selected_inst.max_output_tps" = stats.max_output_tps, "selected_inst.queue_size" = stats.queue_size, "selected_inst.queued_input_size" = stats.queued_input_size, 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 f7385c1a7e..b4539b85e3 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 { + LastMeanInputTps, + #[default] + MaxInputTps, +} + impl LoadBalancerAlgorithm { pub const ALL: [Self; 6] = [ Self::PowerOfN, @@ -296,6 +304,7 @@ pub struct LoadBalancerAlgorithmConfig { pub max_input_work_seconds: Option, pub request_algorithms: HashMap, pub settings: LoadBalancerAlgorithmSettings, + pub(crate) rendezvous_weight: PulsarRendezvousWeight, } impl LoadBalancerAlgorithmConfig { @@ -452,6 +461,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() { @@ -474,6 +484,7 @@ impl RawCommonAlgorithmConfig { max_input_work_seconds: self.max_input_work_seconds, request_algorithms: self.request_algorithms, settings, + rendezvous_weight, }) } } @@ -509,6 +520,8 @@ enum RawLoadBalancerAlgorithmConfig { Pulsar { seed: Option, consider_kv_free_tokens: Option, + #[serde(default)] + rendezvous_weight: PulsarRendezvousWeight, #[serde(flatten)] common: RawCommonAlgorithmConfig, }, @@ -517,6 +530,8 @@ enum RawLoadBalancerAlgorithmConfig { #[serde(flatten)] settings: WaitAndWidenAlgorithmConfig, consider_kv_free_tokens: Option, + #[serde(default)] + rendezvous_weight: PulsarRendezvousWeight, #[serde(flatten)] common: RawCommonAlgorithmConfig, }, @@ -529,43 +544,60 @@ 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(); match &settings { LoadBalancerAlgorithmSettings::PowerOfN(config) => { config.validated_sample_count()?; @@ -578,7 +610,7 @@ impl RawLoadBalancerAlgorithmConfig { | LoadBalancerAlgorithmSettings::Random | LoadBalancerAlgorithmSettings::Pulsar(_) => {} } - 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 f554b61c9b..6433a4be78 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 @@ -46,6 +46,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 d0799661d3..c9dcf03ddb 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)] @@ -37,8 +39,8 @@ pub(super) struct PulsarLoadBalancer { impl PulsarLoadBalancer { pub(super) fn new(config: LoadBalancerAlgorithmConfig) -> Self { Self { + rankings: PulsarRankingStore::new(config.rendezvous_weight), config, - rankings: PulsarRankingStore::default(), } } @@ -116,7 +118,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 +131,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 +190,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 +210,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 +227,8 @@ 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) + && ranking::pulsar_weight(config.rendezvous_weight, candidate).is_some() && 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 65e11df471..43ec56aa2c 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,23 +18,30 @@ 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}; const RANKING_CACHE_LIMIT: usize = 4096; const RANKING_CACHE_PROBATION_LIMIT: usize = 4096; -#[derive(Default)] pub(super) struct PulsarRankingStore { + rendezvous_weight: PulsarRendezvousWeight, cache: RwLock, } impl PulsarRankingStore { + pub(super) fn new(rendezvous_weight: PulsarRendezvousWeight) -> Self { + Self { + rendezvous_weight, + cache: RwLock::default(), + } + } + pub(super) fn lookup_choice( &self, request: &LoadBalancerRequest<'_>, @@ -52,7 +59,7 @@ impl PulsarRankingStore { { let cache = self.cache.read(); - if cache.matches(request.routing_target, candidates) + if cache.matches(request.routing_target, candidates, self.rendezvous_weight) && let Some(ranking) = cache.get_ref(cache_affinity_key) { // Choose under the guard to avoid cloning the cached Arc. @@ -61,7 +68,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, self.rendezvous_weight); 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))); @@ -78,7 +85,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, self.rendezvous_weight); if let Some(cached) = cache.get_ref(cache_affinity_key) { return choose(cached); } @@ -103,7 +110,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 +128,9 @@ impl PulsarRankingCache { &mut self, target: &RoutingTargetKey, candidates: &[RoutedClusterSnapshot], + rendezvous_weight: PulsarRendezvousWeight, ) { - if self.matches(target, candidates) { + if self.matches(target, candidates, rendezvous_weight) { return; } @@ -131,7 +139,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(rendezvous_weight, candidate).map(f64::to_bits), }) .collect(); self.rankings.clear(); @@ -140,7 +148,12 @@ impl PulsarRankingCache { self.probation_order.clear(); } - fn matches(&self, target: &RoutingTargetKey, candidates: &[RoutedClusterSnapshot]) -> bool { + fn matches( + &self, + target: &RoutingTargetKey, + candidates: &[RoutedClusterSnapshot], + rendezvous_weight: PulsarRendezvousWeight, + ) -> bool { self.target.as_ref() == Some(target) && self.candidate_signature.len() == candidates.len() && self @@ -149,8 +162,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(rendezvous_weight, candidate).map(f64::to_bits) }) } @@ -232,10 +245,15 @@ pub(super) fn compare_ranked_candidate( pub(super) struct PulsarScorer { hash_bytes: Vec, prefix_len: usize, + rendezvous_weight: PulsarRendezvousWeight, } impl PulsarScorer { - pub(super) fn new(seed: Option<&str>, request: &LoadBalancerRequest<'_>) -> Self { + pub(super) fn new( + seed: Option<&str>, + request: &LoadBalancerRequest<'_>, + rendezvous_weight: PulsarRendezvousWeight, + ) -> Self { let hash_bytes = pulsar_hash_prefix( seed, &request.routing_target.routing_key, @@ -245,11 +263,12 @@ impl PulsarScorer { Self { prefix_len: hash_bytes.len(), hash_bytes, + rendezvous_weight, } } pub(super) fn score(&mut self, candidate: &RoutedClusterSnapshot) -> Option { - let weight = pulsar_weight(candidate)?; + let weight = pulsar_weight(self.rendezvous_weight, candidate)?; let u = self.hash_to_unit_interval(candidate); let e = -u.ln(); (e.is_finite() && e > 0.0).then(|| weight / e) @@ -269,10 +288,11 @@ impl PulsarScorer { pub(in crate::load_balancer) fn pulsar_ranked_indices( seed: Option<&str>, + rendezvous_weight: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Vec { - let mut scorer = PulsarScorer::new(seed, request); + let mut scorer = PulsarScorer::new(seed, request, rendezvous_weight); 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 +317,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( + rendezvous_weight: PulsarRendezvousWeight, + candidate: &RoutedClusterSnapshot, +) -> Option { + let value = match rendezvous_weight { + 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/pulsar_wait_and_widen.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs index c2b2fcb399..709edc464d 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar_wait_and_widen.rs @@ -203,6 +203,7 @@ mod tests { stats: ModelStats { output_tps: 0.0, last_mean_input_tps: 100.0, + max_input_tps: Some(100.0), max_output_tps: 100.0, queue_size: 0, queued_input_size: 0, 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 23630538a1..908805c6a5 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 @@ -183,6 +183,7 @@ fn candidate(id: &str, kv_cache_free_tokens: u64) -> RoutedClusterSnapshot { cluster_id: id.to_string(), stats: ModelStats { last_mean_input_tps: 100.0, + max_input_tps: Some(100.0), max_output_tps: 100.0, kv_cache_capacity_tokens: 1024, kv_cache_used_tokens: 1024 - kv_cache_free_tokens, @@ -241,6 +242,7 @@ fn work_candidate_with_kv( .with_rtt_ms(rtt_ms) .with_stats(|stats| { stats.last_mean_input_tps = input_tps; + stats.max_input_tps = Some(input_tps); stats.queued_input_size = queued_input; }) } @@ -678,6 +680,60 @@ 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 = seeded_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 pulsar_input_work_capacity_excludes_backends_missing_the_selected_maximum() { + let target = target(); + let request = request(&target, Some("prefix-a"), Some(100)); + let ready = work_candidate("ready", 5, 50.0, 50); + let mut legacy_backend = work_candidate("legacy", 5, 10_000.0, 0); + legacy_backend.stats.max_input_tps = None; + for algorithm in [ + LoadBalancerAlgorithm::Pulsar, + LoadBalancerAlgorithm::PulsarWaitAndWiden, + ] { + let mut config = LoadBalancerAlgorithmConfig::from(algorithm); + assert_eq!( + input_work_seconds_for_request( + &config, + &request, + &[ready.clone(), legacy_backend.clone()] + ), + Some(3.0) + ); + assert_eq!( + input_work_seconds_for_request( + &config, + &request, + std::slice::from_ref(&legacy_backend) + ), + None + ); + config.rendezvous_weight = PulsarRendezvousWeight::LastMeanInputTps; + assert_eq!( + input_work_seconds_for_request( + &config, + &request, + &[ready.clone(), legacy_backend.clone()] + ), + Some(150.0 / 10_050.0) + ); + } +} + #[test] fn invalid_algorithm_name_fails_during_parse() { assert_json_rejected::(r#"{"default":"not-a-real-lb"}"#, "not-a-real-lb"); @@ -723,6 +779,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"}"#, @@ -946,6 +1006,42 @@ 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 + ); + let default: LoadBalancerAlgorithmConfig = + parse_json(&format!(r#"{{"algorithm":"{algorithm}"}}"#)); + assert_eq!( + default.rendezvous_weight, + PulsarRendezvousWeight::MaxInputTps + ); + let legacy: LoadBalancerAlgorithmConfig = parse_json(&format!( + r#"{{"algorithm":"{algorithm}","rendezvous_weight":"last-mean-input-tps"}}"# + )); + assert_eq!( + legacy.rendezvous_weight, + PulsarRendezvousWeight::LastMeanInputTps + ); + } + + for algorithm in [ + LoadBalancerAlgorithm::Pulsar, + LoadBalancerAlgorithm::PulsarWaitAndWiden, + ] { + assert_eq!( + LoadBalancerAlgorithmConfig::from(algorithm).rendezvous_weight, + PulsarRendezvousWeight::MaxInputTps + ); + } +} + #[test] fn unknown_load_balancer_config_fields_are_rejected() { assert_json_rejected::( @@ -2887,7 +2983,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!( @@ -3014,32 +3115,125 @@ fn pulsar_returns_none_when_all_candidates_are_excluded() { } #[test] -fn pulsar_ranking_cache_invalidates_when_capacity_weight_changes() { +fn pulsar_ranking_cache_invalidates_when_the_selected_weight_changes() { + for rendezvous_weight in [ + PulsarRendezvousWeight::MaxInputTps, + PulsarRendezvousWeight::LastMeanInputTps, + ] { + let mut config = seeded_pulsar_algorithm_config("seed-1"); + config.rendezvous_weight = rendezvous_weight; + let pulsar = PulsarLoadBalancer::new(config); + let target = target(); + // Only the selected weight field swaps between the two snapshots. + let with_weights = |a: f64, b: f64| { + let mut candidates = vec![ + work_candidate("inst-a", 5, 100.0, 0), + work_candidate("inst-b", 5, 100.0, 0), + ]; + for (candidate, weight) in candidates.iter_mut().zip([a, b]) { + match rendezvous_weight { + PulsarRendezvousWeight::MaxInputTps => { + candidate.stats.max_input_tps = Some(weight); + } + PulsarRendezvousWeight::LastMeanInputTps => { + candidate.stats.last_mean_input_tps = weight; + } + } + } + candidates + }; + let initial = with_weights(10_000.0, 1.0); + let changed = with_weights(1.0, 10_000.0); + + let changed_key = (0..1024).find_map(|idx| { + let key = format!("affinity-{idx}"); + let request = request(&target, Some(&key), Some(128)); + let first = choose(&pulsar, &request, &initial).candidate.cluster_id; + let second = choose(&pulsar, &request, &changed).candidate.cluster_id; + (first != second).then_some((first, second)) + }); + assert_eq!( + changed_key, + Some(("inst-a".to_string(), "inst-b".to_string())), + "{rendezvous_weight:?}: some key's ranking follows the selected weight" + ); + } +} + +#[test] +fn pulsar_max_ranking_cache_tracks_only_the_selected_weight() { let pulsar = PulsarLoadBalancer::new(seeded_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() + ); - for idx in 0..1024 { - let key = format!("affinity-{idx}"); - let request = request(&target, Some(&key), Some(128)); + 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 both_pulsar_defaults_ignore_mean_changes_and_reject_missing_maxima() { + let target = target(); + for algorithm in [ + LoadBalancerAlgorithm::Pulsar, + LoadBalancerAlgorithm::PulsarWaitAndWiden, + ] { + let config = LoadBalancerAlgorithmConfig::from(algorithm); + let balancer = create_load_balancer_with_config(&config).unwrap(); let initial = vec![ - work_candidate("inst-a", 5, 10_000.0, 0), - work_candidate("inst-b", 5, 1.0, 0), - ]; - let changed = vec![ - work_candidate("inst-a", 5, 1.0, 0), - work_candidate("inst-b", 5, 10_000.0, 0), + work_candidate("a", 5, 100.0, 0), + work_candidate("b", 5, 200.0, 0), ]; - - let first = choose(&pulsar, &request, &initial).candidate.cluster_id; - let second = choose(&pulsar, &request, &changed).candidate.cluster_id; - if first != second { - assert_eq!(first, "inst-a"); - assert_eq!(second, "inst-b"); - return; + let mut changed = initial.clone(); + changed[0].stats.last_mean_input_tps = 100_000.0; + changed[1].stats.last_mean_input_tps = 1.0; + let mut missing = changed.clone(); + for candidate in &mut missing { + candidate.stats.max_input_tps = None; + } + for index in 0..64 { + let key = format!("mean-change-{index}"); + let request = request(&target, Some(&key), Some(128)); + let before = balancer.choose_candidate(&request, &initial).unwrap(); + let after = balancer.choose_candidate(&request, &changed).unwrap(); + assert_eq!(before.candidate_index, after.candidate_index); + assert!(balancer.choose_candidate(&request, &missing).is_none()); } } - - panic!("expected to find an affinity key whose ranking changes after capacity changes"); } #[test] @@ -3071,7 +3265,9 @@ fn pulsar_hash_is_pinned_to_a_fixed_algorithm_and_version() { #[test] fn pulsar_uses_last_mean_input_tps_as_weight() { - let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); + let pulsar = PulsarLoadBalancer::new(parse_json( + r#"{"algorithm":"pulsar","rendezvous_weight":"last-mean-input-tps"}"#, + )); let candidate = work_candidate("inst-a", 5, 123.0, 0); @@ -3079,11 +3275,28 @@ fn pulsar_uses_last_mean_input_tps_as_weight() { } #[test] -fn pulsar_excludes_candidate_with_invalid_last_mean_input_tps() { +fn pulsar_uses_generation_max_input_tps_as_weight_by_default() { + let pulsar = PulsarLoadBalancer::new(seeded_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)); + + candidate.stats.last_mean_input_tps = 123.0; + candidate.stats.max_input_tps = None; + assert_eq!(pulsar.weight(&candidate), None, "no fallback to the mean"); +} + +#[test] +fn pulsar_excludes_candidate_with_invalid_max_input_tps() { let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); let target = target(); - let invalid = work_candidate("inst-a", 5, 0.0, 0); + let mut invalid = work_candidate("inst-a", 5, 100.0, 0); + invalid.stats.max_input_tps = Some(0.0); let valid = candidate("inst-b", 1024); let choice = choose( @@ -3095,12 +3308,15 @@ fn pulsar_excludes_candidate_with_invalid_last_mean_input_tps() { } #[test] -fn pulsar_returns_none_when_all_candidates_lack_valid_last_mean_input_tps() { +fn pulsar_returns_none_when_all_candidates_lack_valid_max_input_tps() { let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); let target = target(); - let invalid_a = work_candidate("inst-a", 5, 0.0, 0); - let invalid_b = work_candidate("inst-b", 5, f64::NAN, 0); + // Valid means do not make up for invalid maxima. + let mut invalid_a = work_candidate("inst-a", 5, 100.0, 0); + invalid_a.stats.max_input_tps = Some(0.0); + let mut invalid_b = work_candidate("inst-b", 5, 100.0, 0); + invalid_b.stats.max_input_tps = Some(f64::NAN); let choice = pulsar.choose_for_test( &request(&target, Some("prefix-1"), Some(128)), 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 0089f65286..aab315af7f 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 @@ -175,6 +175,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_f64); let mut rtt_mean_nanos = 0_u128; let mut rtt_remainder_nanos = 0_u128; @@ -182,6 +183,12 @@ 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(maximum), Some(value)) if value > 0.0 && value.is_finite() => { + Some(maximum.max(value)) + } + _ => None, + }; if valid_output_tps(backend.stats.output_tps) { backend_stats.output_tps += backend.stats.output_tps; } @@ -244,6 +251,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 c5fe310644..7f8f191e14 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, @@ -1718,6 +1720,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, @@ -1737,6 +1740,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, @@ -1758,6 +1762,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(35.0), queue_size: 6, queued_input_size: 60, input_processing_queries: 6, @@ -1775,6 +1780,66 @@ 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())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("complete backend max should aggregate") + .0 + .max_input_tps, + Some(25.0) + ); + + 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 + ); + } + cluster_state.remove_backend(&backend_b.registration); + assert_eq!( + cluster_state.backend_aggregate().unwrap().0.max_input_tps, + Some(15.0) + ); +} + #[tokio::test] async fn list_active_models_filters_by_routing_key() { let scenario = RegistrationScenario::new(None); @@ -2151,6 +2216,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(130.0), output_tps: 7.0, queue_size: 3, queued_input_size: 300, diff --git a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs index b836207cd5..751aa25f2d 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs @@ -402,13 +402,7 @@ impl StargateRuntime { let warmup_config = self.config.warmup.clone(); let shutdown = tasks.shutdown_signal(); tasks.task_tracker().spawn(async move { - run_warmup_stabilization( - warmup_state, - warmup_config, - ready_token, - shutdown, - ) - .await; + run_warmup_stabilization(warmup_state, warmup_config, ready_token, shutdown).await; }); } diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs index ff61991635..b58ae9229c 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs @@ -16,7 +16,9 @@ use std::net::SocketAddr; use std::time::Duration; -use stargate::runtime::{BoundStargateListeners, StargateRuntime, StargateRuntimeConfig, WarmupConfig}; +use stargate::runtime::{ + BoundStargateListeners, StargateRuntime, StargateRuntimeConfig, WarmupConfig, +}; use crate::common::{SelfDiscovery, base_config, init_crypto, make_stargate_runtime}; @@ -27,8 +29,7 @@ fn make_stargate_runtime_with_warmup( let ephemeral: SocketAddr = "127.0.0.1:0".parse().unwrap(); let mut config: StargateRuntimeConfig = base_config(id, ephemeral, ephemeral); config.warmup = warmup; - let listeners = - BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); + let listeners = BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); let grpc_addr = config.grpc_listen_addr; let http_addr = config.http_listen_addr; let discovery = Box::new(SelfDiscovery::new(id, grpc_addr, http_addr)); diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs index b710baf5d3..fa9c1fa83d 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs @@ -1445,6 +1445,7 @@ async fn pulsar_routes_same_affinity_key_consistently() { backend.set_stats(CurrentModelStats { output_tps: 0.0, last_mean_input_tps: *last_mean_input_tps, + max_input_tps: Some(*last_mean_input_tps), max_output_tps: 1000.0, queue_size: 0, queued_input_size: 0, diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs index b485bac172..1603f11e21 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/proxy_contract.rs @@ -216,6 +216,7 @@ fn active_runtime(model: &str) -> PylonRuntimeState { model, CurrentModelStats { last_mean_input_tps: 1000.0, + max_input_tps: Some(1000.0), ..CurrentModelStats::default() }, ); @@ -227,6 +228,7 @@ fn set_model_queue(runtime: &PylonRuntimeState, model: &str, queued_input_size: model, CurrentModelStats { last_mean_input_tps: 1000.0, + max_input_tps: Some(1000.0), queued_input_size, ..CurrentModelStats::default() }, diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/routing_expressions.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/routing_expressions.rs index 913df2ec73..8db8e963d2 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/routing_expressions.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/routing_expressions.rs @@ -61,6 +61,7 @@ async fn routing_expressions_proxy_valid_values_and_reject_before_selection() { model, CurrentModelStats { last_mean_input_tps: 1000.0, + max_input_tps: Some(1000.0), ..CurrentModelStats::default() }, ); diff --git a/src/libraries/rust/stargate/docs/load-balancer-configuration.md b/src/libraries/rust/stargate/docs/load-balancer-configuration.md index 05390a3374..f8853e41aa 100644 --- a/src/libraries/rust/stargate/docs/load-balancer-configuration.md +++ b/src/libraries/rust/stargate/docs/load-balancer-configuration.md @@ -110,7 +110,7 @@ Choose based on the routing goal and available backend statistics: | --- | --- | --- | | Compare a small random sample using TTFT or another load signal. | `power-of-n` | Signals required by the configured comparator. | | Minimize estimated time to first token across heterogeneous or remote clusters while controlling which TTFT bands are eligible. | `wait-and-widen` | Forwarded health RTT and model statistics. Valid `last_mean_input_tps` is needed when queued or request input work is nonzero. | -| Keep the same prefix on a stable, capacity-weighted cluster. | `pulsar` | Positive finite `last_mean_input_tps` for every participating cluster. | +| Keep the same prefix on a stable, capacity-weighted cluster. | `pulsar` | Positive finite `max_input_tps` for every participating cluster by default. | | Keep Pulsar affinity when possible, but escape to lower-latency capacity when the primary cannot meet queue policy. | `pulsar-wait-and-widen` | Pulsar capacity plus the RTT and queue statistics used by `wait-and-widen`. | Use `round-robin` for deterministic cycling and `random` for uniform random @@ -248,8 +248,36 @@ Example with a 200 ms affinity wait and a 10 percent prefill estimate: `pulsar` creates a stable weighted rendezvous ranking from the seed, routing key, model ID, cache affinity key, and cluster ID. The weight is the cluster's -positive finite `last_mean_input_tps`. Transient queue load does not change the -base ranking. +positive finite `max_input_tps` by default. Pylon retains this maximum for the +model generation. Lower live throughput does not reduce the weight or change +the base ranking. A new maximum, membership change, or generation replacement +can change the ranking. + +Both Pulsar variants accept `rendezvous_weight`: + +| Value | Behavior | +|---|---| +| `max-input-tps` | Default. Use the generation maximum. Missing or invalid maximum data makes the cluster ineligible; there is no fallback to the mean. | +| `last-mean-input-tps` | Explicit legacy mode. Use the last valid mean, which can change during traffic. | + +Queue and prefill estimates continue to use `last_mean_input_tps` in either mode. +See [multi-backend clusters](multi-backend-clusters.md) for how a shared-engine +cluster aggregates the maximum. + +Upgrade all Pylons and any registration relays, including +`stargate-k8s-router`, before enabling maximum-based Stargates. Older relays +decode and re-encode registrations without preserving the new field. Older +Pylons omit it and make that cluster ineligible in the new +default mode. Use explicit `last-mean-input-tps` during a mixed-version rollout +if Pylons cannot be upgraded first. + +`rendezvous_weight` is a field of a detailed algorithm configuration object. +The top-level `default`, an algorithm name in `models` or `request_algorithms`, +and the built-in defaults always use `max-input-tps`. For a mixed-version +rollout, give each Pulsar entry in `models` and `request_algorithms` a detailed +object that sets `last-mean-input-tps`. If `default` is `pulsar` or +`pulsar-wait-and-widen`, also list every Pulsar-routed model under `models`, or +move `default` to another algorithm. Retries walk the same ranking after excluding failed clusters. When `consider_kv_free_tokens` is enabled, a ranked candidate is skipped when it diff --git a/src/libraries/rust/stargate/docs/multi-backend-clusters.md b/src/libraries/rust/stargate/docs/multi-backend-clusters.md index c75df11db0..9d1ca3f315 100644 --- a/src/libraries/rust/stargate/docs/multi-backend-clusters.md +++ b/src/libraries/rust/stargate/docs/multi-backend-clusters.md @@ -83,6 +83,15 @@ Cluster latest-wins fields: - `queue_time_estimate_ms_by_priority` - `last_mean_input_tps` is the sum of active backend reports. +- `max_input_tps` is the greatest active backend generation maximum, available + only when every active backend reports a positive finite value. Historical + peaks from Pylons observing the same engine are not summed. Missing maximum + data makes the cluster ineligible for default Pulsar ranking. +- Because the mean is summed and the maximum is not, a shared cluster weighs + less against single-Pylon clusters under maximum weighting than under mean + weighting. This applies to per-Pylon `--initial-input-tps` contributions and + to observed rates when each Pylon measures only its share of the engine's + traffic. `RoutedClusterSnapshot.rtt` is the unweighted arithmetic mean of the latest forwarded `/health` RTT publication from every current active backend in the diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index fed5e21f76..88bb6a8fb7 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -140,6 +140,7 @@ Use it only when the runtime has reliable KV state. Pylon publishes: - sticky completed-request input throughput: `last_mean_input_tps` +- generation maximum input throughput: optional `max_input_tps` - volatile generation throughput: `output_tps` and `max_output_tps` - request phase counts and queue sizes - effective maximum engine concurrency @@ -154,14 +155,29 @@ 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. + +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. +The additive protobuf field preserves older messages as an absent maximum. Shared clusters sum backend-local live load and union labels. Effective input capacity is: ```text -sum(active_runtime_reports) +mean input capacity = sum(active runtime mean reports) +maximum input capacity = max(active generation peak reports) ``` +Maximum capacity is available only when every active backend in the shared +engine cluster reports a valid maximum. Historical per-observer peaks are not +additive. Pulsar uses maximum capacity by default; queue estimates still use +mean capacity. + 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,