From f00fe0179ddc1162c92ed26a58a3c1091f7241a9 Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Tue, 18 Aug 2026 21:48:34 +0000 Subject: [PATCH 1/3] feat(stargate): add configurable Pulsar input TPS weights Signed-off-by: Barry Greengus --- .../crates/proto/proto/stargate.proto | 5 + .../pylon-lib/src/registration/tests.rs | 2 + .../crates/pylon-lib/src/runtime_state.rs | 3 + .../crates/pylon-lib/src/stats/aggregator.rs | 52 ++++-- .../crates/pylon-lib/src/stats/collector.rs | 82 ++++++++-- .../crates/pylon-lib/src/stats/metrics.rs | 15 ++ .../crates/pylon-lib/src/stats/projection.rs | 16 +- .../pylon-lib/src/stats/token_metrics.rs | 153 ++++++++++++++++++ .../rust/stargate/crates/pylon/src/main.rs | 31 ++++ .../rust/stargate/crates/pylon/src/startup.rs | 5 + .../stargate/src/load_balancer/config.rs | 44 ++++- .../crates/stargate/src/load_balancer/mod.rs | 3 +- .../stargate/src/load_balancer/pulsar.rs | 46 ++++-- .../src/load_balancer/pulsar/ranking.rs | 70 +++++--- .../stargate/src/load_balancer/tests.rs | 112 +++++++++++++ .../src/routing_state/cluster_snapshots.rs | 9 ++ .../stargate/src/routing_state/tests.rs | 73 +++++++++ 17 files changed, 651 insertions(+), 70 deletions(-) diff --git a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto index ef4fe5d1f..1ee17256e 100644 --- a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto +++ b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto @@ -126,6 +126,11 @@ message ModelStats { // Sticky source labels for stat observations seen by this backend since the // model metrics state was initialized. repeated string stats_sources = 17; + // Rolling-window maximum input tokens per second observed by this backend. + // Absent after startup or restart until observed, and after the window + // expires. Stargate sums this field only when every active backend in a + // shared cluster reports it. + optional double max_input_tps = 18; } enum InferenceServerStatus { diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs index cba367b0f..3800ba81f 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/registration/tests.rs @@ -450,6 +450,7 @@ fn runtime_snapshot_forwards_bootstrap_and_collected_stats_exactly() { "model-a", CurrentModelStats { last_mean_input_tps: 3.5, + max_input_tps: Some(4.5), output_tps: 2.5, queue_size: 4, queued_input_size: 5, @@ -475,6 +476,7 @@ fn runtime_snapshot_forwards_bootstrap_and_collected_stats_exactly() { assert_eq!(model.status, InferenceServerStatus::Active as i32); let stats = model.stats.as_ref().expect("stats should be present"); assert_eq!(stats.last_mean_input_tps, 3.5); + assert_eq!(stats.max_input_tps, Some(4.5)); assert_eq!(stats.output_tps, 2.5); assert_eq!( stats.queue_time_estimate_ms_by_priority, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index 0cc849340..fe87d09c3 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -31,6 +31,8 @@ use reqwest::header::HeaderMap; pub struct CurrentModelStats { // Sticky runtime-observed mean input TPS for this backend. pub last_mean_input_tps: f64, + // Time-windowed max input TPS for this backend. None until observed. + pub max_input_tps: Option, // Token/sec output rate for streaming generation endpoints. Embeddings item // cardinality is observed separately and is not exported through this field. pub output_tps: f64, @@ -366,6 +368,7 @@ impl PylonRuntimeState { let registration = InferenceServerModelRegistration { stats: Some(ModelStats { last_mean_input_tps: stats.last_mean_input_tps, + max_input_tps: stats.max_input_tps, output_tps: stats.output_tps, max_output_tps: stats.max_output_tps, queue_size: stats.queue_size, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs index c6c680883..d7cabd9b0 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 @@ -28,7 +28,7 @@ use super::collector::{ FinalizeRequestUpdate, RequestCounterUpdate, StatsAggregatorUpdate, StatsCollectorConfig, StatsUpdateSource, }; -use super::token_metrics::TpsDistribution; +use super::token_metrics::{TpsDistribution, WindowedMax}; pub(super) const ENGINE_STATS_SOURCE: &str = "engine_stats_stream"; #[derive(Debug, Default)] @@ -42,6 +42,7 @@ pub(super) struct ModelMetricsState { pub(super) max_embedding_item_tps: f64, pub(super) kv_cache: KvCacheStatsSnapshot, pub(super) input_tps_distribution: TpsDistribution, + pub(super) input_tps_window_max: WindowedMax, aggregate_state_counted: bool, pub(super) counter_output_tps_authoritative: bool, pub(super) chunk_usage_stats_observed: bool, @@ -237,22 +238,20 @@ impl StatsAggregator { if self.per_model.contains_key(generation.model_id()) { return None; } - let (metrics, pinned_input_tps) = match initialization { - super::collector::ModelStatsInitialization::Empty => { - (ModelMetricsState::default(), None) - } + let mut metrics = + ModelMetricsState::with_input_tps_window(self.config.input_tps_capacity_window); + let pinned_input_tps = match initialization { + super::collector::ModelStatsInitialization::Empty => None, super::collector::ModelStatsInitialization::ConfiguredInputTps { input_tps, pin } => { let input_tps_distribution = TpsDistribution::bootstrap(input_tps) .expect("configured input TPS must be positive and finite"); - ( - ModelMetricsState { - last_mean_input_tps: input_tps, - input_tps_distribution, - aggregate_state_counted: true, - ..ModelMetricsState::default() - }, - pin.then_some(input_tps), - ) + metrics.last_mean_input_tps = input_tps; + metrics.input_tps_distribution = input_tps_distribution; + metrics.aggregate_state_counted = true; + metrics + .input_tps_window_max + .observe(TokioInstant::now(), input_tps); + pin.then_some(input_tps) } }; self.aggregate_model_state_count += usize::from(metrics.aggregate_state_counted); @@ -397,6 +396,13 @@ impl StatsAggregator { pub(super) fn sweep_stale(&mut self, now: TokioInstant) -> Vec { let mut dirty_models = Vec::new(); + for (model_id, generation_state) in &mut self.per_model { + if generation_state.pinned_input_tps.is_none() + && generation_state.metrics.input_tps_window_max.expire(now) + { + push_dirty_model(&mut dirty_models, model_id.clone()); + } + } let request_ttl = self.config.engine_stats_request_ttl; if !request_ttl.is_zero() { let metrics = self.runtime_state.metrics(); @@ -650,6 +656,7 @@ impl StatsAggregator { config, model_state, generation_state.pinned_input_tps, + update.observed_at, InputThroughputSample { units, duration, @@ -806,6 +813,7 @@ pub(super) fn apply_input_throughput_sample( config: &StatsCollectorConfig, model_state: &mut ModelMetricsState, pinned_input_tps: Option, + observed_at: TokioInstant, sample: InputThroughputSample, ) -> bool { if sample.units < config.min_input_tokens { @@ -819,16 +827,20 @@ pub(super) fn apply_input_throughput_sample( let Some(input_tps) = tps_for_units(sample.units, duration, config.duration_floor) else { return false; }; + let max_changed = pinned_input_tps.is_none() + && model_state + .input_tps_window_max + .observe(observed_at, input_tps); model_state.input_tps_distribution.update(input_tps); let mean_input_tps = model_state.input_tps_distribution.mean; if !model_state.input_tps_distribution.has_sufficient_data() || !valid_last_mean_input_tps(mean_input_tps) { - return false; + return max_changed; } let last_mean_input_tps = pinned_input_tps.unwrap_or(mean_input_tps); if model_state.last_mean_input_tps == last_mean_input_tps { - return false; + return max_changed; } model_state.last_mean_input_tps = last_mean_input_tps; true @@ -857,6 +869,13 @@ pub(super) struct ModelStatsSnapshotInputs { } impl ModelMetricsState { + fn with_input_tps_window(window: Duration) -> Self { + Self { + input_tps_window_max: WindowedMax::new(window), + ..Self::default() + } + } + pub(super) fn clear_live_output_tps(&mut self) -> bool { self.last_stats_event_at = None; if self.chat_output_tps_samples.is_empty() { @@ -877,6 +896,7 @@ impl ModelMetricsState { }; CurrentModelStats { last_mean_input_tps: self.last_mean_input_tps, + max_input_tps: self.input_tps_window_max.current(), output_tps: active_chat_output_tps.max(average_with_sum( &self.chat_output_tps_samples, self.chat_output_tps_sum, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 1c77c6958..8e787c25d 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -25,6 +25,7 @@ use crate::{CurrentModelStats, PylonRuntimeState, RequestObservationEvent}; use stargate_runtime::OwnedTask; use super::aggregator::{ENGINE_STATS_SOURCE, KvCacheStatsSnapshot, StatsAggregator}; +use super::token_metrics::DEFAULT_INPUT_TPS_CAPACITY_WINDOW; const DEFAULT_OBSERVATION_CHANNEL_CAPACITY: usize = 1024; const DEFAULT_SMOOTHING_WINDOW_SIZE: usize = 8; @@ -41,6 +42,7 @@ const DEFAULT_ENGINE_STATS_SWEEP_INTERVAL: Duration = Duration::from_secs(1); pub struct StatsCollectorConfig { pub observation_channel_capacity: usize, pub smoothing_window_size: usize, + pub input_tps_capacity_window: Duration, pub min_input_tokens: u64, pub min_output_tokens: u64, pub duration_floor: Duration, @@ -58,6 +60,7 @@ impl Default for StatsCollectorConfig { Self { observation_channel_capacity: DEFAULT_OBSERVATION_CHANNEL_CAPACITY, smoothing_window_size: DEFAULT_SMOOTHING_WINDOW_SIZE, + input_tps_capacity_window: DEFAULT_INPUT_TPS_CAPACITY_WINDOW, min_input_tokens: DEFAULT_MIN_INPUT_TOKENS, min_output_tokens: DEFAULT_MIN_OUTPUT_TOKENS, duration_floor: DEFAULT_DURATION_FLOOR, @@ -1239,7 +1242,7 @@ mod tests { aggregator.stream("req-a", (0, 0), false, Duration::ZERO); let updates = aggregator.stream("req-a", (10, 4), false, milliseconds(100)); let stats = published_stats(updates); - assert_stats!(stats; output_tps: 40.0, max_output_tps: 40.0, stats_sources: ["engine_stats_stream"]); + assert_stats!(stats; max_input_tps: Some(100.0), output_tps: 40.0, max_output_tps: 40.0, stats_sources: ["engine_stats_stream"]); for tick in 2..=5 { let updates = aggregator.stream("req-a", (tick * 10, 4), false, milliseconds(tick * 100)); @@ -1254,7 +1257,7 @@ mod tests { #[test] fn pinned_configured_input_tps_is_preserved_across_engine_stats_updates() { let mut aggregator = test_aggregator_with_initialization( - StatsCollectorConfig::default(), + config!(input_tps_capacity_window: seconds(1)), ModelStatsInitialization::ConfiguredInputTps { input_tps: 2_200.0, pin: true, @@ -1262,10 +1265,15 @@ mod tests { ); let stats = aggregator.stream_stats("req-a", (0, 0), false, Duration::ZERO); assert_eq!(stats.last_mean_input_tps, 2_200.0); + assert_eq!(stats.max_input_tps, Some(2_200.0)); for tick in 1..=5 { aggregator.stream("req-a", (tick * 10, 0), false, milliseconds(tick * 100)); } - assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 2_200.0); + let stats = aggregator.snapshot("model-a"); + assert_eq!(stats.last_mean_input_tps, 2_200.0); + assert_eq!(stats.max_input_tps, Some(2_200.0)); + aggregator.sweep(seconds(2)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(2_200.0)); } #[test] @@ -1277,12 +1285,34 @@ mod tests { pin: false, }, ); - assert_eq!(aggregator.snapshot("model-a").last_mean_input_tps, 100.0); + let bootstrapped = aggregator.snapshot("model-a"); + assert_eq!(bootstrapped.last_mean_input_tps, 100.0); + assert_eq!(bootstrapped.max_input_tps, Some(100.0)); aggregator.stream("req-a", (0, 0), false, Duration::ZERO); let stats = published_stats(aggregator.stream("req-a", (20, 0), false, milliseconds(100))); assert!((stats.last_mean_input_tps - (700.0 / 6.0)).abs() < f64::EPSILON); + assert_eq!(stats.max_input_tps, Some(200.0)); + } + + #[test] + fn unpinned_seed_expires_and_empty_restart_has_no_max() { + let config = config!(input_tps_capacity_window: seconds(1)); + let mut initialized = test_aggregator_with_initialization( + config.clone(), + ModelStatsInitialization::ConfiguredInputTps { + input_tps: 100.0, + pin: false, + }, + ); + + let expired = published_stats(initialized.sweep(seconds(2))); + assert_eq!(expired.last_mean_input_tps, 100.0); + assert_eq!(expired.max_input_tps, None); + + let restarted = test_aggregator(config); + assert_eq!(restarted.snapshot("model-a").max_input_tps, None); } #[test] @@ -1336,12 +1366,13 @@ mod tests { .is_empty(), "second output counter is still below the duration floor" ); - let input_only_updates = - aggregator.partial_stream("req-partial", (Some(1), None), false, milliseconds(11)); - assert!( - input_only_updates.is_empty(), - "input-only updates must not publish a stale output TPS sample" - ); + let input_only_stats = aggregator + .partial_stream("req-partial", (Some(1), None), false, milliseconds(11)) + .pop() + .expect("first input counter should publish max input TPS") + .1; + assert_eq!(input_only_stats.output_tps, 100.0); + assert!(input_only_stats.max_input_tps.is_some()); } #[test] @@ -1898,6 +1929,24 @@ mod tests { assert_stats!(stats; last_mean_input_tps: 100.0, output_tps: 0.0, queue_size: 0, queued_input_size: 0, num_running_queries: 0, input_processing_queries: 0, output_generation_queries: 0, stats_sources: ["engine_stats_stream"]); } + #[test] + fn stats_aggregator_publishes_and_expires_windowed_max_input_tps() { + let mut aggregator = test_aggregator(config!( + input_tps_capacity_window: seconds(1), + engine_stats_request_ttl: Duration::ZERO, + engine_stats_model_ttl: Duration::ZERO, + )); + aggregator.stream("req-window-max", (0, 0), false, Duration::ZERO); + + let stats = + published_stats(aggregator.stream("req-window-max", (10, 0), false, milliseconds(100))); + assert_eq!(stats.last_mean_input_tps, 0.0); + assert_eq!(stats.max_input_tps, Some(100.0)); + + let expired = published_stats(aggregator.sweep(seconds(2))); + assert_eq!(expired.max_input_tps, None); + } + #[test] fn stats_aggregator_tombstones_stale_request_before_late_finish() { let config = config!( @@ -2503,7 +2552,11 @@ mod tests { #[test] fn records_metrics_when_configured() { let metrics = PylonMetrics::new().expect("metrics should initialize"); - let config = StatsCollectorConfig::default(); + let config = config!( + input_tps_capacity_window: seconds(1), + engine_stats_request_ttl: Duration::ZERO, + engine_stats_model_ttl: Duration::ZERO, + ); let (runtime_state, _observation_rx) = PylonRuntimeState::observed( stargate_proto::pb::InferenceServerStatus::Unknown, &["model-a".to_string()], @@ -2531,7 +2584,14 @@ mod tests { r#"pylon_requests_total{model="model-a",routing_key="rk-1",status="complete"} 5"# )); assert!(body.contains(r#"pylon_model_last_mean_input_tps{model="model-a"} 10"#)); + assert!(body.contains(r#"pylon_model_max_input_tps{model="model-a"} 10"#)); assert!(body.contains(r#"pylon_model_output_tps{model="model-a"} 5"#)); + + for (model_id, stats) in aggregator.sweep_stale(TokioInstant::now() + seconds(2)) { + publish_model_stats_update(&runtime_state, model_id, stats); + } + 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/metrics.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs index 3ac3727b1..b9c160eeb 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs @@ -151,6 +151,7 @@ metrics! { float_gauge output_tps("model_output_tps", "Current output TPS by model", ["model"]); float_gauge embedding_item_tps("model_embedding_item_tps", "Current embeddings item throughput by model", ["model"]); float_gauge last_mean_input_tps("model_last_mean_input_tps", "Last valid mean input TPS by model", ["model"]); + float_gauge max_input_tps("model_max_input_tps", "Time-windowed max input TPS by model", ["model"]); float_gauge max_output_tps("model_max_output_tps", "Observed max output TPS by model", ["model"]); float_gauge max_embedding_item_tps("model_max_embedding_item_tps", "Observed max embeddings item throughput by model", ["model"]); float_gauge queue_size("model_queue_size", "Current queued request count by model", ["model"]); @@ -311,6 +312,12 @@ impl PylonMetrics { ] { gauge.with_label_values(&[model_id]).set(value); } + match stats.max_input_tps { + Some(value) => self.max_input_tps.with_label_values(&[model_id]).set(value), + None => { + let _ = self.max_input_tps.remove_label_values(&[model_id]); + } + } for (gauge, values) in [ (&self.stats_capability, &stats.stats_capabilities), (&self.stats_source, &stats.stats_sources), @@ -352,6 +359,7 @@ impl PylonMetrics { &self.output_tps, &self.embedding_item_tps, &self.last_mean_input_tps, + &self.max_input_tps, &self.max_output_tps, &self.max_embedding_item_tps, &self.queue_size, @@ -748,6 +756,7 @@ mod tests { let metrics = PylonMetrics::new().expect("metrics should initialize"); let stats = CurrentModelStats { last_mean_input_tps: 12.0, + max_input_tps: Some(13.0), stats_capabilities: vec!["input_tps".to_string()], stats_sources: vec!["engine_stats_stream".to_string()], ..CurrentModelStats::default() @@ -915,6 +924,7 @@ mod tests { output_tps: 20.0, embedding_item_tps: 25.0, last_mean_input_tps: 30.0, + max_input_tps: Some(35.0), max_output_tps: 40.0, max_embedding_item_tps: 45.0, queue_size: 2, @@ -934,6 +944,7 @@ mod tests { r#"pylon_model_output_tps{model="model-a"} 20"#, r#"pylon_model_embedding_item_tps{model="model-a"} 25"#, r#"pylon_model_last_mean_input_tps{model="model-a"} 30"#, + r#"pylon_model_max_input_tps{model="model-a"} 35"#, r#"pylon_model_max_embedding_item_tps{model="model-a"} 45"#, r#"pylon_model_queue_size{model="model-a"} 2"#, r#"pylon_model_queued_input_tokens{model="model-a"} 17"#, @@ -942,6 +953,10 @@ mod tests { r#"pylon_model_stats_source{model="model-a",source="engine_stats_stream"} 1"#, ], ); + + metrics.observe_model_stats("model-a", &CurrentModelStats::default()); + let body = metrics.gather_text().expect("metrics should encode"); + assert!(!body.contains(r#"pylon_model_max_input_tps{model="model-a"}"#)); } #[test] diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs index d416c86e5..4ccd622be 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 @@ -119,7 +119,13 @@ impl StatsAggregator { }; let mut stats = self.per_model.get(model_id).map_or_else( || ModelMetricsState::default().current_stats(inputs), - |state| state.metrics.current_stats(inputs), + |state| { + let mut stats = state.metrics.current_stats(inputs); + if let Some(pinned_input_tps) = state.pinned_input_tps { + stats.max_input_tps = Some(pinned_input_tps); + } + stats + }, ); stats.queue_time_estimate_ms_by_priority = queue.queue_time_estimate_ms_by_priority; stats @@ -201,7 +207,13 @@ impl StatsAggregator { } } if input_sample.is_some_and(|sample| { - apply_input_throughput_sample(&self.config, model_state, pinned_input_tps, sample) + apply_input_throughput_sample( + &self.config, + model_state, + pinned_input_tps, + TokioInstant::now(), + sample, + ) }) { push_changed_model(&mut changed_models, observation.model_id.clone()); } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs index a024bab67..02cd0c0e5 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs @@ -13,7 +13,98 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::VecDeque; +use std::time::Duration; + +use tokio::time::Instant; + pub(crate) const SNAPSHOT_THRESHOLD: usize = 5; +pub(crate) const DEFAULT_INPUT_TPS_CAPACITY_WINDOW: Duration = Duration::from_secs(24 * 60 * 60); +const WINDOW_SEGMENTS: u32 = 1_440; +const MAX_WINDOW_BUCKETS: usize = WINDOW_SEGMENTS as usize + 1; + +#[derive(Debug)] +struct WindowMaxBucket { + end: Instant, + max: f64, +} + +#[derive(Debug)] +pub(crate) struct WindowedMax { + window: Duration, + bucket_width: Duration, + buckets: VecDeque, +} + +impl Default for WindowedMax { + fn default() -> Self { + Self::new(DEFAULT_INPUT_TPS_CAPACITY_WINDOW) + } +} + +impl WindowedMax { + pub(crate) fn new(window: Duration) -> Self { + assert!(!window.is_zero(), "windowed max requires a nonzero window"); + let floor = window / WINDOW_SEGMENTS; + let bucket_width = if floor.saturating_mul(WINDOW_SEGMENTS) < window { + floor + .checked_add(Duration::from_nanos(1)) + .expect("window bucket width should fit Duration") + } else { + floor + }; + Self { + window, + bucket_width, + buckets: VecDeque::new(), + } + } + + pub(crate) fn observe(&mut self, now: Instant, value: f64) -> bool { + if value <= 0.0 || !value.is_finite() { + return false; + } + let previous = self.current(); + self.expire_buckets(now); + match self.buckets.back_mut() { + Some(bucket) if now < bucket.end => bucket.max = bucket.max.max(value), + _ => self.buckets.push_back(WindowMaxBucket { + end: now + .checked_add(self.bucket_width) + .expect("window bucket end should fit Instant"), + max: value, + }), + } + debug_assert!(self.buckets.len() <= MAX_WINDOW_BUCKETS); + self.current() != previous + } + + pub(crate) fn expire(&mut self, now: Instant) -> bool { + let previous = self.current(); + self.expire_buckets(now); + self.current() != previous + } + + pub(crate) fn current(&self) -> Option { + self.buckets + .iter() + .map(|bucket| bucket.max) + .max_by(f64::total_cmp) + } + + fn expire_buckets(&mut self, now: Instant) { + let Some(cutoff) = now.checked_sub(self.window) else { + return; + }; + while self + .buckets + .front() + .is_some_and(|bucket| bucket.end <= cutoff) + { + self.buckets.pop_front(); + } + } +} #[derive(Debug, Clone, Default)] pub(crate) struct TpsDistribution { @@ -69,6 +160,68 @@ impl TpsDistribution { mod tests { use super::*; + fn seconds(value: u64) -> Duration { + Duration::from_secs(value) + } + + #[test] + fn windowed_max_tracks_and_expires_bucket_maxima() { + let start = Instant::now(); + let mut max = WindowedMax::new(seconds(60)); + + assert_eq!(max.current(), None); + assert!(max.observe(start, 10.0)); + assert!(max.observe(start + seconds(1), 20.0)); + assert!(!max.observe(start + seconds(2), 15.0)); + assert_eq!(max.current(), Some(20.0)); + + assert!(!max.observe(start + seconds(31), 12.0)); + assert_eq!(max.current(), Some(20.0)); + assert!(max.expire(start + Duration::from_millis(61_050))); + assert_eq!(max.current(), Some(15.0)); + assert!(max.expire(start + Duration::from_millis(62_050))); + assert_eq!(max.current(), Some(12.0)); + assert!(max.expire(start + seconds(92))); + assert_eq!(max.current(), None); + } + + #[test] + fn windowed_max_rejects_invalid_values() { + let start = Instant::now(); + let mut max = WindowedMax::new(seconds(60)); + + for value in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!(!max.observe(start, value)); + } + assert_eq!(max.current(), None); + } + + #[test] + fn windowed_max_never_expires_early_at_oldest_boundary() { + let start = Instant::now(); + let mut max = WindowedMax::new(Duration::from_millis(1_440)); + assert!(max.observe(start, 10.0)); + + assert!(!max.expire(start + Duration::from_millis(1_440))); + assert_eq!(max.current(), Some(10.0)); + assert!(max.expire(start + Duration::from_millis(1_441))); + assert_eq!(max.current(), None); + } + + #[test] + fn windowed_max_memory_is_bounded() { + let start = Instant::now(); + let mut max = WindowedMax::new(Duration::from_millis(1_440)); + + for offset_ms in 0..10_000 { + max.observe( + start + Duration::from_millis(offset_ms), + offset_ms as f64 + 1.0, + ); + assert!(max.buckets.len() <= MAX_WINDOW_BUCKETS); + } + } + #[test] fn tps_distribution_ignores_non_positive_and_non_finite_samples() { let mut distribution = TpsDistribution::default(); diff --git a/src/libraries/rust/stargate/crates/pylon/src/main.rs b/src/libraries/rust/stargate/crates/pylon/src/main.rs index fd08dc930..d6f08ad35 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/main.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/main.rs @@ -86,6 +86,9 @@ struct Args { /// Bootstrap input TPS for every configured model instead of running calibration #[arg(long, value_name = "TPS")] initial_input_tps: Option, + /// Window for max input TPS capacity in milliseconds + #[arg(long, default_value_t = 86_400_000, value_name = "MS")] + input_tps_capacity_window_ms: u64, /// Interval between active canary requests in milliseconds. `0` disables active canaries #[arg(long, default_value_t = 5000, value_name = "MS")] active_canary_interval_ms: u64, @@ -528,6 +531,34 @@ mod tests { } } + #[test] + fn input_tps_capacity_window_defaults_flow_and_zero_is_rejected() { + let default_args = parse_args(""); + let default_config = + stats_collector_config_from_args(&default_args, &default_args.upstream_http_base_url); + assert_eq!( + default_config.input_tps_capacity_window, + std::time::Duration::from_secs(24 * 60 * 60) + ); + assert_eq!( + default_config.input_tps_capacity_window, + pylon_lib::StatsCollectorConfig::default().input_tps_capacity_window + ); + + let configured_args = parse_args("--input-tps-capacity-window-ms 1234"); + let configured = stats_collector_config_from_args( + &configured_args, + &configured_args.upstream_http_base_url, + ); + assert_eq!( + configured.input_tps_capacity_window, + std::time::Duration::from_millis(1234) + ); + + let zero = parse_args("--input-tps-capacity-window-ms 0"); + assert!(startup::PylonStartupPlan::from_args(&zero).is_err()); + } + #[test] fn calibration_ramp_requires_a_positive_request_increment() { let args = parse_args("--do-calibration --calibration-requests 0"); diff --git a/src/libraries/rust/stargate/crates/pylon/src/startup.rs b/src/libraries/rust/stargate/crates/pylon/src/startup.rs index 9a5196501..dc2fba6ea 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/startup.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/startup.rs @@ -146,6 +146,10 @@ impl BackendTunnelStartup { impl PylonStartupPlan { pub(crate) fn from_args(args: &Args) -> Result { + ensure!( + args.input_tps_capacity_window_ms > 0, + "--input-tps-capacity-window-ms must be greater than zero" + ); let model_initialization = model_initialization_from_args(args)?; let model_source = model_source_from_args(args)?; Ok(Self { @@ -530,6 +534,7 @@ pub(crate) fn stats_collector_config_from_args( upstream: &str, ) -> StatsCollectorConfig { StatsCollectorConfig { + input_tps_capacity_window: Duration::from_millis(args.input_tps_capacity_window_ms), openai_fallback_stats_enabled: args.engine_stats_stream == EngineStatsStreamMode::Off, // Mock benchmark backends can expose live KV-cache occupancy over HTTP; // real upstreams usually do not, so polling is explicit. diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs index 512ef322d..fea4fe1c7 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 enum PulsarRendezvousWeight { + #[default] + LastMeanInputTps, + MaxInputTps, +} + impl LoadBalancerAlgorithm { pub const ALL: [Self; 6] = [ Self::PowerOfN, @@ -280,6 +288,7 @@ pub struct LoadBalancerAlgorithmConfig { pub max_input_work_seconds: Option, pub request_algorithms: HashMap, pub settings: LoadBalancerAlgorithmSettings, + pub(crate) rendezvous_weight: PulsarRendezvousWeight, } impl LoadBalancerAlgorithmConfig { @@ -287,6 +296,10 @@ impl LoadBalancerAlgorithmConfig { self.settings.algorithm() } + pub fn pulsar_rendezvous_weight(&self) -> PulsarRendezvousWeight { + self.rendezvous_weight + } + pub fn requires_cache_affinity_key(&self) -> bool { self.request_policy.require_cache_affinity_key } @@ -436,6 +449,7 @@ impl RawCommonAlgorithmConfig { self, settings: LoadBalancerAlgorithmSettings, consider_kv_free_tokens: Option, + rendezvous_weight: PulsarRendezvousWeight, ) -> Result { let algorithm = settings.algorithm(); if !self.unsupported_fields.is_empty() { @@ -458,6 +472,7 @@ impl RawCommonAlgorithmConfig { max_input_work_seconds: self.max_input_work_seconds, request_algorithms: self.request_algorithms, settings, + rendezvous_weight, }) } } @@ -493,6 +508,8 @@ enum RawLoadBalancerAlgorithmConfig { Pulsar { seed: Option, consider_kv_free_tokens: Option, + #[serde(default)] + rendezvous_weight: PulsarRendezvousWeight, #[serde(flatten)] common: RawCommonAlgorithmConfig, }, @@ -501,6 +518,8 @@ enum RawLoadBalancerAlgorithmConfig { #[serde(flatten)] settings: WaitAndWidenAlgorithmConfig, consider_kv_free_tokens: Option, + #[serde(default)] + rendezvous_weight: PulsarRendezvousWeight, #[serde(flatten)] common: RawCommonAlgorithmConfig, }, @@ -513,47 +532,64 @@ impl RawLoadBalancerAlgorithmConfig { RawCommonAlgorithmConfig, LoadBalancerAlgorithmSettings, Option, + PulsarRendezvousWeight, ) { match self { Self::PowerOfN { settings, common } => ( common, LoadBalancerAlgorithmSettings::PowerOfN(settings), None, + PulsarRendezvousWeight::default(), ), Self::WaitAndWiden { settings, common } => ( common, LoadBalancerAlgorithmSettings::WaitAndWiden(settings), None, + PulsarRendezvousWeight::default(), + ), + Self::RoundRobin(common) => ( + common, + LoadBalancerAlgorithmSettings::RoundRobin, + None, + PulsarRendezvousWeight::default(), + ), + Self::Random(common) => ( + common, + LoadBalancerAlgorithmSettings::Random, + None, + PulsarRendezvousWeight::default(), ), - Self::RoundRobin(common) => (common, LoadBalancerAlgorithmSettings::RoundRobin, None), - Self::Random(common) => (common, LoadBalancerAlgorithmSettings::Random, None), Self::Pulsar { seed, consider_kv_free_tokens, + rendezvous_weight, common, } => ( common, LoadBalancerAlgorithmSettings::Pulsar(seed), consider_kv_free_tokens, + rendezvous_weight, ), Self::PulsarWaitAndWiden { settings, consider_kv_free_tokens, + rendezvous_weight, common, } => ( common, LoadBalancerAlgorithmSettings::PulsarWaitAndWiden(settings), consider_kv_free_tokens, + rendezvous_weight, ), } } fn into_config(self) -> Result { - let (common, settings, consider_kv_free_tokens) = self.normalized(); + let (common, settings, consider_kv_free_tokens, rendezvous_weight) = self.normalized(); if let LoadBalancerAlgorithmSettings::PowerOfN(config) = &settings { config.validated_sample_count()?; } - common.into_config(settings, consider_kv_free_tokens) + common.into_config(settings, consider_kv_free_tokens, rendezvous_weight) } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs index 7e6110f94..02de52bae 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 @@ -48,7 +48,8 @@ pub use config::{ LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, LoadBalancerAlgorithmOverride, LoadBalancerAlgorithmSettings, LoadBalancerConfig, LoadBalancerModelConfig, LoadBalancerRequestPolicy, LoadBalancerRoutingAlgorithmError, LoadBalancerSeedError, - MAX_POWER_OF_N_SAMPLE_COUNT, PowerOfNAlgorithmConfig, WaitAndWidenAlgorithmConfig, + MAX_POWER_OF_N_SAMPLE_COUNT, PowerOfNAlgorithmConfig, PulsarRendezvousWeight, + WaitAndWidenAlgorithmConfig, }; pub use factory::create_load_balancer_with_config; pub use request::{LoadBalancerCandidateChoice, LoadBalancerRequest}; diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs index d0799661d..8bab727fc 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 @@ -23,7 +23,7 @@ use super::{ use crate::routing_state::RoutedClusterSnapshot; use ranking::{ PulsarRankingLookup, PulsarRankingStore, PulsarScorer, ScoredCandidate, - compare_ranked_candidate, has_valid_input_capacity, + compare_ranked_candidate, has_valid_mean_input_tps, }; #[cfg(test)] @@ -56,11 +56,12 @@ impl LoadBalancer for PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Option { - let (lookup, cached_choice) = - self.rankings - .lookup_choice(request, candidates, |ranking| { - self.choose_from_ranked_indices(request, candidates, ranking) - })?; + let (lookup, cached_choice) = self.rankings.lookup_choice( + self.config.pulsar_rendezvous_weight(), + request, + candidates, + |ranking| self.choose_from_ranked_indices(request, candidates, ranking), + )?; match lookup { PulsarRankingLookup::Hit => cached_choice, PulsarRankingLookup::MissCacheable => { @@ -70,9 +71,13 @@ impl LoadBalancer for PulsarLoadBalancer { } let choice = self.choose_from_ranked_indices(request, candidates, &ranking); self.rankings - .insert_or_choose_existing(request, candidates, ranking, |cached| { - self.choose_from_ranked_indices(request, candidates, cached) - }) + .insert_or_choose_existing( + self.config.pulsar_rendezvous_weight(), + request, + candidates, + ranking, + |cached| self.choose_from_ranked_indices(request, candidates, cached), + ) .or(choice) } PulsarRankingLookup::MissBypass => self.choose_by_score_scan(request, candidates), @@ -116,7 +121,12 @@ impl PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Vec { - ranking::pulsar_ranked_indices(self.config.seed(), request, candidates) + ranking::pulsar_ranked_indices_with_weight( + self.config.seed(), + self.config.pulsar_rendezvous_weight(), + request, + candidates, + ) } fn choose_by_score_scan( @@ -124,7 +134,11 @@ 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.pulsar_rendezvous_weight(), + ); let mut best_overall = None; let mut best_feasible = None; for (candidate_index, candidate) in candidates.iter().enumerate() { @@ -182,7 +196,11 @@ 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.pulsar_rendezvous_weight(), + ); let mut rank_depth = 1; let mut skipped_for_kv_free_tokens = false; for candidate in candidates { @@ -201,7 +219,7 @@ impl PulsarLoadBalancer { #[cfg(test)] pub(super) fn weight(&self, candidate: &RoutedClusterSnapshot) -> Option { - ranking::pulsar_weight(candidate) + ranking::pulsar_weight(self.config.pulsar_rendezvous_weight(), candidate) } pub(super) fn feasibility( @@ -218,7 +236,7 @@ pub(super) fn input_work_admission_candidate( request: &LoadBalancerRequest<'_>, candidate: &RoutedClusterSnapshot, ) -> bool { - has_valid_input_capacity(candidate) + has_valid_mean_input_tps(candidate) && candidate_feasibility(config, request, candidate).is_eligible() } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs index 65e11df47..be03082b5 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 @@ -22,7 +22,8 @@ use stargate_protocol::common::valid_last_mean_input_tps; use xxhash_rust::xxh3::xxh3_64; use crate::load_balancer::{ - LoadBalancerCandidateChoice, LoadBalancerRequest, cache_affinity_key_is_cacheable, + LoadBalancerCandidateChoice, LoadBalancerRequest, PulsarRendezvousWeight, + cache_affinity_key_is_cacheable, }; use crate::routing_state::{RoutedClusterSnapshot, RoutingTargetKey}; @@ -37,6 +38,7 @@ pub(super) struct PulsarRankingStore { impl PulsarRankingStore { pub(super) fn lookup_choice( &self, + selector: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], choose: impl Fn(&[usize]) -> Option, @@ -52,7 +54,7 @@ impl PulsarRankingStore { { let cache = self.cache.read(); - if cache.matches(request.routing_target, candidates) + if cache.matches(request.routing_target, candidates, selector) && let Some(ranking) = cache.get_ref(cache_affinity_key) { // Choose under the guard to avoid cloning the cached Arc. @@ -61,7 +63,7 @@ impl PulsarRankingStore { } let mut cache = self.cache.write(); - cache.refresh_if_needed(request.routing_target, candidates); + cache.refresh_if_needed(request.routing_target, candidates, selector); if let Some(ranking) = cache.get_ref(cache_affinity_key) { // Another thread populated the ranking after the read miss. return Some((PulsarRankingLookup::Hit, choose(ranking))); @@ -71,6 +73,7 @@ impl PulsarRankingStore { pub(super) fn insert_or_choose_existing( &self, + selector: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ranking: Arc>, @@ -78,7 +81,7 @@ impl PulsarRankingStore { ) -> Option { let cache_affinity_key = request.cache_affinity_key.unwrap_or(""); let mut cache = self.cache.write(); - cache.refresh_if_needed(request.routing_target, candidates); + cache.refresh_if_needed(request.routing_target, candidates, selector); if let Some(cached) = cache.get_ref(cache_affinity_key) { return choose(cached); } @@ -103,7 +106,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 +124,9 @@ impl PulsarRankingCache { &mut self, target: &RoutingTargetKey, candidates: &[RoutedClusterSnapshot], + selector: PulsarRendezvousWeight, ) { - if self.matches(target, candidates) { + if self.matches(target, candidates, selector) { return; } @@ -131,7 +135,7 @@ impl PulsarRankingCache { .iter() .map(|candidate| PulsarCandidateSignature { cluster_id: candidate.cluster_id.clone(), - last_mean_input_tps_bits: candidate.stats.last_mean_input_tps.to_bits(), + weight_bits: pulsar_weight(selector, candidate).map(f64::to_bits), }) .collect(); self.rankings.clear(); @@ -140,7 +144,12 @@ impl PulsarRankingCache { self.probation_order.clear(); } - fn matches(&self, target: &RoutingTargetKey, candidates: &[RoutedClusterSnapshot]) -> bool { + fn matches( + &self, + target: &RoutingTargetKey, + candidates: &[RoutedClusterSnapshot], + selector: PulsarRendezvousWeight, + ) -> bool { self.target.as_ref() == Some(target) && self.candidate_signature.len() == candidates.len() && self @@ -149,8 +158,8 @@ impl PulsarRankingCache { .zip(candidates) .all(|(cached, candidate)| { cached.cluster_id == candidate.cluster_id - && cached.last_mean_input_tps_bits - == candidate.stats.last_mean_input_tps.to_bits() + && cached.weight_bits + == pulsar_weight(selector, candidate).map(f64::to_bits) }) } @@ -232,10 +241,15 @@ pub(super) fn compare_ranked_candidate( pub(super) struct PulsarScorer { hash_bytes: Vec, prefix_len: usize, + selector: PulsarRendezvousWeight, } impl PulsarScorer { - pub(super) fn new(seed: Option<&str>, request: &LoadBalancerRequest<'_>) -> Self { + pub(super) fn new( + seed: Option<&str>, + request: &LoadBalancerRequest<'_>, + selector: PulsarRendezvousWeight, + ) -> Self { let hash_bytes = pulsar_hash_prefix( seed, &request.routing_target.routing_key, @@ -245,11 +259,12 @@ impl PulsarScorer { Self { prefix_len: hash_bytes.len(), hash_bytes, + selector, } } pub(super) fn score(&mut self, candidate: &RoutedClusterSnapshot) -> Option { - let weight = pulsar_weight(candidate)?; + let weight = pulsar_weight(self.selector, candidate)?; let u = self.hash_to_unit_interval(candidate); let e = -u.ln(); (e.is_finite() && e > 0.0).then(|| weight / e) @@ -267,12 +282,22 @@ impl PulsarScorer { } } +#[cfg(test)] pub(in crate::load_balancer) fn pulsar_ranked_indices( seed: Option<&str>, request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Vec { - let mut scorer = PulsarScorer::new(seed, request); + pulsar_ranked_indices_with_weight(seed, PulsarRendezvousWeight::default(), request, candidates) +} + +pub(super) fn pulsar_ranked_indices_with_weight( + seed: Option<&str>, + selector: PulsarRendezvousWeight, + request: &LoadBalancerRequest<'_>, + candidates: &[RoutedClusterSnapshot], +) -> Vec { + let mut scorer = PulsarScorer::new(seed, request, selector); let mut scored = Vec::with_capacity(candidates.len()); for (candidate_index, candidate) in candidates.iter().enumerate() { if let Some(score) = scorer.score(candidate) { @@ -297,17 +322,18 @@ 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 pulsar_weight( + selector: PulsarRendezvousWeight, + candidate: &RoutedClusterSnapshot, +) -> Option { + let value = match selector { + PulsarRendezvousWeight::LastMeanInputTps => Some(candidate.stats.last_mean_input_tps), + PulsarRendezvousWeight::MaxInputTps => candidate.stats.max_input_tps, + }?; + (value > 0.0 && value.is_finite()).then_some(value) } -pub(super) fn has_valid_input_capacity(candidate: &RoutedClusterSnapshot) -> bool { +pub(super) fn has_valid_mean_input_tps(candidate: &RoutedClusterSnapshot) -> bool { valid_last_mean_input_tps(candidate.stats.last_mean_input_tps) } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs index c55cb9446..9aa90d56f 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs @@ -141,6 +141,12 @@ fn seeded_pulsar_algorithm_config(seed: &str) -> LoadBalancerAlgorithmConfig { config } +fn max_weight_pulsar_algorithm_config(seed: &str) -> LoadBalancerAlgorithmConfig { + parse_json(&format!( + r#"{{"algorithm":"pulsar","seed":"{seed}","rendezvous_weight":"max-input-tps"}}"# + )) +} + #[test] fn set_seed_reports_unsupported_algorithms_without_panicking() { for algorithm in [ @@ -669,6 +675,20 @@ fn input_work_seconds_for_pulsar_excludes_low_free_kv_when_considered() { assert!((seconds - 1.5).abs() < f64::EPSILON); } +#[test] +fn max_rendezvous_weight_keeps_mean_input_work_capacity() { + let config = max_weight_pulsar_algorithm_config("seed-1"); + let target = target(); + let request = request(&target, Some("prefix-a"), Some(100)); + let mut backend = work_candidate("backend", 5, 50.0, 50); + backend.stats.max_input_tps = Some(500.0); + + assert_eq!( + input_work_seconds_for_request(&config, &request, &[backend]), + Some(3.0) + ); +} + #[test] fn invalid_algorithm_name_fails_during_parse() { assert_json_rejected::(r#"{"default":"not-a-real-lb"}"#, "not-a-real-lb"); @@ -714,6 +734,10 @@ fn algorithm_specific_load_balancer_fields_are_rejected_for_other_algorithms() { r#"{"algorithm":"wait-and-widen","consider_kv_free_tokens":true}"#, "consider_kv_free_tokens", ), + ( + r#"{"algorithm":"round-robin","rendezvous_weight":"max-input-tps"}"#, + "rendezvous_weight", + ), (r#"{"algorithm":"random","sample_count":4}"#, "sample_count"), ( r#"{"algorithm":"random","comparator":"ttft"}"#, @@ -937,6 +961,29 @@ fn detailed_algorithm_configs_preserve_all_variant_identities() { } } +#[test] +fn pulsar_rendezvous_weight_parses_for_both_pulsar_variants() { + for algorithm in ["pulsar", "pulsar-wait-and-widen"] { + let config: LoadBalancerAlgorithmConfig = parse_json(&format!( + r#"{{"algorithm":"{algorithm}","rendezvous_weight":"max-input-tps"}}"# + )); + assert_eq!( + config.pulsar_rendezvous_weight(), + PulsarRendezvousWeight::MaxInputTps + ); + } + + for algorithm in [ + LoadBalancerAlgorithm::Pulsar, + LoadBalancerAlgorithm::PulsarWaitAndWiden, + ] { + assert_eq!( + LoadBalancerAlgorithmConfig::from(algorithm).pulsar_rendezvous_weight(), + PulsarRendezvousWeight::LastMeanInputTps + ); + } +} + #[test] fn unknown_load_balancer_config_fields_are_rejected() { assert_json_rejected::( @@ -2560,6 +2607,51 @@ fn pulsar_ranking_cache_invalidates_when_capacity_weight_changes() { panic!("expected to find an affinity key whose ranking changes after capacity changes"); } +#[test] +fn pulsar_max_ranking_cache_tracks_only_the_selected_weight() { + let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); + let target = target(); + let first_key = "max-cache-a"; + let second_key = "max-cache-b"; + let mut snapshots = candidates(&["inst-a", "inst-b"]); + snapshots[0].stats.max_input_tps = Some(100.0); + snapshots[1].stats.max_input_tps = Some(200.0); + choose( + &pulsar, + &request(&target, Some(first_key), Some(128)), + &snapshots, + ); + choose( + &pulsar, + &request(&target, Some(second_key), Some(128)), + &snapshots, + ); + assert_eq!( + pulsar.cached_affinity_key_bytes(), + first_key.len() + second_key.len() + ); + + snapshots[0].stats.last_mean_input_tps = 1.0; + snapshots[1].stats.last_mean_input_tps = 10_000.0; + choose( + &pulsar, + &request(&target, Some(first_key), Some(128)), + &snapshots, + ); + assert_eq!( + pulsar.cached_affinity_key_bytes(), + first_key.len() + second_key.len() + ); + + snapshots[0].stats.max_input_tps = Some(300.0); + choose( + &pulsar, + &request(&target, Some(first_key), Some(128)), + &snapshots, + ); + assert_eq!(pulsar.cached_affinity_key_bytes(), first_key.len()); +} + #[test] fn pulsar_does_not_cache_oversized_affinity_key() { let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); @@ -2596,6 +2688,26 @@ fn pulsar_uses_last_mean_input_tps_as_weight() { assert_eq!(pulsar.weight(&candidate), Some(123.0)); } +#[test] +fn pulsar_can_use_windowed_max_input_tps_as_weight() { + let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); + let mut candidate = work_candidate("inst-a", 5, 123.0, 0); + candidate.stats.max_input_tps = Some(456.0); + + assert_eq!(pulsar.weight(&candidate), Some(456.0)); + + candidate.stats.last_mean_input_tps = 0.0; + assert_eq!(pulsar.weight(&candidate), Some(456.0)); +} + +#[test] +fn pulsar_max_weight_does_not_fall_back_to_mean() { + let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); + let candidate = work_candidate("inst-a", 5, 123.0, 0); + + assert_eq!(pulsar.weight(&candidate), None); +} + #[test] fn pulsar_excludes_candidate_with_invalid_last_mean_input_tps() { let pulsar = PulsarLoadBalancer::new(seeded_pulsar_algorithm_config("seed-1")); diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs index ed781cbe0..2a15fa564 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs @@ -103,6 +103,7 @@ impl ClusterRoutingGeneration { } let backend_count = active_backend_count as u128; let mut backend_stats = ModelStats::default(); + let mut max_input_tps = Some(0.0); let mut rtt_mean_nanos = 0_u128; let mut rtt_remainder_nanos = 0_u128; @@ -111,6 +112,13 @@ impl ClusterRoutingGeneration { if valid_last_mean_input_tps(backend.stats.last_mean_input_tps) { backend_stats.last_mean_input_tps += backend.stats.last_mean_input_tps; } + max_input_tps = match (max_input_tps, backend.stats.max_input_tps) { + (Some(sum), Some(value)) if value > 0.0 && value.is_finite() => { + let total = sum + value; + total.is_finite().then_some(total) + } + _ => None, + }; backend_stats.queue_size += backend.stats.queue_size; backend_stats.queued_input_size += backend.stats.queued_input_size; backend_stats.input_processing_queries += backend.stats.input_processing_queries; @@ -143,6 +151,7 @@ impl ClusterRoutingGeneration { (rtt_mean_nanos / 1_000_000_000) as u64, (rtt_mean_nanos % 1_000_000_000) as u32, ); + backend_stats.max_input_tps = max_input_tps; Some((backend_stats, rtt, active_backend_count)) } diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs index 522f417a1..7735d46cf 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs @@ -307,6 +307,7 @@ fn shared_backend_a_stats() -> ModelStats { ModelStats { output_tps: 2.0, last_mean_input_tps: 100.0, + max_input_tps: Some(110.0), max_output_tps: 50.0, queue_size: 1, queued_input_size: 100, @@ -329,6 +330,7 @@ fn shared_backend_b_stats() -> ModelStats { ModelStats { output_tps: 5.0, last_mean_input_tps: 120.0, + max_input_tps: Some(130.0), max_output_tps: 60.0, queue_size: 2, queued_input_size: 200, @@ -1720,6 +1722,7 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { backend_a.stats = ModelStats { output_tps: 1.5, last_mean_input_tps: 10.0, + max_input_tps: Some(15.0), queue_size: 2, queued_input_size: 20, input_processing_queries: 1, @@ -1739,6 +1742,7 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { backend_b.stats = ModelStats { output_tps: 2.5, last_mean_input_tps: 30.0, + max_input_tps: Some(35.0), queue_size: 4, queued_input_size: 40, input_processing_queries: 5, @@ -1760,6 +1764,7 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { assert_stats!(stats, output_tps: 4.0, last_mean_input_tps: 40.0, + max_input_tps: Some(50.0), queue_size: 6, queued_input_size: 60, input_processing_queries: 6, @@ -1777,6 +1782,73 @@ fn cluster_backend_aggregate_dedupes_sources_and_averages_rtt() { assert_eq!(stats.stats_sources.len(), 3); } +#[test] +fn cluster_backend_aggregate_requires_every_backend_max() { + let mut backend_a = backend!("cluster-max", "backend-a", 1.0, 10.0, 5); + backend_a.stats.max_input_tps = Some(15.0); + let cluster_state = RoutedClusterState::new(backend_a.registration.cluster_generation.clone()); + cluster_state.upsert_backend(Arc::new(backend_a.clone())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("one backend max should aggregate") + .0 + .max_input_tps, + Some(15.0) + ); + + let mut backend_b = backend!( + backend_a.registration.cluster_generation.clone() => + "cluster-max", "backend-b", 1.0, 10.0, 5 + ); + backend_b.stats.max_input_tps = Some(25.0); + cluster_state.upsert_backend(Arc::new(backend_b.clone())); + let complete_bits = cluster_state + .backend_aggregate() + .expect("complete backend max should aggregate") + .0 + .max_input_tps + .map(f64::to_bits); + assert_eq!(complete_bits, Some(40.0_f64.to_bits())); + + let reversed = RoutedClusterState::new(backend_a.registration.cluster_generation.clone()); + reversed.upsert_backend(Arc::new(backend_b.clone())); + reversed.upsert_backend(Arc::new(backend_a)); + assert_eq!( + reversed + .backend_aggregate() + .expect("backend order should not change aggregate") + .0 + .max_input_tps + .map(f64::to_bits), + complete_bits + ); + + backend_b.stats.max_input_tps = None; + cluster_state.upsert_backend(Arc::new(backend_b.clone())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("missing max should not remove the cluster") + .0 + .max_input_tps, + None + ); + + for invalid in [0.0, -1.0, f64::NAN, f64::INFINITY] { + backend_b.stats.max_input_tps = Some(invalid); + cluster_state.upsert_backend(Arc::new(backend_b.clone())); + assert_eq!( + cluster_state + .backend_aggregate() + .expect("invalid max should not remove the cluster") + .0 + .max_input_tps, + None + ); + } +} + #[tokio::test] async fn list_active_models_filters_by_routing_key() { let scenario = RegistrationScenario::new(None); @@ -2153,6 +2225,7 @@ async fn shared_cluster_registration_exposes_one_aggregated_cluster_candidate() assert_eq!(cluster.active_backend_count, 2); assert_stats!(cluster.stats, last_mean_input_tps: 220.0, + max_input_tps: Some(240.0), output_tps: 7.0, queue_size: 3, queued_input_size: 300, From 5944286e692ebfc88828b14a1f4825b98811a305 Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Thu, 20 Aug 2026 16:52:35 +0000 Subject: [PATCH 2/3] refactor(pylon): use generation input TPS maximum Signed-off-by: Barry Greengus --- .../crates/proto/proto/stargate.proto | 8 +- .../crates/pylon-lib/src/runtime_state.rs | 2 +- .../crates/pylon-lib/src/stats/aggregator.rs | 35 +--- .../crates/pylon-lib/src/stats/collector.rs | 50 +++--- .../crates/pylon-lib/src/stats/metrics.rs | 2 +- .../crates/pylon-lib/src/stats/projection.rs | 8 +- .../pylon-lib/src/stats/token_metrics.rs | 153 ------------------ .../rust/stargate/crates/pylon/src/main.rs | 31 ---- .../rust/stargate/crates/pylon/src/startup.rs | 5 - .../stargate/src/load_balancer/tests.rs | 2 +- 10 files changed, 36 insertions(+), 260 deletions(-) diff --git a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto index 1ee17256e..0df216f55 100644 --- a/src/libraries/rust/stargate/crates/proto/proto/stargate.proto +++ b/src/libraries/rust/stargate/crates/proto/proto/stargate.proto @@ -126,10 +126,10 @@ message ModelStats { // Sticky source labels for stat observations seen by this backend since the // model metrics state was initialized. repeated string stats_sources = 17; - // Rolling-window maximum input tokens per second observed by this backend. - // Absent after startup or restart until observed, and after the window - // expires. Stargate sums this field only when every active backend in a - // shared cluster reports it. + // Maximum input tokens per second recorded for the current backend model + // generation. Absent until a valid value is configured or observed. + // Stargate sums this field only when every active backend in a shared + // cluster reports it. optional double max_input_tps = 18; } 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 fe87d09c3..a0b8b3049 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -31,7 +31,7 @@ use reqwest::header::HeaderMap; pub struct CurrentModelStats { // Sticky runtime-observed mean input TPS for this backend. pub last_mean_input_tps: f64, - // Time-windowed max input TPS for this backend. None until observed. + // Max input TPS for this model generation. None until configured or observed. pub max_input_tps: Option, // Token/sec output rate for streaming generation endpoints. Embeddings item // cardinality is observed separately and is not exported through this field. 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 d7cabd9b0..9f9d909ae 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 @@ -28,7 +28,7 @@ use super::collector::{ FinalizeRequestUpdate, RequestCounterUpdate, StatsAggregatorUpdate, StatsCollectorConfig, StatsUpdateSource, }; -use super::token_metrics::{TpsDistribution, WindowedMax}; +use super::token_metrics::TpsDistribution; pub(super) const ENGINE_STATS_SOURCE: &str = "engine_stats_stream"; #[derive(Debug, Default)] @@ -42,7 +42,6 @@ pub(super) struct ModelMetricsState { pub(super) max_embedding_item_tps: f64, pub(super) kv_cache: KvCacheStatsSnapshot, pub(super) input_tps_distribution: TpsDistribution, - pub(super) input_tps_window_max: WindowedMax, aggregate_state_counted: bool, pub(super) counter_output_tps_authoritative: bool, pub(super) chunk_usage_stats_observed: bool, @@ -238,8 +237,7 @@ impl StatsAggregator { if self.per_model.contains_key(generation.model_id()) { return None; } - let mut metrics = - ModelMetricsState::with_input_tps_window(self.config.input_tps_capacity_window); + let mut metrics = ModelMetricsState::default(); let pinned_input_tps = match initialization { super::collector::ModelStatsInitialization::Empty => None, super::collector::ModelStatsInitialization::ConfiguredInputTps { input_tps, pin } => { @@ -248,9 +246,6 @@ impl StatsAggregator { metrics.last_mean_input_tps = input_tps; metrics.input_tps_distribution = input_tps_distribution; metrics.aggregate_state_counted = true; - metrics - .input_tps_window_max - .observe(TokioInstant::now(), input_tps); pin.then_some(input_tps) } }; @@ -396,13 +391,6 @@ impl StatsAggregator { pub(super) fn sweep_stale(&mut self, now: TokioInstant) -> Vec { let mut dirty_models = Vec::new(); - for (model_id, generation_state) in &mut self.per_model { - if generation_state.pinned_input_tps.is_none() - && generation_state.metrics.input_tps_window_max.expire(now) - { - push_dirty_model(&mut dirty_models, model_id.clone()); - } - } let request_ttl = self.config.engine_stats_request_ttl; if !request_ttl.is_zero() { let metrics = self.runtime_state.metrics(); @@ -656,7 +644,6 @@ impl StatsAggregator { config, model_state, generation_state.pinned_input_tps, - update.observed_at, InputThroughputSample { units, duration, @@ -813,7 +800,6 @@ pub(super) fn apply_input_throughput_sample( config: &StatsCollectorConfig, model_state: &mut ModelMetricsState, pinned_input_tps: Option, - observed_at: TokioInstant, sample: InputThroughputSample, ) -> bool { if sample.units < config.min_input_tokens { @@ -827,11 +813,10 @@ pub(super) fn apply_input_throughput_sample( let Some(input_tps) = tps_for_units(sample.units, duration, config.duration_floor) else { return false; }; - let max_changed = pinned_input_tps.is_none() - && model_state - .input_tps_window_max - .observe(observed_at, input_tps); + let previous_max = model_state.input_tps_distribution.max; model_state.input_tps_distribution.update(input_tps); + let max_changed = + pinned_input_tps.is_none() && model_state.input_tps_distribution.max != previous_max; let mean_input_tps = model_state.input_tps_distribution.mean; if !model_state.input_tps_distribution.has_sufficient_data() || !valid_last_mean_input_tps(mean_input_tps) @@ -869,13 +854,6 @@ pub(super) struct ModelStatsSnapshotInputs { } impl ModelMetricsState { - fn with_input_tps_window(window: Duration) -> Self { - Self { - input_tps_window_max: WindowedMax::new(window), - ..Self::default() - } - } - pub(super) fn clear_live_output_tps(&mut self) -> bool { self.last_stats_event_at = None; if self.chat_output_tps_samples.is_empty() { @@ -896,7 +874,8 @@ impl ModelMetricsState { }; CurrentModelStats { last_mean_input_tps: self.last_mean_input_tps, - max_input_tps: self.input_tps_window_max.current(), + max_input_tps: (self.input_tps_distribution.count > 0) + .then_some(self.input_tps_distribution.max), output_tps: active_chat_output_tps.max(average_with_sum( &self.chat_output_tps_samples, self.chat_output_tps_sum, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 8e787c25d..ba5ed70e3 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -25,8 +25,6 @@ use crate::{CurrentModelStats, PylonRuntimeState, RequestObservationEvent}; use stargate_runtime::OwnedTask; use super::aggregator::{ENGINE_STATS_SOURCE, KvCacheStatsSnapshot, StatsAggregator}; -use super::token_metrics::DEFAULT_INPUT_TPS_CAPACITY_WINDOW; - const DEFAULT_OBSERVATION_CHANNEL_CAPACITY: usize = 1024; const DEFAULT_SMOOTHING_WINDOW_SIZE: usize = 8; const DEFAULT_MIN_INPUT_TOKENS: u64 = 1; @@ -42,7 +40,6 @@ const DEFAULT_ENGINE_STATS_SWEEP_INTERVAL: Duration = Duration::from_secs(1); pub struct StatsCollectorConfig { pub observation_channel_capacity: usize, pub smoothing_window_size: usize, - pub input_tps_capacity_window: Duration, pub min_input_tokens: u64, pub min_output_tokens: u64, pub duration_floor: Duration, @@ -60,7 +57,6 @@ impl Default for StatsCollectorConfig { Self { observation_channel_capacity: DEFAULT_OBSERVATION_CHANNEL_CAPACITY, smoothing_window_size: DEFAULT_SMOOTHING_WINDOW_SIZE, - input_tps_capacity_window: DEFAULT_INPUT_TPS_CAPACITY_WINDOW, min_input_tokens: DEFAULT_MIN_INPUT_TOKENS, min_output_tokens: DEFAULT_MIN_OUTPUT_TOKENS, duration_floor: DEFAULT_DURATION_FLOOR, @@ -1257,7 +1253,7 @@ mod tests { #[test] fn pinned_configured_input_tps_is_preserved_across_engine_stats_updates() { let mut aggregator = test_aggregator_with_initialization( - config!(input_tps_capacity_window: seconds(1)), + StatsCollectorConfig::default(), ModelStatsInitialization::ConfiguredInputTps { input_tps: 2_200.0, pin: true, @@ -1297,22 +1293,15 @@ mod tests { } #[test] - fn unpinned_seed_expires_and_empty_restart_has_no_max() { - let config = config!(input_tps_capacity_window: seconds(1)); - let mut initialized = test_aggregator_with_initialization( - config.clone(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 100.0, - pin: false, - }, - ); - - let expired = published_stats(initialized.sweep(seconds(2))); - assert_eq!(expired.last_mean_input_tps, 100.0); - assert_eq!(expired.max_input_tps, None); + fn max_input_tps_persists_for_the_generation() { + let mut aggregator = test_aggregator(StatsCollectorConfig::default()); + aggregator.stream("req-a", (0, 0), false, Duration::ZERO); + aggregator.stream("req-a", (20, 0), false, milliseconds(100)); + aggregator.stream("req-a", (30, 0), false, milliseconds(200)); - let restarted = test_aggregator(config); - assert_eq!(restarted.snapshot("model-a").max_input_tps, None); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(200.0)); + aggregator.sweep(seconds(24 * 60 * 60)); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(200.0)); } #[test] @@ -1930,21 +1919,25 @@ mod tests { } #[test] - fn stats_aggregator_publishes_and_expires_windowed_max_input_tps() { + fn stats_aggregator_publishes_generation_max_input_tps() { let mut aggregator = test_aggregator(config!( - input_tps_capacity_window: seconds(1), engine_stats_request_ttl: Duration::ZERO, engine_stats_model_ttl: Duration::ZERO, )); - aggregator.stream("req-window-max", (0, 0), false, Duration::ZERO); + aggregator.stream("req-generation-max", (0, 0), false, Duration::ZERO); - let stats = - published_stats(aggregator.stream("req-window-max", (10, 0), false, milliseconds(100))); + let stats = published_stats(aggregator.stream( + "req-generation-max", + (10, 0), + false, + milliseconds(100), + )); assert_eq!(stats.last_mean_input_tps, 0.0); assert_eq!(stats.max_input_tps, Some(100.0)); - let expired = published_stats(aggregator.sweep(seconds(2))); - assert_eq!(expired.max_input_tps, None); + let updates = aggregator.stream("req-generation-max", (15, 0), false, milliseconds(200)); + assert!(updates.is_empty()); + assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(100.0)); } #[test] @@ -2553,7 +2546,6 @@ mod tests { fn records_metrics_when_configured() { let metrics = PylonMetrics::new().expect("metrics should initialize"); let config = config!( - input_tps_capacity_window: seconds(1), engine_stats_request_ttl: Duration::ZERO, engine_stats_model_ttl: Duration::ZERO, ); @@ -2591,7 +2583,7 @@ mod tests { publish_model_stats_update(&runtime_state, model_id, stats); } let body = metrics.gather_text().expect("metrics should encode"); - assert!(!body.contains(r#"pylon_model_max_input_tps{model="model-a"}"#)); + assert!(body.contains(r#"pylon_model_max_input_tps{model="model-a"} 10"#)); } #[test] 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 b9c160eeb..cfe392c79 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs @@ -151,7 +151,7 @@ metrics! { float_gauge output_tps("model_output_tps", "Current output TPS by model", ["model"]); float_gauge embedding_item_tps("model_embedding_item_tps", "Current embeddings item throughput by model", ["model"]); float_gauge last_mean_input_tps("model_last_mean_input_tps", "Last valid mean input TPS by model", ["model"]); - float_gauge max_input_tps("model_max_input_tps", "Time-windowed max 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"]); 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 4ccd622be..6f56a5fb8 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 @@ -207,13 +207,7 @@ impl StatsAggregator { } } if input_sample.is_some_and(|sample| { - apply_input_throughput_sample( - &self.config, - model_state, - pinned_input_tps, - TokioInstant::now(), - sample, - ) + apply_input_throughput_sample(&self.config, model_state, pinned_input_tps, sample) }) { push_changed_model(&mut changed_models, observation.model_id.clone()); } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs index 02cd0c0e5..a024bab67 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/token_metrics.rs @@ -13,98 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::VecDeque; -use std::time::Duration; - -use tokio::time::Instant; - pub(crate) const SNAPSHOT_THRESHOLD: usize = 5; -pub(crate) const DEFAULT_INPUT_TPS_CAPACITY_WINDOW: Duration = Duration::from_secs(24 * 60 * 60); -const WINDOW_SEGMENTS: u32 = 1_440; -const MAX_WINDOW_BUCKETS: usize = WINDOW_SEGMENTS as usize + 1; - -#[derive(Debug)] -struct WindowMaxBucket { - end: Instant, - max: f64, -} - -#[derive(Debug)] -pub(crate) struct WindowedMax { - window: Duration, - bucket_width: Duration, - buckets: VecDeque, -} - -impl Default for WindowedMax { - fn default() -> Self { - Self::new(DEFAULT_INPUT_TPS_CAPACITY_WINDOW) - } -} - -impl WindowedMax { - pub(crate) fn new(window: Duration) -> Self { - assert!(!window.is_zero(), "windowed max requires a nonzero window"); - let floor = window / WINDOW_SEGMENTS; - let bucket_width = if floor.saturating_mul(WINDOW_SEGMENTS) < window { - floor - .checked_add(Duration::from_nanos(1)) - .expect("window bucket width should fit Duration") - } else { - floor - }; - Self { - window, - bucket_width, - buckets: VecDeque::new(), - } - } - - pub(crate) fn observe(&mut self, now: Instant, value: f64) -> bool { - if value <= 0.0 || !value.is_finite() { - return false; - } - let previous = self.current(); - self.expire_buckets(now); - match self.buckets.back_mut() { - Some(bucket) if now < bucket.end => bucket.max = bucket.max.max(value), - _ => self.buckets.push_back(WindowMaxBucket { - end: now - .checked_add(self.bucket_width) - .expect("window bucket end should fit Instant"), - max: value, - }), - } - debug_assert!(self.buckets.len() <= MAX_WINDOW_BUCKETS); - self.current() != previous - } - - pub(crate) fn expire(&mut self, now: Instant) -> bool { - let previous = self.current(); - self.expire_buckets(now); - self.current() != previous - } - - pub(crate) fn current(&self) -> Option { - self.buckets - .iter() - .map(|bucket| bucket.max) - .max_by(f64::total_cmp) - } - - fn expire_buckets(&mut self, now: Instant) { - let Some(cutoff) = now.checked_sub(self.window) else { - return; - }; - while self - .buckets - .front() - .is_some_and(|bucket| bucket.end <= cutoff) - { - self.buckets.pop_front(); - } - } -} #[derive(Debug, Clone, Default)] pub(crate) struct TpsDistribution { @@ -160,68 +69,6 @@ impl TpsDistribution { mod tests { use super::*; - fn seconds(value: u64) -> Duration { - Duration::from_secs(value) - } - - #[test] - fn windowed_max_tracks_and_expires_bucket_maxima() { - let start = Instant::now(); - let mut max = WindowedMax::new(seconds(60)); - - assert_eq!(max.current(), None); - assert!(max.observe(start, 10.0)); - assert!(max.observe(start + seconds(1), 20.0)); - assert!(!max.observe(start + seconds(2), 15.0)); - assert_eq!(max.current(), Some(20.0)); - - assert!(!max.observe(start + seconds(31), 12.0)); - assert_eq!(max.current(), Some(20.0)); - assert!(max.expire(start + Duration::from_millis(61_050))); - assert_eq!(max.current(), Some(15.0)); - assert!(max.expire(start + Duration::from_millis(62_050))); - assert_eq!(max.current(), Some(12.0)); - assert!(max.expire(start + seconds(92))); - assert_eq!(max.current(), None); - } - - #[test] - fn windowed_max_rejects_invalid_values() { - let start = Instant::now(); - let mut max = WindowedMax::new(seconds(60)); - - for value in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { - assert!(!max.observe(start, value)); - } - assert_eq!(max.current(), None); - } - - #[test] - fn windowed_max_never_expires_early_at_oldest_boundary() { - let start = Instant::now(); - let mut max = WindowedMax::new(Duration::from_millis(1_440)); - assert!(max.observe(start, 10.0)); - - assert!(!max.expire(start + Duration::from_millis(1_440))); - assert_eq!(max.current(), Some(10.0)); - assert!(max.expire(start + Duration::from_millis(1_441))); - assert_eq!(max.current(), None); - } - - #[test] - fn windowed_max_memory_is_bounded() { - let start = Instant::now(); - let mut max = WindowedMax::new(Duration::from_millis(1_440)); - - for offset_ms in 0..10_000 { - max.observe( - start + Duration::from_millis(offset_ms), - offset_ms as f64 + 1.0, - ); - assert!(max.buckets.len() <= MAX_WINDOW_BUCKETS); - } - } - #[test] fn tps_distribution_ignores_non_positive_and_non_finite_samples() { let mut distribution = TpsDistribution::default(); diff --git a/src/libraries/rust/stargate/crates/pylon/src/main.rs b/src/libraries/rust/stargate/crates/pylon/src/main.rs index d6f08ad35..fd08dc930 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/main.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/main.rs @@ -86,9 +86,6 @@ struct Args { /// Bootstrap input TPS for every configured model instead of running calibration #[arg(long, value_name = "TPS")] initial_input_tps: Option, - /// Window for max input TPS capacity in milliseconds - #[arg(long, default_value_t = 86_400_000, value_name = "MS")] - input_tps_capacity_window_ms: u64, /// Interval between active canary requests in milliseconds. `0` disables active canaries #[arg(long, default_value_t = 5000, value_name = "MS")] active_canary_interval_ms: u64, @@ -531,34 +528,6 @@ mod tests { } } - #[test] - fn input_tps_capacity_window_defaults_flow_and_zero_is_rejected() { - let default_args = parse_args(""); - let default_config = - stats_collector_config_from_args(&default_args, &default_args.upstream_http_base_url); - assert_eq!( - default_config.input_tps_capacity_window, - std::time::Duration::from_secs(24 * 60 * 60) - ); - assert_eq!( - default_config.input_tps_capacity_window, - pylon_lib::StatsCollectorConfig::default().input_tps_capacity_window - ); - - let configured_args = parse_args("--input-tps-capacity-window-ms 1234"); - let configured = stats_collector_config_from_args( - &configured_args, - &configured_args.upstream_http_base_url, - ); - assert_eq!( - configured.input_tps_capacity_window, - std::time::Duration::from_millis(1234) - ); - - let zero = parse_args("--input-tps-capacity-window-ms 0"); - assert!(startup::PylonStartupPlan::from_args(&zero).is_err()); - } - #[test] fn calibration_ramp_requires_a_positive_request_increment() { let args = parse_args("--do-calibration --calibration-requests 0"); diff --git a/src/libraries/rust/stargate/crates/pylon/src/startup.rs b/src/libraries/rust/stargate/crates/pylon/src/startup.rs index dc2fba6ea..9a5196501 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/startup.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/startup.rs @@ -146,10 +146,6 @@ impl BackendTunnelStartup { impl PylonStartupPlan { pub(crate) fn from_args(args: &Args) -> Result { - ensure!( - args.input_tps_capacity_window_ms > 0, - "--input-tps-capacity-window-ms must be greater than zero" - ); let model_initialization = model_initialization_from_args(args)?; let model_source = model_source_from_args(args)?; Ok(Self { @@ -534,7 +530,6 @@ pub(crate) fn stats_collector_config_from_args( upstream: &str, ) -> StatsCollectorConfig { StatsCollectorConfig { - input_tps_capacity_window: Duration::from_millis(args.input_tps_capacity_window_ms), openai_fallback_stats_enabled: args.engine_stats_stream == EngineStatsStreamMode::Off, // Mock benchmark backends can expose live KV-cache occupancy over HTTP; // real upstreams usually do not, so polling is explicit. 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 9aa90d56f..6ce5872be 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 @@ -2689,7 +2689,7 @@ fn pulsar_uses_last_mean_input_tps_as_weight() { } #[test] -fn pulsar_can_use_windowed_max_input_tps_as_weight() { +fn pulsar_can_use_generation_max_input_tps_as_weight() { let pulsar = PulsarLoadBalancer::new(max_weight_pulsar_algorithm_config("seed-1")); let mut candidate = work_candidate("inst-a", 5, 123.0, 0); candidate.stats.max_input_tps = Some(456.0); From 020531174a2cf361deb2cb534c8e44db9825131e Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Thu, 20 Aug 2026 18:52:46 +0000 Subject: [PATCH 3/3] fix(pylon): measure benchmark input throughput --- .../stargate/benches/backend-degradation.yaml | 6 - .../stargate/benches/bursty-8-backends.yaml | 8 -- .../benches/cache-thrash-6-backends.yaml | 2 - .../benches/hotset-8-backends-long.yaml | 8 -- .../stargate/benches/hotset-8-backends.yaml | 8 -- .../benches/lb-balance-bursty-4c2p-2s.yaml | 8 -- .../benches/lb-balance-hotset-8c2p-4s.yaml | 8 -- .../lb-balance-prefix-reuse-4c2p-2s.yaml | 8 -- ...use-pulsar-wait-and-widen-slo-4c2p-1s.yaml | 8 -- ...lb-balance-prefix-reuse-smoke-2c2p-1s.yaml | 2 - .../benches/lb-balance-smoke-2c2p-1s.yaml | 2 - .../stargate/benches/mixed-size-pulsar.yaml | 8 -- .../stargate/benches/overload-6-backends.yaml | 6 - .../benches/queue-mismatch-retry-ab.yaml | 6 +- .../benches/stair-step-2-stargates.yaml | 2 - .../stargate/benches/sticky-hot-prefix.yaml | 2 - .../stargate/benches/uniform-4-backends.yaml | 2 - .../stargate/crates/mock-dynamo/src/openai.rs | 8 +- .../stargate/crates/mock-dynamo/src/tests.rs | 110 ++++++++++++++++++ .../crates/pylon-lib/src/model_lifecycle.rs | 35 ++---- .../crates/pylon-lib/src/stats/aggregator.rs | 29 ++--- .../crates/pylon-lib/src/stats/collector.rs | 55 ++------- .../crates/pylon-lib/src/stats/projection.rs | 15 +-- .../rust/stargate/crates/pylon/src/main.rs | 14 --- .../rust/stargate/crates/pylon/src/startup.rs | 19 +-- .../crates/stargate-bench/src/config.rs | 39 ++----- .../crates/stargate-bench/src/k8s/render.rs | 4 +- .../crates/stargate-bench/src/k8s/tests.rs | 9 +- .../crates/stargate-bench/src/k8s_run.rs | 3 +- .../crates/stargate-bench/src/manifest.rs | 9 +- .../crates/stargate-bench/src/orchestrator.rs | 8 +- .../crates/stargate-bench/src/report.rs | 9 +- .../crates/stargate-bench/src/runtime.rs | 12 +- .../crates/stargate-bench/src/score.rs | 15 +-- .../stargate/src/load_balancer/config.rs | 6 +- .../crates/stargate/src/load_balancer/mod.rs | 4 +- .../stargate/src/load_balancer/pulsar.rs | 30 +++-- .../src/load_balancer/pulsar/ranking.rs | 14 --- .../stargate/src/load_balancer/tests.rs | 11 +- .../crates/stargate/tests/common/mod.rs | 5 +- .../crates/stargate/tests/suite/lifecycle.rs | 5 +- 41 files changed, 211 insertions(+), 351 deletions(-) diff --git a/src/libraries/rust/stargate/benches/backend-degradation.yaml b/src/libraries/rust/stargate/benches/backend-degradation.yaml index c17ff479a..33af7876f 100644 --- a/src/libraries/rust/stargate/benches/backend-degradation.yaml +++ b/src/libraries/rust/stargate/benches/backend-degradation.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7000 decode_tokens_per_s: 140 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 degradation: actions: diff --git a/src/libraries/rust/stargate/benches/bursty-8-backends.yaml b/src/libraries/rust/stargate/benches/bursty-8-backends.yaml index 21d0ce1e9..5a36bc5d5 100644 --- a/src/libraries/rust/stargate/benches/bursty-8-backends.yaml +++ b/src/libraries/rust/stargate/benches/bursty-8-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7200 decode_tokens_per_s: 150 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 2 profile: name: constrained @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 10 prefill_tokens_per_s: 2200 decode_tokens_per_s: 50 - registration: - last_mean_input_tps: 60 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml b/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml index 482895b8f..95e5d6c1b 100644 --- a/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml +++ b/src/libraries/rust/stargate/benches/cache-thrash-6-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 3500 decode_tokens_per_s: 85 - registration: - last_mean_input_tps: 100 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml b/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml index dc588cfb8..759be16fb 100644 --- a/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml +++ b/src/libraries/rust/stargate/benches/hotset-8-backends-long.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7000 decode_tokens_per_s: 140 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 2 profile: name: slow @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 2200 decode_tokens_per_s: 55 - registration: - last_mean_input_tps: 70 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/hotset-8-backends.yaml b/src/libraries/rust/stargate/benches/hotset-8-backends.yaml index 3c1d09041..02af1ec6b 100644 --- a/src/libraries/rust/stargate/benches/hotset-8-backends.yaml +++ b/src/libraries/rust/stargate/benches/hotset-8-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 7000 decode_tokens_per_s: 140 - registration: - last_mean_input_tps: 220 - count: 4 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 2 profile: name: slow @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 2200 decode_tokens_per_s: 55 - registration: - last_mean_input_tps: 70 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml b/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml index a7cbc5df5..9a666d2d9 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-bursty-4c2p-2s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 7000 decode_tokens_per_s: 170 - registration: - last_mean_input_tps: 160 profiles: - count: 2 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 12000 decode_tokens_per_s: 260 - registration: - last_mean_input_tps: 280 - count: 4 profile: name: balanced-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 7000 decode_tokens_per_s: 170 - registration: - last_mean_input_tps: 160 - count: 2 profile: name: constrained-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 7 prefill_tokens_per_s: 3400 decode_tokens_per_s: 85 - registration: - last_mean_input_tps: 80 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml b/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml index 7c28f4a43..0e986f847 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-hotset-8c2p-4s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 8200 decode_tokens_per_s: 190 - registration: - last_mean_input_tps: 185 profiles: - count: 4 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 14000 decode_tokens_per_s: 285 - registration: - last_mean_input_tps: 310 - count: 8 profile: name: balanced-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 8200 decode_tokens_per_s: 190 - registration: - last_mean_input_tps: 185 - count: 4 profile: name: constrained-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 6 prefill_tokens_per_s: 4000 decode_tokens_per_s: 100 - registration: - last_mean_input_tps: 90 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml index b4fdb7d1a..1d85af82a 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-4c2p-2s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 profiles: - count: 2 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 1 prefill_tokens_per_s: 65536 decode_tokens_per_s: 320 - registration: - last_mean_input_tps: 65536 - count: 4 profile: name: balanced-prefix-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 - count: 2 profile: name: constrained-prefix-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 16384 decode_tokens_per_s: 120 - registration: - last_mean_input_tps: 16384 traffic_pattern: kind: prefix_reuse diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml index 447e4db95..1800d1277 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-pulsar-wait-and-widen-slo-4c2p-1s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 profiles: - count: 2 profile: @@ -57,8 +55,6 @@ backends: decode_jitter_ms: 1 prefill_tokens_per_s: 65536 decode_tokens_per_s: 320 - registration: - last_mean_input_tps: 65536 - count: 4 profile: name: balanced-prefix-cluster @@ -71,8 +67,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 32768 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 32768 - count: 2 profile: name: constrained-prefix-cluster @@ -85,8 +79,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 16384 decode_tokens_per_s: 120 - registration: - last_mean_input_tps: 16384 traffic_pattern: kind: prefix_reuse diff --git a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml index 2019da9fa..75d6b10b8 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-prefix-reuse-smoke-2c2p-1s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 1 prefill_tokens_per_s: 32768 decode_tokens_per_s: 512 - registration: - last_mean_input_tps: 32768 traffic_pattern: kind: prefix_reuse diff --git a/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml b/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml index 0e31c2394..b02085a27 100644 --- a/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml +++ b/src/libraries/rust/stargate/benches/lb-balance-smoke-2c2p-1s.yaml @@ -42,8 +42,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 9000 decode_tokens_per_s: 220 - registration: - last_mean_input_tps: 200 traffic_pattern: kind: uniform diff --git a/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml b/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml index 06c36fbbc..9cb7d8d55 100644 --- a/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml +++ b/src/libraries/rust/stargate/benches/mixed-size-pulsar.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 8000 decode_tokens_per_s: 160 - registration: - last_mean_input_tps: 260 - count: 3 profile: name: balanced @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4000 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 120 - count: 1 profile: name: small-cache @@ -83,8 +77,6 @@ backends: decode_jitter_ms: 6 prefill_tokens_per_s: 2600 decode_tokens_per_s: 70 - registration: - last_mean_input_tps: 80 traffic_pattern: kind: mixed_size diff --git a/src/libraries/rust/stargate/benches/overload-6-backends.yaml b/src/libraries/rust/stargate/benches/overload-6-backends.yaml index 6bc9f8206..c61d28485 100644 --- a/src/libraries/rust/stargate/benches/overload-6-backends.yaml +++ b/src/libraries/rust/stargate/benches/overload-6-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 1800 decode_tokens_per_s: 45 - registration: - last_mean_input_tps: 60 profiles: - count: 2 profile: @@ -55,8 +53,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4200 decode_tokens_per_s: 90 - registration: - last_mean_input_tps: 130 - count: 4 profile: name: constrained @@ -69,8 +65,6 @@ backends: decode_jitter_ms: 8 prefill_tokens_per_s: 1800 decode_tokens_per_s: 45 - registration: - last_mean_input_tps: 60 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml b/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml index 713027202..01560fa0e 100644 --- a/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml +++ b/src/libraries/rust/stargate/benches/queue-mismatch-retry-ab.yaml @@ -15,10 +15,10 @@ name: queue-mismatch-retry-ab metadata: - description: A/B diagnostic for pylon rejection of stale Stargate queue estimates under burst load. The scenario runs 2048 requests per arm for benchmark evidence. Mock backends enforce real concurrent-service, prefill, and decode delays; benchmark pylons pin prompt throughput to the matching 2200 tokens/s profile value. Queue admission intentionally models prompt work only because output residency is unknowable at admission time. + description: A/B diagnostic for pylon rejection of stale Stargate queue estimates under burst load. The scenario runs 2048 requests per arm for benchmark evidence. Mock backends enforce real concurrent-service, prefill, and decode delays, and pylons measure the resulting prompt throughput. Queue admission intentionally models prompt work only because output residency is unknowable at admission time. tags: [queue-admission, queue-mismatch, ab, wait-and-widen, diagnostic] expected_runtime: medium - expected_signal: With the same fixed prompt-throughput estimate and real mock service delays in both arms, the enabled row should record pylon rejections and Stargate queue-mismatch retries; the disabled row should record disabled decisions without those retries. + expected_signal: With the same mock service delays in both arms, the enabled row should record pylon rejections and Stargate queue-mismatch retries; the disabled row should record disabled decisions without those retries. model: dummy-model seed: 154 request_count: 2048 @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 2 prefill_tokens_per_s: 2200 decode_tokens_per_s: 100 - registration: - last_mean_input_tps: 2200 traffic_pattern: kind: bursty diff --git a/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml b/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml index c280d752b..6f8c7de71 100644 --- a/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml +++ b/src/libraries/rust/stargate/benches/stair-step-2-stargates.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 4 prefill_tokens_per_s: 4500 decode_tokens_per_s: 95 - registration: - last_mean_input_tps: 130 traffic_pattern: kind: stair_step diff --git a/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml b/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml index 7b3072764..cd264c0e9 100644 --- a/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml +++ b/src/libraries/rust/stargate/benches/sticky-hot-prefix.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 4500 decode_tokens_per_s: 95 - registration: - last_mean_input_tps: 130 traffic_pattern: kind: zipf_hotset diff --git a/src/libraries/rust/stargate/benches/uniform-4-backends.yaml b/src/libraries/rust/stargate/benches/uniform-4-backends.yaml index 419fb13dc..833b72bdc 100644 --- a/src/libraries/rust/stargate/benches/uniform-4-backends.yaml +++ b/src/libraries/rust/stargate/benches/uniform-4-backends.yaml @@ -40,8 +40,6 @@ backends: decode_jitter_ms: 3 prefill_tokens_per_s: 5000 decode_tokens_per_s: 100 - registration: - last_mean_input_tps: 140 traffic_pattern: kind: uniform diff --git a/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs b/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs index 7604e2307..9eb3c80d0 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs @@ -202,12 +202,11 @@ pub(crate) async fn chat_completions( info!(id = %id, model = %model, stream = stream, "received chat/completions request"); let request_id = optional_header(&headers, "x-request-id").unwrap_or_else(|| id.clone()); let cache_affinity_key = optional_header(&headers, "x-cache-affinity-key"); - if stream { - state.emit_counters(&request_id, &model, 0, 0, false); - } + state.emit_counters(&request_id, &model, 0, 0, false); let kv_cache_access = state .process_input_with_cache(cache_affinity_key.as_deref(), input_tokens) .await; + state.emit_counters(&request_id, &model, input_tokens, 0, false); let first_token_delay = state.ttft + Duration::from_millis(jitter_ms(&request_id, "ttft", state.ttft_jitter_ms)); info!( @@ -303,6 +302,7 @@ pub(crate) async fn responses( let kv_cache_access = state .process_input_with_cache(cache_affinity_key.as_deref(), input_tokens) .await; + state.emit_counters(&request_id, &model, input_tokens, 0, false); let first_token_delay = state.ttft + Duration::from_millis(jitter_ms(&request_id, "ttft", state.ttft_jitter_ms)); @@ -514,8 +514,6 @@ fn stream_response(config: StreamResponseConfig) -> Response { } tokio::time::sleep(first_token_delay).await; - state.emit_counters(&request_id, &model, input_tokens, 0, false); - if kind == StreamKind::Chat { yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Role)); } diff --git a/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs b/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs index d4ba751a0..2382cd705 100644 --- a/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs +++ b/src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs @@ -631,6 +631,116 @@ async fn chat_completion_retains_prefix_only_after_modeled_prefill_completes() { ); } +#[tokio::test(start_paused = true)] +async fn non_streaming_chat_reports_input_when_prefill_completes() { + let state = AppState { + prefill_tokens_per_s: 100.0, + ttft: Duration::from_secs(10), + ..test_state() + }; + let mut stats_events = state.stats_events.subscribe(); + let mut headers = HeaderMap::new(); + headers.insert( + "x-request-id", + HeaderValue::from_static("req-prefill-stats"), + ); + headers.insert("x-input-tokens", HeaderValue::from_static("100")); + let mut chat_request = request(Some(1)); + chat_request.stream = Some(false); + + let request = tokio::spawn(chat_completions(State(state), headers, Json(chat_request))); + let StatsStreamEvent::Stats { + tokens_processed, + tokens_generated, + finished, + .. + } = stats_events + .recv() + .await + .expect("baseline event should arrive") + else { + panic!("expected a stats event"); + }; + assert_eq!(tokens_processed, Some(0)); + assert_eq!(tokens_generated, Some(0)); + assert!(!finished); + + tokio::time::advance(Duration::from_millis(999)).await; + assert!(matches!( + stats_events.try_recv(), + Err(broadcast::error::TryRecvError::Empty) + )); + tokio::time::advance(Duration::from_millis(1)).await; + + let StatsStreamEvent::Stats { + tokens_processed, + tokens_generated, + finished, + .. + } = stats_events + .recv() + .await + .expect("prefill event should arrive") + else { + panic!("expected a stats event"); + }; + assert_eq!(tokens_processed, Some(100)); + assert_eq!(tokens_generated, Some(0)); + assert!(!finished); + assert!( + !request.is_finished(), + "TTFT should still delay the response" + ); +} + +#[tokio::test(start_paused = true)] +async fn responses_reports_input_without_waiting_for_stream_polling() { + let state = AppState { + prefill_tokens_per_s: 100.0, + ttft: Duration::from_secs(10), + ..test_state() + }; + let mut stats_events = state.stats_events.subscribe(); + let mut headers = HeaderMap::new(); + headers.insert( + "x-request-id", + HeaderValue::from_static("req-response-prefill"), + ); + headers.insert("x-input-tokens", HeaderValue::from_static("100")); + let response_request = ResponsesRequest { + stream: Some(true), + model: Some("dummy-model".to_string()), + max_output_tokens: Some(1), + input: None, + }; + + let request = tokio::spawn(responses(State(state), headers, Json(response_request))); + stats_events + .recv() + .await + .expect("baseline event should arrive"); + tokio::time::advance(Duration::from_secs(1)).await; + let StatsStreamEvent::Stats { + tokens_processed, + tokens_generated, + finished, + .. + } = stats_events + .recv() + .await + .expect("prefill event should arrive") + else { + panic!("expected a stats event"); + }; + + assert_eq!(tokens_processed, Some(100)); + assert_eq!(tokens_generated, Some(0)); + assert!(!finished); + request + .await + .expect("responses request should return a stream"); +} + #[test] fn kv_cache_evicts_least_recently_used_entry() { let mut cache = KvCacheState::new(300); diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs index 2af6f5c30..58c43cf62 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/model_lifecycle.rs @@ -47,7 +47,7 @@ pub enum ModelSource { #[derive(Clone, Debug)] pub enum ModelInitialization { Calibration(CalibrationConfig), - ConfiguredInputTps { input_tps: f64, pin: bool }, + ConfiguredInputTps { input_tps: f64 }, Uncalibrated, } @@ -348,10 +348,9 @@ impl ModelLifecycleSupervisor { ModelInitialization::Calibration(_) | ModelInitialization::Uncalibrated => { ModelStatsInitialization::Empty } - ModelInitialization::ConfiguredInputTps { input_tps, pin } => { + ModelInitialization::ConfiguredInputTps { input_tps } => { ModelStatsInitialization::ConfiguredInputTps { input_tps: *input_tps, - pin: *pin, } } }; @@ -1014,10 +1013,7 @@ mod tests { poll_interval: Duration::from_millis(10), request_timeout: Duration::from_secs(1), }), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1135,10 +1131,7 @@ mod tests { "model-a".to_string(), "model-b".to_string(), ])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1201,10 +1194,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: upstream.to_string(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { active_canary_interval: Duration::ZERO, ..BringupConfig::default() @@ -1242,10 +1232,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: "http://127.0.0.1:1".to_string(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { active_canary_interval: Duration::ZERO, ..BringupConfig::default() @@ -1317,10 +1304,7 @@ mod tests { config: ModelLifecycleConfig { upstream_http_base_url: "http://127.0.0.1:1".to_string(), source: ModelSource::Static(BTreeSet::new()), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1512,10 +1496,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: upstream.base_url.clone(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 123.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 123.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs index 9f9d909ae..ea19712b8 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/aggregator.rs @@ -55,7 +55,6 @@ pub(super) struct ModelMetricsState { pub(super) struct GenerationMetricsState { pub(super) generation: ModelGeneration, pub(super) metrics: ModelMetricsState, - pub(super) pinned_input_tps: Option, } #[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq)] pub(super) struct KvCacheStatsSnapshot { @@ -237,16 +236,17 @@ impl StatsAggregator { if self.per_model.contains_key(generation.model_id()) { return None; } - let mut metrics = ModelMetricsState::default(); - let pinned_input_tps = match initialization { - super::collector::ModelStatsInitialization::Empty => None, - super::collector::ModelStatsInitialization::ConfiguredInputTps { input_tps, pin } => { + let metrics = match initialization { + super::collector::ModelStatsInitialization::Empty => ModelMetricsState::default(), + super::collector::ModelStatsInitialization::ConfiguredInputTps { input_tps } => { let input_tps_distribution = TpsDistribution::bootstrap(input_tps) .expect("configured input TPS must be positive and finite"); - metrics.last_mean_input_tps = input_tps; - metrics.input_tps_distribution = input_tps_distribution; - metrics.aggregate_state_counted = true; - pin.then_some(input_tps) + ModelMetricsState { + last_mean_input_tps: input_tps, + input_tps_distribution, + aggregate_state_counted: true, + ..ModelMetricsState::default() + } } }; self.aggregate_model_state_count += usize::from(metrics.aggregate_state_counted); @@ -255,7 +255,6 @@ impl StatsAggregator { GenerationMetricsState { generation: generation.clone(), metrics, - pinned_input_tps, }, ); let stats = self.snapshot(generation.model_id()); @@ -643,7 +642,6 @@ impl StatsAggregator { dirty |= apply_input_throughput_sample( config, model_state, - generation_state.pinned_input_tps, InputThroughputSample { units, duration, @@ -799,7 +797,6 @@ fn adjust_live_count(count: &mut usize, delta: isize) { pub(super) fn apply_input_throughput_sample( config: &StatsCollectorConfig, model_state: &mut ModelMetricsState, - pinned_input_tps: Option, sample: InputThroughputSample, ) -> bool { if sample.units < config.min_input_tokens { @@ -815,19 +812,17 @@ pub(super) fn apply_input_throughput_sample( }; let previous_max = model_state.input_tps_distribution.max; model_state.input_tps_distribution.update(input_tps); - let max_changed = - pinned_input_tps.is_none() && model_state.input_tps_distribution.max != previous_max; + let max_changed = model_state.input_tps_distribution.max != previous_max; let mean_input_tps = model_state.input_tps_distribution.mean; if !model_state.input_tps_distribution.has_sufficient_data() || !valid_last_mean_input_tps(mean_input_tps) { return max_changed; } - let last_mean_input_tps = pinned_input_tps.unwrap_or(mean_input_tps); - if model_state.last_mean_input_tps == last_mean_input_tps { + if model_state.last_mean_input_tps == mean_input_tps { return max_changed; } - model_state.last_mean_input_tps = last_mean_input_tps; + model_state.last_mean_input_tps = mean_input_tps; true } 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 ba5ed70e3..2b924af0d 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -109,7 +109,7 @@ impl StatsCollectorHandle { #[derive(Debug, Clone, Copy)] pub(crate) enum ModelStatsInitialization { Empty, - ConfiguredInputTps { input_tps: f64, pin: bool }, + ConfiguredInputTps { input_tps: f64 }, } enum StatsCollectorCommand { @@ -636,7 +636,7 @@ mod tests { } } - async fn begin_configured_model(&self, model_id: &str, input_tps: f64, pin: bool) { + async fn begin_configured_model(&self, model_id: &str, input_tps: f64) { let generation = ModelGeneration::new(model_id, 0); assert!(self.runtime_state.begin_generation(generation.clone())); assert!( @@ -644,7 +644,7 @@ mod tests { .control() .begin_generation( generation, - ModelStatsInitialization::ConfiguredInputTps { input_tps, pin } + ModelStatsInitialization::ConfiguredInputTps { input_tps } ) .await .expect("collector should acknowledge configured generation") @@ -1251,35 +1251,10 @@ mod tests { } #[test] - fn pinned_configured_input_tps_is_preserved_across_engine_stats_updates() { + fn configured_input_tps_moves_with_the_first_real_sample() { let mut aggregator = test_aggregator_with_initialization( StatsCollectorConfig::default(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 2_200.0, - pin: true, - }, - ); - let stats = aggregator.stream_stats("req-a", (0, 0), false, Duration::ZERO); - assert_eq!(stats.last_mean_input_tps, 2_200.0); - assert_eq!(stats.max_input_tps, Some(2_200.0)); - for tick in 1..=5 { - aggregator.stream("req-a", (tick * 10, 0), false, milliseconds(tick * 100)); - } - let stats = aggregator.snapshot("model-a"); - assert_eq!(stats.last_mean_input_tps, 2_200.0); - assert_eq!(stats.max_input_tps, Some(2_200.0)); - aggregator.sweep(seconds(2)); - assert_eq!(aggregator.snapshot("model-a").max_input_tps, Some(2_200.0)); - } - - #[test] - fn unpinned_configured_input_tps_moves_with_the_first_real_sample() { - let mut aggregator = test_aggregator_with_initialization( - StatsCollectorConfig::default(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 100.0, - pin: false, - }, + ModelStatsInitialization::ConfiguredInputTps { input_tps: 100.0 }, ); let bootstrapped = aggregator.snapshot("model-a"); assert_eq!(bootstrapped.last_mean_input_tps, 100.0); @@ -2507,9 +2482,7 @@ mod tests { #[tokio::test] async fn stats_collector_bootstraps_input_tps_for_queue_admission() { let collector = RunningCollector::spawn_empty(StatsCollectorConfig::default(), None, false); - collector - .begin_configured_model("model-a", 2_200.0, false) - .await; + collector.begin_configured_model("model-a", 2_200.0).await; let stats = collector .wait_for_stats("bootstrap TPS stats should be published", |stats| { stats.last_mean_input_tps == 2_200.0 @@ -2545,10 +2518,7 @@ mod tests { #[test] fn records_metrics_when_configured() { let metrics = PylonMetrics::new().expect("metrics should initialize"); - let config = config!( - engine_stats_request_ttl: Duration::ZERO, - engine_stats_model_ttl: Duration::ZERO, - ); + let config = StatsCollectorConfig::default(); let (runtime_state, _observation_rx) = PylonRuntimeState::observed( stargate_proto::pb::InferenceServerStatus::Unknown, &["model-a".to_string()], @@ -2578,12 +2548,6 @@ mod tests { 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"#)); - - for (model_id, stats) in aggregator.sweep_stale(TokioInstant::now() + seconds(2)) { - publish_model_stats_update(&runtime_state, model_id, stats); - } - let body = metrics.gather_text().expect("metrics should encode"); - assert!(body.contains(r#"pylon_model_max_input_tps{model="model-a"} 10"#)); } #[test] @@ -2618,10 +2582,7 @@ mod tests { control .begin_generation( first.clone(), - ModelStatsInitialization::ConfiguredInputTps { - input_tps: 100.0, - pin: false, - }, + ModelStatsInitialization::ConfiguredInputTps { input_tps: 100.0 }, ) .await .expect("stats collector should acknowledge initialization") diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs index 6f56a5fb8..8ced5f997 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs @@ -119,13 +119,7 @@ impl StatsAggregator { }; let mut stats = self.per_model.get(model_id).map_or_else( || ModelMetricsState::default().current_stats(inputs), - |state| { - let mut stats = state.metrics.current_stats(inputs); - if let Some(pinned_input_tps) = state.pinned_input_tps { - stats.max_input_tps = Some(pinned_input_tps); - } - stats - }, + |state| state.metrics.current_stats(inputs), ); stats.queue_time_estimate_ms_by_priority = queue.queue_time_estimate_ms_by_priority; stats @@ -163,7 +157,6 @@ impl StatsAggregator { let Some(generation_state) = self.per_model.get_mut(&observation.model_id) else { return changed_models; }; - let pinned_input_tps = generation_state.pinned_input_tps; let model_state = &mut generation_state.metrics; model_state.chunk_usage_stats_observed |= observation.output_tokens_from_chunk_usage; let record_sample = |samples, sum: &mut f64, max: &mut f64, sample| { @@ -206,9 +199,9 @@ impl StatsAggregator { } } } - if input_sample.is_some_and(|sample| { - apply_input_throughput_sample(&self.config, model_state, pinned_input_tps, sample) - }) { + if input_sample + .is_some_and(|sample| apply_input_throughput_sample(&self.config, model_state, sample)) + { push_changed_model(&mut changed_models, observation.model_id.clone()); } changed_models diff --git a/src/libraries/rust/stargate/crates/pylon/src/main.rs b/src/libraries/rust/stargate/crates/pylon/src/main.rs index fd08dc930..c92dba818 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/main.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/main.rs @@ -116,9 +116,6 @@ struct Args { /// Upstream HTTP path for the engine stats stream #[arg(long, default_value = "/pylon/v1/stats/stream", value_name = "PATH")] engine_stats_stream_path: String, - /// Keep --initial-input-tps fixed for deterministic benchmark/test experiments - #[arg(long, default_value_t = false, hide = true)] - benchmark_pin_input_tps: bool, /// Minimum interval between registration/stat updates to stargate #[arg(long, default_value_t = 1000, value_name = "MS")] min_update_interval_ms: u64, @@ -535,17 +532,6 @@ mod tests { assert!(startup::PylonStartupPlan::from_args(&args).is_err()); } - #[test] - fn benchmark_pin_requires_initial_input_tps() { - let uncalibrated = parse_args("--benchmark-pin-input-tps"); - let calibration = parse_args("--do-calibration --benchmark-pin-input-tps"); - let initial = parse_args("--initial-input-tps 2200 --benchmark-pin-input-tps"); - - assert!(startup::PylonStartupPlan::from_args(&uncalibrated).is_err()); - assert!(startup::PylonStartupPlan::from_args(&calibration).is_err()); - assert!(startup::PylonStartupPlan::from_args(&initial).is_ok()); - } - #[test] fn engine_stats_stream_defaults_to_auto_mode_and_v1_path() { let args = parse_args(""); diff --git a/src/libraries/rust/stargate/crates/pylon/src/startup.rs b/src/libraries/rust/stargate/crates/pylon/src/startup.rs index 9a5196501..edbd9d97e 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/startup.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/startup.rs @@ -598,10 +598,6 @@ fn model_initialization_from_args(args: &Args) -> Result { .is_none_or(|input_tps| input_tps.is_finite() && input_tps > 0.0), "initial input TPS must be finite and positive" ); - ensure!( - !args.benchmark_pin_input_tps || args.initial_input_tps.is_some(), - "--benchmark-pin-input-tps requires --initial-input-tps" - ); if args.do_calibration { ensure!( args.calibration_requests > 0, @@ -617,10 +613,7 @@ fn model_initialization_from_args(args: &Args) -> Result { } Ok(match args.initial_input_tps { - Some(input_tps) => ModelInitialization::ConfiguredInputTps { - input_tps, - pin: args.benchmark_pin_input_tps, - }, + Some(input_tps) => ModelInitialization::ConfiguredInputTps { input_tps }, None => ModelInitialization::Uncalibrated, }) } @@ -1016,10 +1009,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: "http://127.0.0.1:1".to_string(), source: ModelSource::Static(BTreeSet::new()), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() @@ -1373,10 +1363,7 @@ mod tests { ModelLifecycleConfig { upstream_http_base_url: plan.upstream.clone(), source: ModelSource::Static(BTreeSet::from(["model-a".to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1_000.0, - pin: false, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1_000.0 }, bringup: BringupConfig { enabled: false, ..BringupConfig::default() diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs index 683498349..f628c9f82 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/config.rs @@ -337,16 +337,10 @@ fn validate_profile(profile: &BackendProfile) -> anyhow::Result<()> { profile.service_time_ms.decode_tokens_per_s > 0, "backend decode_tokens_per_s must be > 0" ); - if let Some(prefill_tokens_per_s) = profile.service_time_ms.prefill_tokens_per_s { - ensure!( - prefill_tokens_per_s > 0.0 && prefill_tokens_per_s.is_finite(), - "backend prefill_tokens_per_s must be finite and > 0 when set" - ); - } ensure!( - profile.registration.last_mean_input_tps > 0.0 - && profile.registration.last_mean_input_tps.is_finite(), - "backend registration.last_mean_input_tps must be finite and > 0" + profile.service_time_ms.prefill_tokens_per_s > 0.0 + && profile.service_time_ms.prefill_tokens_per_s.is_finite(), + "backend prefill_tokens_per_s must be finite and > 0" ); Ok(()) } @@ -387,7 +381,6 @@ config_struct!(BackendProfile { #[serde(default)] pub kv_cache_capacity_tokens: u64, pub service_time_ms: ServiceTimeConfig, - pub registration: RegistrationConfig, }); fn default_backend_name() -> String { @@ -405,11 +398,7 @@ config_struct!(ServiceTimeConfig { pub decode_tokens_per_s: u64, #[serde(default)] pub decode_jitter_ms: u64, - pub prefill_tokens_per_s: Option, -}); - -config_struct!(RegistrationConfig { - pub last_mean_input_tps: f64, + pub prefill_tokens_per_s: f64, }); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -569,10 +558,7 @@ mod tests { ttft_jitter_ms: 0, decode_tokens_per_s: 100, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 100.0, + prefill_tokens_per_s: 100.0, }, } } @@ -683,7 +669,7 @@ mod tests { ] { remove_yaml_field(&mut yaml["backends"]["profile"], key); } - for key in ["ttft_jitter_ms", "decode_jitter_ms", "prefill_tokens_per_s"] { + for key in ["ttft_jitter_ms", "decode_jitter_ms"] { remove_yaml_field(&mut yaml["backends"]["profile"]["service_time_ms"], key); } @@ -704,10 +690,6 @@ mod tests { assert_eq!(config.backends.profile.kv_cache_capacity_tokens, 0); assert_eq!(config.backends.profile.service_time_ms.ttft_jitter_ms, 0); assert_eq!(config.backends.profile.service_time_ms.decode_jitter_ms, 0); - assert_eq!( - config.backends.profile.service_time_ms.prefill_tokens_per_s, - None - ); } #[test] @@ -865,13 +847,10 @@ action: pause_backend } #[test] - fn rejects_invalid_registered_input_throughput() { + fn rejects_invalid_prefill_throughput() { let mut profile = profile("invalid-throughput"); - profile.registration.last_mean_input_tps = 0.0; - assert_validation_error( - validate_profile(&profile), - "registration.last_mean_input_tps", - ); + profile.service_time_ms.prefill_tokens_per_s = 0.0; + assert_validation_error(validate_profile(&profile), "prefill_tokens_per_s"); } #[test] diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs index 4630a044e..81194f407 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/render.rs @@ -126,14 +126,14 @@ pub(super) fn render_manifest(render: RenderManifestConfig<'_>) -> RenderedManif )); } backends.push_str(&format!( - "apiVersion: apps/v1\nkind: Deployment\nmetadata:\n name: {inference_server_id}-pylon\n namespace: {backends_ns}\nspec:\n replicas: 1\n selector:\n matchLabels:\n app: {inference_server_id}-pylon\n template:\n metadata:\n labels:\n app: {inference_server_id}-pylon\n benchmark.stargate/profile: {profile_name}\n spec:\n containers:\n - name: pylon\n image: {pylon_image}\n imagePullPolicy: IfNotPresent\n args:\n - --upstream-http-base-url=http://{upstream_backend_name}-http.{backends_ns}.svc.cluster.local:8090\n - --model-name={model}\n - --stargate-address=stargate.{stargate_ns}.svc.cluster.local:50071\n - --inference-server-id={inference_server_id}\n{cluster_id_arg} - --backend-connectivity=reverse\n - --quic-insecure\n - --tunnel-protocol={tunnel_protocol}\n - --kv-cache-stats-path=/kv-cache/stats\n - --min-update-interval-ms=100\n - --disable-bringup\n - --active-canary-interval-ms=0\n - --initial-input-tps={last_mean_input_tps}\n - --benchmark-pin-input-tps\n", + "apiVersion: apps/v1\nkind: Deployment\nmetadata:\n name: {inference_server_id}-pylon\n namespace: {backends_ns}\nspec:\n replicas: 1\n selector:\n matchLabels:\n app: {inference_server_id}-pylon\n template:\n metadata:\n labels:\n app: {inference_server_id}-pylon\n benchmark.stargate/profile: {profile_name}\n spec:\n containers:\n - name: pylon\n image: {pylon_image}\n imagePullPolicy: IfNotPresent\n args:\n - --upstream-http-base-url=http://{upstream_backend_name}-http.{backends_ns}.svc.cluster.local:8090\n - --model-name={model}\n - --stargate-address=stargate.{stargate_ns}.svc.cluster.local:50071\n - --inference-server-id={inference_server_id}\n{cluster_id_arg} - --backend-connectivity=reverse\n - --quic-insecure\n - --tunnel-protocol={tunnel_protocol}\n - --kv-cache-stats-path=/kv-cache/stats\n - --min-update-interval-ms=100\n - --disable-bringup\n - --active-canary-interval-ms=0\n - --initial-input-tps={initial_input_tps}\n", upstream_backend_name = pylon.upstream_backend_name, inference_server_id = pylon.inference_server_id, profile_name = pylon.profile_slug, pylon_image = image_refs.pylon, model = config.model, tunnel_protocol = config.tunnel_protocol, - last_mean_input_tps = pylon.last_mean_input_tps, + initial_input_tps = pylon.initial_input_tps, )); backends.push_str(&pylon_queue_admission_args); backends.push_str("---\n"); diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs index e48495859..ff0c6bad6 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s/tests.rs @@ -26,7 +26,7 @@ use super::render::{ use super::run::{BenchmarkK8sRun, prepare_benchmark_k8s_run_with_resolved_dependencies}; use crate::config::{ AlgorithmConfig, ArrivalPatternConfig, BackendConfig, BackendProfile, BenchmarkConfig, - DegradationConfig, RegistrationConfig, ScenarioMetadata, ServiceTimeConfig, StargateConfig, + DegradationConfig, ScenarioMetadata, ServiceTimeConfig, StargateConfig, TokenDistributionConfig, TrafficPatternConfig, UniformTrafficConfig, }; use serde::Deserialize; @@ -57,10 +57,7 @@ fn config() -> BenchmarkConfig { ttft_jitter_ms: 10, decode_tokens_per_s: 50, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 100.0, + prefill_tokens_per_s: 100.0, }, }, }, @@ -725,9 +722,9 @@ fn rendered_pylons_include_per_algorithm_queue_admission_args() { "- --disable-bringup", "- --active-canary-interval-ms=0", "- --initial-input-tps=100", - "- --benchmark-pin-input-tps", "- --pylon-queue-mismatch-tolerance-factor=1", "- --pylon-queue-mismatch-retry-after-ms=5", ], ); + assert!(!rendered.backends.contains("--benchmark-pin-input-tps")); } diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs index a54340e91..71d276968 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/k8s_run.rs @@ -1000,8 +1000,7 @@ backends: count: 1 profile: name: balanced - service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50 } - registration: { last_mean_input_tps: 100.0 } + service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50, prefill_tokens_per_s: 100.0 } traffic_pattern: kind: uniform routing_keys: 1 diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs index 258746d18..d75020bdf 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/manifest.rs @@ -329,8 +329,8 @@ mod tests { use super::*; use crate::config::{ AlgorithmConfig, ArrivalPatternConfig, BackendConfig, BackendProfile, HotsetTrafficConfig, - MixedSizeClassConfig, MixedSizeTrafficConfig, PrefixReuseTrafficConfig, RegistrationConfig, - ServiceTimeConfig, StargateConfig, TokenDistributionConfig, UniformTrafficConfig, + MixedSizeClassConfig, MixedSizeTrafficConfig, PrefixReuseTrafficConfig, ServiceTimeConfig, + StargateConfig, TokenDistributionConfig, UniformTrafficConfig, }; fn base_config() -> BenchmarkConfig { @@ -358,10 +358,7 @@ mod tests { ttft_jitter_ms: 20, decode_tokens_per_s: 80, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 100.0, + prefill_tokens_per_s: 100.0, }, }, }, diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs index 185127b72..ec94be266 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/orchestrator.rs @@ -295,8 +295,7 @@ fn build_compose_spec( "--min-update-interval-ms" => "100", "--disable-bringup", "--active-canary-interval-ms=0", - "--initial-input-tps" => pylon.last_mean_input_tps.to_string(), - "--benchmark-pin-input-tps", + "--initial-input-tps" => pylon.initial_input_tps.to_string(), ]); if let Some(pylon_queue_admission) = &algorithm.pylon_queue_admission { client_command.extend(pylon_queue_admission.pylon_args()); @@ -351,8 +350,7 @@ backends: count: 2 profile: name: balanced - service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50 } - registration: { last_mean_input_tps: 100.0 } + service_time_ms: { ttft_mean: 150, ttft_jitter_ms: 10, decode_tokens_per_s: 50, prefill_tokens_per_s: 100.0 } traffic_pattern: kind: uniform routing_keys: 2 @@ -566,7 +564,7 @@ algorithms: Some("100") ); assert!( - client + !client .command .iter() .any(|candidate| candidate == "--benchmark-pin-input-tps") diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs index 0f2c9dd96..cf7151aeb 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/report.rs @@ -497,8 +497,8 @@ fn render_warnings(out: &mut String, context: &ReportContext, entries: &[ReportE mod tests { use super::*; use crate::config::{ - ArrivalPatternConfig, BackendConfig, BackendProfile, RegistrationConfig, ServiceTimeConfig, - StargateConfig, TokenDistributionConfig, TrafficPatternConfig, UniformTrafficConfig, + ArrivalPatternConfig, BackendConfig, BackendProfile, ServiceTimeConfig, StargateConfig, + TokenDistributionConfig, TrafficPatternConfig, UniformTrafficConfig, }; use crate::score::{ BackendSummary, CacheSummary, FailureSummary, QueueAdmissionSummary, @@ -531,10 +531,7 @@ mod tests { ttft_jitter_ms: 0, decode_tokens_per_s: 1, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps: 1.0, + prefill_tokens_per_s: 1.0, }, }, }, diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs index 04392009a..cc277bd7e 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/runtime.rs @@ -53,7 +53,7 @@ impl BackendRuntimeSpec { decode_jitter_ms: profile.service_time_ms.decode_jitter_ms, ttft_ms: profile.service_time_ms.ttft_mean, ttft_jitter_ms: profile.service_time_ms.ttft_jitter_ms, - prefill_tokens_per_s: profile.service_time_ms.prefill_tokens_per_s.unwrap_or(0.0), + prefill_tokens_per_s: profile.service_time_ms.prefill_tokens_per_s, max_concurrent_requests, kv_cache_capacity_tokens, } @@ -68,7 +68,7 @@ pub(crate) struct PylonRuntimeSpec { pub(crate) inference_server_id: String, pub(crate) cluster_id: Option, pub(crate) profile_slug: String, - pub(crate) last_mean_input_tps: f64, + pub(crate) initial_input_tps: f64, } impl PylonRuntimeSpec { @@ -82,7 +82,7 @@ impl PylonRuntimeSpec { inference_server_id: backend_name(backend_index), cluster_id: config.backends.cluster_id_for_index(backend_index), profile_slug: slugify(&profile.name), - last_mean_input_tps: profile.registration.last_mean_input_tps, + initial_input_tps: profile.service_time_ms.prefill_tokens_per_s, } } @@ -135,8 +135,7 @@ mod tests { "decode_tokens_per_s": 50, "decode_jitter_ms": 2, "prefill_tokens_per_s": 123.0 - }, - "registration": { "last_mean_input_tps": 100.0 } + } } }, "traffic_pattern": { @@ -165,7 +164,7 @@ mod tests { } #[test] - fn pylon_runtime_spec_targets_shared_upstream_and_keeps_registration_identity() { + fn pylon_runtime_spec_targets_shared_upstream() { let config = config(); let spec = PylonRuntimeSpec::for_backend(&config, 1); @@ -175,6 +174,7 @@ mod tests { assert_eq!(spec.upstream_backend_name, "backend-0"); assert_eq!(spec.inference_server_id, "backend-1"); assert_eq!(spec.cluster_id.as_deref(), Some("cluster-0")); + assert_eq!(spec.initial_input_tps, 123.0); assert!(!spec.owns_upstream_backend()); } diff --git a/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs b/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs index 82e54f5e0..c68109123 100644 --- a/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs +++ b/src/libraries/rust/stargate/crates/stargate-bench/src/score.rs @@ -563,8 +563,8 @@ pub fn topology_for(backends: &BackendConfig) -> RoutingTopology { for index in 0..backends.count { let capacity = backends .profile_for_index(index) - .registration - .last_mean_input_tps; + .service_time_ms + .prefill_tokens_per_s; if capacity > 0.0 && capacity.is_finite() { let backend_id = format!("backend-{index}"); let cluster_id = backends.effective_cluster_id_for_index(index); @@ -660,9 +660,7 @@ fn percentile(values: &[u64], q: f64) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::config::{ - BackendProfile, BackendProfileGroup, RegistrationConfig, ServiceTimeConfig, - }; + use crate::config::{BackendProfile, BackendProfileGroup, ServiceTimeConfig}; fn result(backend: &str, ttft: u64, ttlt: u64) -> RequestResult { RequestResult { @@ -705,7 +703,7 @@ mod tests { default_tps: f64, groups: &[(usize, f64)], ) -> BackendConfig { - let profile = |last_mean_input_tps| BackendProfile { + let profile = |prefill_tokens_per_s| BackendProfile { name: "test".to_string(), weight: 1.0, max_concurrent_requests: None, @@ -715,10 +713,7 @@ mod tests { ttft_jitter_ms: 0, decode_tokens_per_s: 1, decode_jitter_ms: 0, - prefill_tokens_per_s: None, - }, - registration: RegistrationConfig { - last_mean_input_tps, + prefill_tokens_per_s, }, }; BackendConfig { diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs index fea4fe1c7..71fd649f1 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/config.rs @@ -70,7 +70,7 @@ pub enum LoadBalancerAlgorithm { #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Deserialize)] #[serde(rename_all = "kebab-case")] -pub enum PulsarRendezvousWeight { +pub(crate) enum PulsarRendezvousWeight { #[default] LastMeanInputTps, MaxInputTps, @@ -296,10 +296,6 @@ impl LoadBalancerAlgorithmConfig { self.settings.algorithm() } - pub fn pulsar_rendezvous_weight(&self) -> PulsarRendezvousWeight { - self.rendezvous_weight - } - pub fn requires_cache_affinity_key(&self) -> bool { self.request_policy.require_cache_affinity_key } 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 02de52bae..f986cb8d2 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/mod.rs @@ -44,12 +44,12 @@ 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, LoadBalancerRequestPolicy, LoadBalancerRoutingAlgorithmError, LoadBalancerSeedError, - MAX_POWER_OF_N_SAMPLE_COUNT, PowerOfNAlgorithmConfig, PulsarRendezvousWeight, - WaitAndWidenAlgorithmConfig, + MAX_POWER_OF_N_SAMPLE_COUNT, PowerOfNAlgorithmConfig, WaitAndWidenAlgorithmConfig, }; pub use factory::create_load_balancer_with_config; pub use request::{LoadBalancerCandidateChoice, LoadBalancerRequest}; 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 8bab727fc..4e1f3d826 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar.rs @@ -17,13 +17,15 @@ mod ranking; use std::sync::Arc; +use stargate_protocol::common::valid_last_mean_input_tps; + use super::{ LoadBalancer, LoadBalancerAlgorithmConfig, LoadBalancerCandidateChoice, LoadBalancerRequest, }; use crate::routing_state::RoutedClusterSnapshot; use ranking::{ PulsarRankingLookup, PulsarRankingStore, PulsarScorer, ScoredCandidate, - compare_ranked_candidate, has_valid_mean_input_tps, + compare_ranked_candidate, }; #[cfg(test)] @@ -57,7 +59,7 @@ impl LoadBalancer for PulsarLoadBalancer { candidates: &[RoutedClusterSnapshot], ) -> Option { let (lookup, cached_choice) = self.rankings.lookup_choice( - self.config.pulsar_rendezvous_weight(), + self.config.rendezvous_weight, request, candidates, |ranking| self.choose_from_ranked_indices(request, candidates, ranking), @@ -72,7 +74,7 @@ impl LoadBalancer for PulsarLoadBalancer { let choice = self.choose_from_ranked_indices(request, candidates, &ranking); self.rankings .insert_or_choose_existing( - self.config.pulsar_rendezvous_weight(), + self.config.rendezvous_weight, request, candidates, ranking, @@ -121,9 +123,9 @@ impl PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Vec { - ranking::pulsar_ranked_indices_with_weight( + ranking::pulsar_ranked_indices( self.config.seed(), - self.config.pulsar_rendezvous_weight(), + self.config.rendezvous_weight, request, candidates, ) @@ -134,11 +136,8 @@ impl PulsarLoadBalancer { request: &LoadBalancerRequest<'_>, candidates: &[RoutedClusterSnapshot], ) -> Option { - let mut scorer = PulsarScorer::new( - self.config.seed(), - request, - self.config.pulsar_rendezvous_weight(), - ); + 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() { @@ -196,11 +195,8 @@ impl PulsarLoadBalancer { chosen: &RoutedClusterSnapshot, chosen_score: f64, ) -> (usize, bool) { - let mut scorer = PulsarScorer::new( - self.config.seed(), - request, - self.config.pulsar_rendezvous_weight(), - ); + 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 { @@ -219,7 +215,7 @@ impl PulsarLoadBalancer { #[cfg(test)] pub(super) fn weight(&self, candidate: &RoutedClusterSnapshot) -> Option { - ranking::pulsar_weight(self.config.pulsar_rendezvous_weight(), candidate) + ranking::pulsar_weight(self.config.rendezvous_weight, candidate) } pub(super) fn feasibility( @@ -236,7 +232,7 @@ pub(super) fn input_work_admission_candidate( request: &LoadBalancerRequest<'_>, candidate: &RoutedClusterSnapshot, ) -> bool { - has_valid_mean_input_tps(candidate) + valid_last_mean_input_tps(candidate.stats.last_mean_input_tps) && candidate_feasibility(config, request, candidate).is_eligible() } diff --git a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs index be03082b5..a68bb11ee 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/pulsar/ranking.rs @@ -18,7 +18,6 @@ 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::{ @@ -282,16 +281,7 @@ impl PulsarScorer { } } -#[cfg(test)] pub(in crate::load_balancer) fn pulsar_ranked_indices( - seed: Option<&str>, - request: &LoadBalancerRequest<'_>, - candidates: &[RoutedClusterSnapshot], -) -> Vec { - pulsar_ranked_indices_with_weight(seed, PulsarRendezvousWeight::default(), request, candidates) -} - -pub(super) fn pulsar_ranked_indices_with_weight( seed: Option<&str>, selector: PulsarRendezvousWeight, request: &LoadBalancerRequest<'_>, @@ -333,10 +323,6 @@ pub(super) fn pulsar_weight( (value > 0.0 && value.is_finite()).then_some(value) } -pub(super) fn has_valid_mean_input_tps(candidate: &RoutedClusterSnapshot) -> bool { - valid_last_mean_input_tps(candidate.stats.last_mean_input_tps) -} - const PULSAR_HASH_VERSION: u8 = 1; #[cfg(test)] 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 6ce5872be..3164f357c 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/load_balancer/tests.rs @@ -968,7 +968,7 @@ fn pulsar_rendezvous_weight_parses_for_both_pulsar_variants() { r#"{{"algorithm":"{algorithm}","rendezvous_weight":"max-input-tps"}}"# )); assert_eq!( - config.pulsar_rendezvous_weight(), + config.rendezvous_weight, PulsarRendezvousWeight::MaxInputTps ); } @@ -978,7 +978,7 @@ fn pulsar_rendezvous_weight_parses_for_both_pulsar_variants() { LoadBalancerAlgorithm::PulsarWaitAndWiden, ] { assert_eq!( - LoadBalancerAlgorithmConfig::from(algorithm).pulsar_rendezvous_weight(), + LoadBalancerAlgorithmConfig::from(algorithm).rendezvous_weight, PulsarRendezvousWeight::LastMeanInputTps ); } @@ -2452,7 +2452,12 @@ fn pulsar_ranking_returns_candidate_slice_indices() { work_candidate("fast", 5, 10.0, 0), ]; - let ranking = pulsar_ranked_indices(config.seed(), &request, &candidates); + let ranking = pulsar_ranked_indices( + config.seed(), + PulsarRendezvousWeight::default(), + &request, + &candidates, + ); assert_eq!(ranking.len(), 2); assert_eq!( diff --git a/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs b/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs index 6a72031a8..06793745d 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs @@ -1269,10 +1269,7 @@ pub async fn start_and_register_backend_with_bringup( ModelLifecycleConfig { upstream_http_base_url: upstream_http_base_url.clone(), source: ModelSource::Static(BTreeSet::from([model.to_string()])), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1.0, - pin: true, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1.0 }, bringup, health_paths: pylon_lib::UpstreamHealthPaths::default(), startup_health_wait: std::time::Duration::ZERO, diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs index d685908a2..082a7bff0 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/lifecycle.rs @@ -1264,10 +1264,7 @@ impl LifecycleBackendOptions<'_> { ModelLifecycleConfig { upstream_http_base_url: upstream_http_base_url.clone(), source: ModelSource::Static(model_ids.into_iter().collect::>()), - initialization: ModelInitialization::ConfiguredInputTps { - input_tps: 1.0, - pin: true, - }, + initialization: ModelInitialization::ConfiguredInputTps { input_tps: 1.0 }, bringup, health_paths: pylon_lib::UpstreamHealthPaths::default(), startup_health_wait: std::time::Duration::ZERO,