diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index 7fc34075b..7488991dc 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -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>>, pub mcp_manager: Arc>>, pub wasm_manager: Option>, + pub worker_service: Arc, } 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, }) } diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs index 01a5b6d7a..ada1d2a2e 100644 --- a/sgl-model-gateway/src/core/mod.rs +++ b/sgl-model-gateway/src/core/mod.rs @@ -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}; diff --git a/sgl-model-gateway/src/core/worker_service.rs b/sgl-model-gateway/src/core/worker_service.rs new file mode 100644 index 000000000..40f7a83a6 --- /dev/null +++ b/sgl-model-gateway/src/core/worker_service.rs @@ -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, + 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, + job_queue: Arc>>, + router_config: RouterConfig, +} + +impl WorkerService { + /// Create a new WorkerService + pub fn new( + worker_registry: Arc, + job_queue: Arc>>, + 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 { + 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, WorkerServiceError> { + self.job_queue + .get() + .ok_or(WorkerServiceError::QueueNotInitialized) + } + + pub async fn create_worker( + &self, + mut config: WorkerConfigRequest, + ) -> Result { + 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 = 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 { + 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 { + 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 { + 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 }) + } +} diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 369e48109..0ffeaeaad 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -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>, Json(config): Json, ) -> 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>) -> Response { - let workers = state.context.worker_registry.get_all_with_ids(); - let worker_infos: Vec = 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 { - 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> { - 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>, Path(worker_id_raw): Path, ) -> 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>, Path(worker_id_raw): Path, ) -> 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, Json(update): Json, ) -> 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(), } } diff --git a/sgl-model-gateway/src/service_discovery.rs b/sgl-model-gateway/src/service_discovery.rs index 5115cee39..7c6fdd271 100644 --- a/sgl-model-gateway/src/service_discovery.rs +++ b/sgl-model-gateway/src/service_discovery.rs @@ -609,19 +609,22 @@ mod tests { } async fn create_test_app_context() -> Arc { - 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, + )), }) }