From 2f283d81525a19b1fd706f38c66a6f25924b0eb3 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Thu, 12 Feb 2026 06:43:24 -0800 Subject: [PATCH] refactor: consolidate gRPC client into shared crate dependency (#18730) --- sgl-model-gateway/Cargo.toml | 2 +- sgl-model-gateway/build.rs | 16 - sgl-model-gateway/src/grpc_client/mod.rs | 7 - .../src/grpc_client/sglang_scheduler.rs | 827 ------------------ .../src/grpc_client/vllm_engine.rs | 747 ---------------- sgl-model-gateway/src/lib.rs | 2 +- .../src/observability/otel_trace.rs | 17 + .../src/proto/sglang_scheduler.proto | 566 ------------ sgl-model-gateway/src/proto/vllm_engine.proto | 195 ----- sgl-model-gateway/src/routers/grpc/client.rs | 12 +- 10 files changed, 29 insertions(+), 2362 deletions(-) delete mode 100644 sgl-model-gateway/src/grpc_client/mod.rs delete mode 100644 sgl-model-gateway/src/grpc_client/sglang_scheduler.rs delete mode 100644 sgl-model-gateway/src/grpc_client/vllm_engine.rs delete mode 100644 sgl-model-gateway/src/proto/sglang_scheduler.proto delete mode 100644 sgl-model-gateway/src/proto/vllm_engine.proto diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index cc35ad5f8..4578a5bb1 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -107,6 +107,7 @@ openmetrics-parser = "0.4.4" arc-swap = "1.7.1" # gRPC and Protobuf dependencies +smg-grpc-client = "1.0.0" tonic = { version = "0.14.2", features = ["gzip", "transport"] } prost = "0.14.1" prost-types = "0.14.1" @@ -124,7 +125,6 @@ sha2 = "0.10" wasmtime = { version = "38.0", features = ["component-model", "async"] } [build-dependencies] -tonic-prost-build = "0.14.2" chrono = { version = "0.4", features = ["clock"] } toml = "0.9" diff --git a/sgl-model-gateway/build.rs b/sgl-model-gateway/build.rs index 57a02656b..90bfb2463 100644 --- a/sgl-model-gateway/build.rs +++ b/sgl-model-gateway/build.rs @@ -12,24 +12,8 @@ macro_rules! set_env { fn main() -> Result<(), Box> { // Rebuild triggers - println!("cargo:rerun-if-changed=src/proto/sglang_scheduler.proto"); - println!("cargo:rerun-if-changed=src/proto/vllm_engine.proto"); println!("cargo:rerun-if-changed=Cargo.toml"); - // Compile protobuf files - tonic_prost_build::configure() - .build_server(true) - .build_client(true) - .type_attribute("GetModelInfoResponse", "#[derive(serde::Serialize)]") - .protoc_arg("--experimental_allow_proto3_optional") - .compile_protos( - &[ - "src/proto/sglang_scheduler.proto", - "src/proto/vllm_engine.proto", - ], - &["src/proto"], - )?; - // Set version info environment variables let version = read_cargo_version().unwrap_or_else(|_| DEFAULT_VERSION.to_string()); let target = std::env::var("TARGET").unwrap_or_else(|_| get_rustc_host().unwrap_or_default()); diff --git a/sgl-model-gateway/src/grpc_client/mod.rs b/sgl-model-gateway/src/grpc_client/mod.rs deleted file mode 100644 index 8edf7ce66..000000000 --- a/sgl-model-gateway/src/grpc_client/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod sglang_scheduler; -pub mod vllm_engine; - -// Export both clients -// Re-export proto modules with explicit names -pub use sglang_scheduler::{proto as sglang_proto, SglangSchedulerClient}; -pub use vllm_engine::{proto as vllm_proto, VllmEngineClient}; diff --git a/sgl-model-gateway/src/grpc_client/sglang_scheduler.rs b/sgl-model-gateway/src/grpc_client/sglang_scheduler.rs deleted file mode 100644 index d14d1abb4..000000000 --- a/sgl-model-gateway/src/grpc_client/sglang_scheduler.rs +++ /dev/null @@ -1,827 +0,0 @@ -use std::{ - convert::TryFrom, - pin::Pin, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, - task::{Context, Poll}, - time::Duration, -}; - -use tonic::{transport::Channel, Request, Streaming}; -use tracing::{debug, warn}; - -use crate::{ - observability::otel_trace::inject_trace_context_grpc, - protocols::{ - chat::ChatCompletionRequest, - common::{ResponseFormat, StringOrArray, ToolChoice, ToolChoiceValue}, - generate::GenerateRequest, - responses::ResponsesRequest, - sampling_params::SamplingParams as GenerateSamplingParams, - }, -}; - -// Include the generated protobuf code -#[allow(clippy::all)] -pub mod proto { - #![allow(clippy::all, unused_qualifications)] - tonic::include_proto!("sglang.grpc.scheduler"); -} - -// The generated module structure depends on the package name in the .proto file -// package sglang.grpc.scheduler; generates a nested module structure - -/// A smart wrapper around Streaming that automatically -/// sends abort when dropped (e.g., due to client disconnection or early termination). -/// -/// This leverages Rust's RAII pattern to ensure cleanup happens automatically, -/// regardless of how the stream is dropped (panic, early return, client disconnect, etc.). -pub struct AbortOnDropStream { - inner: Streaming, - request_id: String, - client: SglangSchedulerClient, - aborted: Arc, -} - -impl AbortOnDropStream { - /// Create a new auto-aborting stream wrapper - pub fn new( - stream: Streaming, - request_id: String, - client: SglangSchedulerClient, - ) -> Self { - debug!("Created AbortOnDropStream for request {}", request_id); - Self { - inner: stream, - request_id, - client, - aborted: Arc::new(AtomicBool::new(false)), - } - } - - /// Manually mark the request as completed to prevent abort on drop. - /// Call this when the request completes successfully to avoid unnecessary abort RPC. - pub fn mark_completed(&self) { - // Use Release ordering to ensure that this write is visible to other threads - // that use Acquire on the same atomic variable - self.aborted.store(true, Ordering::Release); - debug!("Request {} marked as completed", self.request_id); - } -} - -impl Drop for AbortOnDropStream { - fn drop(&mut self) { - // Atomically check and set the aborted flag using compare_exchange. - // If compare_exchange fails, it means the flag was already true (from mark_completed), - // so we don't need to send abort. AcqRel is used for success to synchronize with - // mark_completed's Release, and Acquire for failure to see writes from mark_completed. - if self - .aborted - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_err() - { - return; - } - - let client = self.client.clone(); - let request_id = self.request_id.clone(); - - // Spawn a background task to send abort (since Drop is sync but abort_request is async) - tokio::spawn(async move { - debug!( - "Stream dropped without completion for request {}, sending abort", - request_id - ); - // Clone request_id for the error message since abort_request takes ownership - let request_id_for_log = request_id.clone(); - if let Err(e) = client - .abort_request(request_id, "Stream dropped".to_string()) - .await - { - warn!( - "Failed to send abort on drop for request {}: {}", - request_id_for_log, e - ); - } - }); - } -} - -// Implement Stream trait to make AbortOnDropStream work like the original Streaming -impl futures::Stream for AbortOnDropStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // Delegate to the inner stream - Pin::new(&mut self.inner).poll_next(cx) - } -} - -/// gRPC client for SGLang scheduler -#[derive(Clone)] -pub struct SglangSchedulerClient { - client: proto::sglang_scheduler_client::SglangSchedulerClient, -} - -impl SglangSchedulerClient { - /// Create a new client and connect to the scheduler - pub async fn connect(endpoint: &str) -> Result> { - debug!("Connecting to SGLang scheduler at {}", endpoint); - - // Convert grpc:// to http:// for tonic - let http_endpoint = if let Some(addr) = endpoint.strip_prefix("grpc://") { - format!("http://{}", addr) - } else { - endpoint.to_string() - }; - - let channel = Channel::from_shared(http_endpoint)? - .http2_keep_alive_interval(Duration::from_secs(30)) - .keep_alive_timeout(Duration::from_secs(10)) - .keep_alive_while_idle(true) - .tcp_keepalive(Some(Duration::from_secs(60))) - .tcp_nodelay(true) - .http2_adaptive_window(true) - .initial_stream_window_size(Some(16 * 1024 * 1024)) // 16MB - .initial_connection_window_size(Some(32 * 1024 * 1024)) // 32MB - .connect() - .await?; - - let client = proto::sglang_scheduler_client::SglangSchedulerClient::new(channel); - - Ok(Self { client }) - } - - /// Submit a generation request (returns auto-aborting streaming response) - /// - /// The returned stream automatically sends an abort request when dropped, - /// ensuring proper cleanup even if the HTTP client disconnects or an error occurs. - /// Call `mark_completed()` on the stream after successful completion to prevent - /// unnecessary abort RPCs. - pub async fn generate( - &self, - req: proto::GenerateRequest, - ) -> Result> { - let request_id = req.request_id.clone(); - let mut client = self.client.clone(); - let mut request = Request::new(req); - - // Inject W3C trace context into gRPC metadata for distributed tracing - inject_trace_context_grpc(request.metadata_mut()); - - let response = client.generate(request).await?; - - Ok(AbortOnDropStream::new( - response.into_inner(), - request_id, - self.clone(), - )) - } - - /// Submit an embedding request - pub async fn embed( - &self, - req: proto::EmbedRequest, - ) -> Result> { - let mut client = self.client.clone(); - let mut request = Request::new(req); - - // Inject W3C trace context into gRPC metadata - inject_trace_context_grpc(request.metadata_mut()); - - let response = client.embed(request).await?; - Ok(response.into_inner()) - } - - /// Perform health check - pub async fn health_check( - &self, - ) -> Result> { - debug!("Sending health check request"); - // HealthCheckRequest is now empty - server generates its own health check internally - let request = Request::new(proto::HealthCheckRequest {}); - - let mut client = self.client.clone(); - let response = client.health_check(request).await?; - debug!("Health check response received"); - Ok(response.into_inner()) - } - - /// Abort a request - pub async fn abort_request( - &self, - request_id: String, - reason: String, - ) -> Result<(), Box> { - debug!( - "Sending abort request for {} (reason: {})", - request_id, reason - ); - let request = Request::new(proto::AbortRequest { - request_id: request_id.clone(), - reason, - }); - - let mut client = self.client.clone(); - let response = client.abort(request).await?; - debug!( - "Abort response for {}: success={}, message={}", - request_id, - response.get_ref().success, - response.get_ref().message - ); - Ok(()) - } - - /// Get model information - pub async fn get_model_info( - &self, - ) -> Result> { - debug!("Requesting model info"); - let request = Request::new(proto::GetModelInfoRequest {}); - - let mut client = self.client.clone(); - let response = client.get_model_info(request).await?; - debug!("Model info response received"); - Ok(response.into_inner()) - } - - /// Get server information - pub async fn get_server_info( - &self, - ) -> Result> { - debug!("Requesting server info"); - let request = Request::new(proto::GetServerInfoRequest {}); - - let mut client = self.client.clone(); - let response = client.get_server_info(request).await?; - debug!("Server info response received"); - Ok(response.into_inner()) - } - - /// Build a single SGLang EmbedRequest - pub fn build_embed_request( - &self, - request_id: String, - original_text: Option, - token_ids: Vec, - log_metrics: Option, - ) -> proto::EmbedRequest { - proto::EmbedRequest { - request_id, - tokenized: Some(proto::TokenizedInput { - original_text: original_text.unwrap_or_default(), - input_ids: token_ids, - }), - log_metrics: log_metrics.unwrap_or(false), // Default to false if not specified - ..Default::default() - } - } - - /// Build a single SGLang GenerateRequest from OpenAI ChatCompletionRequest - pub fn build_generate_request_from_chat( - &self, - request_id: String, - body: &ChatCompletionRequest, - processed_text: String, - token_ids: Vec, - multimodal_inputs: Option, - tool_call_constraint: Option<(String, String)>, // (constraint_type, constraint_value) - ) -> Result { - // Build sampling params - let sampling_params = - self.build_grpc_sampling_params_from_chat(body, tool_call_constraint)?; - - let grpc_request = proto::GenerateRequest { - request_id, - tokenized: Some(proto::TokenizedInput { - original_text: processed_text, - input_ids: token_ids, - }), - mm_inputs: multimodal_inputs, - sampling_params: Some(sampling_params), - return_logprob: body.logprobs, - logprob_start_len: -1, - top_logprobs_num: body.top_logprobs.unwrap_or(0) as i32, - return_hidden_states: body.return_hidden_states, - stream: body.stream, - ..Default::default() - }; - - Ok(grpc_request) - } - - /// Build a basic GenerateRequest from the SGLang spec GenerateRequest - pub fn build_plain_generate_request( - &self, - request_id: String, - body: &GenerateRequest, - original_text: Option, - token_ids: Vec, - ) -> Result { - let sampling_params = - Self::build_sampling_params_from_plain(body.sampling_params.as_ref())?; - - let grpc_request = proto::GenerateRequest { - request_id, - tokenized: Some(proto::TokenizedInput { - original_text: original_text.unwrap_or_default(), - input_ids: token_ids, - }), - sampling_params: Some(sampling_params), - return_logprob: body.return_logprob.unwrap_or(false), - logprob_start_len: body.logprob_start_len.unwrap_or(-1), - top_logprobs_num: body.top_logprobs_num.unwrap_or(0), - token_ids_logprob: body.token_ids_logprob.clone().unwrap_or_default(), - return_hidden_states: body.return_hidden_states, - stream: body.stream, - log_metrics: body.log_metrics, - ..Default::default() - }; - - Ok(grpc_request) - } - - /// Build a GenerateRequest from ResponsesRequest (OpenAI Responses API) - /// - /// NOTE: This is used by the Harmony router only. The Regular router uses - /// responses_to_chat() conversion and goes through the chat pipeline. - pub fn build_generate_request_from_responses( - &self, - request_id: String, - body: &ResponsesRequest, - processed_text: String, - token_ids: Vec, - harmony_stop_ids: Option>, - constraint: Option<(String, String)>, - ) -> Result { - // Build sampling params from ResponsesRequest - let mut sampling_params = - self.build_grpc_sampling_params_from_responses(body, constraint)?; - - // Inject Harmony stop token IDs if provided - if let Some(stop_ids) = harmony_stop_ids { - sampling_params.stop_token_ids = stop_ids; - } - - let grpc_request = proto::GenerateRequest { - request_id, - tokenized: Some(proto::TokenizedInput { - original_text: processed_text, - input_ids: token_ids, - }), - mm_inputs: None, // Responses API doesn't support multimodal yet - sampling_params: Some(sampling_params), - return_logprob: false, // Responses API uses top_logprobs field instead - logprob_start_len: -1, - top_logprobs_num: body.top_logprobs.unwrap_or(0) as i32, - return_hidden_states: false, - stream: body.stream.unwrap_or(false), - ..Default::default() - }; - - Ok(grpc_request) - } - - /// Build gRPC SamplingParams from ChatCompletionRequest - fn build_grpc_sampling_params_from_chat( - &self, - request: &ChatCompletionRequest, - tool_call_constraint: Option<(String, String)>, - ) -> Result { - let stop_sequences = self.extract_stop_strings(request); - - let max_new_tokens = request.max_completion_tokens.map(|v| v as i32); - - // Handle skip_special_tokens: set to false if tools are present and tool_choice is not "none" - let skip_special_tokens = if request.tools.is_some() { - match &request.tool_choice { - Some(ToolChoice::Value(ToolChoiceValue::None)) => request.skip_special_tokens, - Some(_) => false, // tool_choice is not "none" - None => false, // TODO: this assumes tool_choice defaults to "auto" when tools present - } - } else { - request.skip_special_tokens - }; - - Ok(proto::SamplingParams { - temperature: request.temperature.unwrap_or(1.0), - top_p: request.top_p.unwrap_or(1.0), - top_k: request.top_k.unwrap_or(-1), - min_p: request.min_p.unwrap_or(0.0), - frequency_penalty: request.frequency_penalty.unwrap_or(0.0), - presence_penalty: request.presence_penalty.unwrap_or(0.0), - repetition_penalty: request.repetition_penalty.unwrap_or(1.0), - max_new_tokens, - stop: stop_sequences, - stop_token_ids: request.stop_token_ids.clone().unwrap_or_default(), - skip_special_tokens, - spaces_between_special_tokens: true, // Default from Python SamplingParams - ignore_eos: request.ignore_eos, - no_stop_trim: request.no_stop_trim, - n: request.n.unwrap_or(1) as i32, - constraint: self.build_constraint_for_chat(request, tool_call_constraint)?, - ..Default::default() - }) - } - - /// Extract stop strings from request - fn extract_stop_strings(&self, request: &ChatCompletionRequest) -> Vec { - match &request.stop { - Some(StringOrArray::String(s)) => vec![s.clone()], - Some(StringOrArray::Array(arr)) => arr.clone(), - None => vec![], - } - } - - /// Build constraint for structured generation - fn build_constraint_for_chat( - &self, - request: &ChatCompletionRequest, - tool_call_constraint: Option<(String, String)>, - ) -> Result, String> { - let mut constraints = Vec::new(); - - // Handle response_format constraints - match &request.response_format { - Some(ResponseFormat::JsonObject) => { - // json_object mode - constrain to valid JSON object - let schema = serde_json::json!({"type": "object"}); - let schema_str = serde_json::to_string(&schema) - .map_err(|e| format!("Failed to serialize JSON schema: {}", e))?; - constraints.push(proto::sampling_params::Constraint::JsonSchema(schema_str)); - } - Some(ResponseFormat::JsonSchema { json_schema }) => { - let schema_str = serde_json::to_string(&json_schema.schema) - .map_err(|e| format!("Failed to serialize JSON schema: {}", e))?; - constraints.push(proto::sampling_params::Constraint::JsonSchema(schema_str)); - } - Some(ResponseFormat::Text) | None => { - // No constraint for text format - } - } - - if let Some(ebnf) = &request.ebnf { - constraints.push(proto::sampling_params::Constraint::EbnfGrammar( - ebnf.clone(), - )); - } - - if let Some(regex) = &request.regex { - constraints.push(proto::sampling_params::Constraint::Regex(regex.clone())); - } - - // Handle tool call constraint from preparation stage - if let Some((constraint_type, constraint_value)) = tool_call_constraint { - if !constraints.is_empty() { - return Err("Constrained decoding is not compatible with tool calls.".to_string()); - } - let tool_constraint = match constraint_type.as_str() { - "structural_tag" => { - proto::sampling_params::Constraint::StructuralTag(constraint_value) - } - "json_schema" => proto::sampling_params::Constraint::JsonSchema(constraint_value), - "ebnf" => proto::sampling_params::Constraint::EbnfGrammar(constraint_value), - "regex" => proto::sampling_params::Constraint::Regex(constraint_value), - _ => return Err(format!("Unknown constraint type: {}", constraint_type)), - }; - constraints.push(tool_constraint); - } - - match constraints.len() { - 0 => Ok(None), - 1 => Ok(constraints.pop()), - _ => Err("Multiple constraints are not allowed.".to_string()), - } - } - - /// Build gRPC SamplingParams from ResponsesRequest - fn build_grpc_sampling_params_from_responses( - &self, - request: &ResponsesRequest, - constraint: Option<(String, String)>, - ) -> Result { - // Used by Harmony models only. Regular models use Chat API path. - // Constraints come from Harmony preparation stage (structural_tag) or tool handling. - - let max_new_tokens = request.max_output_tokens.map(|v| v as i32); - - Ok(proto::SamplingParams { - temperature: request.temperature.unwrap_or(1.0), - top_p: request.top_p.unwrap_or(1.0), - top_k: -1, // ResponsesRequest doesn't expose top_k - min_p: 0.0, // ResponsesRequest doesn't expose min_p - frequency_penalty: 0.0, // ResponsesRequest doesn't expose frequency_penalty - presence_penalty: 0.0, // ResponsesRequest doesn't expose presence_penalty - repetition_penalty: 1.0, // ResponsesRequest doesn't expose repetition_penalty - max_new_tokens, - stop: vec![], // No stop sequences in Responses API - stop_token_ids: vec![], // Handled by Harmony stop tokens - skip_special_tokens: false, // Keep special tokens for Harmony - spaces_between_special_tokens: true, - ignore_eos: false, - no_stop_trim: false, - n: 1, // Responses API doesn't support n>1 - constraint: self.build_constraint_for_responses(constraint)?, - ..Default::default() - }) - } - - /// Build constraint for Responses API - /// - /// Handles constraints from Harmony preparation stage (structural_tag for Harmony models, - /// structured output via text field, or tool call constraints). - /// - /// Note: Regular gRPC models use Chat API path with response_format, not this function. - fn build_constraint_for_responses( - &self, - constraint: Option<(String, String)>, - ) -> Result, String> { - if let Some((constraint_type, constraint_value)) = constraint { - let parsed_constraint = match constraint_type.as_str() { - "structural_tag" => { - // Harmony models: structural tag from preparation stage - proto::sampling_params::Constraint::StructuralTag(constraint_value) - } - "json_schema" => proto::sampling_params::Constraint::JsonSchema(constraint_value), - "ebnf" => proto::sampling_params::Constraint::EbnfGrammar(constraint_value), - "regex" => proto::sampling_params::Constraint::Regex(constraint_value), - _ => return Err(format!("Unknown constraint type: {}", constraint_type)), - }; - Ok(Some(parsed_constraint)) - } else { - Ok(None) - } - } - - fn build_single_constraint_from_plain( - params: &GenerateSamplingParams, - ) -> Result, String> { - let mut constraints = Vec::new(); - if let Some(json_schema) = ¶ms.json_schema { - constraints.push(proto::sampling_params::Constraint::JsonSchema( - json_schema.clone(), - )); - } - if let Some(regex) = ¶ms.regex { - constraints.push(proto::sampling_params::Constraint::Regex(regex.clone())); - } - if let Some(ebnf) = ¶ms.ebnf { - constraints.push(proto::sampling_params::Constraint::EbnfGrammar( - ebnf.clone(), - )); - } - - match constraints.len() { - 0 => Ok(None), - 1 => Ok(constraints.pop()), - _ => Err("Multiple structured constraints are not allowed".to_string()), - } - } - - fn build_sampling_params_from_plain( - params: Option<&GenerateSamplingParams>, - ) -> Result { - let mut sampling = proto::SamplingParams { - temperature: 1.0, - top_p: 1.0, - top_k: -1, - repetition_penalty: 1.0, - n: 1, - skip_special_tokens: true, - spaces_between_special_tokens: true, - ..Default::default() - }; - - let Some(p) = params else { - return Ok(sampling); - }; - - // Simple field mappings using a macro - macro_rules! map_field { - ($field:ident) => { - if let Some(val) = p.$field { - sampling.$field = val; - } - }; - } - - map_field!(temperature); - map_field!(top_p); - map_field!(top_k); - map_field!(frequency_penalty); - map_field!(presence_penalty); - map_field!(repetition_penalty); - map_field!(min_p); - map_field!(ignore_eos); - map_field!(skip_special_tokens); - map_field!(no_stop_trim); - - // Handle stop sequences - if let Some(stop) = &p.stop { - match stop { - StringOrArray::String(s) => sampling.stop.push(s.clone()), - StringOrArray::Array(arr) => sampling.stop.extend(arr.clone()), - } - } - - // Handle stop token IDs - if let Some(stop_token_ids) = &p.stop_token_ids { - sampling.stop_token_ids = stop_token_ids.clone(); - } - - // Handle max_new_tokens with conversion - if let Some(max_new_tokens) = p.max_new_tokens { - sampling.max_new_tokens = - Some(i32::try_from(max_new_tokens).map_err(|_| { - "max_new_tokens must fit into a 32-bit signed integer".to_string() - })?); - } - - // Handle min_new_tokens with conversion - if let Some(min_new_tokens) = p.min_new_tokens { - sampling.min_new_tokens = i32::try_from(min_new_tokens) - .map_err(|_| "min_new_tokens must fit into a 32-bit signed integer".to_string())?; - } - - // Handle n with conversion - if let Some(n) = p.n { - sampling.n = i32::try_from(n) - .map_err(|_| "n must fit into a 32-bit signed integer".to_string())?; - } - - // Handle constraints (exactly one allowed) - sampling.constraint = Self::build_single_constraint_from_plain(p)?; - - Ok(sampling) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_proto_types_compilation() { - let _health_req = proto::HealthCheckRequest {}; - // HealthCheckRequest is now empty - no fields to test - } - - #[test] - fn test_generate_request_construction() { - let sampling_params = proto::SamplingParams { - temperature: 0.7, - max_new_tokens: Some(128), - top_p: 0.9, - top_k: 50, - stop: vec!["".to_string()], - ..Default::default() - }; - - let gen_req = proto::GenerateRequest { - request_id: "test-req-123".to_string(), - tokenized: Some(proto::TokenizedInput { - original_text: "Hello world".to_string(), - input_ids: vec![9906, 1917], // Mock token IDs for "Hello world" - }), - sampling_params: Some(sampling_params), - return_logprob: true, - logprob_start_len: 0, - top_logprobs_num: 5, - ..Default::default() - }; - - assert_eq!(gen_req.request_id, "test-req-123"); - if let Some(ref tokenized) = &gen_req.tokenized { - assert_eq!(tokenized.original_text, "Hello world"); - } - assert!(gen_req.return_logprob); - assert_eq!(gen_req.top_logprobs_num, 5); - - let params = gen_req.sampling_params.unwrap(); - assert_eq!(params.temperature, 0.7); - assert_eq!(params.max_new_tokens, Some(128)); - assert_eq!(params.stop, vec![""]); - } - - #[test] - fn test_health_check_request() { - let _health_req = proto::HealthCheckRequest {}; - // HealthCheckRequest is now empty - server generates its own test internally - } - - #[test] - fn test_abort_request_construction() { - let abort_req = proto::AbortRequest { - request_id: "req-456".to_string(), - reason: "User canceled".to_string(), - }; - assert_eq!(abort_req.request_id, "req-456"); - assert_eq!(abort_req.reason, "User canceled"); - } - - #[test] - fn test_sampling_params_defaults() { - let params = proto::SamplingParams::default(); - // Numeric fields have proto defaults (0) - assert_eq!(params.temperature, 0.0); - assert_eq!(params.top_p, 0.0); - assert_eq!(params.top_k, 0); - assert_eq!(params.repetition_penalty, 0.0); - assert_eq!(params.n, 0); - // Bool fields have proto defaults (false) - assert!(!params.skip_special_tokens); - assert!(!params.spaces_between_special_tokens); - assert!(!params.ignore_eos); - assert!(!params.no_stop_trim); - // Optional int fields should be None - assert_eq!(params.max_new_tokens, None); - assert_eq!(params.stream_interval, None); - // Other non-optional fields - assert_eq!(params.min_p, 0.0); - assert_eq!(params.frequency_penalty, 0.0); - assert_eq!(params.presence_penalty, 0.0); - assert!(params.stop.is_empty()); - } - - #[test] - fn test_multimodal_inputs() { - let mm_inputs = proto::MultimodalInputs { - image_urls: vec!["http://example.com/image.jpg".to_string()], - video_urls: vec![], - audio_urls: vec![], - image_data: vec![], - video_data: vec![], - audio_data: vec![], - modalities: vec!["image".to_string()], - ..Default::default() - }; - - assert_eq!(mm_inputs.image_urls.len(), 1); - assert_eq!(mm_inputs.image_urls[0], "http://example.com/image.jpg"); - assert_eq!(mm_inputs.modalities[0], "image"); - } - - // TODO: SessionParams not in current proto - skip test - - #[test] - fn test_embed_request() { - let embed_req = proto::EmbedRequest { - request_id: "embed-req-202".to_string(), - tokenized: Some(proto::TokenizedInput { - original_text: "This is a test sentence for embedding".to_string(), - input_ids: vec![2028, 374, 264, 1296, 11914, 369, 28537], // Mock token IDs - }), - log_metrics: true, - data_parallel_rank: 0, - ..Default::default() - }; - - assert_eq!(embed_req.request_id, "embed-req-202"); - if let Some(ref tokenized) = &embed_req.tokenized { - assert_eq!( - tokenized.original_text, - "This is a test sentence for embedding" - ); - } - assert!(embed_req.log_metrics); - assert_eq!(embed_req.data_parallel_rank, 0); - } - - #[tokio::test] - async fn test_client_connect_invalid_endpoint() { - let result = SglangSchedulerClient::connect("invalid://endpoint").await; - assert!(result.is_err()); - } - - #[test] - fn test_tokenized_input() { - let tokenized = proto::TokenizedInput { - original_text: "Hello world".to_string(), - input_ids: vec![1, 15043, 1917, 2], - }; - - assert_eq!(tokenized.original_text, "Hello world"); - assert_eq!(tokenized.input_ids, vec![1, 15043, 1917, 2]); - } - - #[test] - fn test_generate_stream_chunk() { - let chunk = proto::GenerateStreamChunk { - token_ids: vec![1234, 5678], - prompt_tokens: 5, - completion_tokens: 2, - cached_tokens: 3, - ..Default::default() - }; - - assert_eq!(chunk.token_ids, vec![1234, 5678]); - assert_eq!(chunk.prompt_tokens, 5); - assert_eq!(chunk.completion_tokens, 2); - assert_eq!(chunk.cached_tokens, 3); - } - - // TODO: ModelInfo not in current proto - skip test -} diff --git a/sgl-model-gateway/src/grpc_client/vllm_engine.rs b/sgl-model-gateway/src/grpc_client/vllm_engine.rs deleted file mode 100644 index d9892e417..000000000 --- a/sgl-model-gateway/src/grpc_client/vllm_engine.rs +++ /dev/null @@ -1,747 +0,0 @@ -use std::{ - pin::Pin, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, - task::{Context, Poll}, - time::Duration, -}; - -use tonic::{transport::Channel, Request, Streaming}; -use tracing::{debug, warn}; - -use crate::{ - observability::otel_trace::inject_trace_context_grpc, - protocols::{ - chat::ChatCompletionRequest, - common::{ResponseFormat, StringOrArray, ToolChoice, ToolChoiceValue}, - generate::GenerateRequest, - responses::ResponsesRequest, - sampling_params::SamplingParams as GenerateSamplingParams, - }, -}; - -// Include the generated protobuf code -#[allow(clippy::all)] -pub mod proto { - #![allow(clippy::all, unused_qualifications)] - tonic::include_proto!("vllm.grpc.engine"); -} - -// The generated module structure depends on the package name in the .proto file -// package vllm.grpc.engine; generates a nested module structure - -/// A smart wrapper around Streaming that automatically -/// sends abort when dropped (e.g., due to client disconnection or early termination). -/// -/// This leverages Rust's RAII pattern to ensure cleanup happens automatically, -/// regardless of how the stream is dropped (panic, early return, client disconnect, etc.). -pub struct AbortOnDropStream { - inner: Streaming, - request_id: String, - client: VllmEngineClient, - aborted: Arc, -} - -impl AbortOnDropStream { - /// Create a new auto-aborting stream wrapper - pub fn new( - stream: Streaming, - request_id: String, - client: VllmEngineClient, - ) -> Self { - debug!("Created AbortOnDropStream for request {}", request_id); - Self { - inner: stream, - request_id, - client, - aborted: Arc::new(AtomicBool::new(false)), - } - } - - /// Manually mark the request as completed to prevent abort on drop. - /// Call this when the request completes successfully to avoid unnecessary abort RPC. - pub fn mark_completed(&self) { - // Use Release ordering to ensure that this write is visible to other threads - // that use Acquire on the same atomic variable - self.aborted.store(true, Ordering::Release); - debug!("Request {} marked as completed", self.request_id); - } -} - -impl Drop for AbortOnDropStream { - fn drop(&mut self) { - // Atomically check and set the aborted flag using compare_exchange. - // If compare_exchange fails, it means the flag was already true (from mark_completed), - // so we don't need to send abort. AcqRel is used for success to synchronize with - // mark_completed's Release, and Acquire for failure to see writes from mark_completed. - if self - .aborted - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_err() - { - return; - } - - let client = self.client.clone(); - let request_id = self.request_id.clone(); - - // Spawn a background task to send abort (since Drop is sync but abort_request is async) - tokio::spawn(async move { - debug!( - "Stream dropped without completion for request {}, sending abort", - request_id - ); - // Clone request_id for the error message since abort_request takes ownership - let request_id_for_log = request_id.clone(); - if let Err(e) = client - .abort_request(request_id, "Stream dropped".to_string()) - .await - { - warn!( - "Failed to send abort on drop for request {}: {}", - request_id_for_log, e - ); - } - }); - } -} - -// Implement Stream trait to make AbortOnDropStream work like the original Streaming -impl futures::Stream for AbortOnDropStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // Delegate to the inner stream - Pin::new(&mut self.inner).poll_next(cx) - } -} - -/// gRPC client for vLLM scheduler -#[derive(Clone)] -pub struct VllmEngineClient { - client: proto::vllm_engine_client::VllmEngineClient, -} - -impl VllmEngineClient { - /// Create a new client and connect to the vLLM server - pub async fn connect(endpoint: &str) -> Result> { - debug!("Connecting to vLLM gRPC server at {}", endpoint); - - // Convert grpc:// to http:// for tonic - let http_endpoint = if let Some(addr) = endpoint.strip_prefix("grpc://") { - format!("http://{}", addr) - } else { - endpoint.to_string() - }; - - let channel = Channel::from_shared(http_endpoint)? - .http2_keep_alive_interval(Duration::from_secs(30)) - .keep_alive_timeout(Duration::from_secs(10)) - .keep_alive_while_idle(true) - .tcp_keepalive(Some(Duration::from_secs(60))) - .tcp_nodelay(true) - .http2_adaptive_window(true) - .initial_stream_window_size(Some(16 * 1024 * 1024)) // 16MB - .initial_connection_window_size(Some(32 * 1024 * 1024)) // 32MB - .connect() - .await?; - - let client = proto::vllm_engine_client::VllmEngineClient::new(channel); - - Ok(Self { client }) - } - - /// Submit a generation request (returns auto-aborting streaming response) - /// - /// The returned stream automatically sends an abort request when dropped, - /// ensuring proper cleanup even if the HTTP client disconnects or an error occurs. - /// Call `mark_completed()` on the stream after successful completion to prevent - /// unnecessary abort RPCs. - pub async fn generate( - &self, - req: proto::GenerateRequest, - ) -> Result> { - let request_id = req.request_id.clone(); - let mut client = self.client.clone(); - let mut request = Request::new(req); - - // Inject W3C trace context into gRPC metadata for distributed tracing - inject_trace_context_grpc(request.metadata_mut()); - - let response = client.generate(request).await?; - - Ok(AbortOnDropStream::new( - response.into_inner(), - request_id, - self.clone(), - )) - } - - /// Perform health check - pub async fn health_check( - &self, - ) -> Result> { - debug!("Sending health check request"); - // HealthCheckRequest is now empty - server generates its own health check internally - let request = Request::new(proto::HealthCheckRequest {}); - - let mut client = self.client.clone(); - let response = client.health_check(request).await?; - debug!("Health check response received"); - Ok(response.into_inner()) - } - - /// Abort a request - pub async fn abort_request( - &self, - request_id: String, - _reason: String, - ) -> Result<(), Box> { - debug!("Sending abort request for {}", request_id); - let request = Request::new(proto::AbortRequest { - request_ids: vec![request_id.clone()], - }); - - let mut client = self.client.clone(); - let _response = client.abort(request).await?; - debug!("Abort response received for {}", request_id); - Ok(()) - } - - /// Get model information - pub async fn get_model_info( - &self, - ) -> Result> { - debug!("Requesting model info"); - let request = Request::new(proto::GetModelInfoRequest {}); - - let mut client = self.client.clone(); - let response = client.get_model_info(request).await?; - debug!("Model info response received"); - Ok(response.into_inner()) - } - - /// Get server information - pub async fn get_server_info( - &self, - ) -> Result> { - debug!("Requesting server info"); - let request = Request::new(proto::GetServerInfoRequest {}); - - let mut client = self.client.clone(); - let response = client.get_server_info(request).await?; - debug!("Server info response received"); - Ok(response.into_inner()) - } - - /// Build a single vLLM GenerateRequest from OpenAI ChatCompletionRequest - pub fn build_generate_request_from_chat( - &self, - request_id: String, - body: &ChatCompletionRequest, - processed_text: String, - token_ids: Vec, - tool_call_constraint: Option<(String, String)>, // (constraint_type, constraint_value) - ) -> Result { - // Build sampling params - let sampling_params = - self.build_grpc_sampling_params_from_chat(body, tool_call_constraint)?; - - let grpc_request = proto::GenerateRequest { - request_id, - input: Some(proto::generate_request::Input::Tokenized( - proto::TokenizedInput { - original_text: processed_text, - input_ids: token_ids, - }, - )), - sampling_params: Some(sampling_params), - stream: body.stream, - }; - - Ok(grpc_request) - } - - /// Build a basic GenerateRequest from the vLLM spec GenerateRequest - pub fn build_plain_generate_request( - &self, - request_id: String, - body: &GenerateRequest, - original_text: Option, - token_ids: Vec, - ) -> Result { - let sampling_params = - Self::build_sampling_params_from_plain(body.sampling_params.as_ref())?; - - let grpc_request = proto::GenerateRequest { - request_id, - input: Some(proto::generate_request::Input::Tokenized( - proto::TokenizedInput { - original_text: original_text.unwrap_or_default(), - input_ids: token_ids, - }, - )), - sampling_params: Some(sampling_params), - stream: body.stream, - }; - - Ok(grpc_request) - } - - /// Build a GenerateRequest from ResponsesRequest (OpenAI Responses API) - /// - /// NOTE: This is used by the Harmony router only. The Regular router uses - /// responses_to_chat() conversion and goes through the chat pipeline. - pub fn build_generate_request_from_responses( - &self, - request_id: String, - body: &ResponsesRequest, - processed_text: String, - token_ids: Vec, - harmony_stop_ids: Option>, - constraint: Option<(String, String)>, - ) -> Result { - // Build sampling params from ResponsesRequest - let mut sampling_params = - self.build_grpc_sampling_params_from_responses(body, constraint)?; - - // Inject Harmony stop token IDs if provided - if let Some(stop_ids) = harmony_stop_ids { - sampling_params.stop_token_ids = stop_ids; - } - - let grpc_request = proto::GenerateRequest { - request_id, - input: Some(proto::generate_request::Input::Tokenized( - proto::TokenizedInput { - original_text: processed_text, - input_ids: token_ids, - }, - )), - sampling_params: Some(sampling_params), - stream: body.stream.unwrap_or(false), - }; - - Ok(grpc_request) - } - - /// Build gRPC SamplingParams from ChatCompletionRequest - fn build_grpc_sampling_params_from_chat( - &self, - request: &ChatCompletionRequest, - tool_call_constraint: Option<(String, String)>, - ) -> Result { - let stop_sequences = self.extract_stop_strings(request); - - let max_tokens = request.max_completion_tokens; - - // Handle skip_special_tokens: set to false if tools are present and tool_choice is not "none" - let skip_special_tokens = if request.tools.is_some() { - match &request.tool_choice { - Some(ToolChoice::Value(ToolChoiceValue::None)) => request.skip_special_tokens, - Some(_) => false, // tool_choice is not "none" - None => false, // TODO: this assumes tool_choice defaults to "auto" when tools present - } - } else { - request.skip_special_tokens - }; - - Ok(proto::SamplingParams { - temperature: request.temperature, - top_p: request.top_p.unwrap_or(1.0), - top_k: request.top_k.map(|v| v.max(0) as u32).unwrap_or(0), // 0 means disabled in vLLM - min_p: request.min_p.unwrap_or(0.0), - frequency_penalty: request.frequency_penalty.unwrap_or(0.0), - presence_penalty: request.presence_penalty.unwrap_or(0.0), - repetition_penalty: request.repetition_penalty.unwrap_or(1.0), - max_tokens, - stop: stop_sequences, - stop_token_ids: request.stop_token_ids.clone().unwrap_or_default(), - skip_special_tokens, - spaces_between_special_tokens: true, // Default from Python SamplingParams - ignore_eos: request.ignore_eos, - n: request.n.unwrap_or(1), - constraint: self.build_constraint_for_chat(request, tool_call_constraint)?, - ..Default::default() - }) - } - - /// Extract stop strings from request - fn extract_stop_strings(&self, request: &ChatCompletionRequest) -> Vec { - match &request.stop { - Some(StringOrArray::String(s)) => vec![s.clone()], - Some(StringOrArray::Array(arr)) => arr.clone(), - None => vec![], - } - } - - /// Build constraint for structured generation - fn build_constraint_for_chat( - &self, - request: &ChatCompletionRequest, - tool_call_constraint: Option<(String, String)>, - ) -> Result, String> { - let mut constraints = Vec::new(); - - // Handle response_format constraints - match &request.response_format { - Some(ResponseFormat::JsonObject) => { - // json_object mode - constrain to valid JSON object - let schema = serde_json::json!({"type": "object"}); - let schema_str = serde_json::to_string(&schema) - .map_err(|e| format!("Failed to serialize JSON schema: {}", e))?; - constraints.push(proto::sampling_params::Constraint::JsonSchema(schema_str)); - } - Some(ResponseFormat::JsonSchema { json_schema }) => { - let schema_str = serde_json::to_string(&json_schema.schema) - .map_err(|e| format!("Failed to serialize JSON schema: {}", e))?; - constraints.push(proto::sampling_params::Constraint::JsonSchema(schema_str)); - } - Some(ResponseFormat::Text) | None => { - // No constraint for text format - } - } - - // vLLM supports: json_schema, regex, grammar, structural_tag, json_object, choice - if let Some(ebnf) = &request.ebnf { - constraints.push(proto::sampling_params::Constraint::Grammar(ebnf.clone())); - } - - if let Some(regex) = &request.regex { - constraints.push(proto::sampling_params::Constraint::Regex(regex.clone())); - } - - // Handle tool call constraint from preparation stage - if let Some((constraint_type, constraint_value)) = tool_call_constraint { - if !constraints.is_empty() { - return Err("Constrained decoding is not compatible with tool calls.".to_string()); - } - let tool_constraint = match constraint_type.as_str() { - "structural_tag" => { - proto::sampling_params::Constraint::StructuralTag(constraint_value) - } - "json_schema" => proto::sampling_params::Constraint::JsonSchema(constraint_value), - "grammar" | "ebnf" => proto::sampling_params::Constraint::Grammar(constraint_value), - "regex" => proto::sampling_params::Constraint::Regex(constraint_value), - _ => return Err(format!("Unknown constraint type: {}", constraint_type)), - }; - constraints.push(tool_constraint); - } - - match constraints.len() { - 0 => Ok(None), - 1 => Ok(constraints.pop()), - _ => Err("Multiple constraints are not allowed.".to_string()), - } - } - - /// Build gRPC SamplingParams from ResponsesRequest - fn build_grpc_sampling_params_from_responses( - &self, - request: &ResponsesRequest, - constraint: Option<(String, String)>, - ) -> Result { - // Used by Harmony models only. Regular models use Chat API path. - // Constraints come from Harmony preparation stage (structural_tag) or tool handling. - - let max_tokens = request.max_output_tokens; - - Ok(proto::SamplingParams { - temperature: request.temperature, - top_p: request.top_p.unwrap_or(1.0), - top_k: 0, // ResponsesRequest doesn't expose top_k (0 means disabled) - min_p: 0.0, // ResponsesRequest doesn't expose min_p - frequency_penalty: 0.0, // ResponsesRequest doesn't expose frequency_penalty - presence_penalty: 0.0, // ResponsesRequest doesn't expose presence_penalty - repetition_penalty: 1.0, // ResponsesRequest doesn't expose repetition_penalty - max_tokens, - stop: vec![], // No stop sequences in Responses API - stop_token_ids: vec![], // Handled by Harmony stop tokens - skip_special_tokens: false, // Keep special tokens for Harmony - spaces_between_special_tokens: true, - ignore_eos: false, - n: 1, // Responses API doesn't support n>1 - constraint: self.build_constraint_for_responses(constraint)?, - ..Default::default() - }) - } - - /// Build constraint for Responses API - /// - /// Handles constraints from Harmony preparation stage (structural_tag for Harmony models, - /// structured output via text field, or tool call constraints). - /// - /// Note: Regular gRPC models use Chat API path with response_format, not this function. - fn build_constraint_for_responses( - &self, - constraint: Option<(String, String)>, - ) -> Result, String> { - if let Some((constraint_type, constraint_value)) = constraint { - let parsed_constraint = match constraint_type.as_str() { - "structural_tag" => { - proto::sampling_params::Constraint::StructuralTag(constraint_value) - } - "json_schema" => proto::sampling_params::Constraint::JsonSchema(constraint_value), - "grammar" | "ebnf" => proto::sampling_params::Constraint::Grammar(constraint_value), - "regex" => proto::sampling_params::Constraint::Regex(constraint_value), - _ => return Err(format!("Unknown constraint type: {}", constraint_type)), - }; - Ok(Some(parsed_constraint)) - } else { - Ok(None) - } - } - - fn build_single_constraint_from_plain( - params: &GenerateSamplingParams, - ) -> Result, String> { - let mut constraints = Vec::new(); - if let Some(json_schema) = ¶ms.json_schema { - constraints.push(proto::sampling_params::Constraint::JsonSchema( - json_schema.clone(), - )); - } - if let Some(regex) = ¶ms.regex { - constraints.push(proto::sampling_params::Constraint::Regex(regex.clone())); - } - if let Some(ebnf) = ¶ms.ebnf { - constraints.push(proto::sampling_params::Constraint::Grammar(ebnf.clone())); - } - - match constraints.len() { - 0 => Ok(None), - 1 => Ok(constraints.pop()), - _ => Err("Multiple structured constraints are not allowed".to_string()), - } - } - - fn build_sampling_params_from_plain( - params: Option<&GenerateSamplingParams>, - ) -> Result { - let mut sampling = proto::SamplingParams { - temperature: Some(1.0), - top_p: 1.0, - top_k: 0, // 0 means disabled in vLLM - repetition_penalty: 1.0, - n: 1, - skip_special_tokens: true, - spaces_between_special_tokens: true, - ..Default::default() - }; - - let Some(p) = params else { - return Ok(sampling); - }; - - // Handle temperature (now optional) - if let Some(val) = p.temperature { - sampling.temperature = Some(val); - } - - // Simple field mappings - if let Some(val) = p.top_p { - sampling.top_p = val; - } - if let Some(val) = p.top_k { - sampling.top_k = val.max(0) as u32; // Clamp negative values to 0 (disabled) - } - if let Some(val) = p.frequency_penalty { - sampling.frequency_penalty = val; - } - if let Some(val) = p.presence_penalty { - sampling.presence_penalty = val; - } - if let Some(val) = p.repetition_penalty { - sampling.repetition_penalty = val; - } - if let Some(val) = p.min_p { - sampling.min_p = val; - } - if let Some(val) = p.ignore_eos { - sampling.ignore_eos = val; - } - if let Some(val) = p.skip_special_tokens { - sampling.skip_special_tokens = val; - } - // Note: no_stop_trim not supported in vLLM - - // Handle stop sequences - if let Some(stop) = &p.stop { - match stop { - StringOrArray::String(s) => sampling.stop.push(s.clone()), - StringOrArray::Array(arr) => sampling.stop.extend(arr.clone()), - } - } - - // Handle stop token IDs - if let Some(stop_token_ids) = &p.stop_token_ids { - sampling.stop_token_ids = stop_token_ids.clone(); - } - - // Handle max_tokens (read from internal max_new_tokens) - if let Some(max_new_tokens) = p.max_new_tokens { - sampling.max_tokens = Some(max_new_tokens); - } - - // Handle min_tokens (read from internal min_new_tokens) - if let Some(min_new_tokens) = p.min_new_tokens { - sampling.min_tokens = min_new_tokens; - } - - // Handle n - if let Some(n) = p.n { - sampling.n = n; - } - - // Handle constraints (exactly one allowed) - sampling.constraint = Self::build_single_constraint_from_plain(p)?; - - Ok(sampling) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_proto_types_compilation() { - let _health_req = proto::HealthCheckRequest {}; - // HealthCheckRequest is now empty - no fields to test - } - - #[test] - fn test_generate_request_construction() { - let sampling_params = proto::SamplingParams { - temperature: Some(0.7), - max_tokens: Some(128), - top_p: 0.9, - top_k: 50, - stop: vec!["".to_string()], - ..Default::default() - }; - - let gen_req = proto::GenerateRequest { - request_id: "test-req-123".to_string(), - input: Some(proto::generate_request::Input::Tokenized( - proto::TokenizedInput { - original_text: "Hello world".to_string(), - input_ids: vec![9906, 1917], // Mock token IDs for "Hello world" - }, - )), - sampling_params: Some(sampling_params), - stream: false, - }; - - assert_eq!(gen_req.request_id, "test-req-123"); - if let Some(proto::generate_request::Input::Tokenized(ref tokenized)) = gen_req.input { - assert_eq!(tokenized.original_text, "Hello world"); - } - // vLLM: logprobs are in SamplingParams, not GenerateRequest - - let params = gen_req.sampling_params.unwrap(); - assert_eq!(params.temperature, Some(0.7)); - assert_eq!(params.max_tokens, Some(128)); - assert_eq!(params.stop, vec![""]); - } - - #[test] - fn test_health_check_request() { - let _health_req = proto::HealthCheckRequest {}; - // HealthCheckRequest is now empty - server generates its own test internally - } - - #[test] - fn test_abort_request_construction() { - let abort_req = proto::AbortRequest { - request_ids: vec!["req-456".to_string(), "req-789".to_string()], - }; - assert_eq!(abort_req.request_ids, vec!["req-456", "req-789"]); - } - - #[test] - fn test_sampling_params_defaults() { - let params = proto::SamplingParams::default(); - // Optional float field defaults to None - assert_eq!(params.temperature, None); - // Non-optional numeric fields have proto defaults (0) - assert_eq!(params.top_p, 0.0); - assert_eq!(params.top_k, 0); - assert_eq!(params.repetition_penalty, 0.0); - assert_eq!(params.n, 0); - // Bool fields have proto defaults (false) - assert!(!params.skip_special_tokens); - assert!(!params.spaces_between_special_tokens); - assert!(!params.ignore_eos); - assert!(!params.include_stop_str_in_output); - // Optional fields should be None - assert_eq!(params.max_tokens, None); - assert_eq!(params.logprobs, None); - // Other non-optional fields - assert_eq!(params.min_p, 0.0); - assert_eq!(params.frequency_penalty, 0.0); - assert_eq!(params.presence_penalty, 0.0); - assert!(params.stop.is_empty()); - } - - // TODO: MultimodalInputs not in vLLM proto - skip test - // vLLM handles multimodal inputs differently than SGLang - - // TODO: SessionParams not in current proto - skip test - - #[test] - fn test_embed_request() { - let embed_req = proto::EmbedRequest { - request_id: "embed-req-202".to_string(), - tokenized: Some(proto::TokenizedInput { - original_text: "This is a test sentence for embedding".to_string(), - input_ids: vec![2028, 374, 264, 1296, 11914, 369, 28537], // Mock token IDs - }), - }; - - assert_eq!(embed_req.request_id, "embed-req-202"); - if let Some(ref tokenized) = &embed_req.tokenized { - assert_eq!( - tokenized.original_text, - "This is a test sentence for embedding" - ); - } - // vLLM: no data_parallel_rank or log_metrics in EmbedRequest - } - - #[tokio::test] - async fn test_client_connect_invalid_endpoint() { - let result = VllmEngineClient::connect("invalid://endpoint").await; - assert!(result.is_err()); - } - - #[test] - fn test_tokenized_input() { - let tokenized = proto::TokenizedInput { - original_text: "Hello world".to_string(), - input_ids: vec![1, 15043, 1917, 2], - }; - - assert_eq!(tokenized.original_text, "Hello world"); - assert_eq!(tokenized.input_ids, vec![1, 15043, 1917, 2]); - } - - #[test] - fn test_generate_stream_chunk() { - let chunk = proto::GenerateStreamChunk { - token_ids: vec![1234, 5678], - prompt_tokens: 5, - completion_tokens: 2, - cached_tokens: 3, - }; - - assert_eq!(chunk.token_ids, vec![1234, 5678]); - assert_eq!(chunk.prompt_tokens, 5); - assert_eq!(chunk.completion_tokens, 2); - assert_eq!(chunk.cached_tokens, 3); - } - - // TODO: ModelInfo not in current proto - skip test -} diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index cd1af3ecd..0166819d0 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -3,7 +3,7 @@ pub use smg_auth as auth; pub mod config; pub mod core; pub use data_connector; -pub mod grpc_client; +pub use smg_grpc_client as grpc_client; pub use smg_mcp as mcp; pub use smg_mesh as mesh; pub mod middleware; diff --git a/sgl-model-gateway/src/observability/otel_trace.rs b/sgl-model-gateway/src/observability/otel_trace.rs index 99ed6b004..cc5b636bc 100644 --- a/sgl-model-gateway/src/observability/otel_trace.rs +++ b/sgl-model-gateway/src/observability/otel_trace.rs @@ -260,3 +260,20 @@ pub fn inject_trace_context_grpc(metadata: &mut MetadataMap) { propagator.inject_context(&context, &mut MetadataInjector(metadata)); }); } + +/// OpenTelemetry trace injector implementing the `smg_grpc_client::TraceInjector` trait. +/// +/// This bridges sglang's OTel integration with the `smg-grpc-client` crate's +/// trace injection interface, enabling distributed tracing across gRPC calls. +#[derive(Clone, Default)] +pub struct OtelTraceInjector; + +impl smg_grpc_client::TraceInjector for OtelTraceInjector { + fn inject( + &self, + metadata: &mut MetadataMap, + ) -> Result<(), Box> { + inject_trace_context_grpc(metadata); + Ok(()) + } +} diff --git a/sgl-model-gateway/src/proto/sglang_scheduler.proto b/sgl-model-gateway/src/proto/sglang_scheduler.proto deleted file mode 100644 index 578fc3317..000000000 --- a/sgl-model-gateway/src/proto/sglang_scheduler.proto +++ /dev/null @@ -1,566 +0,0 @@ -syntax = "proto3"; - -package sglang.grpc.scheduler; - -import "google/protobuf/timestamp.proto"; -import "google/protobuf/struct.proto"; - -// Service definition for SGLang scheduler communication -// This protocol bridges the Rust router and Python scheduler -service SglangScheduler { - // Submit a generation request (supports streaming) - rpc Generate(GenerateRequest) returns (stream GenerateResponse); - - // Submit an embedding request - rpc Embed(EmbedRequest) returns (EmbedResponse); - - // Health check and metrics - rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse); - - // Abort a running request - rpc Abort(AbortRequest) returns (AbortResponse); - - // Get model information - rpc GetModelInfo(GetModelInfoRequest) returns (GetModelInfoResponse); - - // Get server information - rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse); - - // Get comprehensive load metrics - rpc GetLoads(GetLoadsRequest) returns (GetLoadsResponse); - -} - -// ===================== -// Common Types -// ===================== - -// Sampling parameters matching SGLang's SamplingParams -// -// IMPORTANT: Do not use SamplingParams::default() directly! -// The proto3 defaults (0 for numeric fields) do NOT match the semantic defaults -// (temperature=1.0, top_p=1.0, top_k=-1, etc.). Always construct with explicit values -// or use the conversion functions in sglang_scheduler.rs / grpc_server.py. -message SamplingParams { - float temperature = 1; - float top_p = 2; - int32 top_k = 3; - float min_p = 4; - float frequency_penalty = 5; - float presence_penalty = 6; - float repetition_penalty = 7; - - optional int32 max_new_tokens = 8; - repeated string stop = 9; - repeated uint32 stop_token_ids = 10; - bool skip_special_tokens = 11; - bool spaces_between_special_tokens = 12; - - // Structured generation - oneof constraint { - string regex = 13; - string json_schema = 14; - string ebnf_grammar = 15; - string structural_tag = 16; - } - - // Speculative decoding - int32 n = 17; // Number of samples - - // Additional parameters - int32 min_new_tokens = 18; - bool ignore_eos = 19; - bool no_stop_trim = 20; - optional int32 stream_interval = 21; - map logit_bias = 22; - - // Custom parameters for extensibility - google.protobuf.Struct custom_params = 23; -} - - -// Disaggregated serving parameters -message DisaggregatedParams { - string bootstrap_host = 1; - int32 bootstrap_port = 2; - int32 bootstrap_room = 3; -} - -// ===================== -// Generate Request -// ===================== - -message GenerateRequest { - string request_id = 1; - - // Input must be tokenized (no raw text) - TokenizedInput tokenized = 2; - - // Multimodal inputs - MultimodalInputs mm_inputs = 3; - - // Generation parameters - SamplingParams sampling_params = 4; - - // Return options - bool return_logprob = 5; - int32 logprob_start_len = 6; - int32 top_logprobs_num = 7; - repeated uint32 token_ids_logprob = 8; - bool return_hidden_states = 9; - - // For disaggregated serving - DisaggregatedParams disaggregated_params = 10; - - // Custom logit processor (serialized) - string custom_logit_processor = 11; - - // Request metadata - google.protobuf.Timestamp timestamp = 12; - bool log_metrics = 13; - - // Input embeddings (alternative to text/tokens) - repeated float input_embeds = 14; - - // LoRA adapter ID (if pre-loaded) - string lora_id = 15; - - // Data parallel routing - int32 data_parallel_rank = 16; - - // Whether client wants streaming response - bool stream = 17; -} - -message TokenizedInput { - string original_text = 1; // For reference - repeated uint32 input_ids = 2; -} - -message MultimodalInputs { - // Simplified multimodal handling - actual data processed by tokenizer - repeated string image_urls = 1; - repeated string video_urls = 2; - repeated string audio_urls = 3; - - // Pre-processed multimodal features (if available) - google.protobuf.Struct processed_features = 4; - - // Raw data for direct processing - repeated bytes image_data = 5; - repeated bytes video_data = 6; - repeated bytes audio_data = 7; - - // Modality metadata - repeated string modalities = 8; -} - -// ===================== -// Generate Response -// ===================== - -message GenerateResponse { - string request_id = 1; - - // Response type - oneof response { - GenerateStreamChunk chunk = 2; - GenerateComplete complete = 3; - GenerateError error = 4; - } -} - -message GenerateStreamChunk { - // Generated tokens (incremental chunk) - repeated uint32 token_ids = 1; - - // Cumulative counts - int32 prompt_tokens = 2; - int32 completion_tokens = 3; - int32 cached_tokens = 4; - - // Output logprobs (if requested) - incremental for streaming - OutputLogProbs output_logprobs = 5; - - // Hidden states (if requested) - repeated float hidden_states = 6; - - // Input logprobs (if requested) - only in first chunk - InputLogProbs input_logprobs = 7; - - // Index for ordering when n>1 (for parallel request multiplexing) - uint32 index = 8; -} - -message GenerateComplete { - // Final output - repeated uint32 output_ids = 1; - - // Finish reason as OpenAI-compatible string ("stop", "length", "abort") - string finish_reason = 2; - - // Token usage counts - int32 prompt_tokens = 3; - int32 completion_tokens = 4; - int32 cached_tokens = 5; - - // Output logprobs if requested (cumulative) - OutputLogProbs output_logprobs = 6; - - // All hidden states if requested - repeated HiddenStates all_hidden_states = 7; - - // Matched stop information (for stop sequences) - oneof matched_stop { - uint32 matched_token_id = 8; - string matched_stop_str = 9; - } - - // Input logprobs if requested (for prompt tokens) - InputLogProbs input_logprobs = 10; - - // Index for ordering when n>1 (for parallel request multiplexing) - uint32 index = 11; -} - -message GenerateError { - string message = 1; - string http_status_code = 2; - string details = 3; -} - -// Output logprobs - all values are present (no None) -message OutputLogProbs { - repeated float token_logprobs = 1; - repeated int32 token_ids = 2; - - // Top logprobs at each position - repeated TopLogProbs top_logprobs = 3; -} - -// Input logprobs - first token has no logprob (None) -message InputLogProbs { - repeated InputTokenLogProb token_logprobs = 1; - repeated int32 token_ids = 2; - - // Top logprobs at each position - repeated TopLogProbs top_logprobs = 3; -} - -// Wrapper to represent optional logprob (first input token has no logprob) -message InputTokenLogProb { - optional float value = 1; -} - -message TopLogProbs { - repeated float values = 1; - repeated int32 token_ids = 2; -} - -message HiddenStates { - repeated float values = 1; - int32 layer = 2; - int32 position = 3; -} - -// ===================== -// Embedding Request -// ===================== - -message EmbedRequest { - string request_id = 1; - - // Input must be tokenized (no raw text) - TokenizedInput tokenized = 2; - - // Multimodal inputs - MultimodalInputs mm_inputs = 4; - - // Dummy sampling params for compatibility - // EmbedRequest doesn't use sampling_params - SamplingParams sampling_params = 5; - - bool log_metrics = 6; - - // Token type IDs for models that require them - repeated int32 token_type_ids = 7; - - // Data parallel routing - int32 data_parallel_rank = 8; - - // For cross-encoder requests - bool is_cross_encoder = 9; - repeated string texts = 10; // For cross-encoder batch -} - -message EmbedResponse { - string request_id = 1; - - oneof response { - EmbedComplete complete = 2; - EmbedError error = 3; - } -} - -message EmbedComplete { - repeated float embedding = 1; - int32 prompt_tokens = 2; - int32 cached_tokens = 3; - - // Additional metadata - int32 embedding_dim = 4; - - // For batch embeddings - repeated Embedding batch_embeddings = 5; -} - -message Embedding { - repeated float values = 1; - int32 index = 2; -} - -message EmbedError { - string message = 1; - string code = 2; - string details = 3; -} - -// ===================== -// Management Operations -// ===================== - -message HealthCheckRequest {} - -message HealthCheckResponse { - bool healthy = 1; - string message = 2; -} - -message AbortRequest { - string request_id = 1; - string reason = 2; -} - -message AbortResponse { - bool success = 1; - string message = 2; -} - - -// ===================== -// Additional Operations (Future) -// ===================== - -// Load LoRA adapter -message LoadLoRARequest { - string adapter_id = 1; - string adapter_path = 2; - int32 rank = 3; -} - -message LoadLoRAResponse { - bool success = 1; - string adapter_id = 2; - string message = 3; -} - -// Unload LoRA adapter -message UnloadLoRARequest { - string adapter_id = 1; -} - -message UnloadLoRAResponse { - bool success = 1; - string message = 2; -} - -// Update weights -message UpdateWeightsRequest { - oneof source { - string disk_path = 1; - bytes tensor_data = 2; - string remote_url = 3; - } - string weight_name = 4; -} - -message UpdateWeightsResponse { - bool success = 1; - string message = 2; -} - -// Get internal state for debugging -message GetInternalStateRequest { - repeated string state_keys = 1; -} - -message GetInternalStateResponse { - google.protobuf.Struct state = 1; -} - -// Set internal state for testing -message SetInternalStateRequest { - google.protobuf.Struct state = 1; -} - -message SetInternalStateResponse { - bool success = 1; - string message = 2; -} - -// ===================== -// Model and Server Info -// ===================== - -// Get model information -message GetModelInfoRequest {} - -message GetModelInfoResponse { - string model_path = 1; - string tokenizer_path = 2; - bool is_generation = 3; - string preferred_sampling_params = 4; // JSON string or empty - string weight_version = 5; - string served_model_name = 6; - int32 max_context_length = 7; - int32 vocab_size = 8; - bool supports_vision = 9; - string model_type = 10; - repeated int32 eos_token_ids = 11; - int32 pad_token_id = 12; - int32 bos_token_id = 13; - int32 max_req_input_len = 14; - repeated string architectures = 15; - - // Classification model support (from HuggingFace config.json) - // id2label maps class indices to label names, e.g., {"0": "negative", "1": "positive"} - string id2label_json = 16; - // Number of classification labels (0 if not a classifier) - int32 num_labels = 17; -} - -// Get server information -message GetServerInfoRequest {} - -message GetServerInfoResponse { - // Server configuration (as structured data) - google.protobuf.Struct server_args = 1; - - // Scheduler metrics (from scheduler initialization) - google.protobuf.Struct scheduler_info = 2; - - // Runtime state - int32 active_requests = 3; - bool is_paused = 4; - double last_receive_timestamp = 5; - double uptime_seconds = 6; - - // Version info - string sglang_version = 7; - - // Server metadata - string server_type = 8; // "grpc" - google.protobuf.Timestamp start_time = 9; - - // Note: internal_states not provided in gRPC mode - // Scheduler-side metrics (memory usage, throughput) require - // bidirectional communicator infrastructure not available in gRPC. - // Use HTTP /get_server_info if scheduler internal state is needed. -} - -// ===================== -// Load Metrics (v1/loads) -// ===================== - -message GetLoadsRequest { - // Optional: filter to specific DP rank - optional int32 dp_rank = 1; - - // Sections to include: core, memory, spec, lora, disagg, queues, all - repeated string include = 2; -} - -message GetLoadsResponse { - // ISO 8601 timestamp - string timestamp = 1; - - // SGLang version - string version = 2; - - // Number of DP ranks - int32 dp_rank_count = 3; - - // Per-DP-rank load metrics - repeated SchedulerLoad loads = 4; - - // Aggregate metrics across all DP ranks - AggregateMetrics aggregate = 5; -} - -message SchedulerLoad { - int32 dp_rank = 1; - - // Core metrics (always included) - int32 num_running_reqs = 2; - int32 num_waiting_reqs = 3; - int32 num_total_reqs = 4; - int32 num_used_tokens = 5; - int32 max_total_num_tokens = 6; - double token_usage = 7; - double gen_throughput = 8; - double cache_hit_rate = 9; - double utilization = 10; - int32 max_running_requests = 11; - - // Optional sections - optional MemoryMetrics memory = 12; - optional SpeculativeMetrics speculative = 13; - optional LoRAMetrics lora = 14; - optional DisaggregationMetrics disaggregation = 15; - optional QueueMetrics queues = 16; -} - -message MemoryMetrics { - double weight_gb = 1; - double kv_cache_gb = 2; - double graph_gb = 3; - int32 token_capacity = 4; -} - -message SpeculativeMetrics { - double accept_length = 1; - double accept_rate = 2; -} - -message LoRAMetrics { - int32 slots_used = 1; - int32 slots_total = 2; - double utilization = 3; -} - -message DisaggregationMetrics { - string mode = 1; // "prefill", "decode", or "null" - int32 prefill_prealloc_queue_reqs = 2; - int32 prefill_inflight_queue_reqs = 3; - int32 decode_prealloc_queue_reqs = 4; - int32 decode_transfer_queue_reqs = 5; - int32 decode_retracted_queue_reqs = 6; - double kv_transfer_speed_gb_s = 7; - double kv_transfer_latency_ms = 8; -} - -message QueueMetrics { - int32 waiting = 1; - int32 grammar = 2; - int32 paused = 3; - int32 retracted = 4; -} - -message AggregateMetrics { - int32 total_running_reqs = 1; - int32 total_waiting_reqs = 2; - int32 total_reqs = 3; - double avg_token_usage = 4; - double avg_throughput = 5; - double avg_utilization = 6; -} diff --git a/sgl-model-gateway/src/proto/vllm_engine.proto b/sgl-model-gateway/src/proto/vllm_engine.proto deleted file mode 100644 index bbb1b9b00..000000000 --- a/sgl-model-gateway/src/proto/vllm_engine.proto +++ /dev/null @@ -1,195 +0,0 @@ -syntax = "proto3"; - -package vllm.grpc.engine; - -// Service definition for vLLM engine communication -// This protocol is designed for efficient binary communication between -// the Rust router and vLLM Python engine (AsyncLLM). -service VllmEngine { - // Submit a generation request (supports streaming) - rpc Generate(GenerateRequest) returns (stream GenerateResponse); - - // Submit an embedding request - rpc Embed(EmbedRequest) returns (EmbedResponse); - - // Health check - rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse); - - // Abort a running request - rpc Abort(AbortRequest) returns (AbortResponse); - - // Get model information - rpc GetModelInfo(GetModelInfoRequest) returns (GetModelInfoResponse); - - // Get server information - rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse); -} - -// ===================== -// Common Types -// ===================== - -// Sampling parameters for text generation -message SamplingParams { - optional float temperature = 1; - float top_p = 2; - uint32 top_k = 3; - float min_p = 4; - float frequency_penalty = 5; - float presence_penalty = 6; - float repetition_penalty = 7; - - optional uint32 max_tokens = 8; - uint32 min_tokens = 9; - - repeated string stop = 10; - repeated uint32 stop_token_ids = 11; - - bool skip_special_tokens = 12; - bool spaces_between_special_tokens = 13; - bool ignore_eos = 14; - - uint32 n = 15; // Number of parallel samples - - // Logprobs configuration - optional int32 logprobs = 22; // Number of log probabilities per output token (-1 for all) - optional int32 prompt_logprobs = 23; // Number of log probabilities per prompt token (-1 for all) - - // Additional vLLM fields - optional int32 seed = 24; // Random seed for reproducibility - bool include_stop_str_in_output = 25; // Whether to include stop strings in output - map logit_bias = 26; // Token ID to bias mapping (-100 to 100) - optional int32 truncate_prompt_tokens = 27; // Prompt truncation (-1 for model max) - - // Structured outputs (one of) - matches vLLM's StructuredOutputsParams - oneof constraint { - string json_schema = 16; // JSON schema for structured output - string regex = 17; // Regex pattern - string grammar = 18; // Grammar/EBNF for structured output - string structural_tag = 19; // Structural tag (e.g., Harmony models) - bool json_object = 20; // Force JSON object output - ChoiceConstraint choice = 21; // List of allowed choices - } -} - -// Choice constraint for structured outputs -message ChoiceConstraint { - repeated string choices = 1; -} - -// Pre-tokenized input from Rust router -message TokenizedInput { - string original_text = 1; // For reference/debugging - repeated uint32 input_ids = 2; // Actual token IDs to process -} - -// ===================== -// Generate Request -// ===================== - -message GenerateRequest { - string request_id = 1; - - // Prompt input - oneof input { - TokenizedInput tokenized = 2; - string text = 3; - } - - // Generation parameters (includes logprobs config) - SamplingParams sampling_params = 4; - - // Streaming - bool stream = 5; -} - -// ===================== -// Generate Response -// ===================== - -message GenerateResponse { - oneof response { - GenerateStreamChunk chunk = 1; // For streaming - GenerateComplete complete = 2; // For final/non-streaming - } -} - -message GenerateStreamChunk { - repeated uint32 token_ids = 1; // Incremental tokens - uint32 prompt_tokens = 2; - uint32 completion_tokens = 3; - uint32 cached_tokens = 4; - - // Logprobs support (TODO: implement in Phase 4) - // OutputLogProbs output_logprobs = 5; - // InputLogProbs input_logprobs = 6; // Only in first chunk -} - -message GenerateComplete { - repeated uint32 output_ids = 1; // All output tokens - string finish_reason = 2; // "stop", "length", "abort" - uint32 prompt_tokens = 3; - uint32 completion_tokens = 4; - uint32 cached_tokens = 5; - - // Logprobs support (TODO: implement in Phase 4) - // OutputLogProbs output_logprobs = 6; - // InputLogProbs input_logprobs = 7; -} - -// ===================== -// Embedding Request -// ===================== - -message EmbedRequest { - string request_id = 1; - TokenizedInput tokenized = 2; -} - -message EmbedResponse { - repeated float embedding = 1; - uint32 prompt_tokens = 2; - uint32 embedding_dim = 3; -} - -// ===================== -// Management Operations -// ===================== - -message HealthCheckRequest {} - -message HealthCheckResponse { - bool healthy = 1; - string message = 2; -} - -message AbortRequest { - repeated string request_ids = 1; -} - -message AbortResponse { -} - -// ===================== -// Model and Server Info -// ===================== - -message GetModelInfoRequest {} - -message GetModelInfoResponse { - string model_path = 1; - bool is_generation = 2; - uint32 max_context_length = 3; - uint32 vocab_size = 4; - bool supports_vision = 5; -} - -message GetServerInfoRequest {} - -message GetServerInfoResponse { - uint32 active_requests = 1; - bool is_paused = 2; - double last_receive_timestamp = 3; - double uptime_seconds = 4; - string server_type = 5; // "vllm-grpc" -} diff --git a/sgl-model-gateway/src/routers/grpc/client.rs b/sgl-model-gateway/src/routers/grpc/client.rs index eeac4b9bb..75b520acc 100644 --- a/sgl-model-gateway/src/routers/grpc/client.rs +++ b/sgl-model-gateway/src/routers/grpc/client.rs @@ -1,7 +1,10 @@ //! Unified gRPC client wrapper for SGLang and vLLM backends +use std::sync::Arc; + use crate::{ grpc_client::{SglangSchedulerClient, VllmEngineClient}, + observability::otel_trace::OtelTraceInjector, routers::grpc::proto_wrapper::{ ProtoEmbedRequest, ProtoEmbedResponse, ProtoGenerateRequest, ProtoStream, }, @@ -69,9 +72,14 @@ impl GrpcClient { url: &str, runtime_type: &str, ) -> Result> { + let trace_injector = Arc::new(OtelTraceInjector); match runtime_type { - "sglang" => Ok(Self::Sglang(SglangSchedulerClient::connect(url).await?)), - "vllm" => Ok(Self::Vllm(VllmEngineClient::connect(url).await?)), + "sglang" => Ok(Self::Sglang( + SglangSchedulerClient::connect_with_trace_injector(url, trace_injector).await?, + )), + "vllm" => Ok(Self::Vllm( + VllmEngineClient::connect_with_trace_injector(url, trace_injector).await?, + )), _ => Err(format!("Unknown runtime type: {}", runtime_type).into()), } }