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

2025-12-21 14:35:35 -08:00
parent 254de6d2fd
commit 1354063a48
5 changed files with 421 additions and 212 deletions
@@ -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};
@@ -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 })
+    }
+}