diff --git a/crates/protocols/src/worker.rs b/crates/protocols/src/worker.rs index c01f1eeb9..29df13a12 100644 --- a/crates/protocols/src/worker.rs +++ b/crates/protocols/src/worker.rs @@ -1136,6 +1136,10 @@ pub struct FlushCacheResult { pub http_workers: usize, #[serde(default)] pub grpc_workers: usize, + /// Workers skipped because their transport has no cache-flush RPC + /// (direct-ZMQ engines). Keeps `total = http + grpc + zmq` exact. + #[serde(default)] + pub zmq_workers: usize, pub message: String, } @@ -1325,6 +1329,7 @@ impl IntoResponse for FlushCacheResult { "workers_flushed": self.successful.len(), "total_http_workers": self.http_workers, "total_grpc_workers": self.grpc_workers, + "total_zmq_workers_skipped": self.zmq_workers, "total_workers": self.total_workers }); diff --git a/model_gateway/src/routers/grpc/common/stages/encode.rs b/model_gateway/src/routers/grpc/common/stages/encode.rs index 9eeaf7371..0c6bc37ba 100644 --- a/model_gateway/src/routers/grpc/common/stages/encode.rs +++ b/model_gateway/src/routers/grpc/common/stages/encode.rs @@ -320,6 +320,15 @@ fn prepare_items( prefill: BackendClient::Grpc(GrpcClient::TokenSpeed(_)), .. } => prepare_tokenspeed_items(intermediate, workers), + // TokenSpeed supports EPD encode, but only over gRPC — name the + // transport, not the engine, so the error points at the real limit. + ClientSelection::Disaggregated { + prefill: prefill @ BackendClient::Zmq(_), + .. + } if prefill.runtime_type() == RuntimeType::TokenSpeed => Err(anyhow!( + "EPD encode requires a gRPC TokenSpeed prefill worker; the direct-ZMQ backend has \ + no encode dispatch" + )), ClientSelection::Disaggregated { prefill, .. } => Err(anyhow!( "EPD encode is not implemented for {} backend", backend_name(prefill) diff --git a/model_gateway/src/routers/grpc/common/stages/helpers.rs b/model_gateway/src/routers/grpc/common/stages/helpers.rs index fd3fbdb55..070f26dd2 100644 --- a/model_gateway/src/routers/grpc/common/stages/helpers.rs +++ b/model_gateway/src/routers/grpc/common/stages/helpers.rs @@ -452,7 +452,13 @@ fn inject_sglang_bootstrap_metadata( bootstrap_room: room_id, }; - let sglang_request = request.as_sglang_mut(); + // Guarded by the caller's runtime check, but match defensively: a non-SGLang + // proto here (e.g. a ZMQ backend reporting an unexpected runtime) must not + // take down the request task via the panicking accessor. + let ProtoGenerateRequest::Sglang(sglang_request) = request else { + warn!("PD bootstrap metadata requested for a non-SGLang request; skipping injection"); + return; + }; sglang_request.disaggregated_params = Some(disagg_params); debug!( @@ -768,6 +774,39 @@ mod stop_resolution_tests { ); } + #[test] + fn pd_bootstrap_injection_skips_non_sglang_requests() { + use super::{RuntimeType, Worker, WorkerSelection}; + use crate::worker::{BasicWorkerBuilder, WorkerType}; + + // An SGLang-runtime worker selection paired with a non-SGLang proto + // (e.g. a misreporting backend) must skip injection, not panic. + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("grpc://prefill:30000") + .worker_type(WorkerType::Prefill) + .build(), + ); + let selection = WorkerSelection::Disaggregated { + encode_assignments: None, + prefill: worker.clone(), + decode: worker, + runtime_type: RuntimeType::Sglang, + }; + + let mut req = vllm_request(vec!["."], vec![7]); + let before = match &req { + ProtoGenerateRequest::Vllm(inner) => (**inner).clone(), + _ => panic!("vllm_request builds a Vllm variant"), + }; + super::maybe_inject_pd_metadata(&mut req, &selection); + match &req { + ProtoGenerateRequest::Vllm(inner) => { + assert_eq!(**inner, before, "request must be untouched"); + } + _ => panic!("variant must be unchanged"), + } + } + #[test] fn tokenspeed_zmq_multi_token_relies_on_router_decoder() { // "Hello world" => [1, 2]: not a flat stop id, so it must not forward. diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs index c0957be4e..bd49cee7d 100644 --- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs @@ -25,11 +25,29 @@ use crate::{ utils::tonic_ext::{TonicResultExt, TonicStatusExt}, }, }, - worker::{RuntimeType, DEFAULT_BOOTSTRAP_PORT, MOONCAKE_CONNECTOR, NIXL_CONNECTOR}, + worker::{ + ConnectionModeExt, RuntimeType, DEFAULT_BOOTSTRAP_PORT, MOONCAKE_CONNECTOR, NIXL_CONNECTOR, + }, }; type StreamResult = Result; +/// Metric connection labels for the PD legs (a leg can be gRPC or ZMQ). +fn pd_leg_labels(workers: &WorkerSelection) -> (&'static str, &'static str) { + match workers { + WorkerSelection::Disaggregated { + prefill, decode, .. + } => ( + prefill.connection_mode().as_metric_label(), + decode.connection_mode().as_metric_label(), + ), + WorkerSelection::Single { worker } => { + let label = worker.connection_mode().as_metric_label(); + (label, label) + } + } +} + /// KV-transfer params tagged onto the NIXL prefill leg so the engine pins its /// KV blocks and returns the handoff params for the decode worker. const NIXL_PREFILL_KV_PARAMS: &str = r#"{"do_remote_decode":true,"do_remote_prefill":false}"#; @@ -489,11 +507,13 @@ impl RequestExecutionStage { decode_result.cb_status_code(), ); + let (prefill_label, decode_label) = pd_leg_labels(workers); + // Handle prefill result let prefill_stream = prefill_result.map_err(|e| { Metrics::record_worker_error( metrics_labels::WORKER_PREFILL, - metrics_labels::CONNECTION_GRPC, + prefill_label, metrics_labels::ERROR_BACKEND, ); error!(function = "execute_parallel_pd", error = %e, "Prefill worker failed to start"); @@ -507,7 +527,7 @@ impl RequestExecutionStage { let decode_stream = decode_result.map_err(|e| { Metrics::record_worker_error( metrics_labels::WORKER_DECODE, - metrics_labels::CONNECTION_GRPC, + decode_label, metrics_labels::ERROR_BACKEND, ); error!(function = "execute_parallel_pd", error = %e, "Decode worker failed to start"); @@ -648,6 +668,7 @@ impl RequestExecutionStage { ); // Send to prefill, wait for completion + let (prefill_label, decode_label) = pd_leg_labels(workers); let prefill_start = Instant::now(); let mut prefill_stream = prefill_client .generate(prefill_request) @@ -656,7 +677,7 @@ impl RequestExecutionStage { workers.record_outcome_prefill(e.http_status().as_u16()); Metrics::record_worker_error( metrics_labels::WORKER_PREFILL, - metrics_labels::CONNECTION_GRPC, + prefill_label, metrics_labels::ERROR_BACKEND, ); error!(function = "execute_sequential_pd", error = %e, "Prefill worker failed to start"); @@ -678,7 +699,7 @@ impl RequestExecutionStage { workers.record_outcome_prefill(e.http_status().as_u16()); Metrics::record_worker_error( metrics_labels::WORKER_PREFILL, - metrics_labels::CONNECTION_GRPC, + prefill_label, metrics_labels::ERROR_BACKEND, ); error!(function = "execute_sequential_pd", error = %e, "Prefill stream error"); @@ -769,7 +790,7 @@ impl RequestExecutionStage { workers.record_outcome_decode(e.http_status().as_u16()); Metrics::record_worker_error( metrics_labels::WORKER_DECODE, - metrics_labels::CONNECTION_GRPC, + decode_label, metrics_labels::ERROR_BACKEND, ); error!(function = "execute_sequential_pd", error = %e, "Decode worker failed to start"); @@ -803,9 +824,42 @@ impl RequestExecutionStage { #[cfg(test)] mod tests { + use std::sync::Arc; + use smg_grpc_client::vllm_proto as vllm; use super::*; + use crate::worker::{BasicWorkerBuilder, ConnectionMode, Worker, WorkerType}; + + #[test] + fn pd_leg_labels_reflect_each_legs_transport() { + let prefill: Arc = Arc::new( + BasicWorkerBuilder::new("ipc:///tmp/smg-test-prefill") + .worker_type(WorkerType::Prefill) + .connection_mode(ConnectionMode::Zmq) + .build(), + ); + let decode: Arc = Arc::new( + BasicWorkerBuilder::new("grpc://decode:30000") + .worker_type(WorkerType::Decode) + .connection_mode(ConnectionMode::Grpc) + .build(), + ); + let selection = WorkerSelection::Disaggregated { + encode_assignments: None, + prefill, + decode, + runtime_type: RuntimeType::TokenSpeed, + }; + assert_eq!( + pd_leg_labels(&selection), + ( + metrics_labels::CONNECTION_ZMQ, + metrics_labels::CONNECTION_GRPC + ), + "each PD leg must carry its own transport label" + ); + } #[test] fn kv_connector_mode_mooncake_uses_bootstrap_metadata() { diff --git a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs index 9b49ec687..e88293322 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -21,8 +21,8 @@ use crate::{ }, }, worker::{ - ConnectionModeExt, HashRing, RuntimeType, Worker, WorkerRegistry, WorkerType, - UNKNOWN_MODEL_ID, + ConnectionMode, ConnectionModeExt, HashRing, RuntimeType, Worker, WorkerRegistry, + WorkerType, UNKNOWN_MODEL_ID, }, }; @@ -399,13 +399,17 @@ impl WorkerSelectionStage { // Record worker selection metrics for both prefill and decode Metrics::record_worker_selection( metrics_labels::WORKER_PREFILL, - metrics_labels::CONNECTION_GRPC, + available_prefill[prefill_idx] + .connection_mode() + .as_metric_label(), model, prefill_policy.name(), ); Metrics::record_worker_selection( metrics_labels::WORKER_DECODE, - metrics_labels::CONNECTION_GRPC, + available_decode[decode_idx] + .connection_mode() + .as_metric_label(), model, decode_policy.name(), ); @@ -452,7 +456,14 @@ impl WorkerSelectionStage { .fold((Vec::new(), Vec::new(), Vec::new()), |mut acc, w| { if w.connection_mode().uses_grpc_pipeline() && w.is_available() { match w.metadata().spec.worker_type { - WorkerType::Encode => acc.0.push(w), + // Encode dispatch is a gRPC encoder RPC sent to the + // worker's URL; a direct-ZMQ worker has no encode + // path, so only gRPC workers qualify for this pool. + WorkerType::Encode => { + if *w.connection_mode() == ConnectionMode::Grpc { + acc.0.push(w); + } + } WorkerType::Prefill => acc.1.push(w), WorkerType::Decode => acc.2.push(w), WorkerType::Regular => {} @@ -577,13 +588,17 @@ impl WorkerSelectionStage { // recorded in assign_encode_workers. Metrics::record_worker_selection( metrics_labels::WORKER_PREFILL, - metrics_labels::CONNECTION_GRPC, + available_prefill[prefill_idx] + .connection_mode() + .as_metric_label(), model_id, prefill_policy.name(), ); Metrics::record_worker_selection( metrics_labels::WORKER_DECODE, - metrics_labels::CONNECTION_GRPC, + available_decode[decode_idx] + .connection_mode() + .as_metric_label(), model_id, decode_policy.name(), ); diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index bbc42fee6..05a899825 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -1971,6 +1971,12 @@ impl ProtoStream { .next() .await .map(|result| result.map(|r| ProtoGenerateResponse::TokenSpeed(Box::new(r)))), + // Every ZMQ engine (including TokenSpeed) emits vllm-shaped + // responses: the adapter translates wire output into + // `vllm::GenerateResponse`, so variant checks like + // `is_tokenspeed()` on a response are unreliable for ZMQ-backed + // streams. Key response-side engine logic on the worker's + // `runtime_type()`, never on the response variant. Self::Zmq(stream) => stream .next() .await diff --git a/model_gateway/src/routers/grpc/router.rs b/model_gateway/src/routers/grpc/router.rs index 9b84c4d4b..9b8b82ec6 100644 --- a/model_gateway/src/routers/grpc/router.rs +++ b/model_gateway/src/routers/grpc/router.rs @@ -40,7 +40,7 @@ use crate::{ common::retry::{is_retryable_status, RetryExecutor}, error, RouterTrait, }, - worker::{ConnectionMode, WorkerRegistry, WorkerType}, + worker::{WorkerRegistry, WorkerType}, }; const QWEN3_ASR_LANGUAGES: &[(&str, &str)] = &[ @@ -824,23 +824,20 @@ impl std::fmt::Debug for GrpcRouter { .finish() } Mode::PrefillDecode | Mode::EncodePrefillDecode => { - let prefill_workers = self.worker_registry.get_workers_filtered( - None, - Some(WorkerType::Prefill), - Some(ConnectionMode::Grpc), - None, - false, - ); - let decode_workers = self.worker_registry.get_workers_filtered( - None, - Some(WorkerType::Decode), - Some(ConnectionMode::Grpc), - None, - false, - ); + // Count every worker this router can serve (gRPC and ZMQ both + // ride the gRPC pipeline), not just ConnectionMode::Grpc. + let count_pipeline_workers = |worker_type| { + self.worker_registry + .get_workers_filtered(None, Some(worker_type), None, None, false) + .iter() + .filter(|w| w.connection_mode().uses_grpc_pipeline()) + .count() + }; + let prefill_workers = count_pipeline_workers(WorkerType::Prefill); + let decode_workers = count_pipeline_workers(WorkerType::Decode); f.debug_struct("GrpcRouter") - .field("prefill_workers_count", &prefill_workers.len()) - .field("decode_workers_count", &decode_workers.len()) + .field("prefill_workers_count", &prefill_workers) + .field("decode_workers_count", &decode_workers) .finish() } } @@ -1252,7 +1249,7 @@ mod pd_tests { config::{PolicyConfig, RouterConfig, RoutingMode}, policies::PolicyRegistry, tenant::TenantKey, - worker::{BasicWorkerBuilder, WorkerRegistry}, + worker::{BasicWorkerBuilder, ConnectionMode, WorkerRegistry}, }; fn pd_routing_mode() -> RoutingMode { @@ -1354,6 +1351,52 @@ mod pd_tests { /// PD serves /v1/responses: an unknown model is rejected per-request (404), /// not gated behind a blanket 501, and cancel reaches storage. + #[tokio::test] + async fn pd_debug_counts_include_zmq_workers() { + let ctx = grpc_ctx(pd_routing_mode()).await; + for (url, worker_type, mode) in [ + ( + "grpc://prefill:30000", + WorkerType::Prefill, + ConnectionMode::Grpc, + ), + ( + "ipc:///tmp/smg-test-prefill", + WorkerType::Prefill, + ConnectionMode::Zmq, + ), + ( + "ipc:///tmp/smg-test-decode", + WorkerType::Decode, + ConnectionMode::Zmq, + ), + ] { + let worker = BasicWorkerBuilder::new(url) + .worker_type(worker_type) + .connection_mode(mode) + .model(ModelCard::new("m")) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + ctx.worker_registry + .register(Arc::new(worker)) + .expect("register worker"); + } + + let router = GrpcRouter::new(&ctx, Mode::PrefillDecode).expect("pd router"); + let debug = format!("{router:?}"); + assert!( + debug.contains("prefill_workers_count: 2"), + "ZMQ prefill worker missing from debug counts: {debug}" + ); + assert!( + debug.contains("decode_workers_count: 1"), + "ZMQ decode worker missing from debug counts: {debug}" + ); + } + #[tokio::test] async fn pd_router_serves_responses_and_cancel() { let ctx = grpc_ctx(pd_routing_mode()).await; diff --git a/model_gateway/src/worker/manager.rs b/model_gateway/src/worker/manager.rs index 0339d5cd7..1b61df730 100644 --- a/model_gateway/src/worker/manager.rs +++ b/model_gateway/src/worker/manager.rs @@ -864,34 +864,53 @@ impl WorkerManager { } pub async fn flush_cache_all(worker_registry: &WorkerRegistry) -> FlushCacheResult { - let workers = worker_registry.get_all(); - let total_workers = workers.len(); - let http_workers = workers + let all_workers = worker_registry.get_all(); + let total_workers = all_workers.len(); + let http_workers = all_workers .iter() .filter(|w| matches!(w.connection_mode(), ConnectionMode::Http)) .count(); - let grpc_workers = total_workers - http_workers; + // ZMQ engines have no cache-flush RPC; fan out only to workers that can + // succeed instead of reporting every ZMQ worker as failed. + let (workers, zmq_workers): (Vec<_>, Vec<_>) = all_workers + .into_iter() + .partition(|w| !matches!(w.connection_mode(), ConnectionMode::Zmq)); + let zmq_skipped = zmq_workers.len(); + let grpc_workers = total_workers - http_workers - zmq_skipped; if workers.is_empty() { + let message = if zmq_skipped > 0 { + format!( + "No cache-flush-capable workers available \ + ({zmq_skipped} ZMQ workers skipped: no cache-flush RPC)" + ) + } else { + "No workers available for cache flush".to_string() + }; + info!("{}", message); return FlushCacheResult { successful: vec![], failed: vec![], total_workers, http_workers, grpc_workers, - message: "No workers available for cache flush".to_string(), + zmq_workers: zmq_skipped, + message, }; } info!( - "Flushing cache on {} workers ({} HTTP, {} gRPC)", - total_workers, http_workers, grpc_workers + "Flushing cache on {} workers ({} HTTP, {} gRPC, {} ZMQ skipped)", + workers.len(), + http_workers, + grpc_workers, + zmq_skipped ); let (successful, failed) = Self::admin_fan_out(workers, |w| async move { w.flush_cache().await }).await; - let message = if failed.is_empty() { + let mut message = if failed.is_empty() { format!( "Successfully flushed cache on all {} workers", successful.len() @@ -903,6 +922,11 @@ impl WorkerManager { failed.len() ) }; + if zmq_skipped > 0 { + message.push_str(&format!( + " ({zmq_skipped} ZMQ workers skipped: no cache-flush RPC)" + )); + } info!("{}", message); @@ -912,6 +936,7 @@ impl WorkerManager { total_workers, http_workers, grpc_workers, + zmq_workers: zmq_skipped, message, } }