diff --git a/docs/observability/metrics/llm-request-router/metrics.md b/docs/observability/metrics/llm-request-router/metrics.md index 93185db2c6..7bdd683bb1 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` | Local routing reservations not yet released or pruned. A reservation counts until its backend sends a registration update received at or after the soft TTL; time alone never ends it. Release after a queue estimate mismatch and removing or deregistering the backend also end it, so a backend that stops reporting keeps this gauge raised until it reports or deregisters. | +| `stargate_routing_reservation_ttl_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | RTT-based soft reservation TTL in seconds after applying the configured bounds. After it, the next registration update from the backend ends the reservation. | diff --git a/src/libraries/rust/stargate/crates/stargate/src/control_plane/registration.rs b/src/libraries/rust/stargate/crates/stargate/src/control_plane/registration.rs index 0dfddee4dd..3aefdee5e8 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/control_plane/registration.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/control_plane/registration.rs @@ -15,7 +15,7 @@ use std::collections::BTreeSet; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use futures::{Stream, StreamExt}; use tokio_util::sync::CancellationToken; @@ -105,7 +105,7 @@ pub(super) async fn process_registration_stream( stop: CancellationToken, ) { let first_update = match next_registration_update(&mut stream, idle_timeout, &stop).await { - Ok(Some(update)) => update, + Ok(Some(received)) => received, Ok(None) => return, Err(error) => { warn!( @@ -117,7 +117,7 @@ pub(super) async fn process_registration_stream( } }; let session = match RegistrationSession::start( - &first_update, + &first_update.update, state, connection, auth_result.routing_key.as_deref(), @@ -138,7 +138,7 @@ async fn next_registration_update( stream: &mut (impl Stream> + Unpin), idle_timeout: Option, stop: &CancellationToken, -) -> Result, RegistrationError> { +) -> Result, RegistrationError> { let next = if let Some(idle_timeout) = idle_timeout { // Once shutdown begins, cleanup must win over another ready stream update. tokio::select! { @@ -162,13 +162,25 @@ async fn next_registration_update( Some(Err(status)) => return Err(RegistrationError::StreamRead(status)), Some(Ok(update)) => update, }; + // Stamp receipt before any connection or RTT await: routing reservations + // compare it with their soft expiry to decide whether this update can + // include their request. + let received_at = Instant::now(); debug!( inference_server_id = %update.inference_server_id, cluster_id = %update.cluster_id, model_ids = ?sorted_model_ids(&update), "received inference servers update" ); - Ok(Some(update)) + Ok(Some(ReceivedRegistrationUpdate { + update, + received_at, + })) +} + +struct ReceivedRegistrationUpdate { + update: InferenceServerRegistration, + received_at: Instant, } fn log_deregistration( @@ -288,14 +300,14 @@ impl RegistrationSession { async fn run( mut self, - mut update: InferenceServerRegistration, + mut received: ReceivedRegistrationUpdate, stream: &mut (impl Stream> + Unpin), responses: &flume::Sender>, idle_timeout: Option, stop: &CancellationToken, ) { let outcome = loop { - match self.apply_update(&update, stop).await { + match self.apply_update(&received, stop).await { Ok(ApplyUpdateOutcome::Ack(ack)) => { if !send_registration_response(responses, stop, Ok(ack)).await { break if stop.is_cancelled() { @@ -313,7 +325,7 @@ impl RegistrationSession { } } match next_registration_update(stream, idle_timeout, stop).await { - Ok(Some(next_update)) => update = next_update, + Ok(Some(next_update)) => received = next_update, Ok(None) => break Ok(()), Err(error) => break Err(error), } @@ -323,9 +335,13 @@ impl RegistrationSession { async fn apply_update( &mut self, - update: &InferenceServerRegistration, + received: &ReceivedRegistrationUpdate, stop: &CancellationToken, ) -> Result { + let ReceivedRegistrationUpdate { + update, + received_at, + } = received; if let Some(status) = validate_running_update(self.registration.identity(), update) { return Err(status); } @@ -341,7 +357,13 @@ impl RegistrationSession { // Keep the no-ack retry behavior while clearing any stale route // that depended on the now-unavailable direct connection. self.state - .apply_registration_update(&self.registration, update, false, None) + .apply_received_registration_update( + &self.registration, + update, + false, + None, + *received_at, + ) .await; return Ok(ApplyUpdateOutcome::Skip); } @@ -356,7 +378,13 @@ impl RegistrationSession { } self.state - .apply_registration_update(&self.registration, update, reverse_connected, rtt) + .apply_received_registration_update( + &self.registration, + update, + reverse_connected, + rtt, + *received_at, + ) .await; Ok(ApplyUpdateOutcome::Ack(build_registration_ack( @@ -639,7 +667,13 @@ mod tests { assert_eq!(state.candidates_for_target(&target).await.len(), 1); let outcome = session - .apply_update(&update, &CancellationToken::new()) + .apply_update( + &ReceivedRegistrationUpdate { + update, + received_at: Instant::now(), + }, + &CancellationToken::new(), + ) .await .expect("unavailable direct update should be handled"); 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..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; @@ -120,6 +122,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, } @@ -247,7 +251,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()) @@ -286,6 +292,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(), } @@ -308,7 +316,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 002ecce962..13fd844cbd 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,16 @@ 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 { + // 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)] pub(super) struct ProxyAttemptCounters { pub(super) attempt: u32, @@ -112,11 +122,24 @@ impl ProxyRequestRun<'_> { } } - let reservation: Option = selected.cluster.reserve_backend( - &chosen.registration, - self.request.request_inputs.input_tokens, - self.request.request_inputs.priority, - ); + 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); @@ -142,6 +165,13 @@ impl ProxyRequestRun<'_> { &self.request.path_and_query, attempt_headers, || { + // The tunnel calls this after opening the stream and sending + // the request head, so the soft expiry starts only once the + // request has left Stargate. Attempts that end earlier arm + // when the reservation is dropped. + if let Some(reservation) = &reservation { + reservation.arm(); + } let body = self.request.replay_body.body_for_attempt()?; // After handing the body to the tunnel, a transport failure // cannot establish whether the backend applied this POST. @@ -532,19 +562,243 @@ 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 stargate_protocol::tunnel_contract::{ + HEADER_INPUT_TOKENS, HEADER_MODEL, HEADER_REQUEST_ID, HEADER_STARGATE_RETRY_AFTER_MS, + }; use std::time::Duration; use super::super::HEADER_STARGATE_ERROR_CODE; - use super::super::request::ProxyRequestInputs; + use super::super::request::{ProxyRequestInputs, parse_proxy_request_inputs}; 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, }; + fn chat_request_inputs(target: RoutingTargetKey) -> ProxyRequestInputs { + let mut headers = HeaderMap::new(); + headers.insert(HEADER_REQUEST_ID, "req-1".parse().unwrap()); + headers.insert(HEADER_MODEL, target.model_id.parse().unwrap()); + headers.insert(HEADER_INPUT_TOKENS, "10".parse().unwrap()); + let mut inputs = parse_proxy_request_inputs(&headers).unwrap(); + inputs.target = target; + inputs + } + + #[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 + ); + } + + #[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(); + 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: chat_request_inputs(target), + 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" + )); + } + + #[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: chat_request_inputs(target), + 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, @@ -590,6 +844,7 @@ mod tests { stats: ModelStats::default(), rtt: Duration::from_millis(1), snapshot_updated_at: Instant::now(), + update_received_at: Instant::now(), status: InferenceServerStatus::Active, reverse_tunnel: false, }; 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..55ebe6f2f9 --- /dev/null +++ b/src/libraries/rust/stargate/crates/stargate/src/http_proxy/reservation_tests.rs @@ -0,0 +1,641 @@ +// 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 axum::response::IntoResponse; +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 reaches its soft expiry 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 { + Self::start_with_stream_limit(reply, None) + } + + /// Accepts connections but grants no request streams, so Stargate's + /// stream open blocks until it times out or is cancelled. + fn start_without_stream_capacity() -> Self { + Self::start_with_stream_limit(OK_REPLY, Some(0)) + } + + fn start_with_stream_limit(reply: MockReply, max_bidi_streams: Option) -> Self { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let mut server_config = stargate_tls::build_quic_server_config( + &stargate_tls::ServerTlsIdentity::SelfSigned, + TunnelTransportProtocol::RawQuic.alpn_protocols(), + ) + .expect("mock tunnel server config should build"); + if let Some(max_bidi_streams) = max_bidi_streams { + let mut transport = quinn::TransportConfig::default(); + transport.max_concurrent_bidi_streams(max_bidi_streams.into()); + server_config.transport_config(Arc::new(transport)); + } + 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 { + proxy_app_with_request_timeout(ttl_min, ttl_max, Duration::from_secs(5)) +} + +fn proxy_app_with_request_timeout( + ttl_min: Duration, + ttl_max: Duration, + request_timeout: 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, + 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"); + report_idle(app, &running, rtt).await; + running +} + +/// Applies a registration update reporting an empty queue. +async fn report_idle(app: &ProxyAppState, running: &RunningRegistration, rtt: Duration) { + 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; +} + +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"); + let response = + proxy_openai_request(app.clone(), request, OpenAiProxyEndpoint::CHAT_COMPLETIONS) + .await + .into_response(); + if response.status().is_success() { + Ok(response.headers().clone()) + } else { + Err(response.status()) + } +} + +/// 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 +} + +/// Polls until the routed cluster counts `expected` reservations. +async fn wait_for_counted_reservations(app: &ProxyAppState, expected: u64) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while counted_reservations(app).await != expected { + assert!( + tokio::time::Instant::now() < deadline, + "routed cluster never counted {expected} reservations" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } +} + +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 backend reports" + ); + assert_eq!(counted_reservations(&app).await, 3); + + app.state.end_registration(running).await; +} + +#[tokio::test] +async fn reported_reservation_is_excluded_from_expected_queue_header() { + // A zero TTL puts the soft expiry at the dispatch instant. `Instant` is + // monotonic, so the backend update applied after each request is received + // at or past the soft expiry 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"); + report_idle(&app, &running, Duration::from_millis(10)).await; + } + + assert_eq!( + tunnel.expected_queue_ms_headers(), + [Some(BASE_QUEUE_MS), Some(BASE_QUEUE_MS)], + "a reservation settled by its backend's update 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 settling" + ); + + app.state.end_registration(running).await; +} + +#[tokio::test] +async fn soft_expired_reservation_counts_until_the_backend_reports() { + // A zero RTT and TTL put the soft expiry at the dispatch instant, so only + // the missing registration update keeps the reservation counted. + let app = proxy_app(Duration::ZERO, Duration::ZERO); + let tunnel = MockTunnel::start(OK_REPLY); + let running = register_backend(&app, "inst-a", &tunnel, Duration::ZERO).await; + + for request_id in ["req-silent-1", "req-silent-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 + RESERVATION_QUEUE_MS) + ], + "a soft-expired reservation must count until its backend reports" + ); + assert_eq!(counted_reservations(&app).await, 2); + + report_idle(&app, &running, Duration::ZERO).await; + assert_eq!( + counted_reservations(&app).await, + 0, + "an update after the soft expiry replaces reservations with reported stats" + ); + + app.state.end_registration(running).await; +} + +#[tokio::test] +async fn blocked_stream_open_keeps_reservation_unarmed_until_cancelled() { + // A zero TTL armed at dispatch would end on the first update below. + let app = proxy_app(Duration::ZERO, Duration::ZERO); + let tunnel = MockTunnel::start_without_stream_capacity(); + let running = register_backend(&app, "inst-a", &tunnel, Duration::from_millis(10)).await; + + let request = tokio::spawn({ + let app = app.clone(); + async move { send_request(&app, "req-blocked").await } + }); + wait_for_counted_reservations(&app, 1).await; + report_idle(&app, &running, Duration::from_millis(10)).await; + assert_eq!( + counted_reservations(&app).await, + 1, + "an update cannot settle a request still waiting for a stream" + ); + + request.abort(); + assert!( + request + .await + .expect_err("aborted request should not finish") + .is_cancelled() + ); + report_idle(&app, &running, Duration::from_millis(10)).await; + assert_eq!( + counted_reservations(&app).await, + 0, + "cancellation must arm the reservation so the next update settles it" + ); + assert_eq!(tunnel.request_count(), 0); + + app.state.end_registration(running).await; +} + +#[tokio::test] +async fn failed_stream_open_arms_reservation_when_the_attempt_ends() { + let app = + proxy_app_with_request_timeout(Duration::ZERO, Duration::ZERO, Duration::from_millis(100)); + let tunnel = MockTunnel::start_without_stream_capacity(); + let running = register_backend(&app, "inst-a", &tunnel, Duration::from_millis(10)).await; + + send_request(&app, "req-open-timeout") + .await + .expect_err("a request without stream capacity should fail"); + assert!( + counted_reservations(&app).await >= 1, + "failed attempts stay reserved until their backend reports" + ); + report_idle(&app, &running, Duration::from_millis(10)).await; + assert_eq!( + counted_reservations(&app).await, + 0, + "a failed attempt must arm its reservation so the next update settles it" + ); + assert_eq!(tunnel.request_count(), 0); + + 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 backend reports" + ); + 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; +} 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..3df5602947 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, @@ -419,6 +420,7 @@ mod tests { stats: ModelStats::default(), rtt: Duration::from_millis(1), snapshot_updated_at: Instant::now(), + update_received_at: Instant::now(), status: InferenceServerStatus::Active, reverse_tunnel: false, } @@ -517,6 +519,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/main.rs b/src/libraries/rust/stargate/crates/stargate/src/main.rs index 826f4097f3..3fa6ff9ef8 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/main.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/main.rs @@ -161,6 +161,18 @@ struct Args { value_name = "MS" )] registration_update_max_idle_timeout_ms: u64, + /// Minimum soft TTL for a routing reservation. The soft TTL is the + /// selected backend's RTT clamped to the min and max; after it, the + /// backend's next registration update ends the reservation. Settlement + /// already waits for that update, so a larger floor only adds time in + /// which the request is counted twice. + #[arg(long, default_value_t = 1, value_name = "MS")] + routing_reservation_ttl_min_ms: u64, + /// Maximum soft TTL for a routing reservation. It bounds the effect of a + /// bad RTT sample and never ends a reservation on its own; a backend that + /// stops reporting keeps its reservations until it reports or is removed. + #[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 +692,41 @@ 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-max-ms 5"); + let config = runtime_config_from_args(&args, proxy_transport(&args)) + .expect("a maximum above the default minimum should reach runtime config"); + assert_eq!(config.routing_reservation_ttl_min, Duration::from_millis(1)); + assert_eq!(config.routing_reservation_ttl_max, Duration::from_millis(5)); + 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..3b84bdef29 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", "Soft TTL assigned to each routing reservation at dispatch; once that long has passed after the request is sent, a registration update from its backend ends it", ["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", "Routing reservations not yet released or pruned; a reservation counts until its backend reports at or after the soft TTL, or until the backend is removed", ["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..3d8268f680 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 @@ -17,11 +17,16 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; +#[cfg(test)] +use std::time::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, @@ -163,9 +168,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)> { @@ -250,14 +260,20 @@ impl ClusterRoutingGeneration { fn refresh_snapshot(&mut self, updated_backend_id: Option<&str>) { 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; }; if self.snapshot_state.is_none() && updated_backend_id.is_none() { return; } - let source_backend_index = updated_backend_id - .and_then(|backend_id| backend_index(&self.backends, backend_id).ok()) + let updated_backend_index = updated_backend_id + .and_then(|backend_id| backend_index(&self.backends, backend_id).ok()); + let source_backend_index = updated_backend_index .or_else(|| { self.snapshot_state.as_ref().and_then(|state| { backend_index(&self.backends, &state.cluster_stats_source_backend_id).ok() @@ -279,13 +295,26 @@ impl ClusterRoutingGeneration { .map(|state| state.pending_cluster_reservations) .unwrap_or_default(); pending_cluster_reservations.retain(|pending| { + let Some(index) = backend_index(&self.backends, &pending.inference_server_id) + .ok() + .filter(|&index| { + pending.registration.as_ref().is_none_or(|registration| { + Arc::ptr_eq(&self.backends[index].registration, registration) + }) + }) + else { + pending.deactivate(); + return false; + }; + // Only the reserved backend's own update can include the request. + // Updates from other backends in the cluster leave it counting. + // The update's receipt time, not when it was applied here, bounds + // what the backend could have reported. + if updated_backend_index == Some(index) { + pending.settle_on_backend_update(self.backends[index].update_received_at); + } pending.is_active() - && backend_index(&self.backends, &pending.inference_server_id).is_ok() }); - 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 { @@ -306,6 +335,8 @@ impl ClusterRoutingGeneration { fn routing_snapshot(&mut self) -> Option { let snapshot_state = self.snapshot_state.as_mut()?; + // Reservations end only by backend update, release, or removal, so a + // read just drops ones that were released since the last refresh. snapshot_state .pending_cluster_reservations .retain(|pending| pending.is_active()); @@ -322,15 +353,47 @@ impl ClusterRoutingGeneration { registration: &Arc, input_tokens: u64, priority: u32, + 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.inference_server_id().to_string(), + let pending = PendingClusterReservation::new_unarmed( + Some(registration.clone()), input_tokens, priority, + ttl, + metrics, + target, + ); + snapshot_state + .pending_cluster_reservations + .push(pending.clone()); + Some(RoutingReservation(pending)) + } + + #[cfg(test)] + fn reserve_backend_at( + &mut self, + registration: &Arc, + dispatched_at: Instant, + ttl: Duration, + metrics: Arc, + target: &RoutingTargetKey, + ) -> Option { + self.registered_backend(registration)?; + let snapshot_state = self.snapshot_state.as_mut()?; + let pending = PendingClusterReservation::new_at( + Some(registration.clone()), + 1, + 0, + dispatched_at, + ttl, + Some(metrics), + target, ); snapshot_state .pending_cluster_reservations @@ -387,6 +450,28 @@ 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 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 +506,60 @@ impl RoutedClusterState { eligible.nth(index).map(Arc::clone) } + #[cfg(test)] pub(super) fn reserve_backend( &self, registration: &Arc, input_tokens: u64, priority: u32, + ) -> Option { + let reservation = self.generation.lock().reserve_backend( + registration, + input_tokens, + priority, + |_rtt| super::reservations::UNEXPIRING_TEST_RESERVATION_TTL, + None, + &RoutingTargetKey::new(None, ""), + ); + reservation.inspect(RoutingReservation::arm) + } + + /// Reserves one request on `registration` dispatched at `dispatched_at`, + /// so tests can place the soft expiry on an explicit timeline. + #[cfg(test)] + pub(super) fn reserve_backend_at( + &self, + registration: &Arc, + dispatched_at: Instant, + ttl: Duration, + metrics: Arc, + target: &RoutingTargetKey, ) -> Option { self.generation .lock() - .reserve_backend(registration, input_tokens, priority) + .reserve_backend_at(registration, dispatched_at, ttl, metrics, target) + } + + /// Reserves the exact stored registration with a TTL derived from that + /// backend's current RTT. The caller arms the returned reservation when + /// it hands the request to the tunnel. + pub(super) 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.generation.lock().reserve_backend( + registration, + input_tokens, + priority, + ttl_for_rtt, + Some(metrics), + target, + ) } #[cfg(test)] 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/mod.rs b/src/libraries/rust/stargate/crates/stargate/src/routing_state/mod.rs index 060cb7f63c..103e137f63 100644 --- a/src/libraries/rust/stargate/crates/stargate/src/routing_state/mod.rs +++ b/src/libraries/rust/stargate/crates/stargate/src/routing_state/mod.rs @@ -90,12 +90,35 @@ impl StargateState { model_ids } + /// Applies an update as if it had just been received. + #[cfg(test)] pub(crate) async fn apply_registration_update( &self, running: &RunningRegistration, update: &InferenceServerRegistration, reverse_connected: bool, rtt: Option, + ) { + self.apply_received_registration_update( + running, + update, + reverse_connected, + rtt, + Instant::now(), + ) + .await; + } + + /// Applies an update received from the registration stream at + /// `received_at`. Reservations compare that receipt time, not the later + /// apply time, with their soft expiry. + pub(crate) async fn apply_received_registration_update( + &self, + running: &RunningRegistration, + update: &InferenceServerRegistration, + reverse_connected: bool, + rtt: Option, + received_at: Instant, ) { let registration = running.generation(); let identity = ®istration.identity; @@ -156,7 +179,8 @@ impl StargateState { current_rtt, Instant::now(), effective_status, - ), + ) + .with_update_received_at(received_at), ) .await; } else { 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..6aabae08c5 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,39 +14,154 @@ // limitations under the License. use super::*; +use crate::metrics::StargateMetrics; use stargate_protocol::common::queue_time_delta_ms; +use std::sync::OnceLock; use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; + +/// Long enough that no test reservation's soft expiry passes 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, + pub(super) registration: Option>, pub(super) input_tokens: u64, pub(super) priority: u32, + ttl: Duration, + // Soft expiry: a registration update from the reserved backend received + // at or after this instant ends the reservation. Time alone never does, + // so a backend that stops reporting keeps its reservations until it is + // removed. Unset until the request is handed to the tunnel, so no update + // can settle a request that has not left Stargate. `Some(None)` when the + // TTL is too large to represent as an `Instant`; such a reservation only + // ends by release or backend removal. + expires_at: OnceLock>, 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 { + /// Creates a counted reservation whose soft expiry starts when it is + /// armed. + pub(super) fn new_unarmed( + registration: Option>, + input_tokens: u64, + priority: u32, + 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, + ttl, + expires_at: OnceLock::new(), active: AtomicBool::new(true), + metrics, + routing_key: target.routing_key.clone(), + model: target.model_id.clone(), }) } + /// Creates a reservation already armed at `now`. + #[cfg(test)] + pub(super) fn new_at( + registration: Option>, + input_tokens: u64, + priority: u32, + now: Instant, + ttl: Duration, + metrics: Option>, + target: &RoutingTargetKey, + ) -> Arc { + let pending = Self::new_unarmed(registration, input_tokens, priority, ttl, metrics, target); + pending.arm_at(now); + pending + } + pub(super) fn is_active(&self) -> bool { self.active.load(Ordering::Acquire) } + + /// Starts the soft expiry at `now`. Only the first call takes effect. + pub(super) fn arm_at(&self, now: Instant) { + let _ = self.expires_at.set(now.checked_add(self.ttl)); + } + + /// Ends the reservation when a registration update from its own backend, + /// received at `updated_at`, arrives at or after the soft expiry. Receipt + /// time is a settlement threshold, not proof that the update includes the + /// request. Earlier updates, and any update before the reservation is + /// armed, leave it counting. + pub(super) fn settle_on_backend_update(&self, updated_at: Instant) { + if let Some(Some(expires_at)) = self.expires_at.get() + && updated_at >= *expires_at + { + self.deactivate(); + } + } + + 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. +impl Drop for PendingClusterReservation { + fn drop(&mut self) { + self.deactivate(); + } +} + +// Cancellation is explicit because a successful attempt remains pending until +// its backend reports at or after the soft expiry. #[derive(Debug)] pub(crate) struct RoutingReservation(pub(super) Arc); impl RoutingReservation { + /// Starts the soft expiry once the request has been handed to the tunnel. + /// Later calls keep the first arm time. + pub(crate) fn arm(&self) { + self.0.arm_at(Instant::now()); + } + pub(crate) fn release(self) { - self.0.active.store(false, Ordering::Release); + self.0.deactivate(); + } +} + +impl Drop for RoutingReservation { + // An attempt that ends before sending its request, including by + // cancellation, arms the reservation here so it still ends on the + // backend's next update after the TTL. + fn drop(&mut self) { + self.arm(); } } @@ -101,6 +216,8 @@ pub(super) fn apply_pending_cluster_reservations( #[cfg(test)] mod tests { + use prometheus::Encoder; + use std::collections::HashMap; use super::*; @@ -120,7 +237,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 +259,190 @@ mod tests { assert_eq!(estimate(&recovered), Some(0)); } } + + #[test] + fn soft_expiry_waits_for_a_backend_update_at_or_after_it() { + 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, + ); + let at = |millis| dispatched_at + Duration::from_millis(millis); + + pending.settle_on_backend_update(at(59)); + assert!( + pending.is_active(), + "an update received before the soft expiry cannot include the request" + ); + pending.settle_on_backend_update(at(60)); + assert!(!pending.is_active()); + } + + #[test] + fn unarmed_reservation_is_not_settled_by_any_update() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let pending = PendingClusterReservation::new_unarmed( + None, + 10, + 0, + Duration::from_millis(1), + None, + &target, + ); + + pending.settle_on_backend_update(Instant::now() + Duration::from_secs(3600)); + assert!( + pending.is_active(), + "a request that has not left Stargate cannot appear in any update" + ); + } + + #[test] + fn soft_expiry_starts_when_the_reservation_is_armed() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let reserved_at = Instant::now(); + let pending = PendingClusterReservation::new_unarmed( + None, + 10, + 0, + Duration::from_millis(60), + None, + &target, + ); + let armed_at = reserved_at + Duration::from_millis(100); + let after_arm = |millis| armed_at + Duration::from_millis(millis); + + pending.settle_on_backend_update(armed_at); + assert!(pending.is_active()); + pending.arm_at(armed_at); + pending.arm_at(after_arm(1000)); + pending.settle_on_backend_update(after_arm(59)); + assert!( + pending.is_active(), + "an update before arm time plus TTL cannot include the request" + ); + pending.settle_on_backend_update(after_arm(60)); + assert!(!pending.is_active(), "the first arm time must win"); + } + + #[test] + fn dropping_an_unarmed_token_arms_its_reservation() { + let target = RoutingTargetKey::new(Some("routing".to_string()), "model"); + let metrics = StargateMetrics::new().expect("metrics should initialize"); + let ttl = Duration::from_millis(60); + let pending = + PendingClusterReservation::new_unarmed(None, 1, 0, ttl, Some(metrics.clone()), &target); + let before_drop = Instant::now(); + + drop(RoutingReservation(Arc::clone(&pending))); + assert!(active_reservation_gauge(&metrics).ends_with(" 1")); + pending.settle_on_backend_update(before_drop); + assert!(pending.is_active()); + pending.settle_on_backend_update(Instant::now() + ttl); + assert!(!pending.is_active()); + assert!(active_reservation_gauge(&metrics).ends_with(" 0")); + } + + #[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" + )); + } + + 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")); + pending.settle_on_backend_update(dispatched_at + Duration::from_secs(365 * 24 * 60 * 60)); + assert!( + pending.is_active(), + "an unrepresentable soft expiry must not settle on an update" + ); + RoutingReservation(pending).release(); + assert!(active_reservation_gauge(&metrics).ends_with(" 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..8969141369 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,59 @@ impl SelectedRoutedCluster { self.owner .reserve_backend(registration, input_tokens, priority) } + + #[cfg(test)] + pub(crate) fn reserve_backend_with_ttl( + &self, + registration: &Arc, + input_tokens: u64, + priority: u32, + ttl: Duration, + metrics: Arc, + target: &RoutingTargetKey, + ) -> Option { + let reservation = self.reserve_backend_with_rtt_ttl( + registration, + input_tokens, + priority, + |_rtt| ttl, + metrics, + target, + ); + // Tests reserve as if the request were sent at once. + reservation.inspect(RoutingReservation::arm) + } + + /// 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_for_rtt, + metrics, + target, + ) + } + + #[cfg(test)] + pub(super) fn routing_snapshot(&self) -> Option { + self.owner.routing_snapshot() + } + + #[cfg(test)] + pub(super) fn pending_reservation_count(&self) -> usize { + self.owner.pending_reservation_count() + } } #[derive(Debug)] @@ -178,6 +232,9 @@ pub struct RoutedInferenceServerSnapshot { pub stats: ModelStats, pub rtt: Duration, pub snapshot_updated_at: Instant, + // When Stargate received the registration update carrying these stats. + // It can precede `snapshot_updated_at` by connection and RTT awaits. + pub(crate) update_received_at: Instant, pub status: InferenceServerStatus, pub reverse_tunnel: bool, } @@ -197,12 +254,18 @@ impl RoutedInferenceServerSnapshot { stats, rtt, snapshot_updated_at, + update_received_at: snapshot_updated_at, status, reverse_tunnel: registration.identity.reverse_tunnel, registration, } } + pub(super) fn with_update_received_at(mut self, update_received_at: Instant) -> Self { + self.update_received_at = update_received_at; + self + } + pub(super) fn assert_registration_identity(&self) { let identity = &self.registration.identity; assert_eq!( 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..ab2ab51537 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,8 +17,10 @@ use super::registration::{ RegistrationClusterGeneration, test_registration_generation, test_registration_generation_in_cluster, }; -use super::reservations::update_reserved_priority_queue_time; -use super::snapshots::{ClusterBackendUpsert, RoutedClusterState, RoutingTargetGeneration}; +use super::reservations::{UNEXPIRING_TEST_RESERVATION_TTL, update_reserved_priority_queue_time}; +use super::snapshots::{ + ClusterBackendRemoval, ClusterBackendUpsert, RoutedClusterState, RoutingTargetGeneration, +}; use super::*; use crate::load_balancer::{ ClusterComparator, LoadBalancerAlgorithm, LoadBalancerAlgorithmConfig, @@ -26,6 +28,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; @@ -866,7 +869,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 +890,738 @@ 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 registration_update_before_soft_expiry_keeps_reservation_until_release() { + 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"); + // The registration update below is received on the real clock, so the + // soft expiry must outlast any scheduler stall. + let reservation = selected_cluster + .reserve_backend_with_ttl( + &running.generation(), + 37, + 4, + UNEXPIRING_TEST_RESERVATION_TTL, + metrics.clone(), + &target, + ) + .expect("active backend should accept reservation"); + + scenario.publish_connected(&running, &update).await; + assert_eq!( + selected_cluster.pending_reservation_count(), + 1, + "a registration update before the soft expiry must not clear the reservation" + ); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + assert_eq!( + selected_cluster + .routing_snapshot() + .expect("cluster should remain routable") + .stats + .queue_size, + 1 + ); + + reservation.release(); + assert_eq!( + selected_cluster + .routing_snapshot() + .expect("cluster should remain routable") + .stats + .queue_size, + 0 + ); + assert_eq!(selected_cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +/// Matches the default `--routing-reservation-ttl-max-ms`, which only caps +/// the soft TTL; the silent-backend tests report well past it. +const DEFAULT_TTL_MAX: Duration = Duration::from_millis(1000); + +/// Builds one routed cluster whose backends share a cluster generation. Each +/// backend reports an empty queue, so `queue_size` counts live reservations. +fn timeline_cluster( + cluster_id: &str, + backend_ids: &[&str], +) -> (Arc, Vec) { + let first = backend!(cluster_id, backend_ids[0], 1.0, 100.0, 60); + let cluster_generation = first.registration.cluster_generation.clone(); + let mut backends = vec![first]; + for backend_id in &backend_ids[1..] { + backends.push(backend!( + cluster_generation.clone() => cluster_id, backend_id, 1.0, 100.0, 60 + )); + } + let cluster = Arc::new(RoutedClusterState::new(cluster_generation)); + for backend in &backends { + cluster.upsert_backend(Arc::new(backend.clone())); + } + (cluster, backends) +} + +/// Applies a registration update from `backend` received and applied at +/// `received_at`. +fn report_at( + cluster: &RoutedClusterState, + backend: &RoutedInferenceServerSnapshot, + received_at: Instant, + queue_size: u64, +) { + report_received_at(cluster, backend, received_at, received_at, queue_size); +} + +/// Applies at `applied_at` a registration update from `backend` that +/// Stargate received at `received_at`. +fn report_received_at( + cluster: &RoutedClusterState, + backend: &RoutedInferenceServerSnapshot, + received_at: Instant, + applied_at: Instant, + queue_size: u64, +) { + let mut update = backend.clone(); + update.stats.queue_size = queue_size; + update.snapshot_updated_at = applied_at; + update.update_received_at = received_at; + cluster.upsert_backend(Arc::new(update)); +} + +fn queued(cluster: &RoutedClusterState) -> u64 { + cluster + .routing_snapshot() + .expect("cluster should remain routable") + .stats + .queue_size +} + +#[test] +fn reservation_counts_past_soft_expiry_until_its_backend_reports() { + let (cluster, backends) = timeline_cluster("cluster-soft-expiry", &["backend-a"]); + let target = make_target(Some("rk-soft-expiry"), "model-soft-expiry"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let dispatched_at = Instant::now(); + let at = |millis| dispatched_at + Duration::from_millis(millis); + let reservation = cluster + .reserve_backend_at( + &backends[0].registration, + dispatched_at, + Duration::from_millis(60), + metrics.clone(), + &target, + ) + .expect("stored backend should accept reservation"); + assert_eq!(queued(&cluster), 1); + + report_at(&cluster, &backends[0], at(1), 3); + assert_eq!( + queued(&cluster), + 4, + "an update before the soft expiry keeps the reservation on top of its stats" + ); + report_at(&cluster, &backends[0], at(59), 3); + assert_eq!(queued(&cluster), 4); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + + report_at(&cluster, &backends[0], at(60), 3); + assert_eq!( + queued(&cluster), + 3, + "an update at the soft expiry replaces the reservation with reported stats" + ); + assert_eq!(cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); + + reservation.release(); + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[test] +fn settlement_uses_update_receipt_time_not_apply_time() { + let (cluster, backends) = timeline_cluster("cluster-receipt-time", &["backend-a"]); + let target = make_target(Some("rk-receipt-time"), "model-receipt-time"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let received_before_dispatch = Instant::now(); + let at = |millis| received_before_dispatch + Duration::from_millis(millis); + let _reservation = cluster + .reserve_backend_at( + &backends[0].registration, + at(1), + Duration::from_millis(60), + metrics.clone(), + &target, + ) + .expect("stored backend should accept reservation"); + + // Connection and RTT awaits can delay applying an update well past the + // soft expiry; the update still predates the request. + report_received_at(&cluster, &backends[0], at(0), at(100), 3); + assert_eq!( + queued(&cluster), + 4, + "an update received before dispatch must not settle when applied late" + ); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + + report_received_at(&cluster, &backends[0], at(61), at(100), 3); + assert_eq!(queued(&cluster), 3); + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[tokio::test] +async fn registration_update_settles_by_its_receipt_time() { + let scenario = RegistrationScenario::new(Some("rk-receipt-apply")); + let running = scenario.start("inst-receipt-apply", 8888); + let model_id = "model-receipt-apply"; + let update = scenario.update(&running, model_id, Active, ModelStats::default()); + scenario.publish_connected(&running, &update).await; + + let selected_cluster = scenario.selected_cluster(model_id).await; + let target = scenario.target(model_id); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let received_before_dispatch = Instant::now(); + let reservation = selected_cluster + .reserve_backend_with_ttl( + &running.generation(), + 37, + 4, + Duration::from_millis(1), + metrics.clone(), + &target, + ) + .expect("active backend should accept reservation"); + let received_after_expiry = Instant::now() + Duration::from_millis(1); + + let apply_received = |received_at| { + scenario.state.apply_received_registration_update( + &running, + &update, + true, + Some(Duration::from_millis(5)), + received_at, + ) + }; + apply_received(received_before_dispatch).await; + assert_eq!( + scenario.only_cluster(model_id).await.stats.queue_size, + 1, + "an update received before dispatch must not settle the reservation" + ); + apply_received(received_after_expiry).await; + assert_eq!(scenario.only_cluster(model_id).await.stats.queue_size, 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); + + drop(reservation); + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[test] +fn silent_backend_keeps_reservations_until_removed() { + // A hung backend keeps its connection but stops sending updates. Its + // reservations must keep steering traffic away for as long as it stays + // registered, even while another backend in the cluster keeps reporting. + let (cluster, backends) = timeline_cluster("cluster-hung", &["backend-a", "backend-b"]); + let (reporting, silent) = (&backends[0], &backends[1]); + let target = make_target(Some("rk-hung"), "model-hung"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let dispatched_at = Instant::now(); + let reservations: Vec<_> = (0..20) + .map(|_| { + cluster + .reserve_backend_at( + &silent.registration, + dispatched_at, + Duration::from_millis(1), + metrics.clone(), + &target, + ) + .expect("stored backend should accept reservation") + }) + .collect(); + + for reported_after in [ + Duration::from_millis(2), + DEFAULT_TTL_MAX, + 10 * DEFAULT_TTL_MAX, + ] { + report_at(&cluster, reporting, dispatched_at + reported_after, 0); + assert_eq!( + queued(&cluster), + 20, + "the silent backend's reservations must survive {reported_after:?}" + ); + assert_eq!(routing_reservations_active(&metrics, &target), 20); + } + + assert_eq!( + cluster.remove_backend(&silent.registration), + ClusterBackendRemoval::Removed + ); + assert_eq!(queued(&cluster), 0); + assert_eq!(cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); + for reservation in reservations { + reservation.release(); + } + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[test] +fn backend_update_does_not_settle_another_backends_reservation() { + let (cluster, backends) = timeline_cluster("cluster-per-backend", &["backend-a", "backend-b"]); + let (backend_a, backend_b) = (&backends[0], &backends[1]); + let target = make_target(Some("rk-per-backend"), "model-per-backend"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let dispatched_at = Instant::now(); + let at = |millis| dispatched_at + Duration::from_millis(millis); + let reserve = |backend: &RoutedInferenceServerSnapshot| { + cluster + .reserve_backend_at( + &backend.registration, + dispatched_at, + Duration::from_millis(60), + metrics.clone(), + &target, + ) + .expect("stored backend should accept reservation") + }; + let _on_a = reserve(backend_a); + let _on_b = reserve(backend_b); + + report_at(&cluster, backend_a, at(100), 0); + assert_eq!( + queued(&cluster), + 1, + "backend A's update settles only backend A's reservation" + ); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + + report_at(&cluster, backend_b, at(100), 0); + assert_eq!(queued(&cluster), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[test] +fn backend_update_prunes_ten_thousand_soft_expired_reservations() { + let (cluster, backends) = timeline_cluster("cluster-res-prune", &["backend-a"]); + let target = make_target(Some("rk-res-prune"), "model-res-prune"); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + let dispatched_at = Instant::now(); + for _ in 0..10_000 { + let _reservation = cluster + .reserve_backend_at( + &backends[0].registration, + dispatched_at, + Duration::from_millis(1), + metrics.clone(), + &target, + ) + .expect("stored backend should accept reservation"); + } + + assert_eq!(cluster.pending_reservation_count(), 10_000); + assert_eq!(routing_reservations_active(&metrics, &target), 10_000); + assert_eq!( + queued(&cluster), + 10_000, + "a read must not prune soft-expired reservations without a backend update" + ); + report_at( + &cluster, + &backends[0], + dispatched_at + Duration::from_millis(1), + 0, + ); + assert_eq!(cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); +} + +#[tokio::test] +async fn deregistering_silent_backend_drops_soft_expired_reservations() { + let scenario = RegistrationScenario::new(Some("rk-res-deregister")); + let silent = scenario.start_in("inst-silent", "cluster-res-deregister", 1111); + let reporting = scenario.start_in("inst-reporting", "cluster-res-deregister", 2222); + let model_id = "model-res-deregister"; + scenario.activate(&silent, model_id).await; + scenario.activate(&reporting, model_id).await; + + let selected_cluster = scenario.selected_cluster(model_id).await; + let target = scenario.target(model_id); + let metrics = crate::metrics::StargateMetrics::new().expect("metrics should initialize"); + // A zero soft TTL is already past at every later update or read. + let reservation = selected_cluster + .reserve_backend_with_ttl( + &silent.generation(), + 37, + 4, + Duration::ZERO, + metrics.clone(), + &target, + ) + .expect("active backend should accept reservation"); + + scenario.activate(&reporting, model_id).await; + assert_eq!(scenario.only_cluster(model_id).await.stats.queue_size, 1); + assert_eq!(routing_reservations_active(&metrics, &target), 1); + + scenario.state.end_registration(silent).await; + assert_eq!(scenario.only_cluster(model_id).await.stats.queue_size, 0); + assert_eq!(selected_cluster.pending_reservation_count(), 0); + assert_eq!(routing_reservations_active(&metrics, &target), 0); + + drop(reservation); + assert_eq!(routing_reservations_active(&metrics, &target), 0); + scenario.state.end_registration(reporting).await; +} + +#[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, + UNEXPIRING_TEST_RESERVATION_TTL, + 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(); +} + +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; +} + +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"); + 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, + UNEXPIRING_TEST_RESERVATION_TTL, + 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] @@ -915,7 +1646,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; @@ -923,7 +1654,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); diff --git a/src/libraries/rust/stargate/crates/stargate/src/runtime.rs b/src/libraries/rust/stargate/crates/stargate/src/runtime.rs index 5e8785e6c4..0c875058c4 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. @@ -255,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 @@ -373,6 +380,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(), @@ -394,13 +403,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; }); } @@ -621,6 +624,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), @@ -827,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(); 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/crates/stargate/tests/suite/health_lifecycle.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs index ff61991635..b58ae9229c 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/health_lifecycle.rs @@ -16,7 +16,9 @@ use std::net::SocketAddr; use std::time::Duration; -use stargate::runtime::{BoundStargateListeners, StargateRuntime, StargateRuntimeConfig, WarmupConfig}; +use stargate::runtime::{ + BoundStargateListeners, StargateRuntime, StargateRuntimeConfig, WarmupConfig, +}; use crate::common::{SelfDiscovery, base_config, init_crypto, make_stargate_runtime}; @@ -27,8 +29,7 @@ fn make_stargate_runtime_with_warmup( let ephemeral: SocketAddr = "127.0.0.1:0".parse().unwrap(); let mut config: StargateRuntimeConfig = base_config(id, ephemeral, ephemeral); config.warmup = warmup; - let listeners = - BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); + let listeners = BoundStargateListeners::bind(&mut config).expect("test listeners should bind"); let grpc_addr = config.grpc_listen_addr; let http_addr = config.http_listen_addr; let discovery = Box::new(SelfDiscovery::new(id, grpc_addr, http_addr)); diff --git a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs index b710baf5d3..71da4db0f0 100644 --- a/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs +++ b/src/libraries/rust/stargate/crates/stargate/tests/suite/load_balancing.rs @@ -1542,9 +1542,13 @@ async fn wait_and_widen_capacity_rejection_returns_overloaded_error() { routing_key: None, model_id: model.to_string(), }; + // Wait on the backend's own report, not the cluster snapshot. The cluster + // snapshot still counts the reservation from the wait_for_routing probe, + // and a registration update built before set_stats can settle that + // reservation after this loop, leaving a free slot for the next request. let deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { - let candidates = state.cluster_candidates_for_target(&target).await; + let candidates = state.candidates_for_target(&target).await; if candidates.len() == 1 && candidates[0].stats.max_engine_concurrency == 1 && candidates[0].stats.num_running_queries >= 1 diff --git a/src/libraries/rust/stargate/docs/load-balancer-configuration.md b/src/libraries/rust/stargate/docs/load-balancer-configuration.md index 0e37854f6d..2c30926668 100644 --- a/src/libraries/rust/stargate/docs/load-balancer-configuration.md +++ b/src/libraries/rust/stargate/docs/load-balancer-configuration.md @@ -151,8 +151,13 @@ 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. A reservation ends when Stargate +receives a registration update from the selected backend at or after the +reservation's soft expiry (the backend's health-check RTT, bounded by the +minimum and maximum TTL flags, after Stargate sends the request), or when the +backend is removed. See +[Runtime stats interface](runtime-stats-interface.md#registration-updates). 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..fdcfffca43 100644 --- a/src/libraries/rust/stargate/docs/runtime-stats-interface.md +++ b/src/libraries/rust/stargate/docs/runtime-stats-interface.md @@ -40,6 +40,59 @@ 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 for the selected backend +when the backend is selected. The reservation's TTL is the selected backend's +health-check RTT (not the cluster mean), bounded by +`--routing-reservation-ttl-min-ms` (default `1`) and +`--routing-reservation-ttl-max-ms` (default `1000`). The maximum only limits +the effect of a bad RTT sample; it does not end the reservation. The soft +expiry is one TTL after Stargate opens the tunnel stream and sends the request +head. Until then, for example while the stream open waits for tunnel capacity, +no update can end the reservation. An attempt that ends before sending starts +the soft expiry when it ends. + +The reservation counts until Stargate receives a registration update from the +selected backend at or after the soft expiry. Receipt time is the settlement +threshold; it does not show when the update was generated or whether it includes +the request. The request tunnel and registration stream are independent, so an +update generated before the backend receives a request can arrive after the +threshold, while an update that includes the request can arrive before it. This +can briefly leave the request counted in both reported stats and the +reservation, or settle the reservation before the request appears in reported +stats. Updates from other backends in the same cluster never settle it. A queue +estimate mismatch rejection releases the reservation early, and removing or +deregistering the backend drops it. + +Waiting for the backend's own update covers two cases: + +- A backend that hangs but keeps its connection stops sending updates. Its last + reported stats still look idle. With a time-based expiry or cap, its + reservations would all end together and those idle stats would attract a new + burst of requests each time. Held reservations keep steering traffic away + until the backend reports again or deregisters. +- Pylon coalesces change-driven updates (`--stats-update-coalesce-ms`, default + `10`). The reservation uses the selected backend's soft-expiry threshold to + decide when its own updates can settle it. Because the streams are + independent, receipt time is a heuristic and does not prove that an update + includes the request. A larger minimum can lengthen the period in which a + request is counted in both reported stats and its reservation. + +The trade-off is that a reservation on a backend that stops reporting stays +counted, in routing and in `stargate_routing_reservations_active`, until that +backend reports or its registration ends. Only registration liveness bounds +those reservations, so setting `--registration-update-idle-timeout-ms` to `0` +leaves them unbounded for a backend that keeps serving requests without sending +updates. + +Pylons that send only heartbeats keep each reservation until the first +heartbeat received at or after its soft expiry. Receipt time does not prove +that the heartbeat includes the request. The reservation cannot expire while +waiting for that heartbeat, although the reported stats may still lag the +request. +Deploying Pylons with +change-driven stats updates before this Stargate version is recommended for the +most current view but is not required. + ## Stream Events Each non-empty line is JSON with `v: 1` and `type`.