Skip to content

feat(zmq): core per-engine dispatch for direct ZMQ backends (1/7) #2054

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slin1237 merged 1 commit into from
Aug 5, 2026
Merged

feat(zmq): core per-engine dispatch for direct ZMQ backends (1/7) #2054

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 0 additions & 25 deletions crates/grpc_client/src/tokenspeed_scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,12 +178,7 @@ impl TokenSpeedSchedulerClient {

// ── Request builders ──────────────────────────────────────────────

#[expect(
clippy::unused_self,
reason = "receiver kept for API parity with the other engine clients"
)]
pub fn build_generate_request_from_chat(
&self,
request_id: String,
body: &ChatCompletionRequest,
processed_text: String,
Expand All @@ -208,12 +203,7 @@ impl TokenSpeedSchedulerClient {
})
}

#[expect(
clippy::unused_self,
reason = "receiver kept for API parity with the other engine clients"
)]
pub fn build_plain_generate_request(
&self,
request_id: String,
body: &GenerateRequest,
original_text: Option<String>,
Expand Down Expand Up @@ -242,12 +232,7 @@ impl TokenSpeedSchedulerClient {
})
}

#[expect(
clippy::unused_self,
reason = "receiver kept for API parity with the other engine clients"
)]
pub fn build_generate_request_from_responses(
&self,
request_id: String,
body: &ResponsesRequest,
processed_text: String,
Expand All @@ -267,12 +252,7 @@ impl TokenSpeedSchedulerClient {
})
}

#[expect(
clippy::unused_self,
reason = "receiver kept for API parity with the other engine clients"
)]
pub fn build_generate_request_from_messages(
&self,
request_id: String,
body: &CreateMessageRequest,
processed_text: String,
Expand All @@ -295,12 +275,7 @@ impl TokenSpeedSchedulerClient {
})
}

#[expect(
clippy::unused_self,
reason = "receiver kept for API parity with the other engine clients"
)]
pub fn build_generate_request_from_completion(
&self,
request_id: String,
body: &CompletionRequest,
original_text: String,
Expand Down
166 changes: 108 additions & 58 deletions model_gateway/src/routers/grpc/backend_client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,16 +13,22 @@ use openai_protocol::{
messages::CreateMessageRequest, worker::WorkerLoadResponse,
};
use smg_grpc_client::{
common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, VllmEngineClient,
common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, TokenSpeedSchedulerClient,
VllmEngineClient,
};

use crate::routers::grpc::{
client::{GenerateRequestBuildOptions, GrpcClient, HealthCheckResponse, ModelInfo, ServerInfo},
proto_wrapper::{
finish_vllm_request, ProtoEmbedComplete, ProtoEmbedRequest, ProtoGenerateRequest,
ProtoStream,
use crate::{
routers::grpc::{
client::{
GenerateRequestBuildOptions, GrpcClient, HealthCheckResponse, ModelInfo, ServerInfo,
},
proto_wrapper::{
finish_tokenspeed_request, finish_vllm_request, ProtoEmbedComplete, ProtoEmbedRequest,
ProtoGenerateRequest, ProtoStream,
},
zmq_client::ZmqEngineClient,
},
zmq_client::ZmqEngineClient,
worker::RuntimeType,
};

/// A backend connection: gRPC (any engine) or direct ZMQ (vLLM EngineCore or
Expand All @@ -35,7 +41,7 @@ pub enum BackendClient {

impl BackendClient {
/// Runtime type backing this client.
pub fn runtime_type(&self) -> crate::worker::RuntimeType {
pub fn runtime_type(&self) -> RuntimeType {
match self {
Self::Grpc(client) => client.runtime_type(),
Self::Zmq(client) => client.runtime(),
Expand Down Expand Up @@ -88,14 +94,14 @@ impl BackendClient {
pub async fn get_model_info(&self) -> Result<ModelInfo, tonic::Status> {
match self {
Self::Grpc(client) => client.get_model_info().await,
Self::Zmq(client) => Ok(ModelInfo::Vllm(client.get_model_info())),
Self::Zmq(client) => Ok(client.get_model_info()),
}
}

pub async fn get_server_info(&self) -> Result<ServerInfo, tonic::Status> {
match self {
Self::Grpc(client) => client.get_server_info().await,
Self::Zmq(client) => Ok(ServerInfo::Vllm(client.get_server_info())),
Self::Zmq(client) => Ok(client.get_server_info()),
}
}

Expand Down Expand Up @@ -168,14 +174,7 @@ impl BackendClient {
) -> Result<ProtoStream, tonic::Status> {
match self {
Self::Grpc(client) => client.generate(req).await,
Self::Zmq(client) => match req {
ProtoGenerateRequest::Vllm(boxed_req) => {
Ok(ProtoStream::Zmq(client.generate(*boxed_req).await?))
}
_ => Err(tonic::Status::internal(
"ZMQ backend expects a vLLM generate request",
)),
},
Self::Zmq(client) => Ok(ProtoStream::Zmq(client.generate(req).await?)),
}
}

Expand Down Expand Up @@ -203,18 +202,33 @@ impl BackendClient {
Self::Grpc(client) => {
client.build_chat_request(request_id, body, processed_text, token_ids, options)
}
Self::Zmq(_) => {
// A ZMQ backend speaks vLLM EngineCore or TokenSpeed directly; build
// the native request for its runtime, mirroring the gRPC per-engine
// dispatch in `GrpcClient::build_chat_request`.
Self::Zmq(client) => {
reject_zmq_multimodal(&options)?;
finish_vllm_request(None, |mm| {
VllmEngineClient::build_generate_request_from_chat(
request_id,
body,
processed_text,
token_ids,
mm,
options.tool_constraints,
)
})
match client.runtime() {
RuntimeType::TokenSpeed => finish_tokenspeed_request(None, |mm| {
TokenSpeedSchedulerClient::build_generate_request_from_chat(
request_id,
body,
processed_text,
token_ids,
mm,
options.tool_constraints,
)
}),
_ => finish_vllm_request(None, |mm| {
VllmEngineClient::build_generate_request_from_chat(
request_id,
body,
processed_text,
token_ids,
mm,
options.tool_constraints,
)
}),
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}
}
}
Expand All @@ -231,18 +245,32 @@ impl BackendClient {
Self::Grpc(client) => {
client.build_messages_request(request_id, body, processed_text, token_ids, options)
}
Self::Zmq(_) => {
// Mirrors the gRPC per-engine dispatch: build the request natively for
// the ZMQ backend's runtime (vLLM EngineCore or TokenSpeed).
Self::Zmq(client) => {
reject_zmq_multimodal(&options)?;
finish_vllm_request(None, |mm| {
VllmEngineClient::build_generate_request_from_messages(
request_id,
body,
processed_text,
token_ids,
mm,
options.tool_constraints,
)
})
match client.runtime() {
RuntimeType::TokenSpeed => finish_tokenspeed_request(None, |mm| {
TokenSpeedSchedulerClient::build_generate_request_from_messages(
request_id,
body,
processed_text,
token_ids,
mm,
options.tool_constraints,
)
}),
_ => finish_vllm_request(None, |mm| {
VllmEngineClient::build_generate_request_from_messages(
request_id,
body,
processed_text,
token_ids,
mm,
options.tool_constraints,
)
}),
}
}
}
}
Expand All @@ -258,15 +286,26 @@ impl BackendClient {
Self::Grpc(client) => {
client.build_completion_request(request_id, body, original_text, token_ids)
}
Self::Zmq(_) => {
let req = VllmEngineClient::build_generate_request_from_completion(
request_id,
body,
original_text,
token_ids,
)?;
Ok(ProtoGenerateRequest::Vllm(Box::new(req)))
}
Self::Zmq(client) => match client.runtime() {
RuntimeType::TokenSpeed => {
let req = TokenSpeedSchedulerClient::build_generate_request_from_completion(
request_id,
body,
original_text,
token_ids,
)?;
Ok(ProtoGenerateRequest::TokenSpeed(Box::new(req)))
}
_ => {
let req = VllmEngineClient::build_generate_request_from_completion(
request_id,
body,
original_text,
token_ids,
)?;
Ok(ProtoGenerateRequest::Vllm(Box::new(req)))
}
},
}
}

Expand All @@ -281,15 +320,26 @@ impl BackendClient {
Self::Grpc(client) => {
client.build_generate_request(request_id, body, original_text, token_ids)
}
Self::Zmq(_) => {
let req = VllmEngineClient::build_plain_generate_request(
request_id,
body,
original_text,
token_ids,
)?;
Ok(ProtoGenerateRequest::Vllm(Box::new(req)))
}
Self::Zmq(client) => match client.runtime() {
RuntimeType::TokenSpeed => {
let req = TokenSpeedSchedulerClient::build_plain_generate_request(
request_id,
body,
original_text,
token_ids,
)?;
Ok(ProtoGenerateRequest::TokenSpeed(Box::new(req)))
}
_ => {
let req = VllmEngineClient::build_plain_generate_request(
request_id,
body,
original_text,
token_ids,
)?;
Ok(ProtoGenerateRequest::Vllm(Box::new(req)))
}
},
}
}
}
Expand Down
16 changes: 8 additions & 8 deletions model_gateway/src/routers/grpc/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -541,13 +541,13 @@ impl GrpcClient {
)?;
Ok(ProtoGenerateRequest::Mlx(Box::new(req)))
}
Self::TokenSpeed(client) => {
Self::TokenSpeed(_) => {
let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm {
MultimodalData::TokenSpeed(data) => data.into_proto(true),
_ => unreachable!("caller guarantees matching variant"),
});
finish_tokenspeed_request(tokenspeed_mm, |mm| {
client.build_generate_request_from_chat(
TokenSpeedSchedulerClient::build_generate_request_from_chat(
request_id,
body,
processed_text,
Expand Down Expand Up @@ -633,13 +633,13 @@ impl GrpcClient {
)?;
Ok(ProtoGenerateRequest::Mlx(Box::new(req)))
}
Self::TokenSpeed(client) => {
Self::TokenSpeed(_) => {
let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm {
MultimodalData::TokenSpeed(data) => data.into_proto(true),
_ => unreachable!("caller guarantees matching variant"),
});
finish_tokenspeed_request(tokenspeed_mm, |mm| {
client.build_generate_request_from_messages(
TokenSpeedSchedulerClient::build_generate_request_from_messages(
request_id,
body,
processed_text,
Expand Down Expand Up @@ -696,8 +696,8 @@ impl GrpcClient {
)?;
Ok(ProtoGenerateRequest::Mlx(Box::new(req)))
}
Self::TokenSpeed(client) => {
let req = client.build_generate_request_from_completion(
Self::TokenSpeed(_) => {
let req = TokenSpeedSchedulerClient::build_generate_request_from_completion(
request_id,
body,
original_text,
Expand Down Expand Up @@ -752,8 +752,8 @@ impl GrpcClient {
)?;
Ok(ProtoGenerateRequest::Mlx(Box::new(req)))
}
Self::TokenSpeed(client) => {
let req = client.build_plain_generate_request(
Self::TokenSpeed(_) => {
let req = TokenSpeedSchedulerClient::build_plain_generate_request(
request_id,
body,
original_text,
Expand Down
Loading
Loading