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

[model-gateway] add WorkerService abstraction for worker business logic (#15580)

Browse Source
This commit is contained in:
Simo Lin
2025-12-21 12:35:35 -10:00
committed by GitHub
parent 254de6d2fd
commit 1354063a48
5 changed files with 421 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
sgl-model-gateway/src/app_context.rs
View File

@@ -8,7 +8,7 @@ use tracing::debug;
use crate::{
config::RouterConfig,
core::{ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry},
core::{ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry, WorkerService},
data_connector::{
create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage,
},
@@ -60,6 +60,7 @@ pub struct AppContext {
pub workflow_engine: Arc<OnceLock<Arc<WorkflowEngine>>>,
pub mcp_manager: Arc<OnceLock<Arc<McpManager>>>,
pub wasm_manager: Option<Arc<WasmModuleManager>>,
pub worker_service: Arc<WorkerService>,
}
pub struct AppContextBuilder {
@@ -224,6 +225,20 @@ impl AppContextBuilder {
let configured_reasoning_parser = router_config.reasoning_parser.clone();
let configured_tool_parser = router_config.tool_call_parser.clone();
let worker_registry = self
.worker_registry
.ok_or(AppContextBuildError("worker_registry"))?;
let worker_job_queue = self
.worker_job_queue
.ok_or(AppContextBuildError("worker_job_queue"))?;
// Create WorkerService from the already-built components
let worker_service = Arc::new(WorkerService::new(
worker_registry.clone(),
worker_job_queue.clone(),
router_config.clone(),
));
Ok(AppContext {
client: self.client.ok_or(AppContextBuildError("client"))?,
router_config,
@@ -231,9 +246,7 @@ impl AppContextBuilder {
tokenizer: self.tokenizer,
reasoning_parser_factory: self.reasoning_parser_factory,
tool_parser_factory: self.tool_parser_factory,
worker_registry: self
.worker_registry
.ok_or(AppContextBuildError("worker_registry"))?,
worker_registry,
policy_registry: self
.policy_registry
.ok_or(AppContextBuildError("policy_registry"))?,
@@ -250,9 +263,7 @@ impl AppContextBuilder {
load_monitor: self.load_monitor,
configured_reasoning_parser,
configured_tool_parser,
worker_job_queue: self
.worker_job_queue
.ok_or(AppContextBuildError("worker_job_queue"))?,
worker_job_queue,
workflow_engine: self
.workflow_engine
.ok_or(AppContextBuildError("workflow_engine"))?,
@@ -260,6 +271,7 @@ impl AppContextBuilder {
.mcp_manager
.ok_or(AppContextBuildError("mcp_manager"))?,
wasm_manager: self.wasm_manager,
worker_service,
})
}

2
sgl-model-gateway/src/core/mod.rs
View File

@@ -22,6 +22,7 @@ pub mod worker;
pub mod worker_builder;
pub mod worker_manager;
pub mod worker_registry;
pub mod worker_service;
pub use circuit_breaker::{
CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState,
@@ -38,3 +39,4 @@ pub use worker::{
pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder};
pub use worker_manager::{LoadMonitor, WorkerManager};
pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats};
pub use worker_service::{WorkerService, WorkerServiceError};

366
sgl-model-gateway/src/core/worker_service.rs Normal file
View File

@@ -0,0 +1,366 @@
//! Worker Service - Business logic layer for worker operations
//!
//! This module provides a clean separation between HTTP concerns (in routers)
//! and business logic for worker management. The service orchestrates
//! WorkerRegistry and JobQueue operations.
use std::sync::Arc;
use axum::{
http::StatusCode,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use tracing::warn;
use crate::{
config::RouterConfig,
core::{worker_to_info, Job, JobQueue, WorkerId, WorkerRegistry},
protocols::worker_spec::{
WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest,
},
};
/// Error types for worker service operations
#[derive(Debug)]
pub enum WorkerServiceError {
/// Worker with given ID was not found
NotFound { worker_id: String },
/// Invalid worker ID format (expected UUID)
InvalidId { raw: String, message: String },
/// Job queue not initialized
QueueNotInitialized,
/// Failed to submit job to queue
QueueSubmitFailed { message: String },
}
impl WorkerServiceError {
pub fn error_code(&self) -> &'static str {
match self {
Self::NotFound { .. } => "WORKER_NOT_FOUND",
Self::InvalidId { .. } => "BAD_REQUEST",
Self::QueueNotInitialized => "INTERNAL_SERVER_ERROR",
Self::QueueSubmitFailed { .. } => "INTERNAL_SERVER_ERROR",
}
}
pub fn status_code(&self) -> StatusCode {
match self {
Self::NotFound { .. } => StatusCode::NOT_FOUND,
Self::InvalidId { .. } => StatusCode::BAD_REQUEST,
Self::QueueNotInitialized => StatusCode::INTERNAL_SERVER_ERROR,
Self::QueueSubmitFailed { .. } => StatusCode::INTERNAL_SERVER_ERROR,
}
}
}
impl std::fmt::Display for WorkerServiceError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotFound { worker_id } => write!(f, "Worker {} not found", worker_id),
Self::InvalidId { raw, message } => {
write!(
f,
"Invalid worker_id '{}' (expected UUID). Error: {}",
raw, message
)
}
Self::QueueNotInitialized => write!(f, "Job queue not initialized"),
Self::QueueSubmitFailed { message } => write!(f, "{}", message),
}
}
}
impl std::error::Error for WorkerServiceError {}
impl IntoResponse for WorkerServiceError {
fn into_response(self) -> Response {
let error = WorkerErrorResponse {
error: self.to_string(),
code: self.error_code().to_string(),
};
(self.status_code(), Json(error)).into_response()
}
}
/// Result of creating a worker (async job submission)
#[derive(Debug)]
pub struct CreateWorkerResult {
pub worker_id: WorkerId,
pub url: String,
pub location: String,
}
impl IntoResponse for CreateWorkerResult {
fn into_response(self) -> Response {
let response = json!({
"status": "accepted",
"worker_id": self.worker_id.as_str(),
"url": self.url,
"location": self.location,
"message": "Worker addition queued for background processing"
});
(
StatusCode::ACCEPTED,
[(http::header::LOCATION, self.location)],
Json(response),
)
.into_response()
}
}
/// Result of deleting a worker (async job submission)
#[derive(Debug)]
pub struct DeleteWorkerResult {
pub worker_id: WorkerId,
pub url: String,
}
impl IntoResponse for DeleteWorkerResult {
fn into_response(self) -> Response {
let response = json!({
"status": "accepted",
"worker_id": self.worker_id.as_str(),
"message": "Worker removal queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
}
}
/// Result of updating a worker (async job submission)
#[derive(Debug)]
pub struct UpdateWorkerResult {
pub worker_id: WorkerId,
pub url: String,
}
impl IntoResponse for UpdateWorkerResult {
fn into_response(self) -> Response {
let response = json!({
"status": "accepted",
"worker_id": self.worker_id.as_str(),
"message": "Worker update queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
}
}
/// Result of listing workers
#[derive(Debug)]
pub struct ListWorkersResult {
pub workers: Vec<WorkerInfo>,
pub total: usize,
pub prefill_count: usize,
pub decode_count: usize,
pub regular_count: usize,
}
impl IntoResponse for ListWorkersResult {
fn into_response(self) -> Response {
let response = json!({
"workers": self.workers,
"total": self.total,
"stats": {
"prefill_count": self.prefill_count,
"decode_count": self.decode_count,
"regular_count": self.regular_count,
}
});
Json(response).into_response()
}
}
/// Wrapper for WorkerInfo to implement IntoResponse
pub struct GetWorkerResponse(pub WorkerInfo);
impl IntoResponse for GetWorkerResponse {
fn into_response(self) -> Response {
Json(self.0).into_response()
}
}
/// Worker Service - Orchestrates worker business logic
///
/// This service provides a clean API for worker operations, separating
/// business logic from HTTP concerns. Handlers in server.rs become thin
/// wrappers that translate between HTTP and this service.
pub struct WorkerService {
worker_registry: Arc<WorkerRegistry>,
job_queue: Arc<std::sync::OnceLock<Arc<JobQueue>>>,
router_config: RouterConfig,
}
impl WorkerService {
/// Create a new WorkerService
pub fn new(
worker_registry: Arc<WorkerRegistry>,
job_queue: Arc<std::sync::OnceLock<Arc<JobQueue>>>,
router_config: RouterConfig,
) -> Self {
Self {
worker_registry,
job_queue,
router_config,
}
}
/// Parse and validate a worker ID string
pub fn parse_worker_id(&self, raw: &str) -> Result<WorkerId, WorkerServiceError> {
uuid::Uuid::parse_str(raw)
.map(|_| WorkerId::from_string(raw.to_string()))
.map_err(|e| WorkerServiceError::InvalidId {
raw: raw.to_string(),
message: e.to_string(),
})
}
/// Get the job queue, returning an error if not initialized
fn get_job_queue(&self) -> Result<&Arc<JobQueue>, WorkerServiceError> {
self.job_queue
.get()
.ok_or(WorkerServiceError::QueueNotInitialized)
}
pub async fn create_worker(
&self,
mut config: WorkerConfigRequest,
) -> Result<CreateWorkerResult, WorkerServiceError> {
if self.router_config.api_key.is_some() && config.api_key.is_none() {
warn!(
"Adding worker {} without API key while router has API key configured. \
Worker will be accessible without authentication. \
If the worker requires the same API key as the router, please specify it explicitly.",
config.url
);
}
config.dp_aware = self.router_config.dp_aware;
let worker_url = config.url.clone();
let worker_id = self.worker_registry.reserve_id_for_url(&worker_url);
let job = Job::AddWorker {
config: Box::new(config),
};
self.get_job_queue()?
.submit(job)
.await
.map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?;
let location = format!("/workers/{}", worker_id.as_str());
Ok(CreateWorkerResult {
worker_id,
url: worker_url,
location,
})
}
/// List all workers with their info
pub fn list_workers(&self) -> ListWorkersResult {
let workers = self.worker_registry.get_all_with_ids();
let worker_infos: Vec<WorkerInfo> = workers
.iter()
.map(|(worker_id, worker)| {
let mut info = worker_to_info(worker);
info.id = worker_id.as_str().to_string();
info
})
.collect();
let stats = self.worker_registry.stats();
ListWorkersResult {
workers: worker_infos,
total: stats.total_workers,
prefill_count: stats.prefill_workers,
decode_count: stats.decode_workers,
regular_count: stats.regular_workers,
}
}
pub fn get_worker(&self, worker_id_raw: &str) -> Result<GetWorkerResponse, WorkerServiceError> {
let worker_id = self.parse_worker_id(worker_id_raw)?;
let job_queue = self.get_job_queue()?;
if let Some(worker) = self.worker_registry.get(&worker_id) {
let worker_url = worker.url().to_string();
let mut worker_info = worker_to_info(&worker);
worker_info.id = worker_id.as_str().to_string();
if let Some(status) = job_queue.get_status(&worker_url) {
worker_info.job_status = Some(status);
}
return Ok(GetWorkerResponse(worker_info));
}
if let Some(worker_url) = self.worker_registry.get_url_by_id(&worker_id) {
if let Some(status) = job_queue.get_status(&worker_url) {
return Ok(GetWorkerResponse(WorkerInfo::pending(
worker_id.as_str(),
worker_url,
Some(status),
)));
}
}
Err(WorkerServiceError::NotFound {
worker_id: worker_id_raw.to_string(),
})
}
/// Delete a worker by ID (submits async job)
pub async fn delete_worker(
&self,
worker_id_raw: &str,
) -> Result<DeleteWorkerResult, WorkerServiceError> {
let worker_id = self.parse_worker_id(worker_id_raw)?;
let url = self
.worker_registry
.get_url_by_id(&worker_id)
.ok_or_else(|| WorkerServiceError::NotFound {
worker_id: worker_id_raw.to_string(),
})?;
let job = Job::RemoveWorker { url: url.clone() };
let job_queue = self.get_job_queue()?;
job_queue
.submit(job)
.await
.map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?;
Ok(DeleteWorkerResult { worker_id, url })
}
/// Update a worker by ID (submits async job)
pub async fn update_worker(
&self,
worker_id_raw: &str,
update: WorkerUpdateRequest,
) -> Result<UpdateWorkerResult, WorkerServiceError> {
let worker_id = self.parse_worker_id(worker_id_raw)?;
let url = self
.worker_registry
.get_url_by_id(&worker_id)
.ok_or_else(|| WorkerServiceError::NotFound {
worker_id: worker_id_raw.to_string(),
})?;
let job = Job::UpdateWorker {
url: url.clone(),
update: Box::new(update),
};
let job_queue = self.get_job_queue()?;
job_queue
.submit(job)
.await
.map_err(|e| WorkerServiceError::QueueSubmitFailed { message: e })?;
Ok(UpdateWorkerResult { worker_id, url })
}
}

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

@@ -18,7 +18,6 @@ use serde::Deserialize;
use serde_json::{json, Value};
use tokio::{signal, spawn};
use tracing::{debug, error, info, warn, Level};
use uuid::Uuid;
use crate::{
app_context::AppContext,
@@ -30,7 +29,7 @@ use crate::{
create_worker_registration_workflow, create_worker_removal_workflow,
create_worker_update_workflow,
},
worker_to_info, Job, JobQueue, JobQueueConfig, WorkerId, WorkerManager, WorkerType,
Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
},
middleware::{self, AuthConfig, QueuedRequest},
observability::{
@@ -48,7 +47,7 @@ use crate::{
rerank::{RerankRequest, V1RerankReqInput},
responses::{ResponsesGetParams, ResponsesRequest},
validated::ValidatedJson,
worker_spec::{WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest},
worker_spec::{WorkerConfigRequest, WorkerUpdateRequest},
},
routers::{conversations, router_manager::RouterManager, RouterTrait},
service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
@@ -465,185 +464,38 @@ async fn create_worker(
State(state): State<Arc<AppState>>,
Json(config): Json<WorkerConfigRequest>,
) -> Response {
// Warn if router has API key but worker is being added without one
if state.context.router_config.api_key.is_some() && config.api_key.is_none() {
warn!(
"Adding worker {} without API key while router has API key configured. \
Worker will be accessible without authentication. \
If the worker requires the same API key as the router, please specify it explicitly.",
config.url
);
}
// Populate dp_aware from router's configuration
let config = WorkerConfigRequest {
dp_aware: state.context.router_config.dp_aware,
..config
};
// Submit job for async processing
let worker_url = config.url.clone();
let worker_id = state
.context
.worker_registry
.reserve_id_for_url(&worker_url);
let job = Job::AddWorker {
config: Box::new(config),
};
let job_queue = state
.context
.worker_job_queue
.get()
.expect("JobQueue not initialized");
match job_queue.submit(job).await {
Ok(_) => {
let location = format!("/workers/{}", worker_id.as_str());
let response = json!({
"status": "accepted",
"worker_id": worker_id.as_str(),
"url": worker_url,
"location": location,
"message": "Worker addition queued for background processing"
});
(
StatusCode::ACCEPTED,
[(http::header::LOCATION, location)],
Json(response),
)
.into_response()
}
Err(error) => {
let error_response = WorkerErrorResponse {
error,
code: "INTERNAL_SERVER_ERROR".to_string(),
};
(StatusCode::INTERNAL_SERVER_ERROR, Json(error_response)).into_response()
}
match state.context.worker_service.create_worker(config).await {
Ok(result) => result.into_response(),
Err(err) => err.into_response(),
}
}
async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
let workers = state.context.worker_registry.get_all_with_ids();
let worker_infos: Vec<WorkerInfo> = workers
.iter()
.map(|(worker_id, worker)| {
let mut info = worker_to_info(worker);
info.id = worker_id.as_str().to_string();
info
})
.collect();
let response = json!({
"workers": worker_infos,
"total": workers.len(),
"stats": {
"prefill_count": state.context.worker_registry.get_prefill_workers().len(),
"decode_count": state.context.worker_registry.get_decode_workers().len(),
"regular_count": state.context.worker_registry.get_by_type(&WorkerType::Regular).len(),
}
});
Json(response).into_response()
}
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(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())
})
state.context.worker_service.list_workers().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_or_error(&worker_id_raw) {
Ok(id) => id,
Err(resp) => return *resp,
};
let job_queue = state
.context
.worker_job_queue
.get()
.expect("JobQueue not initialized");
if let Some(worker) = state.context.worker_registry.get(&worker_id) {
// Worker exists in registry, get its full info and attach job status if any
let worker_url = worker.url().to_string();
let mut worker_info = worker_to_info(&worker);
worker_info.id = worker_id.as_str().to_string();
if let Some(status) = job_queue.get_status(&worker_url) {
worker_info.job_status = Some(status);
}
return Json(worker_info).into_response();
match state.context.worker_service.get_worker(&worker_id_raw) {
Ok(result) => result.into_response(),
Err(err) => err.into_response(),
}
// 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) {
let worker_info = WorkerInfo::pending(worker_id.as_str(), worker_url, Some(status));
return Json(worker_info).into_response();
}
}
// Worker not found in registry or job queue
let error = WorkerErrorResponse {
error: format!("Worker {worker_id_raw} not found"),
code: "WORKER_NOT_FOUND".to_string(),
};
(StatusCode::NOT_FOUND, Json(error)).into_response()
}
async fn delete_worker(
State(state): State<Arc<AppState>>,
Path(worker_id_raw): Path<String>,
) -> Response {
let worker_id = match parse_worker_id_or_error(&worker_id_raw) {
Ok(id) => id,
Err(resp) => return *resp,
};
let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else {
let error = WorkerErrorResponse {
error: format!("Worker {worker_id_raw} not found"),
code: "WORKER_NOT_FOUND".to_string(),
};
return (StatusCode::NOT_FOUND, Json(error)).into_response();
};
let job = Job::RemoveWorker { url };
let job_queue = state
match state
.context
.worker_job_queue
.get()
.expect("JobQueue not initialized");
match job_queue.submit(job).await {
Ok(_) => {
let response = json!({
"status": "accepted",
"worker_id": worker_id.as_str(),
"message": "Worker removal queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
}
Err(error) => {
let error_response = WorkerErrorResponse {
error,
code: "INTERNAL_SERVER_ERROR".to_string(),
};
(StatusCode::INTERNAL_SERVER_ERROR, Json(error_response)).into_response()
}
.worker_service
.delete_worker(&worker_id_raw)
.await
{
Ok(result) => result.into_response(),
Err(err) => err.into_response(),
}
}
@@ -652,45 +504,14 @@ async fn update_worker(
Path(worker_id_raw): Path<String>,
Json(update): Json<WorkerUpdateRequest>,
) -> Response {
let worker_id = match parse_worker_id_or_error(&worker_id_raw) {
Ok(id) => id,
Err(resp) => return *resp,
};
let Some(url) = state.context.worker_registry.get_url_by_id(&worker_id) else {
let error = WorkerErrorResponse {
error: format!("Worker {worker_id_raw} not found"),
code: "WORKER_NOT_FOUND".to_string(),
};
return (StatusCode::NOT_FOUND, Json(error)).into_response();
};
let job = Job::UpdateWorker {
url,
update: Box::new(update),
};
let job_queue = state
match state
.context
.worker_job_queue
.get()
.expect("JobQueue not initialized");
match job_queue.submit(job).await {
Ok(_) => {
let response = json!({
"status": "accepted",
"worker_id": worker_id.as_str(),
"message": "Worker update queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
}
Err(error) => {
let error_response = WorkerErrorResponse {
error,
code: "INTERNAL_SERVER_ERROR".to_string(),
};
(StatusCode::INTERNAL_SERVER_ERROR, Json(error_response)).into_response()
}
.worker_service
.update_worker(&worker_id_raw, update)
.await
{
Ok(result) => result.into_response(),
Err(err) => err.into_response(),
}
}

14
sgl-model-gateway/src/service_discovery.rs
View File

@@ -609,19 +609,22 @@ mod tests {
}
async fn create_test_app_context() -> Arc<AppContext> {
use crate::{config::RouterConfig, middleware::TokenBucket};
use crate::{config::RouterConfig, core::WorkerService, middleware::TokenBucket};
let router_config = RouterConfig::builder()
.worker_startup_timeout_secs(1)
.build_unchecked();
let worker_registry = Arc::new(crate::core::WorkerRegistry::new());
let worker_job_queue = Arc::new(std::sync::OnceLock::new());
// Note: Using uninitialized queue for tests to avoid spawning background workers
// Jobs submitted during tests will queue but not be processed
Arc::new(AppContext {
client: reqwest::Client::new(),
router_config: router_config.clone(),
rate_limiter: Some(Arc::new(TokenBucket::new(1000, 1000))),
worker_registry: Arc::new(crate::core::WorkerRegistry::new()),
worker_registry: worker_registry.clone(),
policy_registry: Arc::new(crate::policies::PolicyRegistry::new(
router_config.policy.clone(),
)),
@@ -637,10 +640,15 @@ mod tests {
load_monitor: None,
configured_reasoning_parser: None,
configured_tool_parser: None,
worker_job_queue: Arc::new(std::sync::OnceLock::new()),
worker_job_queue: worker_job_queue.clone(),
workflow_engine: Arc::new(std::sync::OnceLock::new()),
mcp_manager: Arc::new(std::sync::OnceLock::new()),
wasm_manager: None,
worker_service: Arc::new(WorkerService::new(
worker_registry,
worker_job_queue,
router_config,
)),
})
}