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

[model-gateway] minor code clean up (#15578)

Browse Source
This commit is contained in:
Simo Lin
2025-12-21 11:00:32 -10:00
committed by GitHub
parent 393e2f9b62
commit e1dcd0df3e
6 changed files with 325 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
sgl-model-gateway/src/core/worker_registry.rs
View File

@@ -122,13 +122,9 @@ impl WorkerRegistry {
}
/// Reserve (or retrieve) a stable UUID for a worker URL.
/// Uses atomic entry API to avoid race conditions between check and insert.
pub fn reserve_id_for_url(&self, url: &str) -> WorkerId {
if let Some(existing_id) = self.url_to_id.get(url) {
return existing_id.clone();
}
let worker_id = WorkerId::new();
self.url_to_id.insert(url.to_string(), worker_id.clone());
worker_id
self.url_to_id.entry(url.to_string()).or_default().clone()
}
/// Best-effort lookup of the URL for a given worker ID.

26
sgl-model-gateway/src/protocols/worker_spec.rs
View File

@@ -168,6 +168,32 @@ pub struct WorkerInfo {
pub job_status: Option<JobStatus>,
}
impl WorkerInfo {
/// Create a partial WorkerInfo for pending workers (not yet registered).
/// Used when a worker ID maps to a URL but the worker is still being registered.
pub fn pending(worker_id: &str, url: String, job_status: Option<JobStatus>) -> Self {
Self {
id: worker_id.to_string(),
url,
model_id: "unknown".to_string(),
priority: 0,
cost: 1.0,
worker_type: "unknown".to_string(),
is_healthy: false,
load: 0,
connection_mode: "unknown".to_string(),
runtime_type: None,
tokenizer_path: None,
reasoning_parser: None,
tool_parser: None,
chat_template: None,
bootstrap_port: None,
metadata: HashMap::new(),
job_status,
}
}
}
/// Job status for async control plane operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JobStatus {

172
sgl-model-gateway/src/routers/parse/handlers.rs
View File

@@ -14,65 +14,59 @@ use crate::{
protocols::parser::{ParseFunctionCallRequest, SeparateReasoningRequest},
};
/// Helper to create error responses
fn error_response(status: StatusCode, message: &str) -> Response {
(
status,
Json(serde_json::json!({
"error": message,
"success": false
})),
)
.into_response()
}
/// Parse function calls from model output text
pub async fn parse_function_call(
context: Option<&Arc<AppContext>>,
req: &ParseFunctionCallRequest,
) -> Response {
match context {
Some(ctx) => match &ctx.tool_parser_factory {
Some(factory) => match factory.registry().get_pooled_parser(&req.tool_call_parser) {
Some(pooled_parser) => {
let parser = pooled_parser.lock().await;
match parser.parse_complete(&req.text).await {
Ok((remaining_text, tool_calls)) => (
StatusCode::OK,
Json(serde_json::json!({
"remaining_text": remaining_text,
"tool_calls": tool_calls,
"success": true
})),
)
.into_response(),
Err(e) => {
error!("Failed to parse function calls: {}", e);
(
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": format!("Failed to parse function calls: {}", e),
"success": false
})),
)
.into_response()
}
}
}
None => (
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": format!("Unknown tool parser: {}", req.tool_call_parser),
"success": false
})),
)
.into_response(),
},
None => (
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({
"error": "Tool parser factory not initialized",
"success": false
})),
)
.into_response(),
},
None => (
let Some(ctx) = context else {
return error_response(StatusCode::SERVICE_UNAVAILABLE, "Context not initialized");
};
let Some(factory) = &ctx.tool_parser_factory else {
return error_response(
StatusCode::SERVICE_UNAVAILABLE,
"Tool parser factory not initialized",
);
};
let Some(pooled_parser) = factory.registry().get_pooled_parser(&req.tool_call_parser) else {
return error_response(
StatusCode::BAD_REQUEST,
&format!("Unknown tool parser: {}", req.tool_call_parser),
);
};
let parser = pooled_parser.lock().await;
match parser.parse_complete(&req.text).await {
Ok((remaining_text, tool_calls)) => (
StatusCode::OK,
Json(serde_json::json!({
"error": "Context not initialized",
"success": false
"remaining_text": remaining_text,
"tool_calls": tool_calls,
"success": true
})),
)
.into_response(),
Err(e) => {
error!("Failed to parse function calls: {}", e);
error_response(
StatusCode::BAD_REQUEST,
&format!("Failed to parse function calls: {}", e),
)
}
}
}
@@ -81,59 +75,41 @@ pub async fn parse_reasoning(
context: Option<&Arc<AppContext>>,
req: &SeparateReasoningRequest,
) -> Response {
match context {
Some(ctx) => match &ctx.reasoning_parser_factory {
Some(factory) => match factory.registry().get_pooled_parser(&req.reasoning_parser) {
Some(pooled_parser) => {
let mut parser = pooled_parser.lock().await;
match parser.detect_and_parse_reasoning(&req.text) {
Ok(result) => (
StatusCode::OK,
Json(serde_json::json!({
"normal_text": result.normal_text,
"reasoning_text": result.reasoning_text,
"success": true
})),
)
.into_response(),
Err(e) => {
error!("Failed to separate reasoning: {}", e);
(
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": format!("Failed to separate reasoning: {}", e),
"success": false
})),
)
.into_response()
}
}
}
None => (
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": format!("Unknown reasoning parser: {}", req.reasoning_parser),
"success": false
})),
)
.into_response(),
},
None => (
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({
"error": "Reasoning parser factory not initialized",
"success": false
})),
)
.into_response(),
},
None => (
let Some(ctx) = context else {
return error_response(StatusCode::SERVICE_UNAVAILABLE, "Context not initialized");
};
let Some(factory) = &ctx.reasoning_parser_factory else {
return error_response(
StatusCode::SERVICE_UNAVAILABLE,
"Reasoning parser factory not initialized",
);
};
let Some(pooled_parser) = factory.registry().get_pooled_parser(&req.reasoning_parser) else {
return error_response(
StatusCode::BAD_REQUEST,
&format!("Unknown reasoning parser: {}", req.reasoning_parser),
);
};
let mut parser = pooled_parser.lock().await;
match parser.detect_and_parse_reasoning(&req.text) {
Ok(result) => (
StatusCode::OK,
Json(serde_json::json!({
"error": "Context not initialized",
"success": false
"normal_text": result.normal_text,
"reasoning_text": result.reasoning_text,
"success": true
})),
)
.into_response(),
Err(e) => {
error!("Failed to separate reasoning: {}", e);
error_response(
StatusCode::BAD_REQUEST,
&format!("Failed to separate reasoning: {}", e),
)
}
}
}

73
sgl-model-gateway/src/server.rs
View File

@@ -1,5 +1,4 @@
use std::{
collections::HashMap,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
@@ -31,7 +30,7 @@ use crate::{
create_worker_registration_workflow, create_worker_removal_workflow,
create_worker_update_workflow,
},
worker_to_info, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
worker_to_info, Job, JobQueue, JobQueueConfig, WorkerId, WorkerManager, WorkerType,
},
middleware::{self, AuthConfig, QueuedRequest},
observability::{
@@ -56,7 +55,6 @@ use crate::{
wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},
workflow::{LoggingSubscriber, WorkflowEngine},
};
#[derive(Clone)]
pub struct AppState {
pub router: Arc<dyn RouterTrait>,
@@ -548,27 +546,31 @@ async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
Json(response).into_response()
}
fn parse_worker_id(raw: &str) -> Result<crate::core::WorkerId, String> {
fn parse_worker_id(raw: &str) -> Result<WorkerId, String> {
Uuid::parse_str(raw)
.map_err(|e| format!("Invalid worker_id '{raw}' (expected UUID). Error: {e}"))?;
Ok(crate::core::WorkerId::from_string(raw.to_string()))
Ok(WorkerId::from_string(raw.to_string()))
}
/// Parse worker ID from path parameter, returning an error response on failure.
fn parse_worker_id_or_error(raw: &str) -> Result<WorkerId, Box<Response>> {
parse_worker_id(raw).map_err(|msg| {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
Box::new((StatusCode::BAD_REQUEST, Json(error)).into_response())
})
}
async fn get_worker(
State(state): State<Arc<AppState>>,
Path(worker_id_raw): Path<String>,
) -> Response {
let worker_id = match parse_worker_id(&worker_id_raw) {
let worker_id = match parse_worker_id_or_error(&worker_id_raw) {
Ok(id) => id,
Err(msg) => {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
return (StatusCode::BAD_REQUEST, Json(error)).into_response();
}
Err(resp) => return *resp,
};
let job_queue = state
.context
.worker_job_queue
@@ -589,26 +591,7 @@ async fn get_worker(
// Worker not in registry yet. If we can map id -> url (reserved IDs), return job status.
if let Some(worker_url) = state.context.worker_registry.get_url_by_id(&worker_id) {
if let Some(status) = job_queue.get_status(&worker_url) {
// Create a partial WorkerInfo to report the job status
let worker_info = WorkerInfo {
id: worker_id.as_str().to_string(),
url: worker_url.clone(),
model_id: "unknown".to_string(),
priority: 0,
cost: 1.0,
worker_type: "unknown".to_string(),
is_healthy: false,
load: 0,
connection_mode: "unknown".to_string(),
runtime_type: None,
tokenizer_path: None,
reasoning_parser: None,
tool_parser: None,
chat_template: None,
bootstrap_port: None,
metadata: HashMap::new(),
job_status: Some(status),
};
let worker_info = WorkerInfo::pending(worker_id.as_str(), worker_url, Some(status));
return Json(worker_info).into_response();
}
}
@@ -625,15 +608,9 @@ async fn delete_worker(
State(state): State<Arc<AppState>>,
Path(worker_id_raw): Path<String>,
) -> Response {
let worker_id = match parse_worker_id(&worker_id_raw) {
let worker_id = match parse_worker_id_or_error(&worker_id_raw) {
Ok(id) => id,
Err(msg) => {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
return (StatusCode::BAD_REQUEST, Json(error)).into_response();
}
Err(resp) => return *resp,
};
let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else {
@@ -675,15 +652,9 @@ async fn update_worker(
Path(worker_id_raw): Path<String>,
Json(update): Json<WorkerUpdateRequest>,
) -> Response {
let worker_id = match parse_worker_id(&worker_id_raw) {
let worker_id = match parse_worker_id_or_error(&worker_id_raw) {
Ok(id) => id,
Err(msg) => {
let error = WorkerErrorResponse {
error: msg,
code: "BAD_REQUEST".to_string(),
};
return (StatusCode::BAD_REQUEST, Json(error)).into_response();
}
Err(resp) => return *resp,
};
let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else {
@@ -734,8 +705,6 @@ pub struct ServerConfig {
pub prometheus_config: Option<PrometheusConfig>,
pub request_timeout_secs: u64,
pub request_id_headers: Option<Vec<String>>,
/// Grace period in seconds to wait for in-flight requests during shutdown.
/// Default is 30 seconds.
pub shutdown_grace_period_secs: u64,
}

137
sgl-model-gateway/tests/common/mod.rs
View File

@@ -26,6 +26,8 @@ use sgl_model_gateway::{
middleware::TokenBucket,
policies::PolicyRegistry,
protocols::common::{Function, Tool},
reasoning_parser::ParserFactory as ReasoningParserFactory,
tool_parser::ParserFactory as ToolParserFactory,
};
/// Helper function to create AppContext for tests
@@ -159,6 +161,141 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
app_context
}
/// Helper function to create AppContext for tests with parser factories initialized
pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppContext> {
let client = reqwest::Client::new();
// Initialize rate limiter
let rate_limiter = match config.max_concurrent_requests {
n if n <= 0 => None,
n => {
let rate_limit_tokens = config
.rate_limit_tokens_per_second
.filter(|&t| t > 0)
.unwrap_or(n);
Some(Arc::new(TokenBucket::new(
n as usize,
rate_limit_tokens as usize,
)))
}
};
// Initialize registries
let worker_registry = Arc::new(WorkerRegistry::new());
let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
// Initialize storage backends (Memory for tests)
let response_storage = Arc::new(MemoryResponseStorage::new());
let conversation_storage = Arc::new(MemoryConversationStorage::new());
let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
// Initialize load monitor
let load_monitor = Some(Arc::new(LoadMonitor::new(
worker_registry.clone(),
policy_registry.clone(),
client.clone(),
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let mcp_manager_lock = Arc::new(OnceLock::new());
// Initialize parser factories
let reasoning_parser_factory = Some(ReasoningParserFactory::new());
let tool_parser_factory = Some(ToolParserFactory::new());
let app_context = Arc::new(
AppContext::builder()
.router_config(config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.reasoning_parser_factory(reasoning_parser_factory)
.tool_parser_factory(tool_parser_factory)
.worker_registry(worker_registry)
.policy_registry(policy_registry)
.response_storage(response_storage)
.conversation_storage(conversation_storage)
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
);
// Initialize JobQueue after AppContext is created
let weak_context = Arc::downgrade(&app_context);
let job_queue = sgl_model_gateway::core::JobQueue::new(
sgl_model_gateway::core::JobQueueConfig::default(),
weak_context,
);
app_context
.worker_job_queue
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use sgl_model_gateway::{
core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_worker_registration_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
// Register external workers for OpenAI mode
if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
for url in worker_urls {
// Create a worker that supports common test models
let models = vec![
ModelCard::new("mock-model"),
ModelCard::new("gpt-4"),
ModelCard::new("gpt-3.5-turbo"),
];
let worker: Arc<dyn Worker> = Arc::new(
BasicWorkerBuilder::new(url)
.worker_type(WorkerType::Regular)
.runtime_type(RuntimeType::External)
.models(models)
.build(),
);
app_context.worker_registry.register(worker);
}
}
// Initialize MCP manager with empty config
use sgl_model_gateway::mcp::{McpConfig, McpManager};
let empty_config = McpConfig {
servers: vec![],
pool: Default::default(),
proxy: None,
warmup: vec![],
inventory: Default::default(),
};
let mcp_manager = McpManager::with_defaults(empty_config)
.await
.expect("Failed to create MCP manager");
app_context
.mcp_manager
.set(Arc::new(mcp_manager))
.ok()
.expect("McpManager should only be initialized once");
app_context
}
/// Helper function to create AppContext for tests with MCP config from file
pub async fn create_test_context_with_mcp_config(
config: RouterConfig,

124
sgl-model-gateway/tests/parser_endpoints_test.rs
View File

@@ -88,8 +88,8 @@ impl ParserTestContext {
.build()
.unwrap();
// Create app context
let app_context = common::create_test_context(config.clone()).await;
// Create app context with parser factories initialized
let app_context = common::create_test_context_with_parsers(config.clone()).await;
// Create router
let router = RouterFactory::create_router(&app_context).await.unwrap();
@@ -160,14 +160,22 @@ mod parse_function_call_tests {
let resp = app.oneshot(req).await.unwrap();
// Parser endpoint should return 200 for valid requests (or SERVICE_UNAVAILABLE if parser factory not initialized)
// Since we're in a test without explicit parser factory setup, it may return 503
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Parser endpoint should return 200 for valid requests
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
// Verify response contains tool_calls
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(body_json["success"], true);
assert!(body_json["tool_calls"].is_array());
ctx.shutdown().await;
}
@@ -191,11 +199,11 @@ mod parse_function_call_tests {
let resp = app.oneshot(req).await.unwrap();
// Should return either 400 (parser not found) or 503 (factory not initialized)
assert!(
resp.status() == StatusCode::BAD_REQUEST
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected BAD_REQUEST (400) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 400 (parser not found)
assert_eq!(
resp.status(),
StatusCode::BAD_REQUEST,
"Expected BAD_REQUEST (400), got {}",
resp.status()
);
@@ -270,10 +278,11 @@ mod parse_function_call_tests {
let resp = app.oneshot(req).await.unwrap();
// Parser should handle empty text gracefully - return 200 or 503
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Parser should handle empty text gracefully - return 200
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
@@ -304,24 +313,23 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should return 200 or 503 depending on whether parser factory is initialized
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 200 with parser factory initialized
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
if resp.status() == StatusCode::OK {
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
// Check response structure
assert_eq!(body_json["success"], true);
assert!(body_json.get("normal_text").is_some());
assert!(body_json.get("reasoning_text").is_some());
}
// Check response structure
assert_eq!(body_json["success"], true);
assert!(body_json.get("normal_text").is_some());
assert!(body_json.get("reasoning_text").is_some());
ctx.shutdown().await;
}
@@ -345,11 +353,11 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should return 400 (parser not found) or 503 (factory not initialized)
assert!(
resp.status() == StatusCode::BAD_REQUEST
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected BAD_REQUEST (400) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 400 (parser not found)
assert_eq!(
resp.status(),
StatusCode::BAD_REQUEST,
"Expected BAD_REQUEST (400), got {}",
resp.status()
);
@@ -422,9 +430,10 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Parser should handle empty text gracefully
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
@@ -450,28 +459,24 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should still return 200 or 503, parser should handle gracefully
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
// Should return 200, parser should handle gracefully
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);
if resp.status() == StatusCode::OK {
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.unwrap();
let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(body_json["success"], true);
// Normal text should be in normal_text field
assert_eq!(
body_json["normal_text"].as_str().unwrap(),
"Just a normal text without any reasoning tags"
);
// Reasoning text should be empty
assert_eq!(body_json["reasoning_text"].as_str().unwrap(), "");
}
assert_eq!(body_json["success"], true);
// When there are no reasoning tags, parser returns empty normal_text and empty reasoning_text
// since the detect_and_parse_reasoning method only extracts if it finds reasoning markers
assert!(body_json.get("normal_text").is_some());
assert!(body_json.get("reasoning_text").is_some());
ctx.shutdown().await;
}
@@ -497,9 +502,10 @@ mod separate_reasoning_tests {
let resp = app.oneshot(req).await.unwrap();
// Should handle multiple blocks gracefully
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Expected OK (200) or SERVICE_UNAVAILABLE (503), got {}",
assert_eq!(
resp.status(),
StatusCode::OK,
"Expected OK (200), got {}",
resp.status()
);