chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

refactor: consolidate gRPC client into shared crate dependency (#18730)

Browse Source
This commit is contained in:
Simo Lin
2026-02-12 06:43:24 -08:00
committed by GitHub
parent 1b8f68af57
commit 2f283d8152
10 changed files with 29 additions and 2362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sgl-model-gateway/Cargo.toml
View File

@@ -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"

16
sgl-model-gateway/build.rs
View File

@@ -12,24 +12,8 @@ macro_rules! set_env {
fn main() -> Result<(), Box<dyn std::error::Error>> {
// 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());

7
sgl-model-gateway/src/grpc_client/mod.rs
View File

@@ -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};

827
sgl-model-gateway/src/grpc_client/sglang_scheduler.rs
View File

@@ -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<GenerateResponse> 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<proto::GenerateResponse>,
request_id: String,
client: SglangSchedulerClient,
aborted: Arc<AtomicBool>,
}
impl AbortOnDropStream {
/// Create a new auto-aborting stream wrapper
pub fn new(
stream: Streaming<proto::GenerateResponse>,
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<proto::GenerateResponse, tonic::Status>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
// 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<Channel>,
}
impl SglangSchedulerClient {
/// Create a new client and connect to the scheduler
pub async fn connect(endpoint: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
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<AbortOnDropStream, Box<dyn std::error::Error + Send + Sync>> {
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<proto::EmbedResponse, Box<dyn std::error::Error + Send + Sync>> {
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<proto::HealthCheckResponse, Box<dyn std::error::Error + Send + Sync>> {
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<dyn std::error::Error + Send + Sync>> {
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<proto::GetModelInfoResponse, Box<dyn std::error::Error + Send + Sync>> {
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<proto::GetServerInfoResponse, Box<dyn std::error::Error + Send + Sync>> {
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<String>,
token_ids: Vec<u32>,
log_metrics: Option<bool>,
) -> 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<u32>,
multimodal_inputs: Option<proto::MultimodalInputs>,
tool_call_constraint: Option<(String, String)>, // (constraint_type, constraint_value)
) -> Result<proto::GenerateRequest, String> {
// 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<String>,
token_ids: Vec<u32>,
) -> Result<proto::GenerateRequest, String> {
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<u32>,
harmony_stop_ids: Option<Vec<u32>>,
constraint: Option<(String, String)>,
) -> Result<proto::GenerateRequest, String> {
// 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<proto::SamplingParams, String> {
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<String> {
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<Option<proto::sampling_params::Constraint>, 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<proto::SamplingParams, String> {
// 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<Option<proto::sampling_params::Constraint>, 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<Option<proto::sampling_params::Constraint>, String> {
let mut constraints = Vec::new();
if let Some(json_schema) = &params.json_schema {
constraints.push(proto::sampling_params::Constraint::JsonSchema(
json_schema.clone(),
));
}
if let Some(regex) = &params.regex {
constraints.push(proto::sampling_params::Constraint::Regex(regex.clone()));
}
if let Some(ebnf) = &params.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<proto::SamplingParams, String> {
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!["</s>".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!["</s>"]);
}
#[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
}

747
sgl-model-gateway/src/grpc_client/vllm_engine.rs
View File

@@ -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<GenerateResponse> 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<proto::GenerateResponse>,
request_id: String,
client: VllmEngineClient,
aborted: Arc<AtomicBool>,
}
impl AbortOnDropStream {
/// Create a new auto-aborting stream wrapper
pub fn new(
stream: Streaming<proto::GenerateResponse>,
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<proto::GenerateResponse, tonic::Status>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
// 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<Channel>,
}
impl VllmEngineClient {
/// Create a new client and connect to the vLLM server
pub async fn connect(endpoint: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
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<AbortOnDropStream, Box<dyn std::error::Error + Send + Sync>> {
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<proto::HealthCheckResponse, Box<dyn std::error::Error + Send + Sync>> {
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<dyn std::error::Error + Send + Sync>> {
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<proto::GetModelInfoResponse, Box<dyn std::error::Error + Send + Sync>> {
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<proto::GetServerInfoResponse, Box<dyn std::error::Error + Send + Sync>> {
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<u32>,
tool_call_constraint: Option<(String, String)>, // (constraint_type, constraint_value)
) -> Result<proto::GenerateRequest, String> {
// 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<String>,
token_ids: Vec<u32>,
) -> Result<proto::GenerateRequest, String> {
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<u32>,
harmony_stop_ids: Option<Vec<u32>>,
constraint: Option<(String, String)>,
) -> Result<proto::GenerateRequest, String> {
// 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<proto::SamplingParams, String> {
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<String> {
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<Option<proto::sampling_params::Constraint>, 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<proto::SamplingParams, String> {
// 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<Option<proto::sampling_params::Constraint>, 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<Option<proto::sampling_params::Constraint>, String> {
let mut constraints = Vec::new();
if let Some(json_schema) = &params.json_schema {
constraints.push(proto::sampling_params::Constraint::JsonSchema(
json_schema.clone(),
));
}
if let Some(regex) = &params.regex {
constraints.push(proto::sampling_params::Constraint::Regex(regex.clone()));
}
if let Some(ebnf) = &params.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<proto::SamplingParams, String> {
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!["</s>".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!["</s>"]);
}
#[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
}

2
sgl-model-gateway/src/lib.rs
View File

@@ -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;

17
sgl-model-gateway/src/observability/otel_trace.rs
View File

@@ -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<dyn std::error::Error + Send + Sync>> {
inject_trace_context_grpc(metadata);
Ok(())
}
}

566
sgl-model-gateway/src/proto/sglang_scheduler.proto
View File

@@ -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<string, float> 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;
}

195
sgl-model-gateway/src/proto/vllm_engine.proto
View File

@@ -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<int32, float> 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"
}

12
sgl-model-gateway/src/routers/grpc/client.rs
View File

@@ -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<Self, Box<dyn std::error::Error + Send + Sync>> {
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()),
}
}