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
sgl-model-gateway/src/core/worker_registry.rs
+2 -6
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.
sgl-model-gateway/src/protocols/worker_spec.rs
+26
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 {
sgl-model-gateway/src/routers/parse/handlers.rs
+74 -98
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),
)
}
}
}
sgl-model-gateway/src/server.rs
+21 -52
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,
}