diff --git a/crates/grpc_client/src/sglang_scheduler.rs b/crates/grpc_client/src/sglang_scheduler.rs index 18ea2618b..60d69ca24 100644 --- a/crates/grpc_client/src/sglang_scheduler.rs +++ b/crates/grpc_client/src/sglang_scheduler.rs @@ -761,6 +761,24 @@ impl From for openai_protocol::worker::SchedulerLoadSnapsh // the two canonical totals: prefill = prealloc + inflight, decode = // prealloc + transfer + retracted. let disagg = load.disaggregation; + let memory = + load.memory.map( + |memory| openai_protocol::worker::EngineMemoryMetricsSnapshot { + weight_gb: memory.weight_gb, + kv_cache_gb: memory.kv_cache_gb, + graph_gb: memory.graph_gb, + token_capacity: memory.token_capacity, + }, + ); + let queues = + load.queues.map( + |queues| openai_protocol::worker::EngineQueueMetricsSnapshot { + waiting: queues.waiting, + grammar: queues.grammar, + paused: queues.paused, + retracted: queues.retracted, + }, + ); Self { dp_rank: load.dp_rank, num_running_reqs: load.num_running_reqs, @@ -774,6 +792,8 @@ impl From for openai_protocol::worker::SchedulerLoadSnapsh cache_hit_rate: load.cache_hit_rate, utilization: load.utilization, max_running_requests: load.max_running_requests, + memory, + queues, kv_transfer_latency_ms: disagg.as_ref().map(|d| d.kv_transfer_latency_ms), kv_transfer_speed_gb_s: disagg.as_ref().map(|d| d.kv_transfer_speed_gb_s), prefill_queue_reqs: disagg.as_ref().map(|d| { @@ -792,10 +812,22 @@ impl From for openai_protocol::worker::SchedulerLoadSnapsh impl From for openai_protocol::worker::WorkerLoadResponse { fn from(resp: proto::GetLoadsResponse) -> Self { + let aggregate = resp.aggregate.map(|aggregate| { + openai_protocol::worker::EngineAggregateMetricsSnapshot { + total_running_reqs: aggregate.total_running_reqs, + total_waiting_reqs: aggregate.total_waiting_reqs, + total_reqs: aggregate.total_reqs, + avg_token_usage: aggregate.avg_token_usage, + avg_throughput: aggregate.avg_throughput, + avg_utilization: aggregate.avg_utilization, + } + }); Self { timestamp: resp.timestamp, + version: resp.version, dp_rank_count: resp.dp_rank_count, loads: resp.loads.into_iter().map(Into::into).collect(), + aggregate, } } } diff --git a/crates/grpc_client/src/tokenspeed_scheduler.rs b/crates/grpc_client/src/tokenspeed_scheduler.rs index 9a543f947..89312ad8a 100644 --- a/crates/grpc_client/src/tokenspeed_scheduler.rs +++ b/crates/grpc_client/src/tokenspeed_scheduler.rs @@ -651,6 +651,24 @@ impl TokenSpeedSchedulerClient { impl From for openai_protocol::worker::SchedulerLoadSnapshot { fn from(load: tokenspeed_proto::SchedulerLoad) -> Self { + let memory = + load.memory.map( + |memory| openai_protocol::worker::EngineMemoryMetricsSnapshot { + weight_gb: memory.weight_gb, + kv_cache_gb: memory.kv_cache_gb, + graph_gb: memory.graph_gb, + token_capacity: memory.token_capacity, + }, + ); + let queues = + load.queues.map( + |queues| openai_protocol::worker::EngineQueueMetricsSnapshot { + waiting: queues.waiting, + grammar: queues.grammar, + paused: queues.paused, + retracted: queues.retracted, + }, + ); Self { dp_rank: load.dp_rank, num_running_reqs: load.num_running_reqs, @@ -664,6 +682,8 @@ impl From for openai_protocol::worker::Schedule cache_hit_rate: load.cache_hit_rate, utilization: load.utilization, max_running_requests: load.max_running_requests, + memory, + queues, // TokenSpeed has no disagg section; canonical PD fields stay None. ..Default::default() } @@ -672,10 +692,101 @@ impl From for openai_protocol::worker::Schedule impl From for openai_protocol::worker::WorkerLoadResponse { fn from(resp: tokenspeed_proto::GetLoadsResponse) -> Self { + let aggregate = resp.aggregate.map(|aggregate| { + openai_protocol::worker::EngineAggregateMetricsSnapshot { + total_running_reqs: aggregate.total_running_reqs, + total_waiting_reqs: aggregate.total_waiting_reqs, + total_reqs: aggregate.total_reqs, + avg_token_usage: aggregate.avg_token_usage, + avg_throughput: aggregate.avg_throughput, + avg_utilization: aggregate.avg_utilization, + } + }); Self { timestamp: resp.timestamp, + version: resp.version, dp_rank_count: resp.dp_rank_count, loads: resp.loads.into_iter().map(Into::into).collect(), + aggregate, } } } + +#[cfg(test)] +mod load_conversion_tests { + use super::*; + + #[test] + fn conversion_preserves_version_sections_and_aggregate() { + let response = tokenspeed_proto::GetLoadsResponse { + timestamp: "2026-08-09T12:34:56Z".to_owned(), + version: "dsv4-engine/0.1.0".to_owned(), + dp_rank_count: 1, + loads: vec![tokenspeed_proto::SchedulerLoad { + dp_rank: 0, + num_running_reqs: 3, + num_waiting_reqs: 2, + num_total_reqs: 5, + num_used_tokens: 98_304, + max_total_num_tokens: 196_608, + max_running_requests: 32, + num_waiting_uncached_tokens: 1_280, + token_usage: 0.5, + gen_throughput: 105.25, + cache_hit_rate: 0.75, + utilization: 0.5, + memory: Some(tokenspeed_proto::MemoryMetrics { + weight_gb: 44.0, + kv_cache_gb: 12.5, + graph_gb: 0.75, + token_capacity: 196_608, + }), + queues: Some(tokenspeed_proto::QueueMetrics { + waiting: 2, + grammar: 0, + paused: 0, + retracted: 0, + }), + }], + aggregate: Some(tokenspeed_proto::AggregateMetrics { + total_running_reqs: 3, + total_waiting_reqs: 2, + total_reqs: 5, + avg_token_usage: 0.5, + avg_throughput: 105.25, + avg_utilization: 0.5, + }), + }; + + let converted = openai_protocol::worker::WorkerLoadResponse::from(response); + assert_eq!(converted.timestamp, "2026-08-09T12:34:56Z"); + assert_eq!(converted.version, "dsv4-engine/0.1.0"); + let load = &converted.loads[0]; + assert!(matches!( + load.memory, + Some(ref memory) + if memory.weight_gb == 44.0 + && memory.kv_cache_gb == 12.5 + && memory.graph_gb == 0.75 + && memory.token_capacity == 196_608 + )); + assert!(matches!( + load.queues, + Some(ref queues) + if queues.waiting == 2 + && queues.grammar == 0 + && queues.paused == 0 + && queues.retracted == 0 + )); + assert!(matches!( + converted.aggregate, + Some(ref aggregate) + if aggregate.total_running_reqs == 3 + && aggregate.total_waiting_reqs == 2 + && aggregate.total_reqs == 5 + && aggregate.avg_token_usage == 0.5 + && aggregate.avg_throughput == 105.25 + && aggregate.avg_utilization == 0.5 + )); + } +} diff --git a/crates/grpc_client/src/vllm_engine.rs b/crates/grpc_client/src/vllm_engine.rs index cf4c44a6e..fc1ddbb25 100644 --- a/crates/grpc_client/src/vllm_engine.rs +++ b/crates/grpc_client/src/vllm_engine.rs @@ -716,8 +716,10 @@ impl From for openai_protocol::worker::WorkerLoadRespon fn from(resp: proto::GetLoadsResponse) -> Self { Self { timestamp: resp.timestamp, + version: resp.version, dp_rank_count: resp.dp_rank_count, loads: resp.loads.into_iter().map(Into::into).collect(), + aggregate: None, } } } diff --git a/crates/protocols/src/common.rs b/crates/protocols/src/common.rs index bf6d7825f..867c7a80d 100644 --- a/crates/protocols/src/common.rs +++ b/crates/protocols/src/common.rs @@ -562,9 +562,10 @@ impl Usage { /// Add cached token details to this Usage pub fn with_cached_tokens(mut self, cached_tokens: u32) -> Self { - if cached_tokens > 0 { - self.prompt_tokens_details = Some(PromptTokenUsageInfo { cached_tokens }); - } + // Calling this builder means the backend supplied cache accounting. + // Zero is therefore evidence of a cold miss, not absence of support, + // and must remain distinguishable from `prompt_tokens_details: None`. + self.prompt_tokens_details = Some(PromptTokenUsageInfo { cached_tokens }); self } @@ -852,6 +853,15 @@ mod tests { assert!(result.is_err()); } + #[test] + fn cached_token_builder_preserves_explicit_zero() { + let usage = Usage::from_counts(16, 1).with_cached_tokens(0); + assert!(matches!( + usage.prompt_tokens_details, + Some(PromptTokenUsageInfo { cached_tokens: 0 }) + )); + } + #[test] fn content_part_deserializes_audio_url() { let value = json!({ diff --git a/crates/protocols/src/worker.rs b/crates/protocols/src/worker.rs index af0301d9c..cd374285c 100644 --- a/crates/protocols/src/worker.rs +++ b/crates/protocols/src/worker.rs @@ -1243,6 +1243,10 @@ pub struct SchedulerLoadSnapshot { pub cache_hit_rate: f64, pub utilization: f64, pub max_running_requests: i32, + #[serde(skip_serializing_if = "Option::is_none")] + pub memory: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub queues: Option, /// PD disaggregation signals, populated only when the backend reports a /// `disagg` section. `None` for HTTP or older engines. Canonical schema /// other engines map into; SGLang derives the queue depths from its @@ -1260,13 +1264,45 @@ pub struct SchedulerLoadSnapshot { pub disagg_mode: Option, } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct EngineMemoryMetricsSnapshot { + pub weight_gb: f64, + pub kv_cache_gb: f64, + pub graph_gb: f64, + pub token_capacity: i32, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct EngineQueueMetricsSnapshot { + pub waiting: i32, + pub grammar: i32, + pub paused: i32, + pub retracted: i32, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct EngineAggregateMetricsSnapshot { + pub total_running_reqs: i32, + pub total_waiting_reqs: i32, + pub total_reqs: i32, + pub avg_token_usage: f64, + pub avg_throughput: f64, + pub avg_utilization: f64, +} + /// Full load response for a single worker across all DP ranks. #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(default)] pub struct WorkerLoadResponse { pub timestamp: String, + pub version: String, pub dp_rank_count: i32, pub loads: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub aggregate: Option, } impl WorkerLoadResponse { diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 13e2740e3..f3744cb21 100755 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -105,9 +105,9 @@ pub struct RouterConfig { /// to 90% of the ceiling. `None`/`0` disables the ceiling (default). #[serde(default, skip_serializing_if = "Option::is_none")] pub kv_indexer_max_entries: Option, - /// Re-export engine `GetLoads` signals as `smg_engine_*` gauges, polling - /// even when no load-aware routing policy is active. Decouples engine - /// observability from routing. + /// Force `GetLoads` polling for `smg_engine_*` gauges even when no + /// load-aware routing policy is active. Successful routing-owned polls are + /// always re-exported without an additional Engine RPC. #[serde(default)] pub engine_metrics: bool, /// Global multimodal tensor transport mode (`inline` | `shm` | `auto` | `rdma`). diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index b160b34f6..d43c469fd 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -461,8 +461,8 @@ struct CliArgs { #[arg(long, default_value_t = 10, help_heading = "Load Monitoring")] load_monitor_interval: u64, - /// Re-export engine GetLoads signals (incl. PD) as smg_engine_* Prometheus - /// gauges, polling even without a load-aware routing policy. + /// Force GetLoads polling for smg_engine_* Prometheus gauges even without + /// a load-aware routing policy. Routing-owned polls are always re-exported. #[arg(long, default_value_t = false, help_heading = "Load Monitoring")] engine_metrics: bool, diff --git a/model_gateway/src/observability/metrics.rs b/model_gateway/src/observability/metrics.rs index 497d75fae..33611d58c 100644 --- a/model_gateway/src/observability/metrics.rs +++ b/model_gateway/src/observability/metrics.rs @@ -1928,6 +1928,7 @@ mod tests { cache_hit_rate: 0.25, ..Default::default() }], + ..Default::default() }; let rendered = render_with_recorder(|| { @@ -1960,6 +1961,7 @@ mod tests { decode_queue_reqs: Some(4), ..Default::default() }], + ..Default::default() }; let rendered = render_with_recorder(|| { diff --git a/model_gateway/src/policies/least_load.rs b/model_gateway/src/policies/least_load.rs index 0003285d4..ee4d3181c 100644 --- a/model_gateway/src/policies/least_load.rs +++ b/model_gateway/src/policies/least_load.rs @@ -460,6 +460,7 @@ mod tests { max_running_requests: 0, ..Default::default() }], + ..Default::default() } } diff --git a/model_gateway/src/policies/power_of_two.rs b/model_gateway/src/policies/power_of_two.rs index 21abf693e..e8952e2c3 100644 --- a/model_gateway/src/policies/power_of_two.rs +++ b/model_gateway/src/policies/power_of_two.rs @@ -108,6 +108,7 @@ mod tests { gen_throughput, ..Default::default() }], + ..Default::default() } } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 689f35ae8..8a15d5304 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -679,6 +679,7 @@ impl ZmqEngineClient { timestamp: String::new(), dp_rank_count: i32::try_from(loads.len()).unwrap_or(i32::MAX), loads, + ..Default::default() } } diff --git a/model_gateway/src/worker/monitor.rs b/model_gateway/src/worker/monitor.rs index d8993550b..a36b7bf79 100644 --- a/model_gateway/src/worker/monitor.rs +++ b/model_gateway/src/worker/monitor.rs @@ -221,8 +221,8 @@ pub struct WorkerMonitor { pub worker_load_manager: Arc, client: reqwest::Client, default_interval: Duration, - /// When set, poll loads and re-export `smg_engine_*` gauges even if no - /// load-aware routing policy is active (`--engine-metrics`). + /// When set, force load polling even if no load-aware routing policy is + /// active (`--engine-metrics`). Every successful poll is re-exported. engine_metrics: bool, load_tx: watch::Sender>, load_rx: watch::Receiver>, @@ -454,16 +454,17 @@ impl WorkerMonitor { /// channel snapshot and the DP cache. Used by the event loop on /// `Removed`, `Replaced`, and `StatusChanged` away from `Ready`. /// - /// Also sentinels the worker's `smg_engine_*` series when engine-metrics - /// re-export is on, since metrics-rs cannot delete series. + /// Also sentinels the worker's `smg_engine_*` series after this monitor has + /// published a load for it, since metrics-rs cannot delete series. fn evict_worker_loads(&self, worker: &Arc) { let url = worker.url(); + let had_published_load = self.load_tx.borrow().contains_key(url); self.load_tx.send_modify(|map| { map.remove(url); }); self.worker_load_manager.remove_worker(url); self.native_loads_absent.remove(url); - if self.engine_metrics { + if had_published_load { // A worker can serve multiple models (one load group per model), // so sentinel every model's series — not just the primary. let dp_size = worker.dp_size().unwrap_or(1); @@ -820,10 +821,7 @@ async fn run_event_loop( } } Ok(WorkerEvent::StatusChanged { - worker, - new_status, - old_status: _, - .. + worker, new_status, .. }) => { if new_status != WorkerStatus::Ready { monitor.evict_worker_loads(&worker); @@ -985,12 +983,12 @@ async fn group_monitor_loop( monitor.worker_load_manager.remove_workers(&dp_evict); } - // Re-export the freshly fetched loads as `smg_engine_*` gauges. Reuses - // this poll's data; no extra fetch. Model label comes from the group. - if monitor.engine_metrics { - for (url, load) in &group_loads { - Metrics::record_engine_load(url, &group_key.model_id, load); - } + // Every successful load poll is also the canonical observability + // sample. Load-aware policies already require this poll; the explicit + // engine-metrics option only forces polling when routing does not. + // Reusing the response avoids a second Engine RPC. + for (url, load) in &group_loads { + Metrics::record_engine_load(url, &group_key.model_id, load); } // Atomically merge into the shared watch channel: clear stale diff --git a/model_gateway/tests/engine_metrics_test.rs b/model_gateway/tests/engine_metrics_test.rs index 589821337..297759ecc 100644 --- a/model_gateway/tests/engine_metrics_test.rs +++ b/model_gateway/tests/engine_metrics_test.rs @@ -41,6 +41,7 @@ async fn engine_pd_gauge_appears_on_metrics_endpoint() { decode_queue_reqs: Some(2), ..Default::default() }], + ..Default::default() }; Metrics::record_engine_load("grpc://prefill-0:30000", "test-model", &response);