From 4cb4ac85a83f52ed84baf691907ae75e0001327b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Mon, 5 Oct 2026 22:42:46 +0300 Subject: [PATCH 01/11] fix(stargate): expire reservations after backend RTT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- .../metrics/llm-request-router/metrics.md | 2 + .../crates/stargate/src/http_proxy.rs | 4 + .../crates/stargate/src/http_proxy/attempt.rs | 36 ++++- .../rust/stargate/crates/stargate/src/main.rs | 36 +++++ .../crates/stargate/src/main/startup.rs | 6 + .../stargate/crates/stargate/src/metrics.rs | 30 ++++ .../src/routing_state/cluster_snapshots.rs | 86 +++++++++-- .../src/routing_state/reservations.rs | 141 +++++++++++++++++- .../stargate/src/routing_state/snapshots.rs | 30 ++++ .../stargate/src/routing_state/tests.rs | 127 +++++++++++++++- .../stargate/crates/stargate/src/runtime.rs | 7 + .../crates/stargate/tests/common/mod.rs | 2 + .../docs/load-balancer-configuration.md | 5 +- .../stargate/docs/runtime-stats-interface.md | 11 ++ 14 files changed, 502 insertions(+), 21 deletions(-) diff --git a/docs/observability/metrics/llm-request-router/metrics.md b/docs/observability/metrics/llm-request-router/metrics.md index 93185db2c6..572049806c 100644 --- a/docs/observability/metrics/llm-request-router/metrics.md +++ b/docs/observability/metrics/llm-request-router/metrics.md @@ -38,3 +38,5 @@ unbounded request fields as metric labels. | `stargate_proxy_duration_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model`, `inference_server_id` | Time to first byte from upstream in seconds. | | `stargate_routing_duration_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | Load-balancer decision time in seconds. | | `stargate_active_inference_servers` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | Currently routable inference servers for a routing target. | +| `stargate_routing_reservations_active` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | In-flight local routing reservations that have not expired or been explicitly released. | +| `stargate_routing_reservation_ttl_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | RTT-based reservation TTL in seconds after applying the configured bounds. | diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs index bf6b4354f0..4b285871c4 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs @@ -120,6 +120,8 @@ pub struct ProxyAppState { pub readiness: ReadinessState, pub lb_router: Arc, pub metrics: Arc, + pub routing_reservation_ttl_min: std::time::Duration, + pub routing_reservation_ttl_max: std::time::Duration, pub retry: ProxyRetryConfig, pub debug_config: DebugConfig, } @@ -286,6 +288,8 @@ mod test_support { .expect("load balancer should initialize"), ), metrics, + routing_reservation_ttl_min: std::time::Duration::from_millis(1), + routing_reservation_ttl_max: std::time::Duration::from_millis(1000), retry: ProxyRetryConfig::default(), debug_config: DebugConfig::default(), } diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs index 002ecce962..b209443d7c 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs @@ -43,6 +43,14 @@ use super::upstream::{ const HEADER_CHOSEN_INFERENCE_SERVER_URL: &str = "x-inference-server-url"; const HEADER_CHOSEN_CLUSTER_ID: &str = "x-stargate-cluster-id"; +fn clamp_routing_reservation_ttl( + rtt: std::time::Duration, + min: std::time::Duration, + max: std::time::Duration, +) -> std::time::Duration { + rtt.clamp(min, max) +} + #[derive(Default)] pub(super) struct ProxyAttemptCounters { pub(super) attempt: u32, @@ -112,10 +120,20 @@ impl ProxyRequestRun<'_> { } } - let reservation: Option = selected.cluster.reserve_backend( + let reservation: Option = selected.cluster.reserve_backend_with_ttl( &chosen.registration, self.request.request_inputs.input_tokens, self.request.request_inputs.priority, + clamp_routing_reservation_ttl( + chosen.rtt, + self.app.routing_reservation_ttl_min, + self.app.routing_reservation_ttl_max, + ), + self.app.metrics.clone(), + &crate::routing_state::RoutingTargetKey::new( + self.routing_key().map(str::to_owned), + self.model_id(), + ), ); record_proxy_attempt_start(self, selected, chosen); @@ -545,6 +563,22 @@ mod tests { RegistrationIdentity, RoutingTargetKey, test_registration_generation, }; + #[test] + fn reservation_ttl_is_clamped_to_the_configured_bounds() { + let min = Duration::from_millis(1); + let max = Duration::from_millis(1000); + + assert_eq!(clamp_routing_reservation_ttl(Duration::ZERO, min, max), min); + assert_eq!( + clamp_routing_reservation_ttl(Duration::from_millis(60), min, max), + Duration::from_millis(60) + ); + assert_eq!( + clamp_routing_reservation_ttl(Duration::from_secs(5), min, max), + max + ); + } + fn finish_response( disposition: FinalRetryDisposition, status: StatusCode, diff --git a/src/libraries/rust/stargate/crates/stargate/src/main.rs b/src/libraries/rust/stargate/crates/stargate/src/main.rs index 826f4097f3..00996aa4de 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/main.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/main.rs @@ -161,6 +161,12 @@ struct Args { value_name = "MS" )] registration_update_max_idle_timeout_ms: u64, + /// Minimum backend RTT-based routing reservation lifetime. + #[arg(long, default_value_t = 1, value_name = "MS")] + routing_reservation_ttl_min_ms: u64, + /// Maximum backend RTT-based routing reservation lifetime. + #[arg(long, default_value_t = 1000, value_name = "MS")] + routing_reservation_ttl_max_ms: u64, /// Grace period for shutdown tasks after Stargate starts draining. #[arg(long, default_value_t = 30000, value_name = "MS")] shutdown_drain_timeout_ms: u64, @@ -680,6 +686,36 @@ mod tests { ); } + #[test] + fn routing_reservation_ttl_defaults_and_bounds_are_validated() { + let args = parse_args(""); + assert_eq!(args.routing_reservation_ttl_min_ms, 1); + assert_eq!(args.routing_reservation_ttl_max_ms, 1000); + let args = + parse_args("--routing-reservation-ttl-min-ms 20 --routing-reservation-ttl-max-ms 200"); + let config = runtime_config_from_args(&args, proxy_transport(&args)) + .expect("valid TTL bounds should reach runtime config"); + assert_eq!( + config.routing_reservation_ttl_min, + Duration::from_millis(20) + ); + assert_eq!( + config.routing_reservation_ttl_max, + Duration::from_millis(200) + ); + + let args = + parse_args("--routing-reservation-ttl-min-ms 20 --routing-reservation-ttl-max-ms 10"); + let error = runtime_config_from_args(&args, proxy_transport(&args)) + .err() + .expect("minimum TTL above maximum should fail startup config"); + assert!( + error + .to_string() + .contains("minimum must not exceed maximum") + ); + } + #[test] fn direct_quic_connections_default_and_override_parse() { assert_eq!(parse_args("").direct_quic_connections, 1); diff --git a/src/libraries/rust/stargate/crates/stargate/src/main/startup.rs b/src/libraries/rust/stargate/crates/stargate/src/main/startup.rs index 2176d54ed8..5e519f361b 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/main/startup.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/main/startup.rs @@ -173,6 +173,10 @@ pub(super) fn runtime_config_from_args( args: &Args, proxy_transport: ProxyTransportConfig, ) -> Result { + ensure!( + args.routing_reservation_ttl_min_ms <= args.routing_reservation_ttl_max_ms, + "routing reservation TTL minimum must not exceed maximum" + ); let millis = Duration::from_millis; Ok(StargateRuntimeConfig { stargate_id: args.stargate_id.clone(), @@ -188,6 +192,8 @@ pub(super) fn runtime_config_from_args( watch_heartbeat_interval: millis(args.watch_heartbeat_ms), registration_update_idle_timeout: millis(args.registration_update_idle_timeout_ms), registration_update_max_idle_timeout: millis(args.registration_update_max_idle_timeout_ms), + routing_reservation_ttl_min: millis(args.routing_reservation_ttl_min_ms), + routing_reservation_ttl_max: millis(args.routing_reservation_ttl_max_ms), proxy_transport, lb_config_path: args.lb_config_path.clone(), metrics_prefix: args.metrics_prefix.clone(), diff --git a/src/libraries/rust/stargate/crates/stargate/src/metrics.rs b/src/libraries/rust/stargate/crates/stargate/src/metrics.rs index 0e2e0d817a..c1d71b9339 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/metrics.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/metrics.rs @@ -111,9 +111,11 @@ define_stargate_metrics! { proxy_replay_buffer_bytes("proxy_replay_buffer_bytes", "Bytes currently retained for proxied request body replay", ["model"], [0.0, 1024.0, 4096.0, 16_384.0, 65_536.0, 262_144.0, 1_048_576.0, 4_194_304.0, 16_777_216.0, 67_108_864.0]); proxy_duration_seconds("proxy_duration_seconds", "Time to first byte from upstream", ["routing_key", "model", "inference_server_id"], [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0]); routing_duration_seconds("routing_duration_seconds", "Time spent selecting an inference server", ["routing_key", "model"], [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0]); + routing_reservation_ttl_seconds("routing_reservation_ttl_seconds", "Configured lifetime of routing reservations", ["routing_key", "model"], [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0]); } gauges { active_inference_servers("active_inference_servers", "Active inference servers available for a routing target", ["routing_key", "model"]); + routing_reservations_active("routing_reservations_active", "Active in-flight routing reservations", ["routing_key", "model"]); tls_certificate_expiry_seconds("tls_certificate_expiry_seconds", "Unix timestamp when the active TLS certificate expires", ["material_type"]); } } @@ -200,6 +202,19 @@ impl StargateMetrics { Histogram, proxy_replay_buffer_bytes(model: &str) => [model]; Histogram, proxy_duration_seconds(routing_key: Option<&str>, model: &str, inference_server_id: &str) => [routing_key.unwrap_or(""), model, inference_server_id]; Histogram, routing_duration_seconds(routing_key: Option<&str>, model: &str) => [routing_key.unwrap_or(""), model]; + Histogram, routing_reservation_ttl_seconds(routing_key: Option<&str>, model: &str) => [routing_key.unwrap_or(""), model]; + } + + pub(crate) fn inc_routing_reservations_active(&self, routing_key: Option<&str>, model: &str) { + self.routing_reservations_active + .with_label_values(&[routing_key.unwrap_or(""), model]) + .inc(); + } + + pub(crate) fn dec_routing_reservations_active(&self, routing_key: Option<&str>, model: &str) { + self.routing_reservations_active + .with_label_values(&[routing_key.unwrap_or(""), model]) + .dec(); } #[inline] @@ -322,6 +337,21 @@ mod tests { ); } + #[test] + fn routing_reservation_metrics_use_routing_target_labels() { + let metrics = StargateMetrics::new().expect("metrics should initialize"); + metrics.inc_routing_reservations_active(Some("routing-a"), "model-a"); + metrics + .routing_reservation_ttl_seconds(Some("routing-a"), "model-a") + .observe(0.06); + + let body = metrics.gather_text().expect("metrics should encode"); + assert!(body.contains( + "stargate_routing_reservations_active{model=\"model-a\",routing_key=\"routing-a\"} 1" + )); + assert!(body.contains("stargate_routing_reservation_ttl_seconds_count{model=\"model-a\",routing_key=\"routing-a\"} 1")); + } + #[test] fn default_metrics_prefix_keeps_stargate_metric_names() { let metrics = StargateMetrics::new().expect("metrics should initialize"); diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/cluster_snapshots.rs index 0089f65286..1355707cf5 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 @@ -16,12 +16,15 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use parking_lot::Mutex; use stargate_proto::pb::ModelStats; use stargate_protocol::common::valid_last_mean_input_tps; +use crate::metrics::StargateMetrics; + +use super::keys::RoutingTargetKey; use super::registration::{RegistrationClusterGeneration, RegistrationGeneration}; use super::reservations::{ PendingClusterReservation, RoutingReservation, apply_pending_cluster_reservations, @@ -249,7 +252,13 @@ impl ClusterRoutingGeneration { } fn refresh_snapshot(&mut self, updated_backend_id: Option<&str>) { + let now = Instant::now(); let Some((backend_stats, rtt, active_backend_count)) = self.backend_aggregate() else { + if let Some(state) = self.snapshot_state.take() { + for pending in state.pending_cluster_reservations { + pending.deactivate(); + } + } self.snapshot_state = None; return; }; @@ -279,13 +288,17 @@ impl ClusterRoutingGeneration { .map(|state| state.pending_cluster_reservations) .unwrap_or_default(); pending_cluster_reservations.retain(|pending| { - pending.is_active() - && backend_index(&self.backends, &pending.inference_server_id).is_ok() + let backend_active = backend_index(&self.backends, &pending.inference_server_id) + .is_ok_and(|index| { + pending.registration.as_ref().is_none_or(|registration| { + Arc::ptr_eq(&self.backends[index].registration, registration) + }) + }); + if !backend_active { + pending.deactivate(); + } + pending.is_active_at(now) && backend_active }); - if let Some(updated_backend_id) = updated_backend_id { - pending_cluster_reservations - .retain(|pending| pending.inference_server_id != updated_backend_id); - } let mut stats = backend_stats; set_shared_engine_stats(&mut stats, &source_backend.stats); let base_snapshot = RoutedClusterSnapshot { @@ -305,10 +318,14 @@ impl ClusterRoutingGeneration { } fn routing_snapshot(&mut self) -> Option { + self.routing_snapshot_at(Instant::now()) + } + + pub(super) fn routing_snapshot_at(&mut self, now: Instant) -> Option { let snapshot_state = self.snapshot_state.as_mut()?; snapshot_state .pending_cluster_reservations - .retain(|pending| pending.is_active()); + .retain(|pending| pending.is_active_at(now)); let mut snapshot = snapshot_state.base_snapshot.clone(); apply_pending_cluster_reservations( &mut snapshot.stats, @@ -322,15 +339,21 @@ impl ClusterRoutingGeneration { registration: &Arc, input_tokens: u64, priority: u32, + ttl: Duration, + metrics: Option>, + target: &RoutingTargetKey, ) -> Option { if !self.contains_registration(registration) { return None; } let snapshot_state = self.snapshot_state.as_mut()?; let pending = PendingClusterReservation::new( - registration.inference_server_id().to_string(), + registration.clone(), input_tokens, priority, + ttl, + metrics, + target, ); snapshot_state .pending_cluster_reservations @@ -387,6 +410,20 @@ impl RoutedClusterState { self.generation.lock().routing_snapshot() } + #[cfg(test)] + pub(super) fn routing_snapshot_at(&self, now: Instant) -> Option { + self.generation.lock().routing_snapshot_at(now) + } + + #[cfg(test)] + pub(super) fn pending_reservation_count(&self) -> usize { + self.generation + .lock() + .snapshot_state + .as_ref() + .map_or(0, |state| state.pending_cluster_reservations.len()) + } + pub(super) fn backend_snapshot_values(&self) -> Vec { let backends = self.generation.lock().backends.clone(); backends @@ -421,15 +458,40 @@ impl RoutedClusterState { eligible.nth(index).map(Arc::clone) } + #[cfg(test)] pub(super) fn reserve_backend( &self, registration: &Arc, input_tokens: u64, priority: u32, ) -> Option { - self.generation - .lock() - .reserve_backend(registration, input_tokens, priority) + self.generation.lock().reserve_backend( + registration, + input_tokens, + priority, + Duration::from_secs(1), + None, + &RoutingTargetKey::new(None, ""), + ) + } + + pub(super) fn reserve_backend_with_ttl( + &self, + registration: &Arc, + input_tokens: u64, + priority: u32, + ttl: Duration, + metrics: Arc, + target: &RoutingTargetKey, + ) -> Option { + self.generation.lock().reserve_backend( + registration, + input_tokens, + priority, + ttl, + Some(metrics), + target, + ) } #[cfg(test)] diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs index 796b04457c..08f1d4ae39 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs @@ -14,30 +14,104 @@ // limitations under the License. use super::*; +use crate::metrics::StargateMetrics; use stargate_protocol::common::queue_time_delta_ms; use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; #[derive(Debug)] pub(super) struct PendingClusterReservation { pub(super) inference_server_id: String, + pub(super) registration: Option>, pub(super) input_tokens: u64, pub(super) priority: u32, + expires_at: Instant, active: AtomicBool, + metrics: Option>, + routing_key: Option, + model: String, } impl PendingClusterReservation { - pub(super) fn new(inference_server_id: String, input_tokens: u64, priority: u32) -> Arc { + pub(super) fn new( + registration: Arc, + input_tokens: u64, + priority: u32, + ttl: Duration, + metrics: Option>, + target: &RoutingTargetKey, + ) -> Arc { + Self::new_at( + Some(registration), + input_tokens, + priority, + Instant::now(), + ttl, + metrics, + target, + ) + } + + fn new_at( + registration: Option>, + input_tokens: u64, + priority: u32, + now: Instant, + ttl: Duration, + metrics: Option>, + target: &RoutingTargetKey, + ) -> Arc { + let inference_server_id = registration + .as_ref() + .map_or_else(String::new, |registration| { + registration.inference_server_id().to_string() + }); + if let Some(metrics) = &metrics { + metrics + .inc_routing_reservations_active(target.routing_key.as_deref(), &target.model_id); + metrics + .routing_reservation_ttl_seconds(target.routing_key.as_deref(), &target.model_id) + .observe(ttl.as_secs_f64()); + } Arc::new(Self { inference_server_id, + registration, input_tokens, priority, + expires_at: now + ttl, active: AtomicBool::new(true), + metrics, + routing_key: target.routing_key.clone(), + model: target.model_id.clone(), }) } pub(super) fn is_active(&self) -> bool { self.active.load(Ordering::Acquire) } + + #[cfg(test)] + pub(super) fn expires_at(&self) -> Instant { + self.expires_at + } + + pub(super) fn is_active_at(&self, now: Instant) -> bool { + if now >= self.expires_at { + self.deactivate(); + } + self.is_active() + } + + pub(super) fn deactivate(&self) { + if self + .active + .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + && let Some(metrics) = &self.metrics + { + metrics.dec_routing_reservations_active(self.routing_key.as_deref(), &self.model); + } + } } // Cancellation is explicit because a successful attempt remains pending until its heartbeat. @@ -46,7 +120,7 @@ pub(crate) struct RoutingReservation(pub(super) Arc); impl RoutingReservation { pub(crate) fn release(self) { - self.0.active.store(false, Ordering::Release); + self.0.deactivate(); } } @@ -101,6 +175,8 @@ pub(super) fn apply_pending_cluster_reservations( #[cfg(test)] mod tests { + use prometheus::Encoder; + use std::collections::HashMap; use super::*; @@ -120,7 +196,16 @@ mod tests { crate::queue_estimate::queue_time_estimate_ms_for_priority(stats, 0) }; assert_eq!(estimate(&base), Some(0)); - let pending = PendingClusterReservation::new("backend".to_string(), 8_000, 0); + let target = RoutingTargetKey::new(None, "model"); + let pending = PendingClusterReservation::new_at( + None, + 8_000, + 0, + Instant::now(), + Duration::from_millis(1), + None, + &target, + ); let reservations = [pending.clone()]; let mut full = base.clone(); apply_pending_cluster_reservations(&mut full, &reservations); @@ -133,4 +218,54 @@ mod tests { assert_eq!(estimate(&recovered), Some(0)); } } + + #[test] + fn reservation_expires_at_its_fixed_ttl_boundary() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let dispatched_at = Instant::now(); + let pending = PendingClusterReservation::new_at( + None, + 10, + 0, + dispatched_at, + Duration::from_millis(60), + None, + &target, + ); + + assert!(pending.is_active_at(dispatched_at + Duration::from_millis(59))); + assert!(!pending.is_active_at(dispatched_at + Duration::from_millis(60))); + assert!(!pending.is_active()); + } + + #[test] + fn reservation_metrics_decrement_once_on_release() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let metrics = StargateMetrics::new().expect("metrics should initialize"); + let pending = PendingClusterReservation::new_at( + None, + 1, + 0, + Instant::now(), + Duration::from_millis(60), + Some(metrics.clone()), + &target, + ); + let active_gauge = || { + let mut encoded = Vec::new(); + prometheus::TextEncoder::new() + .encode(&metrics.registry().gather(), &mut encoded) + .expect("metrics should encode"); + String::from_utf8(encoded).expect("metrics should be UTF-8") + }; + + assert!(active_gauge().contains( + "stargate_routing_reservations_active{model=\"model\",routing_key=\"routing\"} 1" + )); + pending.deactivate(); + pending.deactivate(); + assert!(active_gauge().contains( + "stargate_routing_reservations_active{model=\"model\",routing_key=\"routing\"} 0" + )); + } } diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs index 788296258c..682a26375f 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs @@ -129,6 +129,7 @@ impl SelectedRoutedCluster { self.owner.select_backend(failed_backend_ids) } + #[cfg(test)] pub(crate) fn reserve_backend( &self, registration: &Arc, @@ -138,6 +139,35 @@ impl SelectedRoutedCluster { self.owner .reserve_backend(registration, input_tokens, priority) } + + pub(crate) fn reserve_backend_with_ttl( + &self, + registration: &Arc, + input_tokens: u64, + priority: u32, + ttl: Duration, + metrics: Arc, + target: &RoutingTargetKey, + ) -> Option { + self.owner.reserve_backend_with_ttl( + registration, + input_tokens, + priority, + ttl, + metrics, + target, + ) + } + + #[cfg(test)] + pub(super) fn routing_snapshot_at(&self, now: Instant) -> Option { + self.owner.routing_snapshot_at(now) + } + + #[cfg(test)] + pub(super) fn pending_reservation_count(&self) -> usize { + self.owner.pending_reservation_count() + } } #[derive(Debug)] diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/tests.rs index c5fe310644..4a37cdfdf1 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 @@ -866,7 +866,7 @@ async fn active_registration_keeps_connection_rtt_in_snapshot() { } #[tokio::test] -async fn reservation_updates_local_snapshot_until_next_registration_update() { +async fn reservation_survives_registration_update_until_expiry() { let scenario = RegistrationScenario::new(Some("rk-res")); let running = scenario.start("inst-res", 8888); let mut stats = priority_stats(100.0, [(4, 5)]); @@ -887,10 +887,131 @@ async fn reservation_updates_local_snapshot_until_next_registration_update() { scenario.publish_connected(&running, &update).await; let candidates = scenario.clusters("model-res").await; - assert_queue_stats(&candidates[0].stats, 0, 0, 0, 0, 4, 5); + assert_queue_stats(&candidates[0].stats, 1, 1, 37, 37, 4, 375); let clusters = scenario.clusters("model-res").await; - assert_queue_stats(&clusters[0].stats, 0, 0, 0, 0, 4, 5); + assert_queue_stats(&clusters[0].stats, 1, 1, 37, 37, 4, 375); +} + +#[tokio::test] +async fn reservation_is_in_snapshot_before_expiry_and_pruned_at_expiry() { + let scenario = RegistrationScenario::new(Some("rk-res-expiry")); + let running = scenario.start("inst-res-expiry", 8888); + let update = scenario.update( + &running, + "model-res-expiry", + Active, + priority_stats(100.0, [(4, 5)]), + ); + scenario.publish_connected(&running, &update).await; + + let selected_cluster = scenario.selected_cluster("model-res-expiry").await; + let target = scenario.target("model-res-expiry"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let reservation = selected_cluster + .reserve_backend_with_ttl( + &running.generation(), + 37, + 4, + Duration::from_millis(60), + metrics, + &target, + ) + .expect("active backend should accept reservation"); + let expires_at = reservation.0.expires_at(); + + scenario.publish_connected(&running, &update).await; + assert_eq!( + selected_cluster + .routing_snapshot_at(expires_at - Duration::from_millis(1)) + .expect("cluster should remain routable") + .stats + .queue_size, + 1, + "a registration update must not clear the reservation" + ); + assert_eq!( + selected_cluster + .routing_snapshot_at(expires_at) + .expect("cluster should remain routable") + .stats + .queue_size, + 0, + "reservation should not count at its expiry boundary" + ); + assert_eq!(selected_cluster.pending_reservation_count(), 0); + + reservation.release(); +} + +#[tokio::test] +async fn snapshot_read_prunes_ten_thousand_idle_expired_reservations() { + let scenario = RegistrationScenario::new(Some("rk-res-prune")); + let running = scenario.start("inst-res-prune", 8888); + let update = scenario.update( + &running, + "model-res-prune", + Active, + priority_stats(100.0, [(4, 5)]), + ); + scenario.publish_connected(&running, &update).await; + + let selected_cluster = scenario.selected_cluster("model-res-prune").await; + let target = scenario.target("model-res-prune"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let mut last_expiry = Instant::now(); + for _ in 0..10_000 { + let reservation = selected_cluster + .reserve_backend_with_ttl( + &running.generation(), + 1, + 0, + Duration::from_millis(1), + metrics.clone(), + &target, + ) + .expect("active backend should accept reservation"); + last_expiry = last_expiry.max(reservation.0.expires_at()); + } + + assert_eq!(selected_cluster.pending_reservation_count(), 10_000); + selected_cluster + .routing_snapshot_at(last_expiry + Duration::from_millis(1)) + .expect("cluster should remain routable"); + assert_eq!(selected_cluster.pending_reservation_count(), 0); +} + +#[tokio::test] +async fn removing_backend_drops_its_reservations_immediately() { + let scenario = RegistrationScenario::new(Some("rk-res-remove")); + let running = scenario.start("inst-res-remove", 8888); + let update = scenario.update( + &running, + "model-res-remove", + Active, + priority_stats(100.0, [(4, 5)]), + ); + scenario.publish_connected(&running, &update).await; + + let selected_cluster = scenario.selected_cluster("model-res-remove").await; + let target = scenario.target("model-res-remove"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let reservation = selected_cluster + .reserve_backend_with_ttl( + &running.generation(), + 37, + 4, + Duration::from_secs(1), + metrics, + &target, + ) + .expect("active backend should accept reservation"); + assert_eq!(selected_cluster.pending_reservation_count(), 1); + + scenario.state.end_registration(running).await; + + assert_eq!(selected_cluster.pending_reservation_count(), 0); + reservation.release(); } #[tokio::test] diff --git a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs index 5e8785e6c4..163fb63987 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs @@ -74,6 +74,9 @@ pub struct StargateRuntimeConfig { pub registration_update_idle_timeout: Duration, /// Registration idle-time hard cap and fallback without heartbeat hints; zero disables enforcement. pub registration_update_max_idle_timeout: Duration, + /// Lower and upper bounds for backend RTT-based routing reservation expiry. + pub routing_reservation_ttl_min: Duration, + pub routing_reservation_ttl_max: Duration, /// QUIC, TLS, tunnel-protocol, and retry configuration for backend request forwarding. pub proxy_transport: ProxyTransportConfig, /// Optional JSON load-balancer config path. @@ -373,6 +376,8 @@ impl StargateRuntime { quic_proxy, lb_router, metrics: metrics.clone(), + routing_reservation_ttl_min: self.config.routing_reservation_ttl_min, + routing_reservation_ttl_max: self.config.routing_reservation_ttl_max, retry: self.config.proxy_transport.retry.clone(), debug_config: DebugConfig { stargate_id: self.config.stargate_id.clone(), @@ -621,6 +626,8 @@ mod tests { crate::control_plane::DEFAULT_REGISTRATION_UPDATE_IDLE_TIMEOUT, registration_update_max_idle_timeout: crate::control_plane::DEFAULT_REGISTRATION_UPDATE_MAX_IDLE_TIMEOUT, + routing_reservation_ttl_min: Duration::from_millis(1), + routing_reservation_ttl_max: Duration::from_millis(1000), proxy_transport: ProxyTransportConfig { quic: crate::tunnel::QuicTunnelConfig { connect_timeout: Duration::from_secs(5), 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 1cb0883fb4..42172c6374 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/common/mod.rs @@ -548,6 +548,8 @@ pub fn base_config( stargate::registration::DEFAULT_REGISTRATION_UPDATE_IDLE_TIMEOUT, registration_update_max_idle_timeout: stargate::registration::DEFAULT_REGISTRATION_UPDATE_MAX_IDLE_TIMEOUT, + routing_reservation_ttl_min: Duration::from_millis(1), + routing_reservation_ttl_max: Duration::from_millis(1000), proxy_transport: ProxyTransportConfig { quic: QuicTunnelConfig { connect_timeout: Duration::from_secs(5), diff --git a/src/libraries/rust/stargate/docs/load-balancer-configuration.md b/src/libraries/rust/stargate/docs/load-balancer-configuration.md index 0e37854f6d..fcd1abc2e4 100644 --- a/src/libraries/rust/stargate/docs/load-balancer-configuration.md +++ b/src/libraries/rust/stargate/docs/load-balancer-configuration.md @@ -151,8 +151,9 @@ forwarded health RTT + queue delay + request prefill time Queue delay is zero when the reported engine concurrency limit is positive and the active request count is below that limit. Active requests include prefill, -decode, and pending routing reservations. The new request's own prefill time -still contributes to TTFT. +decode, and pending routing reservations. Reservations expire after the selected +backend's health-check RTT, not on its next registration update. The new +request's own prefill time still contributes to TTFT. At or above the limit, or when the limit is unknown, queue delay uses the backend's priority-aware queue estimate when present. Otherwise it divides diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index fed5e21f76..aa6e497dd5 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -40,6 +40,17 @@ Stargate routes on the most recent update plus its own pending reservations. Change-driven updates keep that view within the coalescing window and network delay of the backend's actual load, instead of up to one heartbeat behind. +Each dispatched request adds a local reservation until one health-check RTT +for the selected backend has elapsed. The TTL uses that backend's RTT (not the +cluster mean), fixed at dispatch, and is bounded by +`--routing-reservation-ttl-min-ms` (default `1`) and +`--routing-reservation-ttl-max-ms` (default `1000`). Registration updates do not +clear reservations. If an update includes the request before its reservation +expires, routing temporarily counts it in both reported stats and the local +reservation. Deploy Pylons with change-driven stats updates before enabling +this Stargate behavior; the RTT-based expiry can otherwise leave a short gap +until the next heartbeat. + ## Stream Events Each non-empty line is JSON with `v: 1` and `type`. From ffc97bbac390286ba3927090b5a69ebaf120908e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Fri, 9 Oct 2026 21:48:53 +0300 Subject: [PATCH 02/11] fix(stargate): clean up retired reservations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- .../crates/stargate/src/http_proxy.rs | 9 +- .../crates/stargate/src/http_proxy/attempt.rs | 108 +++++++++++++++++- .../crates/stargate/src/http_proxy/retry.rs | 4 + .../crates/stargate/src/http_proxy/run.rs | 86 +++++++++++++- .../src/routing_state/reservations.rs | 8 +- .../stargate/src/routing_state/tests.rs | 88 ++++++++++++++ .../stargate/crates/stargate/src/runtime.rs | 8 +- .../stargate/tests/suite/health_lifecycle.rs | 7 +- .../stargate/docs/runtime-stats-interface.md | 4 +- 9 files changed, 305 insertions(+), 17 deletions(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs index 4b285871c4..1b3d18bcd0 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs @@ -249,7 +249,9 @@ mod test_support { use crate::routing_state::StargateState; use crate::tunnel::{QuicHttpProxy, QuicTunnelConfig}; - use super::{DebugConfig, ProxyAppState, ProxyRetryConfig, ProxyTrafficState, ReadinessState, readyz}; + use super::{ + DebugConfig, ProxyAppState, ProxyRetryConfig, ProxyTrafficState, ReadinessState, readyz, + }; pub(super) fn test_proxy_app_state() -> ProxyAppState { test_proxy_app_state_with_lb_config(LoadBalancerConfig::default()) @@ -312,7 +314,10 @@ mod test_support { let ready_token = CancellationToken::new(); app.readiness = ReadinessState::warming_up(ready_token.clone()); - assert_eq!(readyz(State(app.clone())).await, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + readyz(State(app.clone())).await, + StatusCode::SERVICE_UNAVAILABLE + ); ready_token.cancel(); diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs index b209443d7c..fbcbf567fe 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs @@ -550,6 +550,7 @@ fn upstream_status(result: &Result) -> St mod tests { use super::*; use axum::http::{HeaderMap, Method}; + use prometheus::Encoder; use stargate_proto::pb::{InferenceServerStatus, ModelStats}; use stargate_protocol::tunnel_contract::HEADER_STARGATE_RETRY_AFTER_MS; use std::time::Duration; @@ -559,8 +560,9 @@ mod tests { use super::super::retry::ReplayableRequestBody; use super::super::run::PreparedProxyRequest; use super::super::test_support::test_proxy_app_state; + use crate::load_balancer::LoadBalancerCandidateChoice; use crate::routing_state::{ - RegistrationIdentity, RoutingTargetKey, test_registration_generation, + RegistrationIdentity, RoutingTargetKey, StargateState, test_registration_generation, }; #[test] @@ -579,6 +581,110 @@ mod tests { ); } + #[tokio::test] + async fn proxy_reservation_ttl_uses_the_selected_backends_rtt() { + let mut app = test_proxy_app_state(); + app.state = Arc::new(StargateState::new_with_metrics(app.metrics.clone())); + let target = RoutingTargetKey::new(Some("tenant-a".to_string()), "model-a"); + + for (backend_id, port, rtt_ms) in [("fast", 5001, 10), ("slow", 5002, 100)] { + let identity = RegistrationIdentity { + inference_server_id: backend_id.to_string(), + cluster_id: "cluster-a".to_string(), + inference_server_url: format!("quic://127.0.0.1:{port}"), + routing_key: target.routing_key.clone(), + reverse_tunnel: true, + }; + let registration = app.state.begin_registration(&identity).unwrap(); + let update = stargate_proto::pb::InferenceServerRegistration { + inference_server_id: identity.inference_server_id.clone(), + cluster_id: identity.cluster_id.clone(), + inference_server_url: identity.inference_server_url.clone(), + models: std::collections::HashMap::from([( + target.model_id.clone(), + stargate_proto::pb::InferenceServerModelRegistration { + stats: Some(ModelStats { + last_mean_input_tps: 100.0, + queue_time_estimate_ms_by_priority: std::collections::HashMap::from([ + (0, 10), + ]), + ..ModelStats::default() + }), + status: InferenceServerStatus::Active.into(), + }, + )]), + reverse_tunnel: true, + }; + app.state + .apply_registration_update( + ®istration, + &update, + true, + Some(Duration::from_millis(rtt_ms)), + ) + .await; + } + + let candidates = app.state.candidates_for_target(&target).await; + let chosen = Arc::new( + candidates + .into_iter() + .min_by_key(|candidate| candidate.rtt) + .expect("both backends should be routable"), + ); + assert_eq!(chosen.inference_server_id, "fast"); + assert_eq!(chosen.rtt, Duration::from_millis(10)); + + let target_snapshot = app + .state + .routing_target_snapshot(&target) + .await + .expect("registered backends should produce a target snapshot"); + let resolution = app + .lb_router + .resolve_algorithm_override(&target.model_id, None) + .unwrap(); + let selected = super::super::run::SelectedClusterRun::new( + target_snapshot, + resolution.selection(LoadBalancerCandidateChoice::with_rank_depth_1(0)), + 0, + ); + let request = PreparedProxyRequest { + lb_resolution: resolution, + request_inputs: ProxyRequestInputs { + target, + input_tokens: 10, + priority: 0, + max_wait_ms: None, + request_slo_ms: None, + cache_affinity_key: None, + routing_algorithm_override: None, + }, + endpoint_name: "chat_completions", + method: Method::POST, + path_and_query: "/v1/chat/completions".to_string(), + forwarded_headers: HeaderMap::new(), + retry_deadline: None, + request_start: Instant::now(), + replay_body: ReplayableRequestBody::new(&HeaderMap::new(), Body::empty(), 1024) + .unwrap(), + }; + let mut run = ProxyRequestRun::new(&app, request); + let _ = run.run_proxy_attempt(&selected, &chosen).await; + + let mut encoded = Vec::new(); + prometheus::TextEncoder::new() + .encode(&app.metrics.registry().gather(), &mut encoded) + .unwrap(); + let metrics = String::from_utf8(encoded).unwrap(); + assert!(metrics.contains( + "stargate_routing_reservation_ttl_seconds_sum{model=\"model-a\",routing_key=\"tenant-a\"} 0.01" + )); + assert!(metrics.contains( + "stargate_routing_reservations_active{model=\"model-a\",routing_key=\"tenant-a\"} 1" + )); + } + fn finish_response( disposition: FinalRetryDisposition, status: StatusCode, diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/retry.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/retry.rs index f7e74072d7..76f9a9700d 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/retry.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/retry.rs @@ -419,6 +419,10 @@ mod tests { StatusCode::SERVICE_UNAVAILABLE, &headers )); + assert!(!should_release_queue_mismatch_reservation( + StatusCode::INTERNAL_SERVER_ERROR, + &headers + )); headers.insert( HEADER_STARGATE_RETRY_REASON, diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs index e0cf508cd6..63e5cd5bc7 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/run.rs @@ -313,7 +313,7 @@ pub(super) struct SelectedClusterRun { } impl SelectedClusterRun { - fn new( + pub(super) fn new( target_snapshot: RoutingTargetSnapshot, selection: LoadBalancerCandidateSelection, priority: u32, @@ -364,6 +364,7 @@ mod tests { use prometheus::Encoder; use stargate_proto::pb::{InferenceServerStatus, ModelStats}; + use super::super::upstream::headers_for_upstream_attempt; use super::*; use crate::load_balancer::{ LoadBalancerAlgorithm, LoadBalancerCandidateChoice, LoadBalancerCandidateSelection, @@ -517,6 +518,89 @@ mod tests { ); } + #[tokio::test] + async fn upstream_expected_queue_header_includes_an_active_reservation() { + let app = super::super::test_support::test_proxy_app_state(); + let target = target(); + let identity = RegistrationIdentity { + inference_server_id: "inst-a".to_string(), + cluster_id: "cluster-a".to_string(), + inference_server_url: "quic://127.0.0.1:5000".to_string(), + routing_key: target.routing_key.clone(), + reverse_tunnel: false, + }; + let registration = app.state.begin_registration(&identity).unwrap(); + let update = stargate_proto::pb::InferenceServerRegistration { + inference_server_id: identity.inference_server_id.clone(), + cluster_id: identity.cluster_id.clone(), + inference_server_url: identity.inference_server_url.clone(), + models: std::collections::HashMap::from([( + target.model_id.clone(), + stargate_proto::pb::InferenceServerModelRegistration { + stats: Some(ModelStats { + last_mean_input_tps: 100.0, + max_engine_concurrency: 1, + queue_time_estimate_ms_by_priority: std::collections::HashMap::from([( + 0, 0, + )]), + ..ModelStats::default() + }), + status: InferenceServerStatus::Active.into(), + }, + )]), + reverse_tunnel: false, + }; + app.state + .apply_registration_update( + ®istration, + &update, + true, + Some(Duration::from_millis(10)), + ) + .await; + + let before = app + .state + .routing_target_snapshot(&target) + .await + .expect("registered backend should be routable") + .into_selected_cluster(0); + let _reservation = before + .reserve_backend(®istration.generation(), 10, 0) + .expect("active backend should accept reservation"); + let resolution = app + .lb_router + .resolve_algorithm_override(&target.model_id, None) + .unwrap(); + let selected = SelectedClusterRun::new( + app.state + .routing_target_snapshot(&target) + .await + .expect("target should remain routable"), + resolution.selection(LoadBalancerCandidateChoice::with_rank_depth_1(0)), + 0, + ); + let expected_queue_ms = selected + .expected_queue_ms + .expect("reservation should contribute a queue estimate"); + assert!(expected_queue_ms > 0); + let headers = headers_for_upstream_attempt( + &HeaderMap::new(), + &tracing::info_span!("reservation_queue_header_test"), + Some(expected_queue_ms), + ); + + assert_eq!( + headers + .get(stargate_protocol::tunnel_contract::HEADER_STARGATE_EXPECTED_QUEUE_MS) + .unwrap() + .to_str() + .unwrap(), + expected_queue_ms.to_string() + ); + app.state.end_registration(registration).await; + } + #[tokio::test] async fn load_balancer_request_excludes_failed_clusters() { let app = super::super::test_support::test_proxy_app_state(); diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs index 08f1d4ae39..97aaf1cb51 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs @@ -114,7 +114,13 @@ impl PendingClusterReservation { } } -// Cancellation is explicit because a successful attempt remains pending until its heartbeat. +impl Drop for PendingClusterReservation { + fn drop(&mut self) { + self.deactivate(); + } +} + +// Cancellation is explicit because a successful attempt remains pending until its TTL expires. #[derive(Debug)] pub(crate) struct RoutingReservation(pub(super) Arc); 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 4a37cdfdf1..6da6dabb4e 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 @@ -26,6 +26,7 @@ use crate::load_balancer::{ LoadBalancerRequest, LoadBalancerRouter, PowerOfNAlgorithmConfig, WaitAndWidenAlgorithmConfig, }; use InferenceServerStatus::{Active, Inactive}; +use prometheus::Encoder; use stargate_proto::pb::InferenceServerModelRegistration; use std::collections::HashMap; use std::sync::Arc; @@ -1014,6 +1015,93 @@ async fn removing_backend_drops_its_reservations_immediately() { reservation.release(); } +#[tokio::test] +async fn dropping_retired_cluster_generation_decrements_active_reservation_gauge() { + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let state = StargateState::new_with_metrics(metrics.clone()); + let target = make_target(Some("rk-generation-drop"), "model-generation-drop"); + let identity = |port| RegistrationIdentity { + inference_server_id: "inst-generation-drop".to_string(), + cluster_id: "cluster-generation-drop".to_string(), + inference_server_url: format!("quic://127.0.0.1:{port}"), + routing_key: target.routing_key.clone(), + reverse_tunnel: false, + }; + let update = |identity: &RegistrationIdentity| InferenceServerRegistration { + inference_server_id: identity.inference_server_id.clone(), + cluster_id: identity.cluster_id.clone(), + inference_server_url: identity.inference_server_url.clone(), + models: HashMap::from([( + target.model_id.clone(), + InferenceServerModelRegistration { + stats: Some(priority_stats(100.0, [(0, 5)])), + status: Active as i32, + }, + )]), + reverse_tunnel: identity.reverse_tunnel, + }; + + let old_identity = identity(1111); + let old = state.begin_registration(&old_identity).unwrap(); + state + .apply_registration_update( + &old, + &update(&old_identity), + true, + Some(Duration::from_millis(10)), + ) + .await; + let selected = state + .routing_target_snapshot(&target) + .await + .unwrap() + .into_selected_cluster(0); + let reservation = selected + .reserve_backend_with_ttl( + &old.generation(), + 10, + 0, + Duration::from_secs(1), + metrics.clone(), + &target, + ) + .expect("active old backend should accept a reservation"); + let metrics_text = || { + let mut encoded = Vec::new(); + prometheus::TextEncoder::new() + .encode(&metrics.registry().gather(), &mut encoded) + .expect("metrics should encode"); + String::from_utf8(encoded).expect("metrics should be UTF-8") + }; + assert!(metrics_text().contains( + "stargate_routing_reservations_active{model=\"model-generation-drop\",routing_key=\"rk-generation-drop\"} 1" + )); + + state.end_registration(old).await; + let replacement_identity = identity(2222); + let replacement = state.begin_registration(&replacement_identity).unwrap(); + state + .apply_registration_update( + &replacement, + &update(&replacement_identity), + true, + Some(Duration::from_millis(10)), + ) + .await; + assert_eq!( + state.candidates_for_target(&target).await[0].inference_server_url, + "quic://127.0.0.1:2222" + ); + + drop(reservation); + drop(selected); + let metrics_text = metrics_text(); + assert!(metrics_text.contains( + "stargate_routing_reservations_active{model=\"model-generation-drop\",routing_key=\"rk-generation-drop\"} 0" + )); + state.end_registration(replacement).await; +} + #[tokio::test] async fn released_reservation_restores_local_snapshot_before_registration_update() { let scenario = RegistrationScenario::new(Some("rk-release")); diff --git a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs index 163fb63987..addacd535d 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs @@ -399,13 +399,7 @@ impl StargateRuntime { let warmup_config = self.config.warmup.clone(); let shutdown = tasks.shutdown_signal(); tasks.task_tracker().spawn(async move { - run_warmup_stabilization( - warmup_state, - warmup_config, - ready_token, - shutdown, - ) - .await; + run_warmup_stabilization(warmup_state, warmup_config, ready_token, shutdown).await; }); } diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs index ff61991635..b58ae9229c 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs @@ -16,7 +16,9 @@ use std::net::SocketAddr; use std::time::Duration; -use stargate::runtime::{BoundStargateListeners, StargateRuntime, StargateRuntimeConfig, WarmupConfig}; +use stargate::runtime::{ + BoundStargateListeners, StargateRuntime, StargateRuntimeConfig, WarmupConfig, +}; use crate::common::{SelfDiscovery, base_config, init_crypto, make_stargate_runtime}; @@ -27,8 +29,7 @@ fn make_stargate_runtime_with_warmup( let ephemeral: SocketAddr = "127.0.0.1:0".parse().unwrap(); let mut config: StargateRuntimeConfig = base_config(id, ephemeral, ephemeral); config.warmup = warmup; - let listeners = - BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); + let listeners = BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); let grpc_addr = config.grpc_listen_addr; let http_addr = config.http_listen_addr; let discovery = Box::new(SelfDiscovery::new(id, grpc_addr, http_addr)); diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index aa6e497dd5..a8158615d2 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -47,8 +47,8 @@ cluster mean), fixed at dispatch, and is bounded by `--routing-reservation-ttl-max-ms` (default `1000`). Registration updates do not clear reservations. If an update includes the request before its reservation expires, routing temporarily counts it in both reported stats and the local -reservation. Deploy Pylons with change-driven stats updates before enabling -this Stargate behavior; the RTT-based expiry can otherwise leave a short gap +reservation. Deploy Pylons with change-driven stats updates before deploying +this Stargate version; the RTT-based expiry can otherwise leave a short gap until the next heartbeat. ## Stream Events From 0c276fcbc7755524bbbf6c321e96f50f47db8714 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 12:48:26 -0700 Subject: [PATCH 03/11] test(stargate): prove retired cluster drops settle the reservation gauge The existing retired-generation test ends the registration through StargateState, which deactivates reservations on the routing removal path, so it still passes with the Drop impl removed. These tests retire the generation through the registration table only, so the replacement upsert or Inactive removal drops the old cluster state with a live reservation. Both fail without Drop. They also check that an earlier release followed by the drop does not double-decrement. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../stargate/src/routing_state/tests.rs | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) 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 6da6dabb4e..416241fe63 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 @@ -1015,6 +1015,120 @@ async fn removing_backend_drops_its_reservations_immediately() { reservation.release(); } +fn routing_reservations_active( + metrics: &crate::metrics::StargateMetrics, + target: &RoutingTargetKey, +) -> i64 { + use prometheus::Encoder; + + let mut encoded = Vec::new(); + prometheus::TextEncoder::new() + .encode(&metrics.registry().gather(), &mut encoded) + .expect("metrics should encode"); + let series = format!( + "stargate_routing_reservations_active{{model=\"{}\",routing_key=\"{}\"}} ", + target.model_id, + target.routing_key.as_deref().unwrap_or_default() + ); + String::from_utf8(encoded) + .expect("metrics should be UTF-8") + .lines() + .find_map(|line| line.strip_prefix(&series)) + .expect("reservation gauge series should exist") + .parse() + .expect("reservation gauge should be an integer") +} + +#[derive(Clone, Copy)] +enum RetiredClusterReplacement { + ActiveUpsert, + InactiveRemoval, +} + +async fn assert_retired_cluster_drop_settles_reservation_gauge( + name: &str, + replacement_kind: RetiredClusterReplacement, +) { + let scenario = RegistrationScenario::new(Some(&format!("rk-{name}"))); + let model_id = format!("model-{name}"); + let inference_server_id = format!("inst-{name}"); + let cluster_id = format!("cluster-{name}"); + let target = scenario.target(&model_id); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let old = scenario.start_in(&inference_server_id, &cluster_id, 1111); + scenario.activate(&old, &model_id).await; + + let selected_cluster = scenario.selected_cluster(&model_id).await; + let reserve = || { + selected_cluster + .reserve_backend_with_ttl( + &old.generation(), + 37, + 4, + Duration::from_secs(60), + metrics.clone(), + &target, + ) + .expect("active backend should accept reservation") + }; + reserve().release(); + let pending = reserve(); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + drop(pending); + drop(selected_cluster); + assert_eq!( + routing_reservations_active(&metrics, &target), + 1, + "the cluster generation still owns the pending reservation" + ); + + let ended = scenario + .state + .registrations + .end_registration(old) + .expect("exact old registration should be removed"); + assert!(ended.registration.cluster_generation.is_retired()); + let replacement = scenario.start_in(&inference_server_id, &cluster_id, 2222); + let status = match replacement_kind { + RetiredClusterReplacement::ActiveUpsert => Active, + RetiredClusterReplacement::InactiveRemoval => Inactive, + }; + scenario + .publish_default_stats(&replacement, &model_id, status, Some(6)) + .await; + + assert_eq!( + routing_reservations_active(&metrics, &target), + 0, + "dropping a retired cluster generation must settle its reservations once" + ); + scenario + .state + .routing + .remove_inference_server_targets(&ended.registration, &HashSet::from([target.clone()])) + .await; + scenario.state.end_registration(replacement).await; + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[tokio::test] +async fn replacing_retired_cluster_generation_settles_reservation_gauge() { + assert_retired_cluster_drop_settles_reservation_gauge( + "res-retired-upsert", + RetiredClusterReplacement::ActiveUpsert, + ) + .await; +} + +#[tokio::test] +async fn removing_retired_cluster_generation_settles_reservation_gauge() { + assert_retired_cluster_drop_settles_reservation_gauge( + "res-retired-remove", + RetiredClusterReplacement::InactiveRemoval, + ) + .await; +} + #[tokio::test] async fn dropping_retired_cluster_generation_decrements_active_reservation_gauge() { let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); From 215acbbe104238bf4bbd54b4def27654edbb011b Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 13:10:25 -0700 Subject: [PATCH 04/11] test(stargate): cover per-backend reservation TTL and release paths Drive proxied requests through a minimal raw QUIC mock tunnel so the tests observe the expected queue header exactly as Stargate sends it. Cover that the reservation TTL follows the dispatched backend RTT rather than the cluster mean, that unexpired reservations raise the expected queue estimate while expired ones do not, and that only a queue-mismatch 429 releases a reservation early. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/stargate/src/http_proxy.rs | 2 + .../src/http_proxy/reservation_tests.rs | 497 ++++++++++++++++++ 2 files changed, 499 insertions(+) create mode 100644 src/libraries/rust/stargate/crates/stargate/src/http_proxy/reservation_tests.rs diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs index 1b3d18bcd0..fbc723e0cc 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs @@ -33,6 +33,8 @@ use crate::tunnel::{QuicHttpProxy, QuicTunnelConfig}; mod attempt; mod diagnostics; mod request; +#[cfg(test)] +mod reservation_tests; mod retry; mod routing; mod run; diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/reservation_tests.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/reservation_tests.rs new file mode 100644 index 0000000000..725f55c395 --- /dev/null +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/reservation_tests.rs @@ -0,0 +1,497 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Proxy-level coverage for routing reservations created per dispatch attempt. +//! +//! Requests run through `proxy_openai_request` against a minimal raw QUIC +//! tunnel that plays the Pylon side, so tests observe the +//! `x-stargate-expected-queue-ms` header exactly as Stargate sent it and can +//! return arbitrary tunnel retry metadata. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use axum::body::Body; +use axum::extract::Request; +use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode}; +use parking_lot::Mutex; +use stargate_proto::pb::{ + InferenceServerModelRegistration, InferenceServerRegistration, InferenceServerStatus, + ModelStats, +}; +use stargate_protocol::tunnel_contract::{ + HEADER_INFERENCE_SERVER_ID, HEADER_INPUT_TOKENS, HEADER_MODEL, HEADER_REQUEST_ID, + HEADER_STARGATE_EXPECTED_QUEUE_MS, HEADER_STARGATE_RETRY_REASON, HEADER_STARGATE_RETRYABLE, +}; +use stargate_protocol::{RecvStream, SendStream, TunnelTransportProtocol}; + +use super::test_support::test_proxy_app_state; +use super::{OpenAiProxyEndpoint, ProxyAppState, proxy_openai_request}; +use crate::auth::OpenAuthenticator; +use crate::routing_state::{RegistrationIdentity, RoutingTargetKey, RunningRegistration}; +use crate::tunnel::{QuicHttpProxy, QuicTunnelConfig}; + +const MODEL: &str = "model-reservation"; +const CLUSTER: &str = "cluster-reservation"; +const INPUT_TOKENS: u64 = 37; +const BASE_QUEUE_MS: u64 = 5; +const LAST_MEAN_INPUT_TPS: f64 = 100.0; +/// Queue time one reservation adds: 37 tokens at 100 tokens/s. +const RESERVATION_QUEUE_MS: u64 = 370; +/// Long enough that no reservation can expire while a test runs. +const UNEXPIRING_TTL: Duration = Duration::from_secs(600); + +#[derive(Clone, Copy)] +enum MockReply { + Respond { + status: u16, + headers: &'static [(&'static str, &'static str)], + }, + ResetStream, +} + +const OK_REPLY: MockReply = MockReply::Respond { + status: 200, + headers: &[], +}; + +/// Raw QUIC tunnel endpoint standing in for Pylon. It records every request +/// header block Stargate sends and answers with a fixed reply. +struct MockTunnel { + endpoint: quinn::Endpoint, + received: Arc>>, +} + +impl MockTunnel { + fn start(reply: MockReply) -> Self { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let server_config = stargate_tls::build_quic_server_config( + &stargate_tls::ServerTlsIdentity::SelfSigned, + TunnelTransportProtocol::RawQuic.alpn_protocols(), + ) + .expect("mock tunnel server config should build"); + let endpoint = quinn::Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()) + .expect("mock tunnel endpoint should bind"); + let received = Arc::new(Mutex::new(Vec::new())); + let accept_endpoint = endpoint.clone(); + let accept_received = Arc::clone(&received); + tokio::spawn(async move { + while let Some(incoming) = accept_endpoint.accept().await { + let received = Arc::clone(&accept_received); + tokio::spawn(async move { + let Ok(connection) = incoming.await else { + return; + }; + while let Ok((send, recv)) = connection.accept_bi().await { + tokio::spawn(serve_stream(send, recv, reply, Arc::clone(&received))); + } + }); + } + }); + Self { endpoint, received } + } + + fn url(&self) -> String { + format!( + "quic://{}", + self.endpoint + .local_addr() + .expect("mock tunnel should have a local address") + ) + } + + fn request_count(&self) -> usize { + self.received.lock().len() + } + + fn expected_queue_ms_headers(&self) -> Vec> { + self.received + .lock() + .iter() + .map(|headers| { + headers + .get(HEADER_STARGATE_EXPECTED_QUEUE_MS) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()) + }) + .collect() + } +} + +impl Drop for MockTunnel { + fn drop(&mut self) { + self.endpoint.close(0u32.into(), b"test complete"); + } +} + +async fn serve_stream( + mut quinn_send: quinn::SendStream, + quinn_recv: quinn::RecvStream, + reply: MockReply, + received: Arc>>, +) { + let mut recv = RecvStream::new(quinn_recv); + let Ok(request_headers) = recv.recv_header().await else { + return; + }; + received.lock().push(request_headers); + let (status, extra_headers) = match reply { + MockReply::ResetStream => { + let _ = quinn_send.reset(0u32.into()); + return; + } + MockReply::Respond { status, headers } => (status, headers), + }; + while let Ok(frame) = recv.recv_body().await { + if frame.into_body().is_none() { + break; + } + } + let mut response_headers = HeaderMap::new(); + response_headers.insert( + HeaderName::from_static("x-status"), + HeaderValue::from(status), + ); + for (name, value) in extra_headers { + response_headers.insert( + HeaderName::from_static(name), + HeaderValue::from_static(value), + ); + } + let mut send = SendStream::new(quinn_send); + if send.send_header(response_headers).await.is_ok() + && send + .send_body(bytes::Bytes::from_static(b"{}")) + .await + .is_ok() + { + let _ = send.finish(); + } +} + +fn proxy_app(ttl_min: Duration, ttl_max: Duration) -> ProxyAppState { + let mut app = test_proxy_app_state(); + // The shared fixture uses 10 ms tunnel timeouts because it never dials a + // backend; these tests complete real QUIC requests. + app.quic_proxy = Arc::new( + QuicHttpProxy::new( + QuicTunnelConfig { + connect_timeout: Duration::from_secs(5), + request_timeout: Duration::from_secs(5), + direct_quic_connections: 1, + tls_cert_pem: None, + server_tls_identity: stargate_tls::ServerTlsIdentity::SelfSigned, + server_identity_reloader: None, + tls_reload_interval: stargate_tls::DEFAULT_TLS_RELOAD_INTERVAL, + quic_insecure: true, + tunnel_protocol: TunnelTransportProtocol::RawQuic, + }, + Arc::new(OpenAuthenticator), + ) + .expect("quic proxy should initialize"), + ); + app.routing_reservation_ttl_min = ttl_min; + app.routing_reservation_ttl_max = ttl_max; + app +} + +fn target() -> RoutingTargetKey { + RoutingTargetKey::new(None, MODEL) +} + +/// Registers one active direct backend in the shared test cluster. +async fn register_backend( + app: &ProxyAppState, + inference_server_id: &str, + tunnel: &MockTunnel, + rtt: Duration, +) -> RunningRegistration { + let running = app + .state + .begin_registration(&RegistrationIdentity { + inference_server_id: inference_server_id.to_string(), + cluster_id: CLUSTER.to_string(), + inference_server_url: tunnel.url(), + routing_key: None, + reverse_tunnel: false, + }) + .expect("registration should begin"); + let update = InferenceServerRegistration { + models: HashMap::from([( + MODEL.to_string(), + InferenceServerModelRegistration { + status: InferenceServerStatus::Active.into(), + stats: Some(ModelStats { + last_mean_input_tps: LAST_MEAN_INPUT_TPS, + queue_time_estimate_ms_by_priority: HashMap::from([(0, BASE_QUEUE_MS)]), + ..Default::default() + }), + }, + )]), + ..Default::default() + }; + app.state + .apply_registration_update(&running, &update, false, Some(rtt)) + .await; + running +} + +async fn send_request(app: &ProxyAppState, request_id: &str) -> Result { + let request = Request::builder() + .method(Method::POST) + .uri(OpenAiProxyEndpoint::CHAT_COMPLETIONS.path) + .header(HEADER_REQUEST_ID, request_id) + .header(HEADER_MODEL, MODEL) + .header(HEADER_INPUT_TOKENS, INPUT_TOKENS.to_string()) + .body(Body::from("{}")) + .expect("request should build"); + proxy_openai_request(app.clone(), request, OpenAiProxyEndpoint::CHAT_COMPLETIONS) + .await + .map(|response| response.headers().clone()) +} + +/// Reservations currently counted in the routed cluster snapshot. Base stats +/// report an empty queue, so every queued entry is a live reservation. +async fn counted_reservations(app: &ProxyAppState) -> u64 { + let snapshot = app + .state + .routing_target_snapshot(&target()) + .await + .expect("target should stay routable"); + let clusters = snapshot.clusters(); + assert_eq!(clusters.len(), 1); + clusters[0].stats.queue_size +} + +fn observed_ttl_sum_ms(app: &ProxyAppState) -> u64 { + let sum_seconds = app + .metrics + .routing_reservation_ttl_seconds(None, MODEL) + .get_sample_sum(); + (sum_seconds * 1000.0).round() as u64 +} + +#[tokio::test] +async fn reservation_ttl_uses_the_dispatched_backend_rtt_not_the_cluster_mean() { + let app = proxy_app(Duration::from_millis(1), Duration::from_millis(1000)); + let fast_tunnel = MockTunnel::start(OK_REPLY); + let slow_tunnel = MockTunnel::start(OK_REPLY); + let rtts = HashMap::from([ + ("inst-fast", Duration::from_millis(10)), + ("inst-slow", Duration::from_millis(100)), + ]); + let fast = register_backend(&app, "inst-fast", &fast_tunnel, rtts["inst-fast"]).await; + let slow = register_backend(&app, "inst-slow", &slow_tunnel, rtts["inst-slow"]).await; + let cluster_rtt = app + .state + .routing_target_snapshot(&target()) + .await + .expect("target should be routable") + .clusters()[0] + .rtt; + assert_eq!( + cluster_rtt, + Duration::from_millis(55), + "the cluster mean must differ from both backend RTTs" + ); + + // Round robin dispatches one request to each backend. + let mut dispatched = Vec::new(); + for request_id in ["req-ttl-1", "req-ttl-2"] { + let ttl_sum_before_ms = observed_ttl_sum_ms(&app); + let headers = send_request(&app, request_id) + .await + .expect("request should be proxied"); + let backend = headers[HEADER_INFERENCE_SERVER_ID] + .to_str() + .expect("backend id should be ASCII") + .to_string(); + let expected_ttl = rtts[backend.as_str()]; + assert_eq!( + observed_ttl_sum_ms(&app) - ttl_sum_before_ms, + expected_ttl.as_millis() as u64, + "reservation for {backend} must use its own RTT, not the {cluster_rtt:?} cluster mean" + ); + dispatched.push(backend); + } + dispatched.sort(); + assert_eq!(dispatched, ["inst-fast", "inst-slow"]); + assert_eq!( + app.metrics + .routing_reservation_ttl_seconds(None, MODEL) + .get_sample_count(), + 2 + ); + assert_eq!(fast_tunnel.request_count(), 1); + assert_eq!(slow_tunnel.request_count(), 1); + + app.state.end_registration(fast).await; + app.state.end_registration(slow).await; +} + +#[tokio::test] +async fn unexpired_reservation_raises_expected_queue_header_for_later_requests() { + let app = proxy_app(UNEXPIRING_TTL, UNEXPIRING_TTL); + let tunnel = MockTunnel::start(OK_REPLY); + let running = register_backend(&app, "inst-a", &tunnel, Duration::from_millis(10)).await; + + for request_id in ["req-queue-1", "req-queue-2", "req-queue-3"] { + send_request(&app, request_id) + .await + .expect("request should be proxied"); + } + + assert_eq!( + tunnel.expected_queue_ms_headers(), + [ + Some(BASE_QUEUE_MS), + Some(BASE_QUEUE_MS + RESERVATION_QUEUE_MS), + Some(BASE_QUEUE_MS + 2 * RESERVATION_QUEUE_MS), + ], + "each completed dispatch stays reserved until its TTL" + ); + assert_eq!(counted_reservations(&app).await, 3); + + app.state.end_registration(running).await; +} + +#[tokio::test] +async fn expired_reservation_is_excluded_from_expected_queue_header() { + // A zero TTL sets the expiry to the dispatch instant. `Instant` is + // monotonic, so every later snapshot read is at or past the expiry + // boundary without sleeping or controlling the clock. + let app = proxy_app(Duration::ZERO, Duration::ZERO); + let tunnel = MockTunnel::start(OK_REPLY); + let running = register_backend(&app, "inst-a", &tunnel, Duration::from_millis(10)).await; + + for request_id in ["req-expired-1", "req-expired-2"] { + send_request(&app, request_id) + .await + .expect("request should be proxied"); + } + + assert_eq!( + tunnel.expected_queue_ms_headers(), + [Some(BASE_QUEUE_MS), Some(BASE_QUEUE_MS)], + "an expired reservation must not inflate the next estimate" + ); + assert_eq!(counted_reservations(&app).await, 0); + assert_eq!( + app.metrics + .routing_reservation_ttl_seconds(None, MODEL) + .get_sample_count(), + 2, + "both dispatches must have reserved before expiring" + ); + + app.state.end_registration(running).await; +} + +async fn assert_reply_keeps_reservations(reply: MockReply, case: &str) { + let app = proxy_app(UNEXPIRING_TTL, UNEXPIRING_TTL); + let tunnel = MockTunnel::start(reply); + let running = register_backend(&app, "inst-a", &tunnel, Duration::from_millis(10)).await; + + let _ = send_request(&app, "req-failure").await; + + let attempts = tunnel.request_count(); + assert!( + attempts >= 1, + "{case}: the backend should receive a dispatch" + ); + assert_eq!( + counted_reservations(&app).await, + attempts as u64, + "{case}: every failed attempt must stay reserved until its TTL" + ); + let _ = send_request(&app, "req-follow-up").await; + assert_eq!( + tunnel.expected_queue_ms_headers().get(attempts).copied(), + Some(Some(BASE_QUEUE_MS + attempts as u64 * RESERVATION_QUEUE_MS)), + "{case}: the next request must still see every unreleased reservation" + ); + + app.state.end_registration(running).await; +} + +#[tokio::test] +async fn non_queue_mismatch_failures_keep_reservation_until_ttl() { + for (reply, case) in [ + ( + MockReply::Respond { + status: 503, + headers: &[], + }, + "503 without retry metadata", + ), + ( + MockReply::Respond { + status: 503, + headers: &[ + (HEADER_STARGATE_RETRYABLE, "true"), + (HEADER_STARGATE_RETRY_REASON, "upstream_admission_rejected"), + ], + }, + "retryable 503", + ), + ( + MockReply::Respond { + status: 429, + headers: &[], + }, + "429 without retry metadata", + ), + ( + MockReply::Respond { + status: 429, + headers: &[ + (HEADER_STARGATE_RETRYABLE, "true"), + (HEADER_STARGATE_RETRY_REASON, "upstream_admission_rejected"), + ], + }, + "retryable 429 with another reason", + ), + (MockReply::ResetStream, "transport failure"), + ] { + assert_reply_keeps_reservations(reply, case).await; + } +} + +#[tokio::test] +async fn queue_mismatch_rejection_releases_reservation_before_ttl() { + let app = proxy_app(UNEXPIRING_TTL, UNEXPIRING_TTL); + let tunnel = MockTunnel::start(MockReply::Respond { + status: 429, + headers: &[ + (HEADER_STARGATE_RETRYABLE, "true"), + (HEADER_STARGATE_RETRY_REASON, "queue_estimate_mismatch"), + ], + }); + let running = register_backend(&app, "inst-a", &tunnel, Duration::from_millis(10)).await; + + for request_id in ["req-mismatch-1", "req-mismatch-2"] { + let _ = send_request(&app, request_id).await; + } + + assert_eq!(counted_reservations(&app).await, 0); + assert_eq!( + tunnel.expected_queue_ms_headers(), + [Some(BASE_QUEUE_MS), Some(BASE_QUEUE_MS)], + "a released reservation must not inflate the next estimate" + ); + + app.state.end_registration(running).await; +} From db5e2170e185d2a48b8464d4028484ce20962023 Mon Sep 17 00:00:00 2001 From: jcameron Date: Fri, 9 Oct 2026 13:43:24 -0700 Subject: [PATCH 05/11] test(stargate): remove wall-clock dependence from reservation tests - Add a 600 s UNEXPIRING_TEST_RESERVATION_TTL and use it in the cfg(test) reserve_backend helper and the reservation tests that previously used a 1 s TTL, so they no longer depend on finishing within 1 s of wall time. - Reserve with the long TTL in the expiry boundary test. The registration update (which prunes with the real clock) can no longer flake on a slow runner; boundary checks stay relative to expires_at via routing_snapshot_at, and the test also checks the active gauge. - Rename consumed_by_heartbeat to released_after_update, since updates no longer consume reservations. - Assert stargate_routing_reservations_active goes from 10000 to 0 in the idle expired prune test. - Describe lazy expiry and backend removal for the stargate_routing_reservations_active gauge in the metrics doc. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../metrics/llm-request-router/metrics.md | 2 +- .../src/routing_state/cluster_snapshots.rs | 2 +- .../src/routing_state/reservations.rs | 5 ++++ .../stargate/src/routing_state/tests.rs | 28 +++++++++++++------ 4 files changed, 27 insertions(+), 10 deletions(-) diff --git a/docs/observability/metrics/llm-request-router/metrics.md b/docs/observability/metrics/llm-request-router/metrics.md index 572049806c..8943e06edd 100644 --- a/docs/observability/metrics/llm-request-router/metrics.md +++ b/docs/observability/metrics/llm-request-router/metrics.md @@ -38,5 +38,5 @@ unbounded request fields as metric labels. | `stargate_proxy_duration_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model`, `inference_server_id` | Time to first byte from upstream in seconds. | | `stargate_routing_duration_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | Load-balancer decision time in seconds. | | `stargate_active_inference_servers` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | Currently routable inference servers for a routing target. | -| `stargate_routing_reservations_active` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | In-flight local routing reservations that have not expired or been explicitly released. | +| `stargate_routing_reservations_active` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | Local routing reservations not yet released or pruned. Expiry is lazy: an expired reservation stays counted until the next snapshot read or registration update for its cluster prunes it. Removing the backend also drops its reservations. | | `stargate_routing_reservation_ttl_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | RTT-based reservation TTL in seconds after applying the configured bounds. | 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 1355707cf5..817e1af4bc 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 @@ -469,7 +469,7 @@ impl RoutedClusterState { registration, input_tokens, priority, - Duration::from_secs(1), + super::reservations::UNEXPIRING_TEST_RESERVATION_TTL, None, &RoutingTargetKey::new(None, ""), ) diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs index 97aaf1cb51..19425ad6ca 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs @@ -19,6 +19,11 @@ use stargate_protocol::common::queue_time_delta_ms; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Instant; +/// Long enough that no test reservation expires on wall-clock time while a +/// test runs. `Duration::MAX` would overflow `Instant + Duration`. +#[cfg(test)] +pub(super) const UNEXPIRING_TEST_RESERVATION_TTL: Duration = Duration::from_secs(600); + #[derive(Debug)] pub(super) struct PendingClusterReservation { pub(super) inference_server_id: String, 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 416241fe63..7c846430f5 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 @@ -17,7 +17,7 @@ use super::registration::{ RegistrationClusterGeneration, test_registration_generation, test_registration_generation_in_cluster, }; -use super::reservations::update_reserved_priority_queue_time; +use super::reservations::{UNEXPIRING_TEST_RESERVATION_TTL, update_reserved_priority_queue_time}; use super::snapshots::{ClusterBackendUpsert, RoutedClusterState, RoutingTargetGeneration}; use super::*; use crate::load_balancer::{ @@ -909,19 +909,28 @@ async fn reservation_is_in_snapshot_before_expiry_and_pruned_at_expiry() { let selected_cluster = scenario.selected_cluster("model-res-expiry").await; let target = scenario.target("model-res-expiry"); let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + // The registration update below prunes with the real clock, so the TTL must + // outlast any scheduler stall. The boundary checks use explicit instants + // relative to expires_at instead. let reservation = selected_cluster .reserve_backend_with_ttl( &running.generation(), 37, 4, - Duration::from_millis(60), - metrics, + UNEXPIRING_TEST_RESERVATION_TTL, + metrics.clone(), &target, ) .expect("active backend should accept reservation"); let expires_at = reservation.0.expires_at(); scenario.publish_connected(&running, &update).await; + assert_eq!( + selected_cluster.pending_reservation_count(), + 1, + "a registration update must not clear the reservation" + ); + assert_eq!(routing_reservations_active(&metrics, &target), 1); assert_eq!( selected_cluster .routing_snapshot_at(expires_at - Duration::from_millis(1)) @@ -929,7 +938,7 @@ async fn reservation_is_in_snapshot_before_expiry_and_pruned_at_expiry() { .stats .queue_size, 1, - "a registration update must not clear the reservation" + "reservation should count just before its expiry boundary" ); assert_eq!( selected_cluster @@ -941,6 +950,7 @@ async fn reservation_is_in_snapshot_before_expiry_and_pruned_at_expiry() { "reservation should not count at its expiry boundary" ); assert_eq!(selected_cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); reservation.release(); } @@ -976,10 +986,12 @@ async fn snapshot_read_prunes_ten_thousand_idle_expired_reservations() { } assert_eq!(selected_cluster.pending_reservation_count(), 10_000); + assert_eq!(routing_reservations_active(&metrics, &target), 10_000); selected_cluster .routing_snapshot_at(last_expiry + Duration::from_millis(1)) .expect("cluster should remain routable"); assert_eq!(selected_cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); } #[tokio::test] @@ -1002,7 +1014,7 @@ async fn removing_backend_drops_its_reservations_immediately() { &running.generation(), 37, 4, - Duration::from_secs(1), + UNEXPIRING_TEST_RESERVATION_TTL, metrics, &target, ) @@ -1175,7 +1187,7 @@ async fn dropping_retired_cluster_generation_decrements_active_reservation_gauge &old.generation(), 10, 0, - Duration::from_secs(1), + UNEXPIRING_TEST_RESERVATION_TTL, metrics.clone(), &target, ) @@ -1238,7 +1250,7 @@ async fn released_reservation_restores_local_snapshot_before_registration_update let candidates = scenario.clusters("model-release").await; assert_queue_stats(&candidates[0].stats, 0, 0, 0, 0, 4, 5); - let consumed_by_heartbeat = selected_cluster + let released_after_update = selected_cluster .reserve_backend(&running.generation(), 10, 4) .expect("active backend should accept reservation"); scenario.publish_connected(&running, &update).await; @@ -1246,7 +1258,7 @@ async fn released_reservation_restores_local_snapshot_before_registration_update .reserve_backend(&running.generation(), 20, 4) .expect("active backend should accept reservation"); - consumed_by_heartbeat.release(); + released_after_update.release(); let candidates = scenario.clusters("model-release").await; assert_queue_stats(&candidates[0].stats, 1, 1, 20, 20, 4, 205); From c17e51759e8bcd2aa5bb5dd508f872516172a076 Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:04:26 -0300 Subject: [PATCH 06/11] fix(stargate): keep oversized reservation TTLs from panicking PendingClusterReservation incremented the active reservation gauge and then computed `now + ttl`. A programmatically configured TTL bound such as Duration::MAX overflows Instant addition, so every dispatch panicked after the gauge had already been incremented, leaking one active reservation per request. Compute the expiry with checked_add before touching metrics. A TTL that cannot be represented as an Instant now means the reservation never expires on time; it still ends on release or backend removal. Add unit tests for the unrepresentable TTL and for Drop settling the gauge when the last owner goes away. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../src/routing_state/reservations.rs | 77 ++++++++++++++++++- 1 file changed, 74 insertions(+), 3 deletions(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs index 19425ad6ca..2100081c31 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/reservations.rs @@ -30,7 +30,9 @@ pub(super) struct PendingClusterReservation { pub(super) registration: Option>, pub(super) input_tokens: u64, pub(super) priority: u32, - expires_at: Instant, + // `None` when the TTL is too large to represent as an `Instant`; such a + // reservation only ends by release or backend removal. + expires_at: Option, active: AtomicBool, metrics: Option>, routing_key: Option, @@ -71,6 +73,9 @@ impl PendingClusterReservation { .map_or_else(String::new, |registration| { registration.inference_server_id().to_string() }); + // Compute expiry before touching metrics so an unrepresentable TTL + // cannot panic after the active gauge was incremented. + let expires_at = now.checked_add(ttl); if let Some(metrics) = &metrics { metrics .inc_routing_reservations_active(target.routing_key.as_deref(), &target.model_id); @@ -83,7 +88,7 @@ impl PendingClusterReservation { registration, input_tokens, priority, - expires_at: now + ttl, + expires_at, active: AtomicBool::new(true), metrics, routing_key: target.routing_key.clone(), @@ -98,10 +103,11 @@ impl PendingClusterReservation { #[cfg(test)] pub(super) fn expires_at(&self) -> Instant { self.expires_at + .expect("test reservation should have a representable expiry") } pub(super) fn is_active_at(&self, now: Instant) -> bool { - if now >= self.expires_at { + if self.expires_at.is_some_and(|expires_at| now >= expires_at) { self.deactivate(); } self.is_active() @@ -279,4 +285,69 @@ mod tests { "stargate_routing_reservations_active{model=\"model\",routing_key=\"routing\"} 0" )); } + + fn active_reservation_gauge(metrics: &StargateMetrics) -> String { + let mut encoded = Vec::new(); + prometheus::TextEncoder::new() + .encode(&metrics.registry().gather(), &mut encoded) + .expect("metrics should encode"); + String::from_utf8(encoded) + .expect("metrics should be UTF-8") + .lines() + .find(|line| line.starts_with("stargate_routing_reservations_active{")) + .expect("reservation gauge series should exist") + .to_string() + } + + #[test] + fn dropping_last_reservation_owner_decrements_active_gauge() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let metrics = StargateMetrics::new().expect("metrics should initialize"); + let pending = PendingClusterReservation::new_at( + None, + 1, + 0, + Instant::now(), + Duration::from_millis(60), + Some(metrics.clone()), + &target, + ); + let token = RoutingReservation(Arc::clone(&pending)); + assert!(active_reservation_gauge(&metrics).ends_with(" 1")); + + drop(pending); + assert!( + active_reservation_gauge(&metrics).ends_with(" 1"), + "a remaining owner keeps the reservation active" + ); + drop(token); + assert!( + active_reservation_gauge(&metrics).ends_with(" 0"), + "dropping the last owner must settle the gauge" + ); + } + + #[test] + fn unrepresentable_ttl_does_not_panic_or_leak_the_active_gauge() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let metrics = StargateMetrics::new().expect("metrics should initialize"); + let dispatched_at = Instant::now(); + let pending = PendingClusterReservation::new_at( + None, + 1, + 0, + dispatched_at, + Duration::MAX, + Some(metrics.clone()), + &target, + ); + + assert!(active_reservation_gauge(&metrics).ends_with(" 1")); + assert!( + pending.is_active_at(dispatched_at + Duration::from_secs(365 * 24 * 60 * 60)), + "an unrepresentable expiry must not expire early" + ); + RoutingReservation(pending).release(); + assert!(active_reservation_gauge(&metrics).ends_with(" 0")); + } } From b73af5d433a24338636fff262619a8e1cb0df33a Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:05:18 -0300 Subject: [PATCH 07/11] fix(stargate): validate reservation TTL bounds at runtime start The minimum-above-maximum check for routing reservation TTL bounds only ran when converting CLI arguments. StargateRuntimeConfig and ProxyAppState are public and can be built directly, and with min > max Duration::clamp panics on every proxied dispatch. StargateRuntime::start now rejects inverted bounds with the same error the CLI reports. As defense in depth, the per-dispatch clamp no longer panics on inverted bounds; the maximum wins. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/stargate/src/http_proxy/attempt.rs | 14 ++++++- .../stargate/crates/stargate/src/runtime.rs | 39 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs index fbcbf567fe..1dcb1f6763 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs @@ -48,7 +48,9 @@ fn clamp_routing_reservation_ttl( min: std::time::Duration, max: std::time::Duration, ) -> std::time::Duration { - rtt.clamp(min, max) + // Startup rejects min > max; `Duration::clamp` would panic on every + // dispatch if a programmatic config bypassed that check, so max wins. + rtt.max(min).min(max) } #[derive(Default)] @@ -581,6 +583,16 @@ mod tests { ); } + #[test] + fn reservation_ttl_clamp_does_not_panic_on_inverted_bounds() { + let min = Duration::from_millis(20); + let max = Duration::from_millis(10); + + for rtt in [Duration::ZERO, Duration::from_millis(15), Duration::MAX] { + assert_eq!(clamp_routing_reservation_ttl(rtt, min, max), max); + } + } + #[tokio::test] async fn proxy_reservation_ttl_uses_the_selected_backends_rtt() { let mut app = test_proxy_app_state(); diff --git a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs index addacd535d..0c875058c4 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs @@ -258,6 +258,10 @@ impl StargateRuntime { } pub async fn start(self) -> Result { + ensure!( + self.config.routing_reservation_ttl_min <= self.config.routing_reservation_ttl_max, + "routing reservation TTL minimum must not exceed maximum" + ); if let Some(grpc_pylon_dial_addr) = self .config .grpc_pylon_dial_addr @@ -828,6 +832,41 @@ mod tests { ); } + #[tokio::test] + async fn runtime_rejects_routing_reservation_ttl_minimum_above_maximum() { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let mut config = test_runtime_config("test-invalid-reservation-ttl"); + config.routing_reservation_ttl_min = Duration::from_millis(20); + config.routing_reservation_ttl_max = Duration::from_millis(10); + let listeners = + BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); + let runtime = StargateRuntime::new( + config, + Box::new(BlockingDiscovery { + active_calls: Arc::new(AtomicUsize::new(0)), + self_info: StargateInfo::default(), + }), + listeners, + None, + ); + + let error = match runtime.start().await { + Ok(handle) => { + handle.begin_shutdown(); + let _ = handle.wait_for_shutdown(Duration::from_secs(2)).await; + panic!("reservation TTL minimum above maximum should fail startup") + } + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("routing reservation TTL minimum must not exceed maximum"), + "unexpected error: {error}" + ); + } + #[tokio::test] async fn shutdown_cancels_in_flight_discovery_poll() { let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); From 0cd99667ee70d3716ef2ec2f2a30133f8b1a69e2 Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:06:25 -0300 Subject: [PATCH 08/11] fix(stargate): settle reservations when a retired cluster is detached When a replacement registration replaces or removes a retired cluster generation, the old RoutedClusterState leaves the routing target but its pending reservations were only deactivated by Drop. An in-flight request holds the selected cluster and its reservation token until the upstream response arrives, so stargate_routing_reservations_active stayed elevated past the TTL. The old registration's own cleanup cannot help: it finds the replacement generation and returns without touching it. Detach the old generation explicitly at both sites. Detaching deactivates its pending reservations and clears its snapshot, so a stale selection can no longer add reservations that nothing would prune. Drop remains the backstop. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../src/routing_state/cluster_snapshots.rs | 13 +++ .../stargate/src/routing_state/clusters.rs | 12 ++- .../stargate/src/routing_state/tests.rs | 98 +++++++++++++++++++ 3 files changed, 119 insertions(+), 4 deletions(-) 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 817e1af4bc..7690ec6b6a 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 @@ -410,6 +410,19 @@ impl RoutedClusterState { self.generation.lock().routing_snapshot() } + /// Settles pending reservations when a retired generation leaves its + /// routing target. In-flight selections can keep this state alive past + /// the TTL, so clearing the snapshot also refuses new reservations. + pub(super) fn detach(&self) { + let snapshot_state = self.generation.lock().snapshot_state.take(); + for pending in snapshot_state + .into_iter() + .flat_map(|state| state.pending_cluster_reservations) + { + pending.deactivate(); + } + } + #[cfg(test)] pub(super) fn routing_snapshot_at(&self, now: Instant) -> Option { self.generation.lock().routing_snapshot_at(now) diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/clusters.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/clusters.rs index fd761330d3..7e445812de 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/clusters.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/clusters.rs @@ -72,7 +72,8 @@ impl RoutingTargetState { replacement.upsert_backend(backend); clusters .insert(cluster_id, replacement) - .expect("existing routed cluster generation should be replaced"); + .expect("existing routed cluster generation should be replaced") + .detach(); update_active_backend_count(active_backend_count, stale_backend_count, 1); Ok(()) } @@ -103,9 +104,12 @@ impl RoutingTargetState { return; } let backend_count = cluster_state.generation.lock().backends.len(); - clusters.remove(cluster_id).expect( - "retired routed cluster generation should remain target-owned until removal", - ); + clusters + .remove(cluster_id) + .expect( + "retired routed cluster generation should remain target-owned until removal", + ) + .detach(); backend_count } else { match cluster_state.remove_backend(registration) { 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 7c846430f5..c4fd18c3df 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 @@ -1141,6 +1141,104 @@ async fn removing_retired_cluster_generation_settles_reservation_gauge() { .await; } +async fn assert_detaching_retired_cluster_settles_held_reservations( + name: &str, + replacement_kind: RetiredClusterReplacement, +) { + let scenario = RegistrationScenario::new(Some(&format!("rk-{name}"))); + let model_id = format!("model-{name}"); + let inference_server_id = format!("inst-{name}"); + let cluster_id = format!("cluster-{name}"); + let target = scenario.target(&model_id); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let old = scenario.start_in(&inference_server_id, &cluster_id, 1111); + scenario.activate(&old, &model_id).await; + + // An in-flight request keeps both its selected cluster and its + // reservation token until the upstream response arrives. + let selected_cluster = scenario.selected_cluster(&model_id).await; + let reservation = selected_cluster + .reserve_backend_with_ttl( + &old.generation(), + 37, + 4, + UNEXPIRING_TEST_RESERVATION_TTL, + metrics.clone(), + &target, + ) + .expect("active backend should accept reservation"); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + + let ended = scenario + .state + .registrations + .end_registration(old) + .expect("exact old registration should be removed"); + assert!(ended.registration.cluster_generation.is_retired()); + let replacement = scenario.start_in(&inference_server_id, &cluster_id, 2222); + let status = match replacement_kind { + RetiredClusterReplacement::ActiveUpsert => Active, + RetiredClusterReplacement::InactiveRemoval => Inactive, + }; + scenario + .publish_default_stats(&replacement, &model_id, status, Some(6)) + .await; + + assert_eq!( + routing_reservations_active(&metrics, &target), + 0, + "detaching a retired cluster generation must settle reservations still held in flight" + ); + assert!( + selected_cluster + .reserve_backend_with_ttl( + &ended.registration, + 37, + 4, + UNEXPIRING_TEST_RESERVATION_TTL, + metrics.clone(), + &target, + ) + .is_none(), + "a detached cluster generation must not accept new reservations" + ); + assert_eq!(routing_reservations_active(&metrics, &target), 0); + + // The old registration's cleanup finds the replacement generation and + // returns without touching it. + scenario + .state + .routing + .remove_inference_server_targets(&ended.registration, &HashSet::from([target.clone()])) + .await; + drop(reservation); + drop(selected_cluster); + assert_eq!( + routing_reservations_active(&metrics, &target), + 0, + "dropping the last owner must not decrement a settled reservation again" + ); + scenario.state.end_registration(replacement).await; +} + +#[tokio::test] +async fn replacing_retired_cluster_generation_settles_held_reservations() { + assert_detaching_retired_cluster_settles_held_reservations( + "res-held-upsert", + RetiredClusterReplacement::ActiveUpsert, + ) + .await; +} + +#[tokio::test] +async fn removing_retired_cluster_generation_settles_held_reservations() { + assert_detaching_retired_cluster_settles_held_reservations( + "res-held-remove", + RetiredClusterReplacement::InactiveRemoval, + ) + .await; +} + #[tokio::test] async fn dropping_retired_cluster_generation_decrements_active_reservation_gauge() { let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); From 9f6b61c9117a768a9009ee2807d404f53ec1f28b Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:07:48 -0300 Subject: [PATCH 09/11] fix(stargate): use the current backend RTT for retry reservations Each proxy attempt is a new dispatch with its own reservation, but a same-backend retry after a reconnect reused the backend snapshot chosen for the first attempt. If a registration update changed the backend RTT in between (for example 60 ms to 200 ms), the retry's reservation TTL still used the old RTT. The same staleness applied across the hot-path reconnect before the first attempt. Derive the TTL inside the reservation call from the RTT of the exact stored registration, which the call already looks up to verify the registration. Reservations that already exist keep their fixed expiry. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- .../crates/stargate/src/http_proxy/attempt.rs | 129 ++++++++++++++++-- .../src/routing_state/cluster_snapshots.rs | 27 ++-- .../stargate/src/routing_state/snapshots.rs | 26 +++- 3 files changed, 155 insertions(+), 27 deletions(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs index 1dcb1f6763..7b32cae31f 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/attempt.rs @@ -122,21 +122,24 @@ impl ProxyRequestRun<'_> { } } - let reservation: Option = selected.cluster.reserve_backend_with_ttl( - &chosen.registration, - self.request.request_inputs.input_tokens, - self.request.request_inputs.priority, - clamp_routing_reservation_ttl( - chosen.rtt, - self.app.routing_reservation_ttl_min, - self.app.routing_reservation_ttl_max, - ), - self.app.metrics.clone(), - &crate::routing_state::RoutingTargetKey::new( - self.routing_key().map(str::to_owned), - self.model_id(), - ), - ); + let reservation: Option = + selected.cluster.reserve_backend_with_rtt_ttl( + &chosen.registration, + self.request.request_inputs.input_tokens, + self.request.request_inputs.priority, + |rtt| { + clamp_routing_reservation_ttl( + rtt, + self.app.routing_reservation_ttl_min, + self.app.routing_reservation_ttl_max, + ) + }, + self.app.metrics.clone(), + &crate::routing_state::RoutingTargetKey::new( + self.routing_key().map(str::to_owned), + self.model_id(), + ), + ); record_proxy_attempt_start(self, selected, chosen); let upstream_span = proxy_upstream_attempt_span(self, selected, chosen); @@ -697,6 +700,102 @@ mod tests { )); } + #[tokio::test] + async fn same_backend_retry_reserves_with_the_current_backend_rtt() { + let mut app = test_proxy_app_state(); + app.state = Arc::new(StargateState::new_with_metrics(app.metrics.clone())); + let target = RoutingTargetKey::new(Some("tenant-a".to_string()), "model-a"); + let identity = RegistrationIdentity { + inference_server_id: "backend-a".to_string(), + cluster_id: "cluster-a".to_string(), + inference_server_url: "quic://127.0.0.1:5001".to_string(), + routing_key: target.routing_key.clone(), + reverse_tunnel: true, + }; + let registration = app.state.begin_registration(&identity).unwrap(); + let update = stargate_proto::pb::InferenceServerRegistration { + inference_server_id: identity.inference_server_id.clone(), + cluster_id: identity.cluster_id.clone(), + inference_server_url: identity.inference_server_url.clone(), + models: std::collections::HashMap::from([( + target.model_id.clone(), + stargate_proto::pb::InferenceServerModelRegistration { + stats: Some(ModelStats::default()), + status: InferenceServerStatus::Active.into(), + }, + )]), + reverse_tunnel: true, + }; + app.state + .apply_registration_update( + ®istration, + &update, + true, + Some(Duration::from_millis(60)), + ) + .await; + let resolution = app + .lb_router + .resolve_algorithm_override(&target.model_id, None) + .unwrap(); + let selected = super::super::run::SelectedClusterRun::new( + app.state + .routing_target_snapshot(&target) + .await + .expect("registered backend should produce a target snapshot"), + resolution.selection(LoadBalancerCandidateChoice::with_rank_depth_1(0)), + 0, + ); + // A same-backend retry reuses the snapshot chosen for the first attempt. + let retained = selected + .cluster + .select_backend(&std::collections::HashSet::new()) + .expect("registered backend should be selectable"); + assert_eq!(retained.rtt, Duration::from_millis(60)); + + app.state + .apply_registration_update( + ®istration, + &update, + true, + Some(Duration::from_millis(200)), + ) + .await; + let request = PreparedProxyRequest { + lb_resolution: resolution, + request_inputs: ProxyRequestInputs { + target, + input_tokens: 10, + priority: 0, + max_wait_ms: None, + request_slo_ms: None, + cache_affinity_key: None, + routing_algorithm_override: None, + }, + endpoint_name: "chat_completions", + method: Method::POST, + path_and_query: "/v1/chat/completions".to_string(), + forwarded_headers: HeaderMap::new(), + retry_deadline: None, + request_start: Instant::now(), + replay_body: ReplayableRequestBody::new(&HeaderMap::new(), Body::empty(), 1024) + .unwrap(), + }; + let mut run = ProxyRequestRun::new(&app, request); + let _ = run.run_proxy_attempt(&selected, &retained).await; + + let ttl = app + .metrics + .routing_reservation_ttl_seconds(Some("tenant-a"), "model-a"); + assert_eq!(ttl.get_sample_count(), 1); + assert_eq!( + Duration::from_secs_f64(ttl.get_sample_sum()), + Duration::from_millis(200), + "a new dispatch must use the backend's current RTT, not the retained snapshot" + ); + app.state.end_registration(registration).await; + } + fn finish_response( disposition: FinalRetryDisposition, status: StatusCode, 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 7690ec6b6a..41b13cd42a 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 @@ -166,9 +166,14 @@ impl ClusterRoutingGeneration { true } - fn contains_registration(&self, registration: &Arc) -> bool { + fn registered_backend( + &self, + registration: &Arc, + ) -> Option<&Arc> { backend_index(&self.backends, registration.inference_server_id()) - .is_ok_and(|index| Arc::ptr_eq(&self.backends[index].registration, registration)) + .ok() + .map(|index| &self.backends[index]) + .filter(|backend| Arc::ptr_eq(&backend.registration, registration)) } fn backend_aggregate(&self) -> Option<(ModelStats, Duration, usize)> { @@ -339,13 +344,13 @@ impl ClusterRoutingGeneration { registration: &Arc, input_tokens: u64, priority: u32, - ttl: Duration, + ttl_for_rtt: impl FnOnce(Duration) -> Duration, metrics: Option>, target: &RoutingTargetKey, ) -> Option { - if !self.contains_registration(registration) { - return None; - } + // Read the RTT from the stored backend so a caller holding an older + // backend snapshot, such as a same-backend retry, gets the current RTT. + let ttl = ttl_for_rtt(self.registered_backend(registration)?.rtt); let snapshot_state = self.snapshot_state.as_mut()?; let pending = PendingClusterReservation::new( registration.clone(), @@ -482,18 +487,20 @@ impl RoutedClusterState { registration, input_tokens, priority, - super::reservations::UNEXPIRING_TEST_RESERVATION_TTL, + |_rtt| super::reservations::UNEXPIRING_TEST_RESERVATION_TTL, None, &RoutingTargetKey::new(None, ""), ) } - pub(super) fn reserve_backend_with_ttl( + /// Reserves the exact stored registration with a TTL derived from that + /// backend's current RTT. + pub(super) fn reserve_backend_with_rtt_ttl( &self, registration: &Arc, input_tokens: u64, priority: u32, - ttl: Duration, + ttl_for_rtt: impl FnOnce(Duration) -> Duration, metrics: Arc, target: &RoutingTargetKey, ) -> Option { @@ -501,7 +508,7 @@ impl RoutedClusterState { registration, input_tokens, priority, - ttl, + ttl_for_rtt, Some(metrics), target, ) diff --git a/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs index 682a26375f..61eee22bff 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/snapshots.rs @@ -140,6 +140,7 @@ impl SelectedRoutedCluster { .reserve_backend(registration, input_tokens, priority) } + #[cfg(test)] pub(crate) fn reserve_backend_with_ttl( &self, registration: &Arc, @@ -149,11 +150,32 @@ impl SelectedRoutedCluster { metrics: Arc, target: &RoutingTargetKey, ) -> Option { - self.owner.reserve_backend_with_ttl( + self.reserve_backend_with_rtt_ttl( + registration, + input_tokens, + priority, + |_rtt| ttl, + metrics, + target, + ) + } + + /// Reserves `registration` if it is still the stored backend, deriving + /// the TTL from that backend's current RTT rather than a caller snapshot. + pub(crate) fn reserve_backend_with_rtt_ttl( + &self, + registration: &Arc, + input_tokens: u64, + priority: u32, + ttl_for_rtt: impl FnOnce(Duration) -> Duration, + metrics: Arc, + target: &RoutingTargetKey, + ) -> Option { + self.owner.reserve_backend_with_rtt_ttl( registration, input_tokens, priority, - ttl, + ttl_for_rtt, metrics, target, ) From f14aedacdb618610fea3b3715bd7cce664de3b6e Mon Sep 17 00:00:00 2001 From: jcameron Date: Sun, 11 Oct 2026 23:34:51 -0300 Subject: [PATCH 10/11] fix(stargate): describe lazy pruning in reservation metric help Live testing showed stargate_routing_reservations_active holding a nonzero value for an idle cluster after its reservations expired, because expired reservations are pruned lazily on the next snapshot read or registration update. The help text said "Active in-flight", which reads as unexpired. Describe what the gauge counts, and say the TTL histogram records the lifetime assigned at dispatch. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- src/libraries/rust/stargate/crates/stargate/src/metrics.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/metrics.rs b/src/libraries/rust/stargate/crates/stargate/src/metrics.rs index c1d71b9339..db6855da76 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/metrics.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/metrics.rs @@ -111,11 +111,11 @@ define_stargate_metrics! { proxy_replay_buffer_bytes("proxy_replay_buffer_bytes", "Bytes currently retained for proxied request body replay", ["model"], [0.0, 1024.0, 4096.0, 16_384.0, 65_536.0, 262_144.0, 1_048_576.0, 4_194_304.0, 16_777_216.0, 67_108_864.0]); proxy_duration_seconds("proxy_duration_seconds", "Time to first byte from upstream", ["routing_key", "model", "inference_server_id"], [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0]); routing_duration_seconds("routing_duration_seconds", "Time spent selecting an inference server", ["routing_key", "model"], [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0]); - routing_reservation_ttl_seconds("routing_reservation_ttl_seconds", "Configured lifetime of routing reservations", ["routing_key", "model"], [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0]); + routing_reservation_ttl_seconds("routing_reservation_ttl_seconds", "Lifetime assigned to each routing reservation at dispatch", ["routing_key", "model"], [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0]); } gauges { active_inference_servers("active_inference_servers", "Active inference servers available for a routing target", ["routing_key", "model"]); - routing_reservations_active("routing_reservations_active", "Active in-flight routing reservations", ["routing_key", "model"]); + routing_reservations_active("routing_reservations_active", "Routing reservations not yet released or pruned; expired reservations count until the next snapshot read or update prunes them", ["routing_key", "model"]); tls_certificate_expiry_seconds("tls_certificate_expiry_seconds", "Unix timestamp when the active TLS certificate expires", ["material_type"]); } } From 02a3c1d4ca11fbbd708926fb643268a3db94dc2e Mon Sep 17 00:00:00 2001 From: jcameron Date: Mon, 12 Oct 2026 09:26:17 -0300 Subject: [PATCH 11/11] fix(stargate): default reservation TTL floor to Pylon coalescing Live testing with Pylon's default 10 ms stats update coalescing showed the 1 ms floor undercounting in-flight work on low-RTT links: a reservation expired before the change-driven update that included its request arrived. At 48 concurrent requests Stargate's expected queue was 77 ms against Pylon's own 170 ms estimate (164 ms before RTT expiry). Raise the --routing-reservation-ttl-min-ms default to 10 ms to match the Pylon default and document why the floor should stay at or above the coalescing interval. This differs from the 1 ms default proposed in issue #2284. Co-Authored-By: Claude Opus 5.5 Signed-off-by: jcameron --- src/libraries/rust/stargate/crates/stargate/src/main.rs | 8 +++++--- .../rust/stargate/docs/runtime-stats-interface.md | 9 ++++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/libraries/rust/stargate/crates/stargate/src/main.rs b/src/libraries/rust/stargate/crates/stargate/src/main.rs index 00996aa4de..2592bbd0e0 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/main.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/main.rs @@ -161,8 +161,10 @@ struct Args { value_name = "MS" )] registration_update_max_idle_timeout_ms: u64, - /// Minimum backend RTT-based routing reservation lifetime. - #[arg(long, default_value_t = 1, value_name = "MS")] + /// Minimum backend RTT-based routing reservation lifetime. The default + /// matches Pylon's default stats update coalescing interval, so a request + /// stays counted until the first change-driven update that can include it. + #[arg(long, default_value_t = 10, value_name = "MS")] routing_reservation_ttl_min_ms: u64, /// Maximum backend RTT-based routing reservation lifetime. #[arg(long, default_value_t = 1000, value_name = "MS")] @@ -689,7 +691,7 @@ mod tests { #[test] fn routing_reservation_ttl_defaults_and_bounds_are_validated() { let args = parse_args(""); - assert_eq!(args.routing_reservation_ttl_min_ms, 1); + assert_eq!(args.routing_reservation_ttl_min_ms, 10); assert_eq!(args.routing_reservation_ttl_max_ms, 1000); let args = parse_args("--routing-reservation-ttl-min-ms 20 --routing-reservation-ttl-max-ms 200"); diff --git a/src/libraries/rust/stargate/docs/runtime-stats-interface.md b/src/libraries/rust/stargate/docs/runtime-stats-interface.md index a8158615d2..5e6d502e99 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -43,9 +43,12 @@ delay of the backend's actual load, instead of up to one heartbeat behind. Each dispatched request adds a local reservation until one health-check RTT for the selected backend has elapsed. The TTL uses that backend's RTT (not the cluster mean), fixed at dispatch, and is bounded by -`--routing-reservation-ttl-min-ms` (default `1`) and -`--routing-reservation-ttl-max-ms` (default `1000`). Registration updates do not -clear reservations. If an update includes the request before its reservation +`--routing-reservation-ttl-min-ms` (default `10`) and +`--routing-reservation-ttl-max-ms` (default `1000`). The minimum matches Pylon's +default `--stats-update-coalesce-ms`; on low-RTT links a smaller floor lets a +reservation expire before the change-driven update that includes its request +arrives, so routing undercounts in-flight work. Keep the minimum at or above the +Pylon coalescing interval. Registration updates do not clear reservations. If an update includes the request before its reservation expires, routing temporarily counts it in both reported stats and the local reservation. Deploy Pylons with change-driven stats updates before deploying this Stargate version; the RTT-based expiry can otherwise leave a short gap