diff --git a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs index 79999c992..3bc2f0221 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs @@ -118,397 +118,67 @@ impl PipelineStage for HarmonyRequestBuildingStage { // Build gRPC request using token_ids directly (Harmony encoding already handled message rendering) let placeholder_processed_text = "[harmony]".to_string(); - // Build proto request based on backend type and request type - let mut proto_request = match builder_client { - BackendClient::Grpc(GrpcClient::Sglang(sglang_client)) => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - sglang_client - .build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - SglangGenerateRequestOptions { - multimodal_inputs: None, - tool_call_constraint: tool_constraints, - require_reasoning: false, - }, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build SGLang generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => sglang_client - .build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build SGLang generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })?, - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::Sglang(Box::new(req)) - } - BackendClient::Grpc(GrpcClient::Vllm(_)) => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - VllmEngineClient::build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - None, // No multimodal in Harmony pipeline - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build vLLM generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => { - VllmEngineClient::build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build vLLM generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::Vllm(Box::new(req)) - } - BackendClient::Grpc(GrpcClient::Trtllm(trtllm_client)) => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - trtllm_client - .build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - None, // No multimodal in Harmony pipeline - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TensorRT-LLM generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => trtllm_client - .build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TensorRT-LLM generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })?, - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::Trtllm(Box::new(req)) - } - BackendClient::Grpc(GrpcClient::Mlx(mlx_client)) => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - mlx_client - .build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build MLX generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => mlx_client - .build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build MLX generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })?, - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::Mlx(Box::new(req)) - } - BackendClient::Grpc(GrpcClient::TokenSpeed(_)) => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - TokenSpeedSchedulerClient::build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - None, // Harmony path: multimodal not yet wired - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => TokenSpeedSchedulerClient::build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })?, - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::TokenSpeed(Box::new(req)) - } - // A ZMQ worker speaks vLLM EngineCore or TokenSpeed directly; build the - // request natively for its runtime, mirroring the gRPC per-engine - // dispatch above. Both support the Harmony request types (Chat + - // Responses). - BackendClient::Zmq(zmq_client) if zmq_client.runtime() == RuntimeType::TokenSpeed => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request - .as_deref() - .unwrap_or_else(|| request.as_ref()); - TokenSpeedSchedulerClient::build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - None, // Harmony path: multimodal not yet wired - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed ZMQ generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => { - TokenSpeedSchedulerClient::build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed ZMQ generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::TokenSpeed(Box::new(req)) - } - BackendClient::Zmq(zmq_client) if zmq_client.runtime() == RuntimeType::Vllm => { - let req = match &ctx.input.request_type { - RequestType::Chat(request) => { - let body = modified_request - .as_deref() - .unwrap_or_else(|| request.as_ref()); - VllmEngineClient::build_generate_request_from_chat( - request_id, - body, - placeholder_processed_text, - token_ids, - None, // No multimodal in Harmony pipeline - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build ZMQ generate request"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Responses(request) => { - VllmEngineClient::build_generate_request_from_responses( - request_id, - request.as_ref(), - placeholder_processed_text, - token_ids, - tool_constraints, - ) - .map_err(|e| { - error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build ZMQ generate request from responses"); - error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) - })? - } - RequestType::Embedding(_) => { - return Err(error::bad_request( - "harmony_embedding_not_supported", - "Embedding requests are not supported with Harmony models".to_string(), - )); - } - _ => { - return Err(error::bad_request( - "unsupported_request_type", - "Unsupported request type for Harmony models".to_string(), - )); - } - }; - ProtoGenerateRequest::Vllm(Box::new(req)) + // Resolve the request kind once; the backend dispatch below is a single + // compact match with one shared error path. (Non-Chat/Responses kinds + // were already rejected by the request-id match above; the reject arms + // here keep this exhaustive without a panic.) + let body = match &ctx.input.request_type { + RequestType::Chat(request) => HarmonyBody::Chat( + modified_request + .as_deref() + .unwrap_or_else(|| request.as_ref()), + ), + RequestType::Responses(request) => HarmonyBody::Responses(request.as_ref()), + RequestType::Embedding(_) => { + return Err(error::bad_request( + "harmony_embedding_not_supported", + "Embedding requests are not supported with Harmony models".to_string(), + )); } - // connect() admits only vLLM/TokenSpeed runtimes over ZMQ; a client - // reporting anything else is a wiring bug, not a request to serve. - BackendClient::Zmq(zmq_client) => { - return Err(error::internal_error( - "unsupported_zmq_runtime", - format!( - "ZMQ backend reports unsupported runtime {:?} for Harmony requests", - zmq_client.runtime() - ), + _ => { + return Err(error::bad_request( + "unsupported_request_type", + "Unsupported request type for Harmony models".to_string(), )); } }; - // Inject Harmony stop token IDs into sampling params for ALL Harmony requests - // These stop tokens (<|return|> and <|call|>) prevent the model from generating - // malformed Harmony sequences + let mut proto_request = build_harmony_proto( + builder_client, + body, + request_id, + placeholder_processed_text, + token_ids, + tool_constraints, + ) + .map_err(|e| match e { + HarmonyBuildError::Request(e) => { + error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build Harmony generate request"); + error::bad_request( + "invalid_request_parameters", + format!("Invalid request parameters: {e}"), + ) + } + HarmonyBuildError::Wiring(e) => { + error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Harmony backend wiring bug"); + error::internal_error("unsupported_backend_runtime", e) + } + })?; + + // Inject the Harmony stop ids (<|return|> and <|call|>) so the model + // cannot generate past a channel boundary. if !harmony_stop_ids.is_empty() { - match &mut proto_request { - ProtoGenerateRequest::Sglang(req) => { - if let Some(params) = req.sampling_params.as_mut() { - params.stop_token_ids.extend_from_slice(&harmony_stop_ids); - debug!( - stop_token_count = harmony_stop_ids.len(), - "Injected Harmony stop tokens into SGLang sampling params" - ); - } - } - ProtoGenerateRequest::Vllm(req) => { - if let Some(params) = req.sampling_params.as_mut() { - params.stop_token_ids.extend_from_slice(&harmony_stop_ids); - debug!( - stop_token_count = harmony_stop_ids.len(), - "Injected Harmony stop tokens into vLLM sampling params" - ); - } - } - ProtoGenerateRequest::Trtllm(req) => { - req.stop_token_ids.extend_from_slice(&harmony_stop_ids); - // TRT-LLM strips stop tokens from output by default, but - // the Harmony parser needs them to detect channel boundaries - // (e.g. <|call|> marks the tool-call channel transition). - req.include_stop_token_in_output = true; - debug!( - stop_token_count = harmony_stop_ids.len(), - "Injected Harmony stop tokens into TensorRT-LLM stop_token_ids" - ); - } - ProtoGenerateRequest::Mlx(req) => { - if let Some(ref mut params) = req.sampling_params { - params.stop_token_ids.extend_from_slice(&harmony_stop_ids); - debug!( - stop_token_count = harmony_stop_ids.len(), - "Injected Harmony stop tokens into MLX sampling params" - ); - } - } - ProtoGenerateRequest::TokenSpeed(req) => { - if let Some(params) = req.sampling_params.as_mut() { - params.stop_token_ids.extend_from_slice(&harmony_stop_ids); - debug!( - stop_token_count = harmony_stop_ids.len(), - "Injected Harmony stop tokens into TokenSpeed sampling params" - ); - } - } + proto_request.extend_stop_token_ids(&harmony_stop_ids); + if let ProtoGenerateRequest::Trtllm(req) = &mut proto_request { + // TRT-LLM strips stop tokens from output by default, but the + // Harmony parser needs them to detect channel boundaries + // (e.g. <|call|> marks the tool-call channel transition). + req.include_stop_token_in_output = true; } + debug!( + stop_token_count = harmony_stop_ids.len(), + "Injected Harmony stop tokens" + ); } // The client resolves string `stop`s its engine can't match and @@ -539,3 +209,156 @@ impl PipelineStage for HarmonyRequestBuildingStage { ) } } + +/// The two request kinds Harmony serves, with the Chat body already resolved +/// to the (possibly response_format-modified) request. +enum HarmonyBody<'a> { + Chat(&'a openai_protocol::chat::ChatCompletionRequest), + Responses(&'a openai_protocol::responses::ResponsesRequest), +} + +/// Build failure classes: a bad request (builder rejected the parameters, +/// HTTP 400) vs. a wiring bug (a backend/runtime pairing that cannot exist, +/// HTTP 500). +enum HarmonyBuildError { + Request(String), + Wiring(String), +} + +impl From for HarmonyBuildError { + fn from(reason: String) -> Self { + Self::Request(reason) + } +} + +/// One (backend x request-kind) dispatch for Harmony request building: every +/// arm is just the engine's builder call. vLLM and TokenSpeed build through +/// static translators, so one arm each covers both gRPC and direct-ZMQ. +fn build_harmony_proto( + client: &BackendClient, + body: HarmonyBody<'_>, + request_id: String, + text: String, + token_ids: Vec, + tool_constraints: Option<(String, String)>, +) -> Result { + use HarmonyBody::{Chat, Responses}; + let runtime = client.runtime_type(); + Ok(match (client, body) { + (BackendClient::Grpc(GrpcClient::Sglang(c)), Chat(b)) => { + ProtoGenerateRequest::Sglang(Box::new(c.build_generate_request_from_chat( + request_id, + b, + text, + token_ids, + SglangGenerateRequestOptions { + multimodal_inputs: None, + tool_call_constraint: tool_constraints, + require_reasoning: false, + }, + )?)) + } + (BackendClient::Grpc(GrpcClient::Sglang(c)), Responses(b)) => { + ProtoGenerateRequest::Sglang(Box::new(c.build_generate_request_from_responses( + request_id, + b, + text, + token_ids, + tool_constraints, + )?)) + } + (BackendClient::Grpc(GrpcClient::Trtllm(c)), Chat(b)) => { + ProtoGenerateRequest::Trtllm(Box::new(c.build_generate_request_from_chat( + request_id, + b, + text, + token_ids, + None, // No multimodal in the Harmony pipeline + tool_constraints, + )?)) + } + (BackendClient::Grpc(GrpcClient::Trtllm(c)), Responses(b)) => { + ProtoGenerateRequest::Trtllm(Box::new(c.build_generate_request_from_responses( + request_id, + b, + text, + token_ids, + tool_constraints, + )?)) + } + (BackendClient::Grpc(GrpcClient::Mlx(c)), Chat(b)) => ProtoGenerateRequest::Mlx(Box::new( + c.build_generate_request_from_chat(request_id, b, text, token_ids, tool_constraints)?, + )), + (BackendClient::Grpc(GrpcClient::Mlx(c)), Responses(b)) => { + ProtoGenerateRequest::Mlx(Box::new(c.build_generate_request_from_responses( + request_id, + b, + text, + token_ids, + tool_constraints, + )?)) + } + (BackendClient::Grpc(GrpcClient::Vllm(_)) | BackendClient::Zmq(_), Chat(b)) + if runtime == RuntimeType::Vllm => + { + ProtoGenerateRequest::Vllm(Box::new( + VllmEngineClient::build_generate_request_from_chat( + request_id, + b, + text, + token_ids, + None, // No multimodal in the Harmony pipeline + tool_constraints, + )?, + )) + } + (BackendClient::Grpc(GrpcClient::Vllm(_)) | BackendClient::Zmq(_), Responses(b)) + if runtime == RuntimeType::Vllm => + { + ProtoGenerateRequest::Vllm(Box::new( + VllmEngineClient::build_generate_request_from_responses( + request_id, + b, + text, + token_ids, + tool_constraints, + )?, + )) + } + (BackendClient::Grpc(GrpcClient::TokenSpeed(_)) | BackendClient::Zmq(_), Chat(b)) + if runtime == RuntimeType::TokenSpeed => + { + ProtoGenerateRequest::TokenSpeed(Box::new( + TokenSpeedSchedulerClient::build_generate_request_from_chat( + request_id, + b, + text, + token_ids, + None, // Harmony path: multimodal not yet wired + tool_constraints, + )?, + )) + } + (BackendClient::Grpc(GrpcClient::TokenSpeed(_)) | BackendClient::Zmq(_), Responses(b)) + if runtime == RuntimeType::TokenSpeed => + { + ProtoGenerateRequest::TokenSpeed(Box::new( + TokenSpeedSchedulerClient::build_generate_request_from_responses( + request_id, + b, + text, + token_ids, + tool_constraints, + )?, + )) + } + // Guards above keep the match non-exhaustive to the compiler; the only + // real way here is a ZMQ client reporting a runtime it cannot have + // (connect() admits vLLM/TokenSpeed only) - a wiring bug, so error out. + _ => { + return Err(HarmonyBuildError::Wiring(format!( + "unsupported backend runtime {runtime:?} for Harmony requests" + ))) + } + }) +} diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index 05a899825..b44103f89 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -1091,6 +1091,35 @@ pub enum ProtoGenerateRequest { } impl ProtoGenerateRequest { + /// Append stop token ids to the request's sampling params (TRT-LLM keeps + /// them on the request itself). Requests without sampling params are left + /// unchanged, matching the per-engine injection this replaces. + pub fn extend_stop_token_ids(&mut self, ids: &[u32]) { + match self { + Self::Sglang(req) => { + if let Some(params) = req.sampling_params.as_mut() { + params.stop_token_ids.extend_from_slice(ids); + } + } + Self::Vllm(req) => { + if let Some(params) = req.sampling_params.as_mut() { + params.stop_token_ids.extend_from_slice(ids); + } + } + Self::Mlx(req) => { + if let Some(params) = req.sampling_params.as_mut() { + params.stop_token_ids.extend_from_slice(ids); + } + } + Self::TokenSpeed(req) => { + if let Some(params) = req.sampling_params.as_mut() { + params.stop_token_ids.extend_from_slice(ids); + } + } + Self::Trtllm(req) => req.stop_token_ids.extend_from_slice(ids), + } + } + /// Get SGLang variant (panics if not SGLang) #[expect( clippy::panic, @@ -2345,4 +2374,31 @@ mod tests { let image = vllm_mm_data(common::Modality::Image).into_proto(); assert_eq!(image.modality, common::Modality::Image as i32); } + #[test] + fn extend_stop_token_ids_reaches_every_variant() { + let ids = [7u32, 8]; + let mut req = ProtoGenerateRequest::Vllm(Box::new(vllm::GenerateRequest { + sampling_params: Some(vllm::SamplingParams::default()), + ..Default::default() + })); + req.extend_stop_token_ids(&ids); + match &req { + ProtoGenerateRequest::Vllm(r) => { + assert_eq!(r.sampling_params.as_ref().unwrap().stop_token_ids, ids); + } + _ => panic!("variant changed"), + } + + // TRT-LLM keeps ids on the request itself. + let mut req = ProtoGenerateRequest::Trtllm(Box::default()); + req.extend_stop_token_ids(&ids); + match &req { + ProtoGenerateRequest::Trtllm(r) => assert_eq!(r.stop_token_ids, ids), + _ => panic!("variant changed"), + } + + // Missing sampling params: untouched, no panic. + let mut req = ProtoGenerateRequest::TokenSpeed(Box::default()); + req.extend_stop_token_ids(&ids); + } } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index f27ccb855..71167fed0 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -159,6 +159,169 @@ pub(crate) fn fold_tokenizer_eos_backstop( } } +/// Time to wait for a ZMQ engine to complete the startup handshake. Generous: +/// the engine loads the model and profiles KV cache between INIT and READY. +const ZMQ_CONNECT_TIMEOUT: Duration = Duration::from_secs(600); + +/// Derive a deterministic TCP handshake port from the ipc data-plane path. +/// +/// vLLM's headless engine dials a *TCP* handshake (`--data-parallel-address` + +/// `--data-parallel-rpc-port`); making the port a pure function of the worker +/// URL lets the operator compute the same `--data-parallel-rpc-port` without a +/// side channel. FNV-1a keeps it stable across processes and builds. Mapped +/// into 20000..=29999 to avoid well-known and typical ephemeral ranges. +/// +/// `_zmq_handshake_port` in `bindings/python/src/smg/serve.py` mirrors this +/// function — keep them in sync. +fn derive_handshake_port(path: &str) -> u16 { + let mut hash: u64 = 0xcbf2_9ce4_8422_2325; + for b in path.as_bytes() { + hash ^= u64::from(*b); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + // Map into 20000..=29999: below the Linux default ephemeral range + // (`net.ipv4.ip_local_port_range` = 32768..60999) so an outbound socket + // can't already hold the port. `hash % 10000` always fits u16. + 20000 + (hash % 10000) as u16 +} + +/// Derive the ZMQ socket addresses for a worker from its base URL. +/// +/// Mirrors vLLM's headless topology: the **handshake is TCP** (the engine dials +/// it, so it matches `vllm serve --headless --data-parallel-rpc-port`), while +/// the **data plane is `ipc://`** for the same-host fast path (SMG chooses these +/// and hands them to the engine during the handshake INIT). The operator gives a +/// single `ipc://` base; SMG binds the ipc input/output at +/// `-in.sock` / `-out.sock` and derives the TCP handshake port from the +/// path. A `WorkerSpec.zmq_handshake_address` override replaces the derived +/// handshake address verbatim (it must be `tcp://`), for engines that dial a +/// fixed, pre-agreed address — e.g. TokenSpeed's default dial target is +/// `tcp://127.0.0.1:30500` (its `--data-parallel-address`/ +/// `--data-parallel-rpc-port` defaults, outside the derived 20000..=29999 +/// band), so setting the override to that value pairs a bare +/// `ts serve --headless` with a manually registered worker. +/// Returns `(handshake, input, output)`. +fn zmq_socket_addresses( + base_url: &str, + handshake_override: Option<&str>, +) -> Result<(String, String, String), String> { + let path = base_url + .strip_prefix("ipc://") + .ok_or_else(|| format!("ZMQ worker URL must be ipc://, got '{base_url}'"))?; + let handshake = match handshake_override { + Some(address) => { + if !address.starts_with("tcp://") { + return Err(format!( + "zmq_handshake_address must be a tcp:// address \ + (the engine dials a TCP handshake), got '{address}'" + )); + } + address.to_string() + } + None => format!("tcp://{ZMQ_LOOPBACK_HOST}:{}", derive_handshake_port(path)), + }; + let input = format!("ipc://{path}-in.sock"); + let output = format!("ipc://{path}-out.sock"); + Ok((handshake, input, output)) +} + +/// Create the parent directory for a worker's `ipc://` sockets. Kept off the +/// address computation (which is pure) and async so it doesn't block a runtime +/// thread. +/// +/// The ipc:// data-plane sockets SMG binds here carry no authentication, so the +/// directory must be owner-controlled: when this call creates it, it is created +/// 0700 (mode applied at mkdir time — no chmod window); when it already exists, +/// its permissions are left untouched (never chmod a shared dir like `/tmp`) +/// and it is rejected unless it is a real directory owned by the current user. +async fn ensure_ipc_socket_dir(base_url: &str) -> Result<(), String> { + let path = base_url.strip_prefix("ipc://").unwrap_or(base_url); + let Some(parent) = Path::new(path).parent() else { + return Ok(()); + }; + // symlink_metadata: a symlinked parent must not redirect the checks (or the + // sockets) into a directory we did not verify. + let meta = match tokio::fs::symlink_metadata(parent).await { + Ok(meta) => meta, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + let mut builder = tokio::fs::DirBuilder::new(); + builder.recursive(true); + #[cfg(unix)] + builder.mode(0o700); + builder + .create(parent) + .await + .map_err(|e| format!("failed to create ipc socket dir for {path}: {e}"))?; + tokio::fs::symlink_metadata(parent) + .await + .map_err(|e| format!("failed to stat ipc socket dir for {path}: {e}"))? + } + Err(e) => return Err(format!("failed to stat ipc socket dir for {path}: {e}")), + }; + if !meta.is_dir() { + return Err(format!( + "ipc socket dir {} exists but is not a directory", + parent.display() + )); + } + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + let uid = rustix::process::geteuid().as_raw(); + if meta.uid() != uid { + return Err(format!( + "ipc socket dir {} is owned by uid {} (expected {uid}); refusing to bind \ + unauthenticated ZMQ sockets in a directory owned by another user", + parent.display(), + meta.uid() + )); + } + } + Ok(()) +} + +/// Bind the SMG-side ZMQ sockets and complete the handshake with the engine: +/// the single connect path for a worker's `ipc://` URL, shared by the lazy +/// client accessor and the background handshake driver. `model_id` is the +/// config-resolved served model (EngineCore reports none). Errors are plain +/// reasons; the worker layer wraps them in its own error type. +pub(crate) async fn connect_for_worker( + base_url: &str, + model_id: String, + runtime: RuntimeType, + handshake_override: Option<&str>, +) -> Result { + let (handshake, input, output) = zmq_socket_addresses(base_url, handshake_override)?; + ensure_ipc_socket_dir(base_url).await?; + // The engine can't stop at EOS on its own (it has no tokenizer or model + // config); resolve the EOS ids from the local model dir so every request + // carries them. + let model_dir = Path::new(&model_id); + let eos = if model_dir.is_dir() { + EosTokenIds::from_model_dir(model_dir) + } else { + tracing::warn!( + "ZMQ worker model id '{model_id}' is not a local model directory; connect-time \ + EOS ids unavailable — relying on the tokenizer's EOS set, folded into stop \ + tokens at request time" + ); + EosTokenIds::default() + }; + tracing::info!("Binding ZMQ client for worker {base_url} (handshake={handshake})"); + ZmqEngineClient::connect( + &handshake, + &input, + &output, + 1, + model_id, + eos, + runtime, + ZMQ_CONNECT_TIMEOUT, + ) + .await + .map_err(|e| format!("Failed to connect ZMQ engine: {e}")) +} + /// Direct ZMQ connection to a same-host engine (vLLM EngineCore or TokenSpeed), /// presented behind the vLLM gRPC client surface. #[derive(Clone)] @@ -506,6 +669,60 @@ impl StreamState { top_logprobs: std::mem::take(&mut self.output_top_logprobs), }) } + /// Emit one engine tick as vLLM-proto responses. On a finish tick the + /// `Complete` (with the engine-specific finish reason and matched stop) is + /// returned directly — unless the tick also carried new tokens, in which + /// case a `Chunk` goes out first and the `Complete` is parked in `pending` + /// for the next poll. Non-finish ticks emit a plain `Chunk`. + fn emit_tick( + &mut self, + index: u32, + token_ids: Vec, + chunk_logprobs: Option, + finish: Option<(String, Option)>, + pending: &mut Option, + ) -> vllm::GenerateResponse { + let chunk = |state: &Self, token_ids, chunk_logprobs| { + vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { + token_ids, + prompt_tokens: state.prompt_tokens, + completion_tokens: state.completion_tokens, + cached_tokens: state.cached_tokens, + output_logprobs: chunk_logprobs, + index, + ..Default::default() + }) + }; + let response = match finish { + Some((finish_reason, matched_stop)) => { + let complete = vllm::GenerateResponse { + response: Some(vllm::generate_response::Response::Complete( + vllm::GenerateComplete { + output_ids: std::mem::take(&mut self.output_ids), + finish_reason, + prompt_tokens: self.prompt_tokens, + completion_tokens: self.completion_tokens, + cached_tokens: self.cached_tokens, + matched_stop, + output_logprobs: self.take_complete_logprobs(), + index, + ..Default::default() + }, + )), + }; + if token_ids.is_empty() { + return complete; + } + let chunk = chunk(self, token_ids, chunk_logprobs); + *pending = Some(complete); + chunk + } + None => chunk(self, token_ids, chunk_logprobs), + }; + vllm::GenerateResponse { + response: Some(response), + } + } } /// Streaming generate output for one vLLM EngineCore sub-request, mapping each @@ -592,61 +809,27 @@ impl VllmGenerateStream { state.output_logprobs_idx.extend(tick_logprobs_idx); state.output_top_logprobs.extend(tick_top_logprobs); - let response = match output.finish_reason { - // An engine-side request failure (e.g. grammar compilation) must - // surface as an error, not as a normal completion with empty - // output — that would produce a 200 with no content. - Some(EngineCoreFinishReason::Error) => { - return Err(tonic::Status::internal( - "engine finished the request with an error (see engine logs)", - )); - } - Some(reason) => { - let complete = vllm::GenerateResponse { - response: Some(vllm::generate_response::Response::Complete( - vllm::GenerateComplete { - output_ids: std::mem::take(&mut state.output_ids), - finish_reason: finish_reason_str(reason).to_string(), - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - matched_stop: output.stop_reason.map(map_matched_stop), - output_logprobs: state.take_complete_logprobs(), - index: self.index, - ..Default::default() - }, - )), - }; - if token_ids.is_empty() { - return Ok(complete); - } - // The finish tick carried new tokens: emit them as a `Chunk` - // first and hold the `Complete` for the next poll. - let chunk = vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { - token_ids, - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: chunk_logprobs, - index: self.index, - ..Default::default() - }); - self.pending = Some(complete); - chunk - } - None => vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { - token_ids, - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: chunk_logprobs, - index: self.index, - ..Default::default() - }), - }; - Ok(vllm::GenerateResponse { - response: Some(response), - }) + // An engine-side request failure (e.g. grammar compilation) must + // surface as an error, not as a normal completion with empty output — + // that would produce a 200 with no content. + if matches!(output.finish_reason, Some(EngineCoreFinishReason::Error)) { + return Err(tonic::Status::internal( + "engine finished the request with an error (see engine logs)", + )); + } + let finish = output.finish_reason.map(|reason| { + ( + finish_reason_str(reason).to_string(), + output.stop_reason.map(map_matched_stop), + ) + }); + Ok(state.emit_tick( + self.index, + token_ids, + chunk_logprobs, + finish, + &mut self.pending, + )) } } @@ -696,7 +879,10 @@ impl TokenSpeedGenerateStream { } } - fn map_output(&mut self, output: TokenSpeedOutput) -> vllm::GenerateResponse { + fn map_output( + &mut self, + output: TokenSpeedOutput, + ) -> Result { let state = &mut self.state; // TokenSpeed reports per-request token counts directly (cumulative for // completions), rather than vLLM's per-output prefill-stats deltas. @@ -730,52 +916,25 @@ impl TokenSpeedGenerateStream { .output_logprobs_idx .extend(output.output_logprobs_idx.iter().copied()); - let response = match output.finish_reason { - Some(reason) => { - let complete = vllm::GenerateResponse { - response: Some(vllm::generate_response::Response::Complete( - vllm::GenerateComplete { - output_ids: std::mem::take(&mut state.output_ids), - finish_reason: normalize_finish_reason(&reason).to_string(), - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: state.take_complete_logprobs(), - index: self.index, - ..Default::default() - }, - )), - }; - if output.output_ids.is_empty() { - return complete; - } - // The finish tick carried new tokens: emit them as a `Chunk` - // first and hold the `Complete` for the next poll. - let chunk = vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { - token_ids: output.output_ids, - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: chunk_logprobs, - index: self.index, - ..Default::default() - }); - self.pending = Some(complete); - chunk - } - None => vllm::generate_response::Response::Chunk(vllm::GenerateStreamChunk { - token_ids: output.output_ids, - prompt_tokens: state.prompt_tokens, - completion_tokens: state.completion_tokens, - cached_tokens: state.cached_tokens, - output_logprobs: chunk_logprobs, - index: self.index, - ..Default::default() - }), - }; - vllm::GenerateResponse { - response: Some(response), + // An engine-side failure must surface as an error, not a normal + // completion with empty output — mirroring the vLLM stream's guard. + if output.finish_reason.as_deref() == Some("error") { + return Err(tonic::Status::internal( + "engine finished the request with an error (see engine logs)", + )); } + // No matched_stop on this wire: TokenSpeed reports the finish reason + // only, and the router-side stop machinery owns string matching. + let finish = output + .finish_reason + .map(|reason| (normalize_finish_reason(&reason).to_string(), None)); + Ok(state.emit_tick( + self.index, + output.output_ids, + chunk_logprobs, + finish, + &mut self.pending, + )) } } @@ -792,7 +951,7 @@ impl Stream for TokenSpeedGenerateStream { return Poll::Ready(Some(Ok(pending))); } match std::pin::Pin::new(&mut this.inner).poll_next(cx) { - Poll::Ready(Some(Ok(output))) => Poll::Ready(Some(Ok(this.map_output(output)))), + Poll::Ready(Some(Ok(output))) => Poll::Ready(Some(this.map_output(output))), Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(zmq_status(error)))), Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, @@ -818,17 +977,27 @@ impl Stream for TokenSpeedGenerateStream { /// pipeline de-duplicates (max per prompt), so nothing is counted n times. fn fan_out_requests(req: vllm::GenerateRequest) -> Vec { let n = req.sampling_params.as_ref().map_or(1, |sp| sp.n.max(1)); + fan_out_n(req, n, |sub, i| { + sub.request_id = format!("{}-{i}", sub.request_id); + if let Some(sp) = sub.sampling_params.as_mut() { + sp.n = 1; + // An explicit seed must still yield distinct samples per sub. + sp.seed = sp.seed.map(|seed| seed.wrapping_add(i as i32)); + } + }) +} + +/// Shared n>1 fan-out scaffolding: clone the request into `n` subs and let +/// `per_sub` apply the engine-specific rid suffix and sampling tweaks. An +/// `n <= 1` request passes through untouched. +fn fan_out_n(req: R, n: u32, mut per_sub: impl FnMut(&mut R, u32)) -> Vec { if n <= 1 { return vec![req]; } (0..n) .map(|i| { let mut sub = req.clone(); - sub.request_id = format!("{}-{i}", req.request_id); - if let Some(sp) = sub.sampling_params.as_mut() { - sp.n = 1; - sp.seed = sp.seed.map(|seed| seed.wrapping_add(i as i32)); - } + per_sub(&mut sub, i); sub }) .collect() @@ -844,19 +1013,12 @@ fn fan_out_tokenspeed_requests( req: tokenspeed_proto::GenerateRequest, ) -> Vec { let n = req.sampling_params.as_ref().map_or(1, |sp| sp.n.max(1)); - if n <= 1 { - return vec![req]; - } - (0..n) - .map(|i| { - let mut sub = req.clone(); - sub.request_id = format!("{}-{i}", req.request_id); - if let Some(sp) = sub.sampling_params.as_mut() { - sp.n = 1; - } - sub - }) - .collect() + fan_out_n(req, n, |sub, i| { + sub.request_id = format!("{}-{i}", sub.request_id); + if let Some(sp) = sub.sampling_params.as_mut() { + sp.n = 1; + } + }) } /// Translate a TokenSpeed proto `GenerateRequest` into the wire @@ -1171,6 +1333,91 @@ mod tests { use super::*; + #[test] + fn derive_handshake_port_matches_pinned_vectors() { + // Fixed vectors shared with `_zmq_handshake_port` in + // bindings/python/src/smg/serve.py — a change on either side breaks the + // engine/router port agreement, so these must stay in sync. + assert_eq!(derive_handshake_port("/tmp/smg-zmq/ts0.ipc"), 25152); + assert_eq!(derive_handshake_port("/tmp/smg-zmq/engine-31000"), 22714); + // Range invariant: every path maps into 20000..=29999. + for p in ["", "a", "/x/y/z.ipc", "very/long/path/with/segments.sock"] { + let port = derive_handshake_port(p); + assert!( + (20000..=29999).contains(&port), + "port {port} out of band for {p:?}" + ); + } + } + + #[test] + fn zmq_socket_addresses_derive_handshake_by_default() { + let (handshake, input, output) = + zmq_socket_addresses("ipc:///tmp/smg-zmq/ts0.ipc", None).unwrap(); + assert_eq!(handshake, "tcp://127.0.0.1:25152"); + assert_eq!(input, "ipc:///tmp/smg-zmq/ts0.ipc-in.sock"); + assert_eq!(output, "ipc:///tmp/smg-zmq/ts0.ipc-out.sock"); + } + + #[test] + fn zmq_socket_addresses_honor_handshake_override() { + // TokenSpeed's default dial target — outside the derived band; the + // override must be bound verbatim while the data plane stays derived. + let (handshake, input, output) = + zmq_socket_addresses("ipc:///tmp/smg-zmq/ts0.ipc", Some("tcp://127.0.0.1:30500")) + .unwrap(); + assert_eq!(handshake, "tcp://127.0.0.1:30500"); + assert_eq!(input, "ipc:///tmp/smg-zmq/ts0.ipc-in.sock"); + assert_eq!(output, "ipc:///tmp/smg-zmq/ts0.ipc-out.sock"); + } + + #[test] + fn zmq_socket_addresses_reject_non_tcp_override() { + // The engine dials a TCP handshake; a non-tcp override is a config + // error and must fail loudly rather than bind something unexpected. + let err = zmq_socket_addresses("ipc:///tmp/smg-zmq/ts0.ipc", Some("ipc:///tmp/hs.sock")) + .unwrap_err(); + assert!( + err.contains("tcp://"), + "error must name the required scheme: {err}" + ); + } + + #[tokio::test] + async fn ensure_ipc_socket_dir_creates_a_private_owner_only_dir() { + let base = tempfile::tempdir().unwrap(); + let dir = base.path().join("sockets"); + let url = format!("ipc://{}/x.ipc", dir.display()); + ensure_ipc_socket_dir(&url).await.unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = std::fs::metadata(&dir).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o700, "created socket dir must be 0700"); + } + } + + #[cfg(unix)] + #[tokio::test] + async fn ensure_ipc_socket_dir_leaves_an_existing_owned_dir_untouched() { + use std::os::unix::fs::PermissionsExt; + let dir = tempfile::tempdir().unwrap(); + std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap(); + let url = format!("ipc://{}/x.ipc", dir.path().display()); + ensure_ipc_socket_dir(&url).await.unwrap(); + let mode = std::fs::metadata(dir.path()).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o755, "an existing dir must not be chmod'd"); + } + + #[tokio::test] + async fn ensure_ipc_socket_dir_rejects_a_non_directory_parent() { + let base = tempfile::tempdir().unwrap(); + let file = base.path().join("not-a-dir"); + std::fs::write(&file, b"x").unwrap(); + let url = format!("ipc://{}/x.ipc", file.display()); + assert!(ensure_ipc_socket_dir(&url).await.is_err()); + } + fn eos_request(stop_token_ids: Vec, ignore_eos: bool) -> ProtoGenerateRequest { ProtoGenerateRequest::Vllm(Box::new(vllm::GenerateRequest { sampling_params: Some(vllm::SamplingParams { diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index d34b9a4e3..4cc7bd553 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -32,11 +32,7 @@ use crate::{ observability::metrics::{metrics_labels, Metrics}, routers::{ common::header_utils::extract_routing_key, - grpc::{ - backend_client::BackendClient, - client::GrpcClient, - zmq_client::{EosTokenIds, ZmqEngineClient, ZMQ_LOOPBACK_HOST}, - }, + grpc::{backend_client::BackendClient, client::GrpcClient, zmq_client}, }, }; @@ -52,186 +48,23 @@ const FLUSH_HTTP_TIMEOUT: Duration = Duration::from_secs(45); /// gRPC client's profile deadline. const PROFILE_HTTP_TIMEOUT: Duration = Duration::from_secs(630); -/// Time to wait for a ZMQ engine to complete the startup handshake. Generous: -/// the engine loads the model and profiles KV cache between INIT and READY. -const ZMQ_CONNECT_TIMEOUT: Duration = Duration::from_secs(600); - -/// Derive a deterministic TCP handshake port from the ipc data-plane path. -/// -/// vLLM's headless engine dials a *TCP* handshake (`--data-parallel-address` + -/// `--data-parallel-rpc-port`); making the port a pure function of the worker -/// URL lets the operator compute the same `--data-parallel-rpc-port` without a -/// side channel. FNV-1a keeps it stable across processes and builds. Mapped -/// into 20000..=29999 to avoid well-known and typical ephemeral ranges. -/// -/// `_zmq_handshake_port` in `bindings/python/src/smg/serve.py` mirrors this -/// function — keep them in sync. -fn derive_handshake_port(path: &str) -> u16 { - let mut hash: u64 = 0xcbf2_9ce4_8422_2325; - for b in path.as_bytes() { - hash ^= u64::from(*b); - hash = hash.wrapping_mul(0x0000_0100_0000_01b3); - } - // Map into 20000..=29999: below the Linux default ephemeral range - // (`net.ipv4.ip_local_port_range` = 32768..60999) so an outbound socket - // can't already hold the port. `hash % 10000` always fits u16. - 20000 + (hash % 10000) as u16 -} - -/// Derive the ZMQ socket addresses for a worker from its base URL. -/// -/// Mirrors vLLM's headless topology: the **handshake is TCP** (the engine dials -/// it, so it matches `vllm serve --headless --data-parallel-rpc-port`), while -/// the **data plane is `ipc://`** for the same-host fast path (SMG chooses these -/// and hands them to the engine during the handshake INIT). The operator gives a -/// single `ipc://` base; SMG binds the ipc input/output at -/// `-in.sock` / `-out.sock` and derives the TCP handshake port from the -/// path. A `WorkerSpec.zmq_handshake_address` override replaces the derived -/// handshake address verbatim (it must be `tcp://`), for engines that dial a -/// fixed, pre-agreed address — e.g. TokenSpeed's default dial target is -/// `tcp://127.0.0.1:30500` (its `--data-parallel-address`/ -/// `--data-parallel-rpc-port` defaults, outside the derived 20000..=29999 -/// band), so setting the override to that value pairs a bare -/// `ts serve --headless` with a manually registered worker. -/// Returns `(handshake, input, output)`. -fn zmq_socket_addresses( - base_url: &str, - handshake_override: Option<&str>, -) -> WorkerResult<(String, String, String)> { - let path = base_url - .strip_prefix("ipc://") - .ok_or_else(|| WorkerError::ConnectionFailed { - url: base_url.to_string(), - reason: "ZMQ worker URL must be ipc://".to_string(), - })?; - let handshake = match handshake_override { - Some(address) => { - if !address.starts_with("tcp://") { - return Err(WorkerError::ConnectionFailed { - url: base_url.to_string(), - reason: format!( - "zmq_handshake_address must be a tcp:// address \ - (the engine dials a TCP handshake), got '{address}'" - ), - }); - } - address.to_string() - } - None => format!("tcp://{ZMQ_LOOPBACK_HOST}:{}", derive_handshake_port(path)), - }; - let input = format!("ipc://{path}-in.sock"); - let output = format!("ipc://{path}-out.sock"); - Ok((handshake, input, output)) -} - -/// Create the parent directory for a worker's `ipc://` sockets. Kept off the -/// address computation (which is pure) and async so it doesn't block a runtime -/// thread. -/// -/// The ipc:// data-plane sockets SMG binds here carry no authentication, so the -/// directory must be owner-controlled: when this call creates it, it is created -/// 0700 (mode applied at mkdir time — no chmod window); when it already exists, -/// its permissions are left untouched (never chmod a shared dir like `/tmp`) -/// and it is rejected unless it is a real directory owned by the current user. -async fn ensure_ipc_socket_dir(base_url: &str) -> WorkerResult<()> { - let path = base_url.strip_prefix("ipc://").unwrap_or(base_url); - let Some(parent) = std::path::Path::new(path).parent() else { - return Ok(()); - }; - let fail = |reason: String| WorkerError::ConnectionFailed { - url: base_url.to_string(), - reason, - }; - // symlink_metadata: a symlinked parent must not redirect the checks (or the - // sockets) into a directory we did not verify. - let meta = match tokio::fs::symlink_metadata(parent).await { - Ok(meta) => meta, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => { - let mut builder = tokio::fs::DirBuilder::new(); - builder.recursive(true); - #[cfg(unix)] - builder.mode(0o700); - builder - .create(parent) - .await - .map_err(|e| fail(format!("failed to create ipc socket dir for {path}: {e}")))?; - tokio::fs::symlink_metadata(parent) - .await - .map_err(|e| fail(format!("failed to stat ipc socket dir for {path}: {e}")))? - } - Err(e) => { - return Err(fail(format!( - "failed to stat ipc socket dir for {path}: {e}" - ))) - } - }; - if !meta.is_dir() { - return Err(fail(format!( - "ipc socket dir {} exists but is not a directory", - parent.display() - ))); - } - #[cfg(unix)] - { - use std::os::unix::fs::MetadataExt; - let uid = rustix::process::geteuid().as_raw(); - if meta.uid() != uid { - return Err(fail(format!( - "ipc socket dir {} is owned by uid {} (expected {uid}); refusing to bind \ - unauthenticated ZMQ sockets in a directory owned by another user", - parent.display(), - meta.uid() - ))); - } - } - Ok(()) -} - -/// Bind the SMG-side ZMQ sockets and complete the handshake with the engine. -/// Shared by the lazy client accessor and the background handshake driver so -/// both go through the exact same connect path. `base_url` is the `ipc://` URL; -/// `model_id` is the config-resolved served model (EngineCore reports none). +/// Connect the ZMQ backend for this worker URL. All connect mechanics — +/// address derivation, socket-dir prep, EOS resolution, handshake — live in +/// the client layer ([`zmq_client::connect_for_worker`]); the worker layer +/// only wraps the client into its slot and error type. async fn connect_zmq_backend( base_url: String, model_id: String, runtime: RuntimeType, handshake_override: Option, ) -> WorkerResult> { - let (handshake, input, output) = - zmq_socket_addresses(&base_url, handshake_override.as_deref())?; - ensure_ipc_socket_dir(&base_url).await?; - // The engine can't stop at EOS on its own (it has no tokenizer or model - // config); resolve the EOS ids from the local model dir so every request - // carries them. - let model_dir = std::path::Path::new(&model_id); - let eos = if model_dir.is_dir() { - EosTokenIds::from_model_dir(model_dir) - } else { - tracing::warn!( - "ZMQ worker model id '{model_id}' is not a local model directory; EOS ids \ - unavailable — generation stops only at max_tokens or explicit stops" - ); - EosTokenIds::default() - }; - tracing::info!("Binding ZMQ client for worker {base_url} (handshake={handshake})"); - match ZmqEngineClient::connect( - &handshake, - &input, - &output, - 1, - model_id, - eos, - runtime, - ZMQ_CONNECT_TIMEOUT, - ) - .await - { - Ok(client) => Ok(Arc::new(BackendClient::Zmq(client))), - Err(e) => Err(WorkerError::ConnectionFailed { + zmq_client::connect_for_worker(&base_url, model_id, runtime, handshake_override.as_deref()) + .await + .map(|client| Arc::new(BackendClient::Zmq(client))) + .map_err(|reason| WorkerError::ConnectionFailed { url: base_url, - reason: format!("Failed to connect ZMQ engine: {e}"), - }), - } + reason, + }) } /// Default bootstrap port for PD disaggregation (used by SGLang and vLLM Mooncake) @@ -1791,56 +1624,6 @@ mod tests { assert_eq!(WorkerType::Decode.to_string(), "decode"); } - #[test] - fn derive_handshake_port_matches_pinned_vectors() { - // Fixed vectors shared with `_zmq_handshake_port` in - // bindings/python/src/smg/serve.py — a change on either side breaks the - // engine/router port agreement, so these must stay in sync. - assert_eq!(derive_handshake_port("/tmp/smg-zmq/ts0.ipc"), 25152); - assert_eq!(derive_handshake_port("/tmp/smg-zmq/engine-31000"), 22714); - // Range invariant: every path maps into 20000..=29999. - for p in ["", "a", "/x/y/z.ipc", "very/long/path/with/segments.sock"] { - let port = derive_handshake_port(p); - assert!( - (20000..=29999).contains(&port), - "port {port} out of band for {p:?}" - ); - } - } - - #[test] - fn zmq_socket_addresses_derive_handshake_by_default() { - let (handshake, input, output) = - zmq_socket_addresses("ipc:///tmp/smg-zmq/ts0.ipc", None).unwrap(); - assert_eq!(handshake, "tcp://127.0.0.1:25152"); - assert_eq!(input, "ipc:///tmp/smg-zmq/ts0.ipc-in.sock"); - assert_eq!(output, "ipc:///tmp/smg-zmq/ts0.ipc-out.sock"); - } - - #[test] - fn zmq_socket_addresses_honor_handshake_override() { - // TokenSpeed's default dial target — outside the derived band; the - // override must be bound verbatim while the data plane stays derived. - let (handshake, input, output) = - zmq_socket_addresses("ipc:///tmp/smg-zmq/ts0.ipc", Some("tcp://127.0.0.1:30500")) - .unwrap(); - assert_eq!(handshake, "tcp://127.0.0.1:30500"); - assert_eq!(input, "ipc:///tmp/smg-zmq/ts0.ipc-in.sock"); - assert_eq!(output, "ipc:///tmp/smg-zmq/ts0.ipc-out.sock"); - } - - #[test] - fn zmq_socket_addresses_reject_non_tcp_override() { - // The engine dials a TCP handshake; a non-tcp override is a config - // error and must fail loudly rather than bind something unexpected. - let err = zmq_socket_addresses("ipc:///tmp/smg-zmq/ts0.ipc", Some("ipc:///tmp/hs.sock")) - .unwrap_err(); - assert!( - err.to_string().contains("tcp://"), - "error must name the required scheme: {err}" - ); - } - #[test] fn test_worker_type_equality() { assert_eq!(WorkerType::Regular, WorkerType::Regular); @@ -2783,41 +2566,6 @@ mod tests { assert!(worker.has_models_discovered()); } - #[tokio::test] - async fn ensure_ipc_socket_dir_creates_a_private_owner_only_dir() { - let base = tempfile::tempdir().unwrap(); - let dir = base.path().join("sockets"); - let url = format!("ipc://{}/x.ipc", dir.display()); - ensure_ipc_socket_dir(&url).await.unwrap(); - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let mode = std::fs::metadata(&dir).unwrap().permissions().mode() & 0o777; - assert_eq!(mode, 0o700, "created socket dir must be 0700"); - } - } - - #[cfg(unix)] - #[tokio::test] - async fn ensure_ipc_socket_dir_leaves_an_existing_owned_dir_untouched() { - use std::os::unix::fs::PermissionsExt; - let dir = tempfile::tempdir().unwrap(); - std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap(); - let url = format!("ipc://{}/x.ipc", dir.path().display()); - ensure_ipc_socket_dir(&url).await.unwrap(); - let mode = std::fs::metadata(dir.path()).unwrap().permissions().mode() & 0o777; - assert_eq!(mode, 0o755, "an existing dir must not be chmod'd"); - } - - #[tokio::test] - async fn ensure_ipc_socket_dir_rejects_a_non_directory_parent() { - let base = tempfile::tempdir().unwrap(); - let file = base.path().join("not-a-dir"); - std::fs::write(&file, b"x").unwrap(); - let url = format!("ipc://{}/x.ipc", file.display()); - assert!(ensure_ipc_socket_dir(&url).await.is_err()); - } - /// A ZMQ client whose engine dies must be evicted by the health probe (the /// connection can't reconnect in place — liveness is latched), and the /// handshake guard reset so a later probe rebinds the sockets for a @@ -2829,6 +2577,8 @@ mod tests { EngineId, ENGINE_CORE_DEAD_SENTINEL, }; + use crate::routers::grpc::zmq_client::{EosTokenIds, ZmqEngineClient}; + let base = tempfile::tempdir().unwrap(); let ep = |name: &str| format!("ipc://{}", base.path().join(name).display()); let (handshake, input, output) = (ep("hs.sock"), ep("in.sock"), ep("out.sock"));