diff --git a/sgl-model-gateway/src/core/worker_manager.rs b/sgl-model-gateway/src/core/worker_manager.rs index 3a32fe418..4845c7f80 100644 --- a/sgl-model-gateway/src/core/worker_manager.rs +++ b/sgl-model-gateway/src/core/worker_manager.rs @@ -1,25 +1,27 @@ -//! Unified Worker Management Module +//! Worker Management Module //! -//! Handles all aspects of worker lifecycle including discovery, initialization, -//! runtime management, and health monitoring. +//! Provides worker lifecycle operations and fan-out request utilities. use std::{collections::HashMap, sync::Arc, time::Duration}; -use axum::response::{IntoResponse, Response}; +use axum::{ + response::{IntoResponse, Response}, + Json, +}; use futures::{ future, stream::{self, StreamExt}, }; -use http::{Method, StatusCode}; -use serde_json::Value; +use http::StatusCode; +use serde_json::{json, Value}; use tokio::{ sync::{watch, Mutex}, task::JoinHandle, }; -use tracing::{debug, error, info, warn}; +use tracing::{debug, info, warn}; use crate::{ - core::{metrics_aggregator::MetricPack, ConnectionMode, WorkerRegistry, WorkerType}, + core::{metrics_aggregator::MetricPack, ConnectionMode, Worker, WorkerRegistry, WorkerType}, policies::PolicyRegistry, protocols::worker_spec::{FlushCacheResult, WorkerLoadInfo, WorkerLoadsResult}, }; @@ -27,7 +29,101 @@ use crate::{ const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); const MAX_CONCURRENT: usize = 32; -/// Unified worker management +/// Result of a fan-out request to a single worker +struct WorkerResponse { + url: String, + result: Result, +} + +/// Fan out requests to workers in parallel +async fn fan_out( + workers: &[Arc], + client: &reqwest::Client, + endpoint: &str, + method: reqwest::Method, +) -> Vec { + let futures: Vec<_> = workers + .iter() + .map(|worker| { + let client = client.clone(); + let url = worker.url().to_string(); + let full_url = format!("{}/{}", url, endpoint); + let api_key = worker.api_key().clone(); + let method = method.clone(); + + async move { + let mut req = client.request(method, &full_url).timeout(REQUEST_TIMEOUT); + if let Some(key) = api_key { + req = req.bearer_auth(key); + } + WorkerResponse { + url, + result: req.send().await, + } + } + }) + .collect(); + + stream::iter(futures) + .buffer_unordered(MAX_CONCURRENT) + .collect() + .await +} + +impl IntoResponse for FlushCacheResult { + fn into_response(self) -> Response { + let status = if self.failed.is_empty() { + StatusCode::OK + } else { + StatusCode::PARTIAL_CONTENT + }; + + let mut body = json!({ + "status": if self.failed.is_empty() { "success" } else { "partial_success" }, + "message": self.message, + "workers_flushed": self.successful.len(), + "total_http_workers": self.http_workers, + "total_workers": self.total_workers + }); + + if !self.failed.is_empty() { + body["successful"] = json!(self.successful); + body["failed"] = json!(self + .failed + .into_iter() + .map(|(url, err)| json!({"worker": url, "error": err})) + .collect::>()); + } + + (status, Json(body)).into_response() + } +} + +impl IntoResponse for WorkerLoadsResult { + fn into_response(self) -> Response { + let loads: Vec = self + .loads + .iter() + .map(|info| json!({"worker": &info.worker, "load": info.load})) + .collect(); + Json(json!({"workers": loads})).into_response() + } +} + +pub enum EngineMetricsResult { + Ok(String), + Err(String), +} + +impl IntoResponse for EngineMetricsResult { + fn into_response(self) -> Response { + match self { + Self::Ok(text) => (StatusCode::OK, text).into_response(), + Self::Err(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response(), + } + } +} + pub struct WorkerManager; impl WorkerManager { @@ -39,77 +135,44 @@ impl WorkerManager { .collect() } - /// Flush cache on all workers - /// - /// Sends a POST request to /flush_cache endpoint on all HTTP workers. - /// Returns detailed results showing which workers succeeded and which failed. pub async fn flush_cache_all( worker_registry: &WorkerRegistry, client: &reqwest::Client, - ) -> Result { - warn!("Flushing cache for ALL workers - this may impact performance temporarily"); - + ) -> FlushCacheResult { let workers = worker_registry.get_all(); + let total_workers = workers.len(); let http_workers: Vec<_> = workers - .iter() + .into_iter() .filter(|w| matches!(w.connection_mode(), ConnectionMode::Http)) .collect(); if http_workers.is_empty() { - return Ok(FlushCacheResult { + return FlushCacheResult { successful: vec![], failed: vec![], - total_workers: workers.len(), + total_workers, http_workers: 0, message: "No HTTP workers available for cache flush".to_string(), - }); + }; } info!( - "Flushing cache on {} HTTP workers (out of {} total workers)", + "Flushing cache on {} HTTP workers (out of {} total)", http_workers.len(), - workers.len() + total_workers ); - let mut tasks = Vec::with_capacity(http_workers.len()); - for worker in &http_workers { - let url = worker.url().to_string(); - let flush_url = format!("{}/flush_cache", url); - let mut request = client.post(&flush_url); - - if let Some(api_key) = worker.api_key() { - request = request.header("Authorization", format!("Bearer {}", api_key)); - } - - let worker_url = url.clone(); - tasks.push(async move { - let result = request.send().await; - (worker_url, result) - }); - } - - let results = future::join_all(tasks).await; + let responses = fan_out(&http_workers, client, "flush_cache", reqwest::Method::POST).await; let mut successful = Vec::new(); let mut failed = Vec::new(); - for (url, result) in results { - match result { - Ok(response) if response.status().is_success() => { - debug!("Successfully flushed cache on worker: {}", url); - successful.push(url); - } - Ok(response) => { - let error = format!("HTTP {}", response.status()); - warn!("Failed to flush cache on worker {}: {}", url, error); - failed.push((url, error)); - } - Err(e) => { - let error = e.to_string(); - error!("Failed to connect to worker {}: {}", url, error); - failed.push((url, error)); - } + for resp in responses { + match resp.result { + Ok(r) if r.status().is_success() => successful.push(resp.url), + Ok(r) => failed.push((resp.url, format!("HTTP {}", r.status()))), + Err(e) => failed.push((resp.url, e.to_string())), } } @@ -120,77 +183,20 @@ impl WorkerManager { ) } else { format!( - "Cache flush completed: {} succeeded, {} failed (out of {} HTTP workers)", + "Cache flush: {} succeeded, {} failed", successful.len(), - failed.len(), - http_workers.len() + failed.len() ) }; info!("{}", message); - Ok(FlushCacheResult { + FlushCacheResult { successful, failed, - total_workers: workers.len(), + total_workers, http_workers: http_workers.len(), message, - }) - } - - pub async fn get_worker_load( - url: &str, - api_key: Option<&str>, - client: &reqwest::Client, - ) -> Option { - let load_url = format!("{}/get_load", url); - let mut request = client.get(&load_url); - - if let Some(key) = api_key { - request = request.bearer_auth(key); - } - - match request.send().await { - Ok(response) if response.status().is_success() => { - match response.json::().await { - Ok(json) => { - // The /get_load endpoint returns an array of load info objects (one per DP rank) - // Each object has: {dp_rank, num_reqs, num_waiting_reqs, num_tokens} - if let Some(array) = json.as_array() { - let total_tokens: i64 = array - .iter() - .filter_map(|entry| { - entry.get("num_tokens").and_then(|v| v.as_i64()) - }) - .sum(); - debug!("Worker {} load (total tokens): {}", url, total_tokens); - Some(total_tokens as isize) - } else { - warn!( - "Invalid load response from {}: expected array, got {:?}", - url, json - ); - None - } - } - Err(e) => { - warn!("Failed to parse load response from {}: {}", url, e); - None - } - } - } - Ok(response) => { - warn!( - "Failed to get load from {}: HTTP {}", - url, - response.status() - ); - None - } - Err(e) => { - warn!("Failed to connect to {} for load check: {}", url, e); - None - } } } @@ -201,38 +207,35 @@ impl WorkerManager { let workers = worker_registry.get_all(); let total_workers = workers.len(); - // Prepare tasks for parallel execution - let mut tasks = Vec::with_capacity(workers.len()); - for worker in &workers { - let url = worker.url().to_string(); - let api_key = worker.api_key().clone(); - let worker_type = match worker.worker_type() { - WorkerType::Regular => None, - WorkerType::Prefill { .. } => Some("prefill".to_string()), - WorkerType::Decode => Some("decode".to_string()), - }; - let is_http = matches!(worker.connection_mode(), ConnectionMode::Http); - let client = client.clone(); - - tasks.push(async move { - let load = if is_http { - Self::get_worker_load(&url, api_key.as_deref(), &client) - .await - .unwrap_or(-1) - } else { - -1 + let futures: Vec<_> = workers + .iter() + .map(|worker| { + let url = worker.url().to_string(); + let api_key = worker.api_key().clone(); + let worker_type = match worker.worker_type() { + WorkerType::Regular => None, + WorkerType::Prefill { .. } => Some("prefill".to_string()), + WorkerType::Decode => Some("decode".to_string()), }; + let is_http = matches!(worker.connection_mode(), ConnectionMode::Http); + let client = client.clone(); - WorkerLoadInfo { - worker: url, - worker_type, - load, + async move { + let load = if is_http { + Self::parse_load_response(&client, &url, api_key.as_deref()).await + } else { + -1 + }; + WorkerLoadInfo { + worker: url, + worker_type, + load, + } } - }); - } - - let loads = future::join_all(tasks).await; + }) + .collect(); + let loads = future::join_all(futures).await; let successful = loads.iter().filter(|l| l.load >= 0).count(); let failed = loads.iter().filter(|l| l.load < 0).count(); @@ -244,95 +247,65 @@ impl WorkerManager { } } + async fn parse_load_response( + client: &reqwest::Client, + url: &str, + api_key: Option<&str>, + ) -> isize { + let load_url = format!("{}/get_load", url); + let mut req = client.get(&load_url).timeout(REQUEST_TIMEOUT); + if let Some(key) = api_key { + req = req.bearer_auth(key); + } + + match req.send().await { + Ok(r) if r.status().is_success() => match r.json::().await { + Ok(json) if json.is_array() => json + .as_array() + .unwrap() + .iter() + .filter_map(|e| e.get("num_tokens").and_then(|v| v.as_i64())) + .sum::() as isize, + _ => -1, + }, + _ => -1, + } + } + pub async fn get_engine_metrics( worker_registry: &WorkerRegistry, client: &reqwest::Client, - ) -> Response { - let engine_responses = - match Self::fan_out_simple_request(worker_registry, client, "metrics", Method::GET) - .await - { - Ok(x) => x, - Err(e) => return e, - }; - let engine_responses = engine_responses - .into_iter() - .map(|(worker_base_url, metrics_text)| MetricPack { - labels: vec![("worker_addr".into(), worker_base_url)], - metrics_text, - }) - .collect(); - let text = match crate::core::metrics_aggregator::aggregate_metrics(engine_responses) { - Ok(x) => x, - Err(e) => { - let error_msg = format!("Failed to aggregate metrics: {}", e); - return (StatusCode::INTERNAL_SERVER_ERROR, error_msg).into_response(); - } - }; - (StatusCode::OK, text).into_response() - } - - async fn fan_out_simple_request( - worker_registry: &WorkerRegistry, - client: &reqwest::Client, - endpoint: &str, - method: Method, - ) -> Result, Response> { + ) -> EngineMetricsResult { let workers = worker_registry.get_all(); - let worker_count = workers.len(); if workers.is_empty() { - return Err((StatusCode::SERVICE_UNAVAILABLE, "No available workers").into_response()); + return EngineMetricsResult::Err("No available workers".to_string()); } - let futures: Vec<_> = workers - .into_iter() - .map(|worker| { - let client = client.clone(); - let worker_url = worker.url().to_string(); - let url = format!("{worker_url}/{endpoint}"); - let api_key = worker.api_key().clone(); - let method = method.clone(); + let responses = fan_out(&workers, client, "metrics", reqwest::Method::GET).await; - async move { - let mut req = client.request(method, &url).timeout(REQUEST_TIMEOUT); - - if let Some(key) = api_key { - req = req.bearer_auth(key); - } - - match req.send().await { - Ok(res) if res.status().is_success() => match res.text().await { - Ok(body) => Some((worker_url, body)), - Err(e) => { - warn!("Failed reading response from {url}: {e}"); - None - } - }, - Ok(res) => { - warn!("Request to {url} failed: {}", res.status()); - None - } - Err(e) => { - warn!("Request to {url} failed: {e}"); - None - } + let mut metric_packs = Vec::new(); + for resp in responses { + if let Ok(r) = resp.result { + if r.status().is_success() { + if let Ok(text) = r.text().await { + metric_packs.push(MetricPack { + labels: vec![("worker_addr".into(), resp.url)], + metrics_text: text, + }); } } - }) - .collect(); - - let responses: Vec<_> = stream::iter(futures) - .buffer_unordered(MAX_CONCURRENT) - .filter_map(|r| async { r }) - .collect() - .await; - - if responses.is_empty() && worker_count > 0 { - return Err((StatusCode::BAD_GATEWAY, "All backend requests failed").into_response()); + } } - Ok(responses) + if metric_packs.is_empty() { + return EngineMetricsResult::Err("All backend requests failed".to_string()); + } + + match crate::core::metrics_aggregator::aggregate_metrics(metric_packs) { + Ok(text) => EngineMetricsResult::Ok(text), + Err(e) => EngineMetricsResult::Err(format!("Failed to aggregate metrics: {}", e)), + } } } @@ -348,7 +321,6 @@ pub struct LoadMonitor { } impl LoadMonitor { - /// Create a new load monitor pub fn new( worker_registry: Arc, policy_registry: Arc, @@ -368,7 +340,6 @@ impl LoadMonitor { } } - /// Start monitoring worker loads pub async fn start(&self) { let mut handle_guard = self.monitor_handle.lock().await; if handle_guard.is_some() { @@ -394,7 +365,6 @@ impl LoadMonitor { *handle_guard = Some(handle); } - /// Stop monitoring worker loads pub async fn stop(&self) { let mut handle_guard = self.monitor_handle.lock().await; if let Some(handle) = handle_guard.take() { @@ -404,12 +374,10 @@ impl LoadMonitor { } } - /// Get a receiver for load updates pub fn subscribe(&self) -> watch::Receiver> { self.rx.clone() } - /// The main monitoring loop async fn monitor_loop( worker_registry: Arc, policy_registry: Arc, @@ -452,7 +420,6 @@ impl LoadMonitor { } } - /// Check if monitoring is currently active pub async fn is_running(&self) -> bool { let handle_guard = self.monitor_handle.lock().await; handle_guard.is_some() diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 0ffeaeaad..96a5b478f 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -137,7 +137,9 @@ async fn health_generate(State(state): State>, req: Request) -> Re } async fn engine_metrics(State(state): State>) -> Response { - WorkerManager::get_engine_metrics(&state.context.worker_registry, &state.context.client).await + WorkerManager::get_engine_metrics(&state.context.worker_registry, &state.context.client) + .await + .into_response() } async fn get_server_info(State(state): State>, req: Request) -> Response { @@ -393,71 +395,15 @@ async fn v1_conversations_delete_item( } async fn flush_cache(State(state): State>, _req: Request) -> Response { - match WorkerManager::flush_cache_all(&state.context.worker_registry, &state.context.client) + WorkerManager::flush_cache_all(&state.context.worker_registry, &state.context.client) .await - { - Ok(result) => { - if result.failed.is_empty() { - ( - StatusCode::OK, - Json(json!({ - "status": "success", - "message": result.message, - "workers_flushed": result.successful.len(), - "total_http_workers": result.http_workers, - "total_workers": result.total_workers - })), - ) - .into_response() - } else { - ( - StatusCode::PARTIAL_CONTENT, - Json(json!({ - "status": "partial_success", - "message": result.message, - "successful": result.successful, - "failed": result.failed.into_iter().map(|(url, err)| json!({ - "worker": url, - "error": err - })).collect::>(), - "total_http_workers": result.http_workers, - "total_workers": result.total_workers - })), - ) - .into_response() - } - } - Err(e) => { - error!("Failed to flush cache: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({ - "status": "error", - "message": format!("Failed to flush cache: {}", e) - })), - ) - .into_response() - } - } + .into_response() } async fn get_loads(State(state): State>, _req: Request) -> Response { - let result = - WorkerManager::get_all_worker_loads(&state.context.worker_registry, &state.context.client) - .await; - - let loads: Vec = result - .loads - .iter() - .map(|info| { - json!({ - "worker": &info.worker, - "load": info.load - }) - }) - .collect(); - - (StatusCode::OK, Json(json!({ "workers": loads }))).into_response() + WorkerManager::get_all_worker_loads(&state.context.worker_registry, &state.context.client) + .await + .into_response() } async fn create_worker(