From 53594a5570c70213703834e698b36fa0bb1f877c Mon Sep 17 00:00:00 2001 From: YzXiao101 Date: Sun, 31 May 2026 01:49:19 -0400 Subject: [PATCH] feat(grpc): bridge gRPC get_server_info Signed-off-by: YzXiao101 --- docs/reference/api/admin.md | 32 ++- model_gateway/src/routers/grpc/client.rs | 279 +++++++++++++++++++- model_gateway/src/routers/grpc/pd_router.rs | 84 +++++- model_gateway/src/routers/grpc/router.rs | 78 +++++- 4 files changed, 458 insertions(+), 15 deletions(-) diff --git a/docs/reference/api/admin.md b/docs/reference/api/admin.md index 9a6dbe9e57..9743da3bd9 100644 --- a/docs/reference/api/admin.md +++ b/docs/reference/api/admin.md @@ -382,14 +382,38 @@ Returns detailed model information (proxied to workers). GET /get_server_info ``` -Returns server information (proxied to workers). +Returns server information proxied from a backend worker. + +The response shape is backend-specific: + +- HTTP workers return the worker's native JSON response. +- gRPC-backed structured runtimes such as SGLang and TokenSpeed return the top-level fields from the backend `GetServerInfoResponse` protobuf. +- For those gRPC structured runtimes, `server_args` is intentionally filtered to a curated subset of high-signal fields, while `scheduler_info` is returned as-is. **Response:** `200 OK` + ```json { - "version": "0.1.0", - "backend": "vllm", - "gpu_count": 8 + "server_args": { + "model": "meta-llama/Llama-3.1-8B-Instruct", + "tokenizer": "meta-llama/Llama-3.1-8B-Instruct", + "served_model_name": "meta-llama/Llama-3.1-8B-Instruct", + "tp_size": 1, + "dp_size": 1 + }, + "scheduler_info": { + "max_total_num_tokens": 131072, + "max_req_input_len": 32768 + }, + "active_requests": 0, + "is_paused": false, + "uptime_seconds": 12.34, + "tokenspeed_version": "0.1.0", + "start_time": { + "seconds": 1748660000, + "nanos": 0 + }, + "max_total_num_tokens": 131072 } ``` diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index 465dc2e58f..1b86829701 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -6,6 +6,7 @@ use openai_protocol::{ chat::ChatCompletionRequest, completion::CompletionRequest, generate::GenerateRequest, messages::CreateMessageRequest, worker::WorkerLoadResponse, }; +use serde_json::Value; use smg_grpc_client::{ tokenizer_bundle, tokenizer_bundle::StreamBundle, MlxEngineClient, SglangSchedulerClient, TokenSpeedSchedulerClient, TrtllmServiceClient, VllmEngineClient, @@ -727,6 +728,161 @@ impl ServerInfo { } } } + + /// Convert to JSON. For the potentially very large `server_args` protobuf + /// `Struct`: filter it down to the curated `*_GRPC_KEYS` subset + pub fn to_public_json(&self) -> Value { + fn prost_number_to_json(number: f64) -> Value { + const MAX_SAFE_INTEGER_F64: f64 = 9_007_199_254_740_991.0; + const MIN_SAFE_INTEGER_F64: f64 = -9_007_199_254_740_991.0; + + if number.is_finite() + && number.fract() == 0.0 + && (MIN_SAFE_INTEGER_F64..=MAX_SAFE_INTEGER_F64).contains(&number) + { + return Value::from(number as i64); + } + + serde_json::Number::from_f64(number) + .map(Value::Number) + .unwrap_or(Value::Null) + } + + fn prost_value_to_json(value: &prost_types::Value) -> Value { + match value.kind.as_ref() { + Some(prost_types::value::Kind::NullValue(_)) | None => Value::Null, + Some(prost_types::value::Kind::NumberValue(number)) => prost_number_to_json(*number), + Some(prost_types::value::Kind::StringValue(text)) => Value::String(text.clone()), + Some(prost_types::value::Kind::BoolValue(flag)) => Value::Bool(*flag), + Some(prost_types::value::Kind::StructValue(struct_value)) => { + Value::Object(struct_to_json_map(struct_value)) + } + Some(prost_types::value::Kind::ListValue(list_value)) => Value::Array( + list_value + .values + .iter() + .map(prost_value_to_json) + .collect::>(), + ), + } + } + + fn struct_to_json_map(struct_value: &prost_types::Struct) -> serde_json::Map { + struct_value + .fields + .iter() + .map(|(key, value)| (key.clone(), prost_value_to_json(value))) + .collect() + } + + fn filter_struct_json( + struct_value: Option<&prost_types::Struct>, + keys: &[&str], + ) -> Value { + let Some(struct_value) = struct_value else { + return Value::Null; + }; + + let mut map = serde_json::Map::new(); + for key in keys { + if let Some(value) = struct_value.fields.get(*key) { + map.insert((*key).to_string(), prost_value_to_json(value)); + } + } + Value::Object(map) + } + + fn optional_struct_json(struct_value: Option<&prost_types::Struct>) -> Value { + struct_value + .map(|struct_value| Value::Object(struct_to_json_map(struct_value))) + .unwrap_or(Value::Null) + } + + fn optional_timestamp_json(timestamp: Option<&prost_types::Timestamp>) -> Value { + timestamp + .map(|timestamp| { + serde_json::json!({ + "seconds": timestamp.seconds, + "nanos": timestamp.nanos, + }) + }) + .unwrap_or(Value::Null) + } + + match self { + ServerInfo::Sglang(info) => { + let mut payload = serde_json::Map::new(); + payload.insert( + "server_args".to_string(), + filter_struct_json(info.server_args.as_ref(), SGLANG_GRPC_KEYS), + ); + payload.insert( + "scheduler_info".to_string(), + optional_struct_json(info.scheduler_info.as_ref()), + ); + payload.insert("active_requests".to_string(), Value::from(info.active_requests)); + payload.insert("is_paused".to_string(), Value::Bool(info.is_paused)); + if let Some(number) = serde_json::Number::from_f64(info.last_receive_timestamp) { + payload.insert("last_receive_timestamp".to_string(), Value::Number(number)); + } + if let Some(number) = serde_json::Number::from_f64(info.uptime_seconds) { + payload.insert("uptime_seconds".to_string(), Value::Number(number)); + } + payload.insert( + "max_total_num_tokens".to_string(), + Value::from(info.max_total_num_tokens), + ); + payload.insert( + "sglang_version".to_string(), + Value::String(info.sglang_version.clone()), + ); + payload.insert( + "server_type".to_string(), + Value::String(info.server_type.clone()), + ); + payload.insert( + "start_time".to_string(), + optional_timestamp_json(info.start_time.as_ref()), + ); + Value::Object(payload) + } + ServerInfo::Vllm(info) => serde_json::to_value(info) + .unwrap_or_else(|_| Value::Object(serde_json::Map::new())), + ServerInfo::Trtllm(info) => serde_json::to_value(info) + .unwrap_or_else(|_| Value::Object(serde_json::Map::new())), + ServerInfo::Mlx(info) => serde_json::to_value(info) + .unwrap_or_else(|_| Value::Object(serde_json::Map::new())), + ServerInfo::TokenSpeed(info) => { + let mut payload = serde_json::Map::new(); + payload.insert( + "server_args".to_string(), + filter_struct_json(info.server_args.as_ref(), TOKENSPEED_GRPC_KEYS), + ); + payload.insert( + "scheduler_info".to_string(), + optional_struct_json(info.scheduler_info.as_ref()), + ); + payload.insert("active_requests".to_string(), Value::from(info.active_requests)); + payload.insert("is_paused".to_string(), Value::Bool(info.is_paused)); + if let Some(number) = serde_json::Number::from_f64(info.uptime_seconds) { + payload.insert("uptime_seconds".to_string(), Value::Number(number)); + } + payload.insert( + "max_total_num_tokens".to_string(), + Value::from(info.max_total_num_tokens), + ); + payload.insert( + "tokenspeed_version".to_string(), + Value::String(info.tokenspeed_version.clone()), + ); + payload.insert( + "start_time".to_string(), + optional_timestamp_json(info.start_time.as_ref()), + ); + Value::Object(payload) + } + } + } } /// Keys worth extracting from SGLang gRPC `server_args` (which contains the full config). @@ -775,13 +931,13 @@ const TOKENSPEED_GRPC_KEYS: &[&str] = &[ /// (e.g. `is_generation == "false"` for embedding detection) work correctly. pub(crate) fn flat_labels(value: &T) -> HashMap { let mut labels = HashMap::new(); - if let Ok(serde_json::Value::Object(obj)) = serde_json::to_value(value) { + if let Ok(Value::Object(obj)) = serde_json::to_value(value) { for (key, val) in obj { match val { - serde_json::Value::String(s) if !s.is_empty() && s != "null" => { + Value::String(s) if !s.is_empty() && s != "null" => { labels.insert(key, s); } - serde_json::Value::Number(n) if n.as_f64().is_some_and(|v| v != 0.0) => { + Value::Number(n) if n.as_f64().is_some_and(|v| v != 0.0) => { // Format integers without decimal point let formatted = n .as_i64() @@ -789,10 +945,10 @@ pub(crate) fn flat_labels(value: &T) -> HashMap { + Value::Bool(b) => { labels.insert(key, b.to_string()); } - serde_json::Value::Array(arr) if !arr.is_empty() => { + Value::Array(arr) if !arr.is_empty() => { if let Ok(s) = serde_json::to_string(&arr) { labels.insert(key, s); } @@ -830,3 +986,116 @@ fn pick_prost_fields(labels: &mut HashMap, s: &prost_types::Stru } } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use smg_grpc_client::{sglang_proto, tokenspeed_proto}; + + use super::ServerInfo; + + fn string_value(value: &str) -> prost_types::Value { + prost_types::Value { + kind: Some(prost_types::value::Kind::StringValue(value.to_string())), + } + } + + fn number_value(value: f64) -> prost_types::Value { + prost_types::Value { + kind: Some(prost_types::value::Kind::NumberValue(value)), + } + } + + #[test] + fn to_public_json_sglang_keeps_pb_shape_and_filters_server_args() { + let server_info = ServerInfo::Sglang(Box::new(sglang_proto::GetServerInfoResponse { + server_args: Some(prost_types::Struct { + fields: BTreeMap::from([ + ("model_path".to_string(), string_value("Qwen/Qwen3-8B")), + ("tp_size".to_string(), number_value(2.0)), + ("host".to_string(), string_value("127.0.0.1")), + ]), + }), + scheduler_info: Some(prost_types::Struct { + fields: BTreeMap::from([ + ("max_total_num_tokens".to_string(), number_value(16384.0)), + ("policy".to_string(), string_value("lpm")), + ]), + }), + active_requests: 3, + is_paused: true, + last_receive_timestamp: 12.5, + uptime_seconds: 7.25, + sglang_version: "0.4.0".to_string(), + server_type: "grpc".to_string(), + start_time: Some(prost_types::Timestamp { + seconds: 100, + nanos: 5, + }), + max_total_num_tokens: 8192, + ..Default::default() + })); + + let json = server_info.to_public_json(); + + assert_eq!(json["server_args"]["model_path"], "Qwen/Qwen3-8B"); + assert_eq!(json["server_args"]["tp_size"], 2); + assert!(json["server_args"].get("host").is_none()); + assert_eq!(json["scheduler_info"]["policy"], "lpm"); + assert_eq!(json["scheduler_info"]["max_total_num_tokens"], 16384); + assert_eq!(json["active_requests"], 3); + assert_eq!(json["is_paused"], true); + assert_eq!(json["last_receive_timestamp"], 12.5); + assert_eq!(json["uptime_seconds"], 7.25); + assert_eq!(json["sglang_version"], "0.4.0"); + assert_eq!(json["server_type"], "grpc"); + assert_eq!(json["start_time"]["seconds"], 100); + assert_eq!(json["max_total_num_tokens"], 8192); + } + + #[test] + fn to_public_json_tokenspeed_keeps_pb_shape_and_filters_server_args() { + let server_info = + ServerInfo::TokenSpeed(Box::new(tokenspeed_proto::GetServerInfoResponse { + server_args: Some(prost_types::Struct { + fields: BTreeMap::from([ + ("model".to_string(), string_value("Qwen/Qwen3-8B")), + ("tokenizer".to_string(), string_value("tok-path")), + ("dp_size".to_string(), number_value(4.0)), + ("host".to_string(), string_value("127.0.0.1")), + ]), + }), + scheduler_info: Some(prost_types::Struct { + fields: BTreeMap::from([( + "max_total_num_tokens".to_string(), + number_value(32768.0), + )]), + }), + active_requests: 1, + is_paused: false, + uptime_seconds: 4.5, + max_total_num_tokens: 4096, + tokenspeed_version: "0.1.0".to_string(), + start_time: Some(prost_types::Timestamp { + seconds: 200, + nanos: 9, + }), + ..Default::default() + })); + + let json = server_info.to_public_json(); + + assert_eq!(json["server_args"]["model"], "Qwen/Qwen3-8B"); + assert_eq!(json["server_args"]["tokenizer"], "tok-path"); + assert_eq!(json["server_args"]["dp_size"], 4); + assert!(json["server_args"].get("host").is_none()); + assert_eq!(json["scheduler_info"]["max_total_num_tokens"], 32768); + assert_eq!(json["active_requests"], 1); + assert_eq!(json["is_paused"], false); + assert_eq!(json["uptime_seconds"], 4.5); + assert_eq!(json["tokenspeed_version"], "0.1.0"); + assert_eq!(json["start_time"]["seconds"], 200); + assert_eq!(json["max_total_num_tokens"], 4096); + } +} diff --git a/model_gateway/src/routers/grpc/pd_router.rs b/model_gateway/src/routers/grpc/pd_router.rs index ab96d4101c..adc0f26bec 100644 --- a/model_gateway/src/routers/grpc/pd_router.rs +++ b/model_gateway/src/routers/grpc/pd_router.rs @@ -1,12 +1,18 @@ use std::sync::Arc; use async_trait::async_trait; -use axum::{http::HeaderMap, response::Response}; +use axum::{ + body::Body, + extract::Request, + http::HeaderMap, + response::{IntoResponse, Response}, + Json, +}; use openai_protocol::{ chat::ChatCompletionRequest, completion::CompletionRequest, generate::GenerateRequest, messages::CreateMessageRequest, }; -use tracing::debug; +use tracing::{debug, error}; use super::{ context::SharedComponents, multimodal::MultimodalComponents, pipeline::RequestPipeline, @@ -18,9 +24,11 @@ use crate::{ observability::metrics::{metrics_labels, Metrics}, routers::{ common::retry::{is_retryable_status, RetryExecutor}, + error, + grpc::utils::tonic_ext::TonicStatusExt, RouterTrait, }, - worker::{ConnectionMode, WorkerRegistry, WorkerType}, + worker::{ConnectionMode, Worker, WorkerRegistry, WorkerType}, }; /// gRPC PD (Prefill-Decode) router implementation for SGLang @@ -107,6 +115,72 @@ impl GrpcPDRouter { }) } + fn select_first_prefill_worker(&self) -> Option> { + self.worker_registry + .get_workers_filtered( + None, + Some(WorkerType::Prefill), + Some(ConnectionMode::Grpc), + None, + true, + ) + .into_iter() + .next() + } + + async fn get_server_info_from_first_prefill_worker(&self) -> Response { + let Some(worker) = self.select_first_prefill_worker() else { + return error::service_unavailable( + "no_prefill_servers", + "No prefill servers available", + ); + }; + + let client = match worker.get_grpc_client().await { + Ok(Some(client)) => client, + Ok(None) => { + error!( + function = "GrpcPDRouter::get_server_info_from_first_prefill_worker", + worker = %worker.url(), + "Selected prefill worker is not configured for gRPC", + ); + return error::internal_error( + "worker_not_configured_for_grpc", + "Selected worker is not configured for gRPC", + ); + } + Err(err) => { + error!( + function = "GrpcPDRouter::get_server_info_from_first_prefill_worker", + worker = %worker.url(), + error = %err, + "Failed to acquire gRPC client for prefill worker", + ); + return error::internal_error( + "get_grpc_client_failed", + format!("Failed to get gRPC client: {err}"), + ); + } + }; + + match client.get_server_info().await { + Ok(server_info) => Json(server_info.to_public_json()).into_response(), + Err(err) => { + error!( + function = "GrpcPDRouter::get_server_info_from_first_prefill_worker", + worker = %worker.url(), + grpc_code = ?err.code(), + error = %err, + "Failed to fetch server info from prefill worker", + ); + err.to_http_error( + "get_server_info_failed", + format!("Failed to get server info: {}", err.message()), + ) + } + } + } + /// Main route_generate implementation with PD dual dispatch async fn route_generate_impl( &self, @@ -404,6 +478,10 @@ impl RouterTrait for GrpcPDRouter { self } + async fn get_server_info(&self, _req: Request) -> Response { + self.get_server_info_from_first_prefill_worker().await + } + async fn route_generate( &self, headers: Option<&HeaderMap>, diff --git a/model_gateway/src/routers/grpc/router.rs b/model_gateway/src/routers/grpc/router.rs index c1cb982b50..f1ce7157c9 100644 --- a/model_gateway/src/routers/grpc/router.rs +++ b/model_gateway/src/routers/grpc/router.rs @@ -2,15 +2,18 @@ use std::sync::Arc; use async_trait::async_trait; use axum::{ + body::Body, + extract::Request, http::HeaderMap, response::{IntoResponse, Response}, + Json, }; use openai_protocol::{ chat::ChatCompletionRequest, classify::ClassifyRequest, completion::CompletionRequest, embedding::EmbeddingRequest, generate::GenerateRequest, messages::CreateMessageRequest, responses::ResponsesRequest, }; -use tracing::debug; +use tracing::{debug, error}; use super::{ common::responses::{ @@ -29,9 +32,11 @@ use crate::{ observability::metrics::{metrics_labels, Metrics}, routers::{ common::retry::{is_retryable_status, RetryExecutor}, + error, + grpc::utils::tonic_ext::TonicStatusExt, RouterTrait, }, - worker::WorkerRegistry, + worker::{ConnectionMode, Worker, WorkerRegistry, WorkerType}, }; /// gRPC router implementation for SGLang @@ -174,6 +179,69 @@ impl GrpcRouter { }) } + fn select_first_worker(&self) -> Option> { + self.worker_registry + .get_workers_filtered( + None, + Some(WorkerType::Regular), + Some(ConnectionMode::Grpc), + None, + true, + ) + .into_iter() + .next() + } + + async fn get_server_info_impl(&self) -> Response { + let Some(worker) = self.select_first_worker() else { + return error::service_unavailable("no_workers", "No workers are available"); + }; + + let client = match worker.get_grpc_client().await { + Ok(Some(client)) => client, + Ok(None) => { + error!( + function = "GrpcRouter::get_server_info_impl", + worker = %worker.url(), + "Selected worker is not configured for gRPC", + ); + return error::internal_error( + "worker_not_configured_for_grpc", + "Selected worker is not configured for gRPC", + ); + } + Err(err) => { + error!( + function = "GrpcRouter::get_server_info_impl", + worker = %worker.url(), + error = %err, + "Failed to acquire gRPC client for worker", + ); + return error::internal_error( + "get_grpc_client_failed", + format!("Failed to get gRPC client: {err}"), + ); + } + }; + + match client.get_server_info().await { + Ok(server_info) => Json(server_info.to_public_json()).into_response(), + Err(err) => { + error!( + function = "GrpcRouter::get_server_info_impl", + worker = %worker.url(), + grpc_code = ?err.code(), + error = %err, + "Failed to fetch server info from worker", + ); + err.to_http_error( + "get_server_info_failed", + format!("Failed to get server info: {}", err.message()), + ) + } + } + } + /// Main route_chat implementation async fn route_chat_impl( &self, @@ -354,7 +422,7 @@ impl GrpcRouter { } else { match serve_harmony_responses(&harmony_ctx, body.clone(), tenant_meta.clone()).await { - Ok(response) => axum::Json(response).into_response(), + Ok(response) => Json(response).into_response(), Err(error_response) => error_response, } } @@ -544,6 +612,10 @@ impl RouterTrait for GrpcRouter { self } + async fn get_server_info(&self, _req: Request) -> Response { + self.get_server_info_impl().await + } + async fn route_generate( &self, headers: Option<&HeaderMap>,