From 29b6b0ac2e6c0a344192eceeeaf55c0cd2f166e8 Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Fri, 11 Sep 2026 21:22:52 +0000 Subject: [PATCH 1/5] fix(pylon): bind request state to individual lifetimes Fence queue guards and observations with an internal request identity so a reused correlation ID cannot let earlier work change the replacement. Own queue cleanup across body validation and rejection. Refs: NVIDIA/nvcf#1817 --- .../crates/pylon-lib/src/bringup/upstream.rs | 1 + .../crates/pylon-lib/src/queue_admission.rs | 189 ++++++++++++++---- .../pylon-lib/src/quic_http_tunnel/core.rs | 22 +- .../pylon-lib/src/quic_http_tunnel/tests.rs | 1 + .../crates/pylon-lib/src/request_observer.rs | 36 ++++ .../pylon-lib/src/request_observer/headers.rs | 2 + .../crates/pylon-lib/src/runtime_state.rs | 62 ++++-- .../crates/pylon-lib/src/stats/collector.rs | 3 + 8 files changed, 257 insertions(+), 59 deletions(-) diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs index 4e11e47edb..eb07499282 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs @@ -212,6 +212,7 @@ pub(super) async fn send_completion_request( TunnelRequestObserver::accepted( RequestObservationEndpoint::ChatCompletions, RequiredTunnelHeaders { + request_instance: Default::default(), request_id: request_id.clone(), routing_key: None, model_id: model_id.to_string(), diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs index dc1ff4752b..2c24a32050 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs @@ -24,7 +24,7 @@ use stargate_protocol::common::{ use stargate_protocol::tunnel_contract::HEADER_STARGATE_EXPECTED_QUEUE_MS; use crate::request_observer::{RequestObservation, RequestObservationState, RequiredTunnelHeaders}; -use crate::runtime_state::ModelGeneration; +use crate::runtime_state::{ModelGeneration, RequestInstance}; pub(crate) const RETRY_REASON_QUEUE_ESTIMATE_MISMATCH: &str = "queue_estimate_mismatch"; @@ -55,10 +55,16 @@ pub(crate) struct LiveRequestState { #[derive(Debug, Default)] struct QueueAdmissionState { - requests: HashMap, + requests: HashMap, models: HashMap, } +#[derive(Debug)] +struct LiveRequestRecord { + instance: Option, + request: LiveRequest, +} + #[derive(Debug, Default)] struct QueueModelState { last_mean_input_tps: Option, @@ -155,6 +161,7 @@ pub(crate) struct QueueModelSnapshot { pub(crate) struct QueueTrackedRequestGuard { live_requests: LiveRequestState, request_id: String, + instance: RequestInstance, finished: bool, } @@ -221,7 +228,7 @@ impl LiveRequestState { .lock() .requests .get(request_id) - .map(|request| request.generation().clone()) + .map(|record| record.request.generation().clone()) } pub(crate) fn update_generation_throughput( @@ -247,7 +254,7 @@ impl LiveRequestState { let request_ids = state .requests .iter() - .filter(|(_, request)| request.generation() == generation) + .filter(|(_, record)| record.request.generation() == generation) .map(|(request_id, _)| request_id.clone()) .collect::>(); for request_id in request_ids { @@ -302,7 +309,7 @@ impl LiveRequestState { let excluded_request = state .requests .get(&required.request_id) - .and_then(|request| match request { + .and_then(|record| match &record.request { LiveRequest::Queue(queue, _) => Some(queue), LiveRequest::Observed(_) => None, }) @@ -343,11 +350,11 @@ impl LiveRequestState { self.track_generation_request(required, ModelGeneration::new(required.model_id.clone(), 0)) } - pub(crate) fn track_generation_request( + pub(crate) fn begin_request( &self, required: &RequiredTunnelHeaders, generation: ModelGeneration, - ) -> QueueTrackedRequestGuard { + ) { let request_id = required.request_id.clone(); let request = TrackedPromptRequest { generation, @@ -360,12 +367,40 @@ impl LiveRequestState { let mut state = self.inner.lock(); let observed = state .remove_request(&request_id) - .and_then(|(_, request)| request.into_observed()); - state.insert_request(request_id.clone(), LiveRequest::Queue(request, observed)); + .and_then(|(_, request)| request.request.into_observed()); + state.insert_request( + request_id, + LiveRequest::Queue(request, observed), + Some(required.request_instance.clone()), + ); } + } + + #[cfg(test)] + pub(crate) fn track_generation_request( + &self, + required: &RequiredTunnelHeaders, + generation: ModelGeneration, + ) -> QueueTrackedRequestGuard { + self.begin_request(required, generation); + self.request_guard(required) + } + + pub(crate) fn track_existing_request( + &self, + required: &RequiredTunnelHeaders, + ) -> Option { + self.inner + .lock() + .owns_instance(&required.request_id, &required.request_instance) + .then(|| self.request_guard(required)) + } + + fn request_guard(&self, required: &RequiredTunnelHeaders) -> QueueTrackedRequestGuard { QueueTrackedRequestGuard { live_requests: self.clone(), - request_id, + request_id: required.request_id.clone(), + instance: required.request_instance.clone(), finished: false, } } @@ -390,22 +425,30 @@ impl LiveRequestState { observe: impl FnOnce(&RequestObservationTransition), ) -> RequestObservationTransition { let generation = ModelGeneration::new(observation.model_id.clone(), 0); - self.transition_generation_observation_with(observation, Some(&generation), observe) + self.transition_generation_observation_with(observation, Some(&generation), None, observe) + .unwrap() } pub(crate) fn transition_generation_observation_with( &self, observation: &RequestObservation, generation: Option<&ModelGeneration>, + instance: Option<&RequestInstance>, observe: impl FnOnce(&RequestObservationTransition), - ) -> RequestObservationTransition { + ) -> Option { let _order = self.observation_order.lock(); - let transition = self - .inner - .lock() - .transition_observation(observation, generation); + let transition = { + let mut state = self.inner.lock(); + if let Some(instance) = instance + && generation.is_some() + && !state.owns_instance(&observation.request_id, instance) + { + return None; + } + state.transition_observation(observation, generation, instance) + }; observe(&transition); - transition + Some(transition) } pub(crate) fn update_active_output_tps( @@ -418,24 +461,44 @@ impl LiveRequestState { .update_active_output_tps(request_id, active_chat_output_tps) } - pub(crate) fn finish_queue_request(&self, request_id: &str) { + pub(crate) fn finish_queue_request( + &self, + request_id: &str, + instance: Option<&RequestInstance>, + ) { let mut state = self.inner.lock(); - if let Some((request_id, request)) = state.remove_request(request_id) - && let Some(observed) = request.into_observed() + if instance.is_some_and(|instance| !state.owns_instance(request_id, instance)) { + return; + } + if let Some((request_id, record)) = state.remove_request(request_id) + && let Some(observed) = record.request.into_observed() { - state.insert_request(request_id, LiveRequest::Observed(observed)); + state.insert_request(request_id, LiveRequest::Observed(observed), record.instance); } } } impl QueueAdmissionState { + fn owns_instance(&self, request_id: &str, instance: &RequestInstance) -> bool { + self.requests + .get(request_id) + .is_some_and(|record| record.instance.as_ref() == Some(instance)) + } + fn transition_observation( &mut self, observation: &RequestObservation, generation: Option<&ModelGeneration>, + instance: Option<&RequestInstance>, ) -> RequestObservationTransition { + let prior = self.remove_request(&observation.request_id); + let instance = instance.cloned().or_else(|| { + prior + .as_ref() + .and_then(|(_, record)| record.instance.clone()) + }); let (request_id, prior_queue, prior_observed) = - match self.remove_request(&observation.request_id) { + match prior.map(|(id, record)| (id, record.request)) { Some((request_id, LiveRequest::Queue(queue, observed))) => { (request_id, Some(queue), observed) } @@ -493,6 +556,7 @@ impl QueueAdmissionState { }, Some(current.clone()), ), + instance, ); Some(current) } @@ -513,14 +577,28 @@ impl QueueAdmissionState { request_id: &str, active_chat_output_tps: Option, ) -> Option { - let (request_id, mut request) = self.remove_request(request_id)?; - let model_id = request.update_active_output_tps(active_chat_output_tps); - self.insert_request(request_id, request); + let (request_id, mut record) = self.remove_request(request_id)?; + let model_id = record + .request + .update_active_output_tps(active_chat_output_tps); + self.insert_request(request_id, record.request, record.instance); model_id } - fn advance_request_phase(&mut self, request_id: &str, next_phase: TrackedPromptPhase) { - let Some(LiveRequest::Queue(request, _)) = self.requests.get_mut(request_id) else { + fn advance_request_phase( + &mut self, + request_id: &str, + instance: &RequestInstance, + next_phase: TrackedPromptPhase, + ) { + if !self.owns_instance(request_id, instance) { + return; + } + let Some(LiveRequest::Queue(request, _)) = self + .requests + .get_mut(request_id) + .map(|record| &mut record.request) + else { return; }; if next_phase <= request.phase { @@ -535,15 +613,21 @@ impl QueueAdmissionState { request.phase = next_phase; } - fn remove_request(&mut self, request_id: &str) -> Option<(String, LiveRequest)> { - let (request_id, request) = self.requests.remove_entry(request_id)?; - self.adjust_live_request(&request, -1); - Some((request_id, request)) + fn remove_request(&mut self, request_id: &str) -> Option<(String, LiveRequestRecord)> { + let (request_id, record) = self.requests.remove_entry(request_id)?; + self.adjust_live_request(&record.request, -1); + Some((request_id, record)) } - fn insert_request(&mut self, request_id: String, request: LiveRequest) { + fn insert_request( + &mut self, + request_id: String, + request: LiveRequest, + instance: Option, + ) { self.adjust_live_request(&request, 1); - self.requests.insert(request_id, request); + self.requests + .insert(request_id, LiveRequestRecord { instance, request }); } fn adjust_live_request(&mut self, request: &LiveRequest, delta: i8) { @@ -786,17 +870,26 @@ impl TrackedPromptPhase { impl QueueTrackedRequestGuard { pub(crate) fn on_backend_submission(&mut self) { let mut state = self.live_requests.inner.lock(); - state.advance_request_phase(&self.request_id, TrackedPromptPhase::InputProcessing); + state.advance_request_phase( + &self.request_id, + &self.instance, + TrackedPromptPhase::InputProcessing, + ); } pub(crate) fn observe_output(&mut self) { let mut state = self.live_requests.inner.lock(); - state.advance_request_phase(&self.request_id, TrackedPromptPhase::OutputGeneration); + state.advance_request_phase( + &self.request_id, + &self.instance, + TrackedPromptPhase::OutputGeneration, + ); } pub(crate) fn finish(&mut self) { if !self.finished { - self.live_requests.finish_queue_request(&self.request_id); + self.live_requests + .finish_queue_request(&self.request_id, Some(&self.instance)); self.finished = true; } } @@ -840,7 +933,7 @@ mod tests { .lock() .requests .values() - .filter(|request| matches!(request, LiveRequest::Queue(..))) + .filter(|record| matches!(record.request, LiveRequest::Queue(..))) .count() } } @@ -856,6 +949,7 @@ mod tests { input_tokens: u64, ) -> RequiredTunnelHeaders { RequiredTunnelHeaders { + request_instance: Default::default(), request_id: request_id.to_string(), routing_key: None, model_id: model_id.to_string(), @@ -910,6 +1004,27 @@ mod tests { } } + #[test] + fn stale_guards_cannot_advance_or_remove_reused_request_ids() { + for replacement_model in ["model-a", "model-b"] { + let live = LiveRequestState::default(); + let mut first = live.track_request(&required_for_model("same-id", "model-a", 0, 100)); + let replacement = + live.track_request(&required_for_model("same-id", replacement_model, 1, 20)); + let expected = live.snapshot_model(replacement_model); + first.on_backend_submission(); + first.observe_output(); + drop(first); + assert_eq!(live.snapshot_model(replacement_model), expected); + assert_eq!(expected.num_running_queries, 1); + drop(replacement); + assert_eq!( + live.snapshot_model(replacement_model).num_running_queries, + 0 + ); + } + } + #[test] fn one_live_request_transition_updates_queue_and_active_output_load() { let live_requests = LiveRequestState::default(); diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs index 7fbde8e567..1b61f8b9f0 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs @@ -314,6 +314,10 @@ impl TunnelRequestLifecycle { required: RequiredTunnelHeaders, generation: Option, ) -> Self { + if observation_endpoint.is_none() { + app.runtime_state + .begin_request(&required, generation.as_ref()); + } let observer = observation_endpoint.map(|endpoint| { TunnelRequestObserver::accepted( endpoint, @@ -332,11 +336,14 @@ impl TunnelRequestLifecycle { model_label: request_headers[HEADER_MODEL].clone(), }); + let queue_request = app + .runtime_state + .track_generation_request(&required, generation.as_ref()); Self { required, generation, observer, - queue_request: None, + queue_request, quality_check, } } @@ -374,16 +381,14 @@ impl TunnelRequestLifecycle { "evaluated local queue mismatch admission" ); if !matches!(decision, QueueAdmissionDecision::Rejected { .. }) { - self.queue_request = app - .runtime_state - .track_generation_request(required, self.generation.as_ref()); return None; } - // Observers are created before admission so body validation and terminal - // accounting keep their existing order. Remove the queue projection before - // sending the rejection; fail() clears the observed lifecycle projection. - app.runtime_state.finish_queue_request(&required.request_id); + // Release only this request's queue projection before reporting rejection. + // The observer's terminal transition clears its observed projection. + if let Some(guard) = self.queue_request.as_mut() { + guard.finish(); + } self.fail(); Some(decision) } @@ -1629,6 +1634,7 @@ mod tests { app.force_chat_completions_include_usage = true; app.runtime_state.update_model_throughput("model-a", 100.0); let _queued_request = app.runtime_state.track_request(&RequiredTunnelHeaders { + request_instance: Default::default(), request_id: "req-already-queued".to_string(), routing_key: None, model_id: "model-a".to_string(), diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs index 70ba319fd2..2db2952a64 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs @@ -775,6 +775,7 @@ async fn start_queue_mismatch_test_tunnel( .forwarding .runtime_state .track_request(&RequiredTunnelHeaders { + request_instance: Default::default(), request_id: "req-already-queued".to_string(), routing_key: Some("rk-1".to_string()), model_id: "model-a".to_string(), diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index 5cca12e87e..3853b5ac06 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs @@ -165,6 +165,7 @@ enum ChatChoiceState { pub(crate) struct RequestObserver { endpoint: RequestObservationEndpoint, request_id: String, + request_instance: crate::runtime_state::RequestInstance, started_at: Instant, routing_key: Option, model_id: String, @@ -186,9 +187,11 @@ impl RequestObserver { generation: Option, runtime_state: PylonRuntimeState, ) -> Self { + runtime_state.begin_request(&required, generation.as_ref()); let priority = required.queue_priority(); let RequiredTunnelHeaders { request_id, + request_instance, routing_key, model_id, priority: _, @@ -199,6 +202,7 @@ impl RequestObserver { let mut observer = Self { endpoint, request_id, + request_instance, started_at: accepted_at, routing_key, model_id, @@ -541,6 +545,7 @@ impl RequestObserver { self.runtime_state.observe_request_for_generation( RequestObservationEvent { observation, + request_instance: Some(self.request_instance.clone()), generation: self.generation.clone(), changed_generations: Vec::new(), input_interval, @@ -712,6 +717,35 @@ mod tests { RequestObserver::new(&request_headers(request_id, 42), runtime_state).unwrap() } + #[test] + fn replaced_observer_cannot_change_the_current_request_instance() { + let runtime = PylonRuntimeState::new(InferenceServerStatus::Active, &["model-a".into()]); + let generation = runtime.current_generation("model-a"); + let make_observer = |input_tokens| { + RequestObserver::from_required( + RequestObservationEndpoint::ChatCompletions, + validate_required_tunnel_headers(&request_headers("reused-id", input_tokens)) + .unwrap(), + generation.clone(), + runtime.clone(), + ) + }; + let mut first = make_observer(100); + let mut replacement = make_observer(20); + replacement.submit_now(); + let expected = runtime.snapshot_live_model("model-a"); + first.submit_now(); + first.observe_output_message(); + first.cancel(); + assert_eq!(runtime.snapshot_live_model("model-a"), expected); + assert_eq!(expected.num_running_queries, 1); + replacement.cancel(); + assert_eq!( + runtime.snapshot_live_model("model-a").num_running_queries, + 0 + ); + } + #[test] fn chat_calibration_state_exists_only_when_enabled() { for (enabled, expected_state) in [(false, false), (true, true)] { @@ -898,6 +932,7 @@ mod tests { fn embeddings_required_headers() -> RequiredTunnelHeaders { RequiredTunnelHeaders { + request_instance: Default::default(), request_id: "req-embeddings-terminal".to_string(), routing_key: Some("rk-1".to_string()), model_id: "model-embed".to_string(), @@ -1146,6 +1181,7 @@ mod tests { .checked_sub(Duration::from_secs(10)) .expect("test acceptance time should be representable"); let required = RequiredTunnelHeaders { + request_instance: Default::default(), request_id: "req-upstream-completion".to_string(), routing_key: Some("rk-1".to_string()), model_id: "model-a".to_string(), diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs index 2b11e8f4c6..a62cf13178 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs @@ -60,6 +60,7 @@ impl MissingRequiredHeaderError { #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct RequiredTunnelHeaders { pub request_id: String, + pub request_instance: crate::runtime_state::RequestInstance, pub routing_key: Option, pub model_id: String, pub priority: Option, @@ -88,6 +89,7 @@ pub(crate) fn validate_required_tunnel_headers( let priority = parse_optional_numeric_header(request_headers, HEADER_PRIORITY)?; Ok(RequiredTunnelHeaders { request_id, + request_instance: Default::default(), routing_key, model_id, priority, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index 8558a8b81a..01bdfd35fb 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -82,6 +82,18 @@ impl ModelGeneration { } } +/// Identifies one local request lifetime independently of the caller's request ID. +#[derive(Clone, Debug, Default)] +pub(crate) struct RequestInstance(Arc<()>); + +impl PartialEq for RequestInstance { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for RequestInstance {} + #[derive(Clone, Debug, Default)] pub struct PylonRuntimeState { advertised: Arc>, @@ -94,6 +106,7 @@ pub struct PylonRuntimeState { #[derive(Clone, Debug)] pub struct RequestObservationEvent { pub(crate) observation: RequestObservation, + pub(crate) request_instance: Option, pub(crate) generation: Option, pub(crate) changed_generations: Vec, pub(crate) input_interval: Option, @@ -443,6 +456,7 @@ impl PylonRuntimeState { let request_input_tokens = observation.input_tokens; self.observe_request_for_generation( RequestObservationEvent { + request_instance: None, observation, generation, changed_generations: Vec::new(), @@ -460,7 +474,11 @@ impl PylonRuntimeState { event: RequestObservationEvent, request_input_tokens: u64, ) { - let event = self.transition_request_observation_for_generation(event, request_input_tokens); + let Some(event) = + self.transition_request_observation_for_generation(event, request_input_tokens) + else { + return; + }; if let Some(tx) = &self.observation_tx && let Err(error) = tx.try_send(event) { @@ -502,6 +520,7 @@ impl PylonRuntimeState { let request_input_tokens = observation.input_tokens; self.transition_request_observation_for_generation( RequestObservationEvent { + request_instance: None, observation, generation: Some(generation), changed_generations: Vec::new(), @@ -512,13 +531,14 @@ impl PylonRuntimeState { }, request_input_tokens, ) + .expect("test observation must target a current generation") } fn transition_request_observation_for_generation( &self, mut event: RequestObservationEvent, request_input_tokens: u64, - ) -> RequestObservationEvent { + ) -> Option { // Held across the queue transition below: retire_generation() purges // live-request state under this lock, so releasing it after the // currency check would let a retired generation reinsert queue state. @@ -536,7 +556,7 @@ impl PylonRuntimeState { current_generation = ?current_generation, "dropping request observation from a retired model generation" ); - return event; + return None; } } let mut live_observation = event.observation.clone(); @@ -544,14 +564,15 @@ impl PylonRuntimeState { let transition = self.live_requests.transition_generation_observation_with( &live_observation, event.generation.as_ref(), + event.request_instance.as_ref(), |transition| { if let Some(metrics) = &self.metrics { metrics.observe_request_transition(&event.observation, transition); } }, - ); + )?; event.changed_generations = transition.changed_generations; - event + Some(event) } pub(crate) fn update_request_active_output_tps( @@ -608,6 +629,19 @@ impl PylonRuntimeState { .track_generation_request(required, generation) } + pub(crate) fn begin_request( + &self, + required: &RequiredTunnelHeaders, + generation: Option<&ModelGeneration>, + ) { + let Some(generation) = generation else { return }; + let advertised = self.advertised.lock(); + if advertised.current(generation).is_some() { + self.live_requests + .begin_request(required, generation.clone()); + } + } + pub(crate) fn track_generation_request( &self, required: &RequiredTunnelHeaders, @@ -616,14 +650,12 @@ impl PylonRuntimeState { let generation = generation?; let advertised = self.advertised.lock(); advertised.current(generation)?; - Some( - self.live_requests - .track_generation_request(required, generation.clone()), - ) + self.live_requests.track_existing_request(required) } + #[cfg(test)] pub(crate) fn finish_queue_request(&self, request_id: &str) { - self.live_requests.finish_queue_request(request_id); + self.live_requests.finish_queue_request(request_id, None); } #[cfg(test)] @@ -692,10 +724,11 @@ mod tests { runtime_state: &PylonRuntimeState, observation: RequestObservation, generation: ModelGeneration, - ) -> RequestObservationEvent { + ) -> Option { let request_input_tokens = observation.input_tokens; runtime_state.transition_request_observation_for_generation( RequestObservationEvent { + request_instance: None, observation, generation: Some(generation), changed_generations: Vec::new(), @@ -942,14 +975,15 @@ mod tests { assert!(runtime_state.publish_generation(&first)); let mut first_observation = observation("req-first", "model-a", None); - transition_for_generation(&runtime_state, first_observation.clone(), first.clone()); + transition_for_generation(&runtime_state, first_observation.clone(), first.clone()) + .unwrap(); assert_eq!(runtime_state.snapshot_live_model("model-a").queue_size, 1); assert!(runtime_state.retire_generation(&first).is_some()); assert!(runtime_state.begin_generation(replacement.clone())); assert!(runtime_state.publish_generation(&replacement)); first_observation.state = RequestObservationState::Complete; - transition_for_generation(&runtime_state, first_observation, first); + assert!(transition_for_generation(&runtime_state, first_observation, first).is_none()); assert_eq!( runtime_state.snapshot_live_model("model-a"), @@ -974,7 +1008,7 @@ mod tests { let mut stale = observation("req-first", "model-a", Some("rk-a")); stale.state = RequestObservationState::Failed; - transition_for_generation(&runtime_state, stale, first); + assert!(transition_for_generation(&runtime_state, stale, first).is_none()); let body = metrics.gather_text().expect("metrics should encode"); assert!( diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 35abda75fd..5000b7216b 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -802,6 +802,7 @@ mod tests { let request_input_tokens = observation.input_tokens; runtime_state.observe_request_for_generation( crate::runtime_state::RequestObservationEvent { + request_instance: None, observation, generation, changed_generations: Vec::new(), @@ -2252,6 +2253,7 @@ mod tests { ); let _queued = runtime_state.track_request(&crate::request_observer::RequiredTunnelHeaders { + request_instance: Default::default(), request_id: "req-queued-after-fallback-samples".to_string(), routing_key: None, model_id: "model-a".to_string(), @@ -3674,6 +3676,7 @@ mod tests { assert_eq!(stats.last_mean_input_tps, 2_200.0); let _queued = collector.runtime_state.track_request( &crate::request_observer::RequiredTunnelHeaders { + request_instance: Default::default(), request_id: "req-queued".to_string(), routing_key: None, model_id: "model-a".to_string(), From 018708b67e877ae290c271b38c5db2ce73b76876 Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Fri, 11 Sep 2026 22:04:19 +0000 Subject: [PATCH 2/5] fix(pylon): fence admission exclusions by request instance Preserve the current engine-concurrency behavior while preventing stale requests from excluding replacement reservations. Update the new capacity regression fixture for request identities. Refs: NVIDIA/nvcf#1817 --- .../crates/pylon-lib/src/queue_admission.rs | 25 +++++++++++++++++++ .../crates/pylon-lib/src/stats/collector.rs | 1 + 2 files changed, 26 insertions(+) diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs index 2c24a32050..37d4695002 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs @@ -309,6 +309,9 @@ impl LiveRequestState { let excluded_request = state .requests .get(&required.request_id) + .filter(|record| { + record.instance.as_ref() == Some(&required.request_instance) + }) .and_then(|record| match &record.request { LiveRequest::Queue(queue, _) => Some(queue), LiveRequest::Observed(_) => None, @@ -1004,6 +1007,28 @@ mod tests { } } + #[test] + fn admission_excludes_only_its_own_request_instance() { + let live = LiveRequestState::default(); + live.update_model_throughput("model-a", 100.0); + let stale = required("reused-id", 0, 100); + let _first = live.track_request(&stale); + let current = required("reused-id", 0, 200); + let _replacement = live.track_request(¤t); + let generation = ModelGeneration::new("model-a", 0); + let evaluate = |request| { + live.evaluate_generation( + &PylonQueueMismatchRetryConfig::default(), + request, + Some(&generation), + &headers_with_expected("0"), + Some(1), + ) + }; + assert_eq!(evaluate(&stale).actual_ms(), Some(2000)); + assert_eq!(evaluate(¤t).actual_ms(), Some(0)); + } + #[test] fn stale_guards_cannot_advance_or_remove_reused_request_ids() { for replacement_model in ["model-a", "model-b"] { diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 5000b7216b..79b2b279d0 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -3724,6 +3724,7 @@ mod tests { 2 ); let required = |request_id: &str| crate::request_observer::RequiredTunnelHeaders { + request_instance: Default::default(), request_id: request_id.to_string(), routing_key: None, model_id: "model-a".to_string(), From 7fb1079f003434f8f2aee0c4d04f5ae430ab675e Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Mon, 14 Sep 2026 16:39:56 +0000 Subject: [PATCH 3/5] fix(pylon): preserve request ownership through observation cleanup Fence queued collector updates by request instance and retire inherited observation gauges when a replacement has no observer. Cover delayed events and same-model and cross-model replacements. Refs: NVIDIA/nvcf#1817 --- .../crates/pylon-lib/src/queue_admission.rs | 36 +++++++-- .../pylon-lib/src/quic_http_tunnel/core.rs | 2 +- .../crates/pylon-lib/src/request_observer.rs | 50 ++++++++++++- .../crates/pylon-lib/src/runtime_state.rs | 22 ++++-- .../crates/pylon-lib/src/stats/collector.rs | 74 +++++++++++++++++++ .../crates/pylon-lib/src/stats/metrics.rs | 13 +++- .../crates/pylon-lib/src/stats/projection.rs | 2 +- 7 files changed, 180 insertions(+), 19 deletions(-) diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs index 37d4695002..f077a51a28 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs @@ -357,6 +357,8 @@ impl LiveRequestState { &self, required: &RequiredTunnelHeaders, generation: ModelGeneration, + has_observer: bool, + observe: impl FnOnce(&RequestObservationTransition), ) { let request_id = required.request_id.clone(); let request = TrackedPromptRequest { @@ -366,17 +368,32 @@ impl LiveRequestState { phase: TrackedPromptPhase::Pending, active_chat_output_tps: None, }; - { + let _order = self.observation_order.lock(); + let transition = { let mut state = self.inner.lock(); let observed = state .remove_request(&request_id) .and_then(|(_, request)| request.request.into_observed()); + // The first observer event replaces the inherited metric projection. + // A request without an observer must retire that projection now. + let (inherited, retired) = if has_observer { + (observed, None) + } else { + (None, observed) + }; state.insert_request( request_id, - LiveRequest::Queue(request, observed), + LiveRequest::Queue(request, inherited), Some(required.request_instance.clone()), ); - } + RequestObservationTransition { + changed_generations: Vec::new(), + input_token_totals: state.input_token_totals([retired.as_ref(), None]), + prior: retired, + current: None, + } + }; + observe(&transition); } #[cfg(test)] @@ -385,7 +402,7 @@ impl LiveRequestState { required: &RequiredTunnelHeaders, generation: ModelGeneration, ) -> QueueTrackedRequestGuard { - self.begin_request(required, generation); + self.begin_request(required, generation, true, |_| {}); self.request_guard(required) } @@ -457,11 +474,14 @@ impl LiveRequestState { pub(crate) fn update_active_output_tps( &self, request_id: &str, + instance: Option<&RequestInstance>, active_chat_output_tps: Option, ) -> Option { - self.inner - .lock() - .update_active_output_tps(request_id, active_chat_output_tps) + let mut state = self.inner.lock(); + if instance.is_some_and(|instance| !state.owns_instance(request_id, instance)) { + return None; + } + state.update_active_output_tps(request_id, active_chat_output_tps) } pub(crate) fn finish_queue_request( @@ -1061,7 +1081,7 @@ mod tests { request_id, RequestObservationState::OutputGeneration, )); - live_requests.update_active_output_tps(request_id, Some(output_tps)); + live_requests.update_active_output_tps(request_id, None, Some(output_tps)); } let active = live_requests.snapshot_model("model-a"); diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs index 1b61f8b9f0..2bd34d28de 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs @@ -316,7 +316,7 @@ impl TunnelRequestLifecycle { ) -> Self { if observation_endpoint.is_none() { app.runtime_state - .begin_request(&required, generation.as_ref()); + .begin_request(&required, generation.as_ref(), false); } let observer = observation_endpoint.map(|endpoint| { TunnelRequestObserver::accepted( diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index 3853b5ac06..cb89eeaa90 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs @@ -187,7 +187,7 @@ impl RequestObserver { generation: Option, runtime_state: PylonRuntimeState, ) -> Self { - runtime_state.begin_request(&required, generation.as_ref()); + runtime_state.begin_request(&required, generation.as_ref(), true); let priority = required.queue_priority(); let RequiredTunnelHeaders { request_id, @@ -746,6 +746,54 @@ mod tests { ); } + #[test] + fn unobserved_replacement_retires_prior_observation_and_gauges() { + for replacement_model in ["model-a", "model-b"] { + let metrics = crate::PylonMetrics::new().unwrap(); + let (runtime, _events) = PylonRuntimeState::observed( + InferenceServerStatus::Active, + &["model-a".into(), "model-b".into()], + 16, + Some(metrics.clone()), + ); + let mut first = RequestObserver::from_required( + RequestObservationEndpoint::ChatCompletions, + validate_required_tunnel_headers(&request_headers("reused-id", 100)).unwrap(), + runtime.current_generation("model-a"), + runtime.clone(), + ); + first.submit_now(); + assert!(metrics.gather_text().unwrap().contains( + r#"pylon_requests_state_input_tokens{model="model-a",state="input_processing"} 100"# + )); + + let mut headers = request_headers("reused-id", 20); + headers.insert(HEADER_MODEL, replacement_model.parse().unwrap()); + let required = validate_required_tunnel_headers(&headers).unwrap(); + let generation = runtime.current_generation(replacement_model); + runtime.begin_request(&required, generation.as_ref(), false); + let guard = runtime + .track_generation_request(&required, generation.as_ref()) + .unwrap(); + assert_eq!( + runtime + .snapshot_live_model(replacement_model) + .num_running_queries, + 1 + ); + drop(guard); + first.cancel(); + + assert!(runtime.request_generation("reused-id").is_none()); + for line in metrics.gather_text().unwrap().lines().filter(|line| { + line.starts_with("pylon_requests_state") + || line.starts_with("pylon_requests_inflight") + }) { + assert!(line.ends_with(" 0"), "retained request gauge: {line}"); + } + } + } + #[test] fn chat_calibration_state_exists_only_when_enabled() { for (enabled, expected_state) in [(false, false), (true, true)] { diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index 01bdfd35fb..66bd367921 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -577,11 +577,14 @@ impl PylonRuntimeState { pub(crate) fn update_request_active_output_tps( &self, - request_id: &str, + event: &RequestObservationEvent, active_chat_output_tps: Option, ) -> Option { - self.live_requests - .update_active_output_tps(request_id, active_chat_output_tps) + self.live_requests.update_active_output_tps( + &event.observation.request_id, + event.request_instance.as_ref(), + active_chat_output_tps, + ) } pub(crate) fn request_generation(&self, request_id: &str) -> Option { @@ -633,12 +636,21 @@ impl PylonRuntimeState { &self, required: &RequiredTunnelHeaders, generation: Option<&ModelGeneration>, + has_observer: bool, ) { let Some(generation) = generation else { return }; let advertised = self.advertised.lock(); if advertised.current(generation).is_some() { - self.live_requests - .begin_request(required, generation.clone()); + self.live_requests.begin_request( + required, + generation.clone(), + has_observer, + |transition| { + if let Some(metrics) = &self.metrics { + metrics.observe_request_state_transition(transition); + } + }, + ); } } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 79b2b279d0..83129f6536 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -4069,4 +4069,78 @@ mod tests { assert!(body.contains(r#"pylon_model_last_mean_input_tps{model="model-a"} 10"#)); assert!(!body.contains(r#"pylon_model_last_mean_input_tps{model="model-a"} 999"#)); } + #[test] + fn queued_observations_cannot_change_replacement_output_rate() { + use crate::request_observer::{RequestObserver, RequiredTunnelHeaders}; + use stargate_proto::pb::InferenceServerStatus; + + let (runtime, events) = PylonRuntimeState::observed( + InferenceServerStatus::Active, + &["model-a".into()], + 16, + None, + ); + let generation = runtime.current_generation("model-a").unwrap(); + let required = || RequiredTunnelHeaders { + request_instance: Default::default(), + request_id: "reused-id".into(), + routing_key: None, + model_id: "model-a".into(), + priority: None, + input_tokens: 100, + accepted_at: std::time::Instant::now() - Duration::from_secs(2), + }; + let mut first = RequestObserver::from_required( + RequestObservationEndpoint::ChatCompletions, + required(), + Some(generation.clone()), + runtime.clone(), + ); + first.on_backend_submission(std::time::Instant::now() - Duration::from_secs(1)); + first.observe_generated_output( + std::time::Instant::now() - Duration::from_millis(500), + true, + 100, + false, + ); + first.observe_estimated_output_tokens_total(100); + let old_events: Vec<_> = events.try_iter().collect(); + let mut second = RequestObserver::from_required( + RequestObservationEndpoint::ChatCompletions, + required(), + Some(generation.clone()), + runtime.clone(), + ); + second.on_backend_submission(std::time::Instant::now() - Duration::from_secs(1)); + second.observe_generated_output( + std::time::Instant::now() - Duration::from_millis(500), + true, + 50, + false, + ); + second.observe_estimated_output_tokens_total(50); + let before = runtime.snapshot_live_model("model-a"); + assert_eq!(before.active_chat_output_tps, 0.0); + let mut collector = StatsAggregator::new(StatsCollectorConfig::default(), runtime.clone()); + collector + .begin_generation(generation, ModelStatsInitialization::Empty) + .unwrap(); + for event in old_events { + collector.apply_fallback_observation(&event); + } + assert_eq!( + runtime.snapshot_live_model("model-a"), + before, + "an old queued event changed the replacement's output rate" + ); + for event in events.try_iter() { + collector.apply_fallback_observation(&event); + } + assert!( + runtime + .snapshot_live_model("model-a") + .active_chat_output_tps + > 0.0 + ); + } } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs index 54d1c84a4f..2facf2a5f2 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/metrics.rs @@ -263,12 +263,19 @@ impl PylonMetrics { observation: &RequestObservation, transition: &RequestObservationTransition, ) { - if let Some(prior) = &transition.prior { - self.adjust_observed_request(prior, -1); - } + self.observe_request_state_transition(transition); if observation.is_terminal() { self.record_terminal_observation(observation, request_state_label(observation.state)); } + } + + pub(crate) fn observe_request_state_transition( + &self, + transition: &RequestObservationTransition, + ) { + if let Some(prior) = &transition.prior { + self.adjust_observed_request(prior, -1); + } if let Some(current) = &transition.current { self.adjust_observed_request(current, 1); } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs index c62023dd51..55631379c4 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs @@ -283,7 +283,7 @@ impl StatsAggregator { } if let Some(model_id) = self .runtime_state - .update_request_active_output_tps(&observation.request_id, active_chat_output_tps) + .update_request_active_output_tps(event, active_chat_output_tps) { push_changed_model(&mut changed_models, model_id); } From 38176dea3789e6d9f9e041617bb675d04834a0e0 Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Mon, 14 Sep 2026 17:00:47 +0000 Subject: [PATCH 4/5] fix(pylon): start observed requests in one state transition Replace queue and observation ownership atomically so model retirement cannot encounter an inherited projection from another generation. Start replacements pending and retain guarded cleanup for unobserved requests. Refs: NVIDIA/nvcf#1817 --- .../crates/pylon-lib/src/queue_admission.rs | 54 ++++++++++++------ .../pylon-lib/src/quic_http_tunnel/core.rs | 2 +- .../crates/pylon-lib/src/request_observer.rs | 55 +++++++++++-------- .../crates/pylon-lib/src/runtime_state.rs | 20 ++++--- .../crates/pylon-lib/src/stats/collector.rs | 2 + 5 files changed, 84 insertions(+), 49 deletions(-) diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs index f077a51a28..4178258963 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs @@ -353,11 +353,10 @@ impl LiveRequestState { self.track_generation_request(required, ModelGeneration::new(required.model_id.clone(), 0)) } - pub(crate) fn begin_request( + pub(crate) fn begin_unobserved_request( &self, required: &RequiredTunnelHeaders, generation: ModelGeneration, - has_observer: bool, observe: impl FnOnce(&RequestObservationTransition), ) { let request_id = required.request_id.clone(); @@ -371,19 +370,12 @@ impl LiveRequestState { let _order = self.observation_order.lock(); let transition = { let mut state = self.inner.lock(); - let observed = state + let retired = state .remove_request(&request_id) .and_then(|(_, request)| request.request.into_observed()); - // The first observer event replaces the inherited metric projection. - // A request without an observer must retire that projection now. - let (inherited, retired) = if has_observer { - (observed, None) - } else { - (None, observed) - }; state.insert_request( request_id, - LiveRequest::Queue(request, inherited), + LiveRequest::Queue(request, None), Some(required.request_instance.clone()), ); RequestObservationTransition { @@ -402,7 +394,7 @@ impl LiveRequestState { required: &RequiredTunnelHeaders, generation: ModelGeneration, ) -> QueueTrackedRequestGuard { - self.begin_request(required, generation, true, |_| {}); + self.begin_unobserved_request(required, generation, |_| {}); self.request_guard(required) } @@ -445,8 +437,14 @@ impl LiveRequestState { observe: impl FnOnce(&RequestObservationTransition), ) -> RequestObservationTransition { let generation = ModelGeneration::new(observation.model_id.clone(), 0); - self.transition_generation_observation_with(observation, Some(&generation), None, observe) - .unwrap() + self.transition_generation_observation_with( + observation, + Some(&generation), + None, + false, + observe, + ) + .unwrap() } pub(crate) fn transition_generation_observation_with( @@ -454,18 +452,22 @@ impl LiveRequestState { observation: &RequestObservation, generation: Option<&ModelGeneration>, instance: Option<&RequestInstance>, + begin_request: bool, observe: impl FnOnce(&RequestObservationTransition), ) -> Option { let _order = self.observation_order.lock(); let transition = { let mut state = self.inner.lock(); - if let Some(instance) = instance + // Starting an observed request replaces both projections under this + // lock, so generation retirement cannot see different owners. + if !begin_request + && let Some(instance) = instance && generation.is_some() && !state.owns_instance(&observation.request_id, instance) { return None; } - state.transition_observation(observation, generation, instance) + state.transition_observation(observation, generation, instance, begin_request) }; observe(&transition); Some(transition) @@ -513,6 +515,7 @@ impl QueueAdmissionState { observation: &RequestObservation, generation: Option<&ModelGeneration>, instance: Option<&RequestInstance>, + begin_request: bool, ) -> RequestObservationTransition { let prior = self.remove_request(&observation.request_id); let instance = instance.cloned().or_else(|| { @@ -520,7 +523,7 @@ impl QueueAdmissionState { .as_ref() .and_then(|(_, record)| record.instance.clone()) }); - let (request_id, prior_queue, prior_observed) = + let (request_id, mut prior_queue, prior_observed) = match prior.map(|(id, record)| (id, record.request)) { Some((request_id, LiveRequest::Queue(queue, observed))) => { (request_id, Some(queue), observed) @@ -547,6 +550,9 @@ impl QueueAdmissionState { prior_observed.as_ref().map(|request| &request.generation), ], ); + if begin_request { + prior_queue = None; + } let current = match (observation.is_terminal(), generation) { (false, Some(generation)) => { let phase = [ @@ -1518,4 +1524,18 @@ mod tests { ); } } + + #[test] + fn retiring_a_generation_preserves_replacement_accounting() { + let live = LiveRequestState::default(); + let _first = live.track_request(&required_for_model("reused-id", "model-a", 0, 100)); + live.transition_observation(&observation( + "reused-id", + RequestObservationState::InputProcessing, + )); + let replacement = live.track_request(&required_for_model("reused-id", "model-b", 0, 20)); + live.retire_generation(&ModelGeneration::new("model-a", 0)); + drop(replacement); + assert_eq!(live.snapshot_model("model-b").num_running_queries, 0); + } } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs index 2bd34d28de..692b916d05 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs @@ -316,7 +316,7 @@ impl TunnelRequestLifecycle { ) -> Self { if observation_endpoint.is_none() { app.runtime_state - .begin_request(&required, generation.as_ref(), false); + .begin_unobserved_request(&required, generation.as_ref()); } let observer = observation_endpoint.map(|endpoint| { TunnelRequestObserver::accepted( diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index cb89eeaa90..0db08f899a 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs @@ -187,7 +187,6 @@ impl RequestObserver { generation: Option, runtime_state: PylonRuntimeState, ) -> Self { - runtime_state.begin_request(&required, generation.as_ref(), true); let priority = required.queue_priority(); let RequiredTunnelHeaders { request_id, @@ -199,7 +198,7 @@ impl RequestObserver { accepted_at, } = required; let output_token_calibration_enabled = runtime_state.output_token_calibration_enabled(); - let mut observer = Self { + let observer = Self { endpoint, request_id, request_instance, @@ -218,7 +217,13 @@ impl RequestObserver { state: RequestLifecycleState::UpstreamConnecting, runtime_state, }; - observer.emit(); + let event = observer.observation_event(); + log_observation(&event.observation); + observer.runtime_state.observe_request_for_generation( + event, + observer.request_input_tokens, + true, + ); observer } @@ -501,7 +506,14 @@ impl RequestObserver { ), } } - fn emit(&mut self) { + fn emit(&self) { + let event = self.observation_event(); + log_observation(&event.observation); + self.runtime_state + .observe_request_for_generation(event, self.request_input_tokens, false); + } + + fn observation_event(&self) -> RequestObservationEvent { let backend = self.state.backend(); let input_interval = backend.and_then(|backend| { backend @@ -541,25 +553,20 @@ impl RequestObserver { .map(|instant| instant.saturating_duration_since(self.started_at)), total_duration: self.started_at.elapsed(), }; - log_observation(&observation); - self.runtime_state.observe_request_for_generation( - RequestObservationEvent { - observation, - request_instance: Some(self.request_instance.clone()), - generation: self.generation.clone(), - changed_generations: Vec::new(), - input_interval, - input_tokens_explicit: self.input_tokens_explicit, - output_calibration: backend - .map_or_else(OutputCalibrationFacts::default, |backend| { - backend.output_calibration - }), - upstream_duration: backend - .and_then(|backend| backend.last_upstream_event_at) - .map(|instant| instant.saturating_duration_since(self.started_at)), - }, - self.request_input_tokens, - ); + RequestObservationEvent { + observation, + request_instance: Some(self.request_instance.clone()), + generation: self.generation.clone(), + changed_generations: Vec::new(), + input_interval, + input_tokens_explicit: self.input_tokens_explicit, + output_calibration: backend.map_or_else(OutputCalibrationFacts::default, |backend| { + backend.output_calibration + }), + upstream_duration: backend + .and_then(|backend| backend.last_upstream_event_at) + .map(|instant| instant.saturating_duration_since(self.started_at)), + } } } @@ -771,7 +778,7 @@ mod tests { headers.insert(HEADER_MODEL, replacement_model.parse().unwrap()); let required = validate_required_tunnel_headers(&headers).unwrap(); let generation = runtime.current_generation(replacement_model); - runtime.begin_request(&required, generation.as_ref(), false); + runtime.begin_unobserved_request(&required, generation.as_ref()); let guard = runtime .track_generation_request(&required, generation.as_ref()) .unwrap(); diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index 66bd367921..c40040d077 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -466,6 +466,7 @@ impl PylonRuntimeState { upstream_duration: None, }, request_input_tokens, + false, ); } @@ -473,10 +474,13 @@ impl PylonRuntimeState { &self, event: RequestObservationEvent, request_input_tokens: u64, + begin_request: bool, ) { - let Some(event) = - self.transition_request_observation_for_generation(event, request_input_tokens) - else { + let Some(event) = self.transition_request_observation_for_generation( + event, + request_input_tokens, + begin_request, + ) else { return; }; if let Some(tx) = &self.observation_tx @@ -530,6 +534,7 @@ impl PylonRuntimeState { upstream_duration: None, }, request_input_tokens, + false, ) .expect("test observation must target a current generation") } @@ -538,6 +543,7 @@ impl PylonRuntimeState { &self, mut event: RequestObservationEvent, request_input_tokens: u64, + begin_request: bool, ) -> Option { // Held across the queue transition below: retire_generation() purges // live-request state under this lock, so releasing it after the @@ -565,6 +571,7 @@ impl PylonRuntimeState { &live_observation, event.generation.as_ref(), event.request_instance.as_ref(), + begin_request, |transition| { if let Some(metrics) = &self.metrics { metrics.observe_request_transition(&event.observation, transition); @@ -632,19 +639,17 @@ impl PylonRuntimeState { .track_generation_request(required, generation) } - pub(crate) fn begin_request( + pub(crate) fn begin_unobserved_request( &self, required: &RequiredTunnelHeaders, generation: Option<&ModelGeneration>, - has_observer: bool, ) { let Some(generation) = generation else { return }; let advertised = self.advertised.lock(); if advertised.current(generation).is_some() { - self.live_requests.begin_request( + self.live_requests.begin_unobserved_request( required, generation.clone(), - has_observer, |transition| { if let Some(metrics) = &self.metrics { metrics.observe_request_state_transition(transition); @@ -750,6 +755,7 @@ mod tests { upstream_duration: None, }, request_input_tokens, + false, ) } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 83129f6536..2a779e6456 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -812,6 +812,7 @@ mod tests { upstream_duration: None, }, request_input_tokens, + false, ); } @@ -4111,6 +4112,7 @@ mod tests { Some(generation.clone()), runtime.clone(), ); + assert_eq!(runtime.snapshot_live_model("model-a").queue_size, 1); second.on_backend_submission(std::time::Instant::now() - Duration::from_secs(1)); second.observe_generated_output( std::time::Instant::now() - Duration::from_millis(500), From e128ffd6dc0107dbcaa4b8f7c3280561a20fe51f Mon Sep 17 00:00:00 2001 From: Barry Greengus Date: Mon, 14 Sep 2026 21:20:47 +0000 Subject: [PATCH 5/5] fix(pylon): publish unobserved load changes with explicit ownership Wake the stats collector for coalesced model load changes, including unobserved request completion. Require request instances in production records and move synthetic identity handling into test adapters. Cover immediate and delayed completion across models. Refs: NVIDIA/nvcf#1817 --- .../crates/pylon-lib/src/queue_admission.rs | 128 ++++++++++++------ .../crates/pylon-lib/src/request_observer.rs | 2 +- .../crates/pylon-lib/src/runtime_state.rs | 41 ++++-- .../crates/pylon-lib/src/stats/collector.rs | 94 ++++++++++++- .../crates/pylon-lib/src/stats/projection.rs | 16 +++ 5 files changed, 225 insertions(+), 56 deletions(-) diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs index 4178258963..19f0d01500 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs @@ -13,7 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; use parking_lot::Mutex; @@ -51,17 +51,19 @@ impl Default for PylonQueueMismatchRetryConfig { pub(crate) struct LiveRequestState { inner: Arc>, observation_order: Arc>, + changed: Arc, } #[derive(Debug, Default)] struct QueueAdmissionState { requests: HashMap, models: HashMap, + pending_publications: BTreeSet, } #[derive(Debug)] struct LiveRequestRecord { - instance: Option, + instance: RequestInstance, request: LiveRequest, } @@ -223,6 +225,24 @@ impl QueueAdmissionDecision { } impl LiveRequestState { + pub(crate) async fn changed(&self) { + self.changed.notified().await; + } + + pub(crate) fn take_pending_publications(&self) -> Vec { + std::mem::take(&mut self.inner.lock().pending_publications) + .into_iter() + .collect() + } + + #[cfg(any(test, feature = "test-support"))] + pub(crate) fn request_instance_for_test(&self, request_id: &str) -> (RequestInstance, bool) { + self.inner.lock().requests.get(request_id).map_or_else( + || (RequestInstance::default(), true), + |record| (record.instance.clone(), false), + ) + } + pub(crate) fn request_generation(&self, request_id: &str) -> Option { self.inner .lock() @@ -261,6 +281,7 @@ impl LiveRequestState { state.remove_request(&request_id); } state.models.remove(generation); + state.pending_publications.remove(generation); } #[cfg(test)] @@ -309,9 +330,7 @@ impl LiveRequestState { let excluded_request = state .requests .get(&required.request_id) - .filter(|record| { - record.instance.as_ref() == Some(&required.request_instance) - }) + .filter(|record| record.instance == required.request_instance) .and_then(|record| match &record.request { LiveRequest::Queue(queue, _) => Some(queue), LiveRequest::Observed(_) => None, @@ -370,13 +389,20 @@ impl LiveRequestState { let _order = self.observation_order.lock(); let transition = { let mut state = self.inner.lock(); - let retired = state - .remove_request(&request_id) - .and_then(|(_, request)| request.request.into_observed()); + let prior = state.remove_request(&request_id); + if let Some((_, record)) = &prior { + state + .pending_publications + .insert(record.request.generation().clone()); + } + state + .pending_publications + .insert(request.generation.clone()); + let retired = prior.and_then(|(_, request)| request.request.into_observed()); state.insert_request( request_id, LiveRequest::Queue(request, None), - Some(required.request_instance.clone()), + required.request_instance.clone(), ); RequestObservationTransition { changed_generations: Vec::new(), @@ -386,6 +412,7 @@ impl LiveRequestState { } }; observe(&transition); + self.changed.notify_one(); } #[cfg(test)] @@ -437,11 +464,12 @@ impl LiveRequestState { observe: impl FnOnce(&RequestObservationTransition), ) -> RequestObservationTransition { let generation = ModelGeneration::new(observation.model_id.clone(), 0); + let (instance, begin_request) = self.request_instance_for_test(&observation.request_id); self.transition_generation_observation_with( observation, Some(&generation), - None, - false, + &instance, + begin_request, observe, ) .unwrap() @@ -451,7 +479,7 @@ impl LiveRequestState { &self, observation: &RequestObservation, generation: Option<&ModelGeneration>, - instance: Option<&RequestInstance>, + instance: &RequestInstance, begin_request: bool, observe: impl FnOnce(&RequestObservationTransition), ) -> Option { @@ -461,7 +489,6 @@ impl LiveRequestState { // Starting an observed request replaces both projections under this // lock, so generation retirement cannot see different owners. if !begin_request - && let Some(instance) = instance && generation.is_some() && !state.owns_instance(&observation.request_id, instance) { @@ -476,29 +503,48 @@ impl LiveRequestState { pub(crate) fn update_active_output_tps( &self, request_id: &str, - instance: Option<&RequestInstance>, + instance: &RequestInstance, active_chat_output_tps: Option, ) -> Option { let mut state = self.inner.lock(); - if instance.is_some_and(|instance| !state.owns_instance(request_id, instance)) { + if !state.owns_instance(request_id, instance) { return None; } state.update_active_output_tps(request_id, active_chat_output_tps) } - pub(crate) fn finish_queue_request( + pub(crate) fn finish_queue_request(&self, request_id: &str, instance: &RequestInstance) { + let mut state = self.inner.lock(); + if !state.owns_instance(request_id, instance) { + return; + } + if let Some((request_id, record)) = state.remove_request(request_id) { + match record.request { + LiveRequest::Queue(_, Some(observed)) | LiveRequest::Observed(observed) => { + state.insert_request( + request_id, + LiveRequest::Observed(observed), + record.instance, + ); + } + LiveRequest::Queue(queue, None) => { + state.pending_publications.insert(queue.generation); + self.changed.notify_one(); + } + } + } + } + + fn advance_request_phase( &self, request_id: &str, - instance: Option<&RequestInstance>, + instance: &RequestInstance, + phase: TrackedPromptPhase, ) { let mut state = self.inner.lock(); - if instance.is_some_and(|instance| !state.owns_instance(request_id, instance)) { - return; - } - if let Some((request_id, record)) = state.remove_request(request_id) - && let Some(observed) = record.request.into_observed() - { - state.insert_request(request_id, LiveRequest::Observed(observed), record.instance); + if let Some(generation) = state.advance_request_phase(request_id, instance, phase) { + state.pending_publications.insert(generation); + self.changed.notify_one(); } } } @@ -507,22 +553,18 @@ impl QueueAdmissionState { fn owns_instance(&self, request_id: &str, instance: &RequestInstance) -> bool { self.requests .get(request_id) - .is_some_and(|record| record.instance.as_ref() == Some(instance)) + .is_some_and(|record| &record.instance == instance) } fn transition_observation( &mut self, observation: &RequestObservation, generation: Option<&ModelGeneration>, - instance: Option<&RequestInstance>, + instance: &RequestInstance, begin_request: bool, ) -> RequestObservationTransition { let prior = self.remove_request(&observation.request_id); - let instance = instance.cloned().or_else(|| { - prior - .as_ref() - .and_then(|(_, record)| record.instance.clone()) - }); + let instance = instance.clone(); let (request_id, mut prior_queue, prior_observed) = match prior.map(|(id, record)| (id, record.request)) { Some((request_id, LiveRequest::Queue(queue, observed))) => { @@ -619,19 +661,19 @@ impl QueueAdmissionState { request_id: &str, instance: &RequestInstance, next_phase: TrackedPromptPhase, - ) { + ) -> Option { if !self.owns_instance(request_id, instance) { - return; + return None; } - let Some(LiveRequest::Queue(request, _)) = self + let Some(LiveRequest::Queue(request, observed)) = self .requests .get_mut(request_id) .map(|record| &mut record.request) else { - return; + return None; }; if next_phase <= request.phase { - return; + return None; } let model = self .models @@ -640,6 +682,7 @@ impl QueueAdmissionState { model.adjust_phase(request.phase, request.priority, request.input_tokens, -1); model.adjust_phase(next_phase, request.priority, request.input_tokens, 1); request.phase = next_phase; + observed.is_none().then(|| request.generation.clone()) } fn remove_request(&mut self, request_id: &str) -> Option<(String, LiveRequestRecord)> { @@ -652,7 +695,7 @@ impl QueueAdmissionState { &mut self, request_id: String, request: LiveRequest, - instance: Option, + instance: RequestInstance, ) { self.adjust_live_request(&request, 1); self.requests @@ -898,8 +941,7 @@ impl TrackedPromptPhase { impl QueueTrackedRequestGuard { pub(crate) fn on_backend_submission(&mut self) { - let mut state = self.live_requests.inner.lock(); - state.advance_request_phase( + self.live_requests.advance_request_phase( &self.request_id, &self.instance, TrackedPromptPhase::InputProcessing, @@ -907,8 +949,7 @@ impl QueueTrackedRequestGuard { } pub(crate) fn observe_output(&mut self) { - let mut state = self.live_requests.inner.lock(); - state.advance_request_phase( + self.live_requests.advance_request_phase( &self.request_id, &self.instance, TrackedPromptPhase::OutputGeneration, @@ -918,7 +959,7 @@ impl QueueTrackedRequestGuard { pub(crate) fn finish(&mut self) { if !self.finished { self.live_requests - .finish_queue_request(&self.request_id, Some(&self.instance)); + .finish_queue_request(&self.request_id, &self.instance); self.finished = true; } } @@ -1087,7 +1128,8 @@ mod tests { request_id, RequestObservationState::OutputGeneration, )); - live_requests.update_active_output_tps(request_id, None, Some(output_tps)); + let (instance, _) = live_requests.request_instance_for_test(request_id); + live_requests.update_active_output_tps(request_id, &instance, Some(output_tps)); } let active = live_requests.snapshot_model("model-a"); diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index 0db08f899a..60593dc2b3 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs @@ -555,7 +555,7 @@ impl RequestObserver { }; RequestObservationEvent { observation, - request_instance: Some(self.request_instance.clone()), + request_instance: self.request_instance.clone(), generation: self.generation.clone(), changed_generations: Vec::new(), input_interval, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs index c40040d077..292cf76e8c 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/runtime_state.rs @@ -106,7 +106,7 @@ pub struct PylonRuntimeState { #[derive(Clone, Debug)] pub struct RequestObservationEvent { pub(crate) observation: RequestObservation, - pub(crate) request_instance: Option, + pub(crate) request_instance: RequestInstance, pub(crate) generation: Option, pub(crate) changed_generations: Vec, pub(crate) input_interval: Option, @@ -450,13 +450,20 @@ impl PylonRuntimeState { model_ids } + #[cfg(any(test, feature = "test-support"))] + pub(crate) fn request_instance_for_test(&self, request_id: &str) -> (RequestInstance, bool) { + self.live_requests.request_instance_for_test(request_id) + } + #[cfg(any(test, feature = "test-support"))] pub fn observe_request_for_test(&self, observation: RequestObservation) { let generation = self.current_generation(&observation.model_id); let request_input_tokens = observation.input_tokens; + let (request_instance, begin_request) = + self.request_instance_for_test(&observation.request_id); self.observe_request_for_generation( RequestObservationEvent { - request_instance: None, + request_instance, observation, generation, changed_generations: Vec::new(), @@ -466,7 +473,7 @@ impl PylonRuntimeState { upstream_duration: None, }, request_input_tokens, - false, + begin_request, ); } @@ -522,9 +529,11 @@ impl PylonRuntimeState { .current_generation(&observation.model_id) .expect("test model generation should already exist"); let request_input_tokens = observation.input_tokens; + let (request_instance, begin_request) = + self.request_instance_for_test(&observation.request_id); self.transition_request_observation_for_generation( RequestObservationEvent { - request_instance: None, + request_instance, observation, generation: Some(generation), changed_generations: Vec::new(), @@ -534,7 +543,7 @@ impl PylonRuntimeState { upstream_duration: None, }, request_input_tokens, - false, + begin_request, ) .expect("test observation must target a current generation") } @@ -570,7 +579,7 @@ impl PylonRuntimeState { let transition = self.live_requests.transition_generation_observation_with( &live_observation, event.generation.as_ref(), - event.request_instance.as_ref(), + &event.request_instance, begin_request, |transition| { if let Some(metrics) = &self.metrics { @@ -589,7 +598,7 @@ impl PylonRuntimeState { ) -> Option { self.live_requests.update_active_output_tps( &event.observation.request_id, - event.request_instance.as_ref(), + &event.request_instance, active_chat_output_tps, ) } @@ -598,6 +607,14 @@ impl PylonRuntimeState { self.live_requests.request_generation(request_id) } + pub(crate) async fn live_requests_changed(&self) { + self.live_requests.changed().await; + } + + pub(crate) fn take_live_request_changes(&self) -> Vec { + self.live_requests.take_pending_publications() + } + pub(crate) fn snapshot_live_model(&self, model_id: &str) -> QueueModelSnapshot { self.current_generation(model_id) .map_or_else(QueueModelSnapshot::default, |generation| { @@ -672,7 +689,9 @@ impl PylonRuntimeState { #[cfg(test)] pub(crate) fn finish_queue_request(&self, request_id: &str) { - self.live_requests.finish_queue_request(request_id, None); + let (instance, _) = self.request_instance_for_test(request_id); + self.live_requests + .finish_queue_request(request_id, &instance); } #[cfg(test)] @@ -743,9 +762,11 @@ mod tests { generation: ModelGeneration, ) -> Option { let request_input_tokens = observation.input_tokens; + let (request_instance, begin_request) = + runtime_state.request_instance_for_test(&observation.request_id); runtime_state.transition_request_observation_for_generation( RequestObservationEvent { - request_instance: None, + request_instance, observation, generation: Some(generation), changed_generations: Vec::new(), @@ -755,7 +776,7 @@ mod tests { upstream_duration: None, }, request_input_tokens, - false, + begin_request, ) } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 2a779e6456..5872c5812b 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -350,6 +350,10 @@ async fn run_stats_collector( let mut stats_aggregator_latest_models = IndexMap::with_capacity(2); 'collector: loop { + // Drain coalesced queue changes even while observation traffic is busy. + let updates = + aggregator.apply_live_request_changes(runtime_state.take_live_request_changes()); + publish_model_stats_updates(&runtime_state, updates); tokio::select! { biased; _ = stop.cancelled() => break 'collector, @@ -382,6 +386,7 @@ async fn run_stats_collector( }; publish_observation_event(&mut aggregator, &runtime_state, event); } + _ = runtime_state.live_requests_changed() => {} update = async { match &stats_update_rx { Some(rx) => rx.recv_async().await.ok(), @@ -800,9 +805,11 @@ mod tests { ..crate::runtime_state::OutputCalibrationFacts::default() }; let request_input_tokens = observation.input_tokens; + let (request_instance, begin_request) = + runtime_state.request_instance_for_test(&observation.request_id); runtime_state.observe_request_for_generation( crate::runtime_state::RequestObservationEvent { - request_instance: None, + request_instance, observation, generation, changed_generations: Vec::new(), @@ -812,7 +819,7 @@ mod tests { upstream_duration: None, }, request_input_tokens, - false, + begin_request, ); } @@ -4145,4 +4152,87 @@ mod tests { > 0.0 ); } + + #[tokio::test] + async fn unobserved_replacements_publish_load_through_completion() { + use crate::request_observer::{RequestObserver, RequiredTunnelHeaders}; + + for replacement_model in ["model-a", "model-b"] { + for wait_for_submission in [false, true] { + let collector = RunningCollector::spawn_with_models( + StatsCollectorConfig::default(), + None, + false, + &["model-a".into(), "model-b".into()], + ); + let runtime = &collector.runtime_state; + let required = |model: &str, input_tokens| RequiredTunnelHeaders { + request_id: "reused-id".into(), + request_instance: Default::default(), + routing_key: None, + model_id: model.into(), + priority: None, + input_tokens, + accepted_at: std::time::Instant::now(), + }; + let mut first = RequestObserver::from_required( + RequestObservationEndpoint::ChatCompletions, + required("model-a", 100), + runtime.current_generation("model-a"), + runtime.clone(), + ); + first.on_backend_submission(std::time::Instant::now()); + collector + .wait_for_stats("first request should be advertised", |stats| { + stats.num_running_queries == 1 + }) + .await; + + let replacement = required(replacement_model, 20); + let generation = runtime.current_generation(replacement_model); + runtime.begin_unobserved_request(&replacement, generation.as_ref()); + let mut guard = runtime + .track_generation_request(&replacement, generation.as_ref()) + .unwrap(); + guard.on_backend_submission(); + if wait_for_submission { + wait_for_model_stats( + runtime, + replacement_model, + "replacement submission should be advertised", + |stats| { + stats.num_running_queries == 1 + && stats.total_query_input_size == 20 + && stats.input_processing_queries == 1 + }, + ) + .await; + } + drop(guard); + first.cancel(); + for model in ["model-a", "model-b"] { + wait_for_model_stats( + runtime, + model, + "completed requests should leave no advertised load", + |stats| { + stats.num_running_queries == 0 + && stats.total_query_input_size == 0 + && stats.input_processing_queries == 0 + }, + ) + .await; + assert_eq!( + runtime.advertised_models()[model] + .stats + .as_ref() + .unwrap() + .num_running_queries, + 0 + ); + } + collector.handle.shutdown().await; + } + } + } } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs index 55631379c4..2675b84924 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/projection.rs @@ -27,6 +27,22 @@ use super::aggregator::{ use super::collector::StatsCollectorConfig; impl StatsAggregator { + pub(super) fn apply_live_request_changes( + &mut self, + generations: Vec, + ) -> Vec { + let mut changed_models = Vec::with_capacity(generations.len()); + for generation in generations { + if let Some(model) = self.per_model.get_mut(generation.model_id()) + && model.generation == generation + { + model.metrics.stats_observed_at_unix_ms = current_unix_millis(); + changed_models.push(generation.model_id().to_string()); + } + } + self.snapshots(changed_models) + } + pub(super) fn apply_fallback_observation( &mut self, event: &RequestObservationEvent,