EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[router] refactor router and worker management 3/n (#10727)

Browse Source
This commit is contained in:
Simo Lin
2025-09-22 15:17:50 -04:00
committed by GitHub
parent 60dbbd086a
commit 97c3823931
25 changed files with 1427 additions and 2540 deletions
Download Patch File Download Diff File Expand all files Collapse all files

208
sgl-router/src/core/worker.rs
View File

@@ -11,10 +11,9 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, LazyLock};
use tokio::sync::Mutex;
// Shared HTTP client for worker operations (health checks, server info, etc.)
static WORKER_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30)) // Default timeout, overridden per request
.timeout(std::time::Duration::from_secs(30))
.build()
.expect("Failed to create worker HTTP client")
});
@@ -43,7 +42,6 @@ pub trait Worker: Send + Sync + fmt::Debug {
/// Synchronous health check wrapper (for compatibility)
fn check_health(&self) -> WorkerResult<()> {
// Use a small runtime for synchronous contexts
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
@@ -64,10 +62,7 @@ pub trait Worker: Send + Sync + fmt::Debug {
fn decrement_load(&self);
/// Reset the load counter to 0 (for sync/recovery)
fn reset_load(&self) {
// Default implementation - does nothing
// Workers that track load should override this
}
fn reset_load(&self) {}
/// Get the number of processed requests
fn processed_requests(&self) -> usize;
@@ -88,11 +83,9 @@ pub trait Worker: Send + Sync + fmt::Debug {
/// Record the outcome of a request to this worker
fn record_outcome(&self, success: bool) {
// Record outcome-level metric with worker label
let outcome_str = if success { "success" } else { "failure" };
RouterMetrics::record_cb_outcome(self.url(), outcome_str);
// Record into circuit breaker and infer state change for metrics
let before = self.circuit_breaker().state();
self.circuit_breaker().record_outcome(success);
let after = self.circuit_breaker().state();
@@ -119,8 +112,6 @@ pub trait Worker: Send + Sync + fmt::Debug {
RouterMetrics::set_cb_state(self.url(), state_code);
}
// === DP-aware methods ===
/// Check if this worker is DP-aware
fn is_dp_aware(&self) -> bool {
false
@@ -156,8 +147,6 @@ pub trait Worker: Send + Sync + fmt::Debug {
true
}
// === Multi-router support ===
// TODO: - Enhanced Worker Discovery
// The Worker trait should handle async discovery of metadata from the worker itself
// rather than having service discovery or other components query /get_server_info.
@@ -356,14 +345,12 @@ impl fmt::Debug for BasicWorker {
impl BasicWorker {
pub fn normalised_url(&self) -> WorkerResult<&str> {
if self.url().contains("@") {
// Need to extract the URL from "http://host:port@dp_rank"
let parts: Vec<&str> = self.url().split('@').collect();
if parts.len() != 2 {
return Err(WorkerError::InvalidUrl {
url: self.url().to_string(),
});
}
// Ensure the second part (the dp_rank) can be parsed as an integer
match parts[1].parse::<usize>() {
Ok(_) => Ok(parts[0]),
Err(_) => Err(WorkerError::InvalidUrl {
@@ -408,19 +395,22 @@ impl Worker for BasicWorker {
let health_result = match &self.metadata.connection_mode {
ConnectionMode::Http => {
// Perform HTTP health check
let url = self.normalised_url()?;
let health_url = format!("{}{}", url, self.metadata.health_config.endpoint);
let timeout = Duration::from_secs(self.metadata.health_config.timeout_secs);
// Use the shared client with a custom timeout for this request
match WORKER_CLIENT.get(&health_url).timeout(timeout).send().await {
let mut request = WORKER_CLIENT.get(&health_url).timeout(timeout);
if let Some(ref api_key) = self.metadata.api_key {
request = request.header("Authorization", format!("Bearer {}", api_key));
}
match request.send().await {
Ok(response) => response.status().is_success(),
Err(_) => false,
}
}
ConnectionMode::Grpc { .. } => {
// Perform gRPC health check
if let Some(grpc_client) = &self.grpc_client {
let mut client = grpc_client.lock().await;
match client.health_check().await {
@@ -449,11 +439,9 @@ impl Worker for BasicWorker {
};
if health_result {
// Health check succeeded
self.consecutive_failures.store(0, Ordering::Release);
let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1;
// Mark healthy if we've reached the success threshold
if !self.is_healthy()
&& successes >= self.metadata.health_config.success_threshold as usize
{
@@ -462,11 +450,9 @@ impl Worker for BasicWorker {
}
Ok(())
} else {
// Health check failed
self.consecutive_successes.store(0, Ordering::Release);
let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1;
// Mark unhealthy if we've reached the failure threshold
if self.is_healthy()
&& failures >= self.metadata.health_config.failure_threshold as usize
{
@@ -576,7 +562,6 @@ impl Worker for DPAwareWorker {
}
async fn check_health_async(&self) -> WorkerResult<()> {
// Delegate to the base worker's health check logic
self.base_worker.check_health_async().await
}
@@ -612,8 +597,6 @@ impl Worker for DPAwareWorker {
self.base_worker.circuit_breaker()
}
// DP-aware specific implementations
fn is_dp_aware(&self) -> bool {
true
}
@@ -631,7 +614,6 @@ impl Worker for DPAwareWorker {
}
async fn prepare_request(&self, mut req: serde_json::Value) -> WorkerResult<serde_json::Value> {
// Inject data_parallel_rank into the request
if let Some(map) = req.as_object_mut() {
map.insert(
"data_parallel_rank".to_string(),
@@ -646,7 +628,6 @@ impl Worker for DPAwareWorker {
}
fn endpoint_url(&self, route: &str) -> String {
// Use base URL for actual requests
format!("{}{}", self.base_url, route)
}
}
@@ -670,53 +651,52 @@ impl WorkerFactory {
}
Box::new(builder.build())
}
#[allow(dead_code)]
/// Get DP size from a worker
async fn get_worker_dp_size(url: &str, api_key: &Option<String>) -> WorkerResult<usize> {
let mut req_builder = WORKER_CLIENT.get(format!("{}/get_server_info", url));
if let Some(key) = &api_key {
req_builder = req_builder.bearer_auth(key);
}
/// Static health validation before creating a worker
/// This replaces wait_for_worker_health in handlers
pub async fn validate_health(url: &str, timeout_secs: u64) -> WorkerResult<()> {
use std::time::Instant;
let response = req_builder
.send()
.await
.map_err(|e| WorkerError::NetworkError {
url: url.to_string(),
error: e.to_string(),
})?;
let start_time = Instant::now();
let timeout = std::time::Duration::from_secs(timeout_secs);
if !response.status().is_success() {
return Err(WorkerError::NetworkError {
url: url.to_string(),
error: format!("Server returned: {}", response.status()),
});
}
let info: serde_json::Value =
response
.json()
.await
.map_err(|e| WorkerError::NetworkError {
loop {
if start_time.elapsed() > timeout {
return Err(WorkerError::HealthCheckFailed {
url: url.to_string(),
error: format!("Failed to parse JSON: {}", e),
})?;
reason: format!(
"Timeout {}s waiting for worker to become healthy",
timeout_secs
),
});
}
let dp_size = info
.get("dp_size")
.and_then(|v| v.as_u64())
.ok_or_else(|| WorkerError::InvalidConfiguration {
message: "dp_size not found in server info".to_string(),
})?;
// Note: This static function doesn't have access to worker's API key
// API key authentication is handled in the worker instance's check_health_async method
match WORKER_CLIENT
.get(format!("{}/health", url))
.timeout(std::time::Duration::from_secs(5))
.send()
.await
{
Ok(res) if res.status().is_success() => {
tracing::info!("Worker {} is healthy", url);
return Ok(());
}
Ok(res) => {
tracing::warn!(
"Worker {} health check failed with status: {}",
url,
res.status()
);
}
Err(e) => {
tracing::warn!("Failed to contact worker {}: {}", url, e);
}
}
if dp_size > usize::MAX as u64 {
return Err(WorkerError::InvalidConfiguration {
message: format!("dp_size is too large: {}", dp_size),
});
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
Ok(dp_size as usize)
}
}
@@ -893,7 +873,6 @@ mod tests {
use std::thread;
use std::time::Duration;
// Test WorkerType
#[test]
fn test_worker_type_display() {
assert_eq!(WorkerType::Regular.to_string(), "Regular");
@@ -945,7 +924,6 @@ mod tests {
assert_eq!(original, cloned);
}
// Test HealthConfig
#[test]
fn test_health_config_default() {
let config = HealthConfig::default();
@@ -972,13 +950,11 @@ mod tests {
assert_eq!(config.success_threshold, 3);
}
// Test BasicWorker
#[test]
fn test_basic_worker_creation() {
use crate::core::BasicWorkerBuilder;
let worker = BasicWorkerBuilder::new("http://test:8080")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.build();
assert_eq!(worker.url(), "http://test:8080");
assert_eq!(worker.worker_type(), WorkerType::Regular);
@@ -1016,7 +992,6 @@ mod tests {
let worker = BasicWorkerBuilder::new("http://test:8080")
.worker_type(WorkerType::Regular)
.health_config(custom_config.clone())
.api_key("test_api_key")
.build();
assert_eq!(worker.metadata().health_config.timeout_secs, 15);
@@ -1024,13 +999,11 @@ mod tests {
assert_eq!(worker.metadata().health_config.endpoint, "/custom-health");
}
// Test Worker trait implementation
#[test]
fn test_worker_url() {
use crate::core::BasicWorkerBuilder;
let worker = BasicWorkerBuilder::new("http://worker1:8080")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.build();
assert_eq!(worker.url(), "http://worker1:8080");
}
@@ -1040,7 +1013,6 @@ mod tests {
use crate::core::BasicWorkerBuilder;
let regular = BasicWorkerBuilder::new("http://test:8080")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.build();
assert_eq!(regular.worker_type(), WorkerType::Regular);
@@ -1048,7 +1020,6 @@ mod tests {
.worker_type(WorkerType::Prefill {
bootstrap_port: Some(9090),
})
.api_key("test_api_key")
.build();
assert_eq!(
prefill.worker_type(),
@@ -1059,7 +1030,6 @@ mod tests {
let decode = BasicWorkerBuilder::new("http://test:8080")
.worker_type(WorkerType::Decode)
.api_key("test_api_key")
.build();
assert_eq!(decode.worker_type(), WorkerType::Decode);
}
@@ -1071,14 +1041,11 @@ mod tests {
.worker_type(WorkerType::Regular)
.build();
// Initial state is healthy
assert!(worker.is_healthy());
// Set unhealthy
worker.set_healthy(false);
assert!(!worker.is_healthy());
// Set healthy again
worker.set_healthy(true);
assert!(worker.is_healthy());
}
@@ -1088,31 +1055,24 @@ mod tests {
use crate::core::BasicWorkerBuilder;
let worker = BasicWorkerBuilder::new("http://test:8080")
.worker_type(WorkerType::Regular)
.api_key("test_api_key")
.build();
// Initial load is 0
assert_eq!(worker.load(), 0);
// Increment once
worker.increment_load();
assert_eq!(worker.load(), 1);
// Increment twice more
worker.increment_load();
worker.increment_load();
assert_eq!(worker.load(), 3);
// Decrement once
worker.decrement_load();
assert_eq!(worker.load(), 2);
// Decrement to 0
worker.decrement_load();
worker.decrement_load();
assert_eq!(worker.load(), 0);
// Decrement below 0 should stay at 0
worker.decrement_load();
assert_eq!(worker.load(), 0);
}
@@ -1124,17 +1084,14 @@ mod tests {
.worker_type(WorkerType::Regular)
.build();
// Initial count is 0
assert_eq!(worker.processed_requests(), 0);
// Increment multiple times
for i in 1..=100 {
worker.increment_processed();
assert_eq!(worker.processed_requests(), i);
}
}
// Test concurrent operations
#[tokio::test]
async fn test_concurrent_load_increments() {
use crate::core::BasicWorkerBuilder;
@@ -1146,7 +1103,6 @@ mod tests {
let mut handles = vec![];
// Spawn 100 tasks incrementing load
for _ in 0..100 {
let worker_clone = Arc::clone(&worker);
let handle = tokio::spawn(async move {
@@ -1155,12 +1111,10 @@ mod tests {
handles.push(handle);
}
// Wait for all tasks
for handle in handles {
handle.await.unwrap();
}
// Final count should be 100
assert_eq!(worker.load(), 100);
}
@@ -1173,7 +1127,6 @@ mod tests {
.build(),
);
// Set initial load to 100
for _ in 0..100 {
worker.increment_load();
}
@@ -1181,7 +1134,6 @@ mod tests {
let mut handles = vec![];
// Spawn 100 tasks decrementing load
for _ in 0..100 {
let worker_clone = Arc::clone(&worker);
let handle = tokio::spawn(async move {
@@ -1190,12 +1142,10 @@ mod tests {
handles.push(handle);
}
// Wait for all tasks
for handle in handles {
handle.await.unwrap();
}
// Final count should be 0
assert_eq!(worker.load(), 0);
}
@@ -1210,7 +1160,6 @@ mod tests {
let mut handles = vec![];
// Spawn threads randomly setting health status
for i in 0..100 {
let worker_clone = Arc::clone(&worker);
let handle = tokio::spawn(async move {
@@ -1220,13 +1169,11 @@ mod tests {
handles.push(handle);
}
// Wait for all tasks
for handle in handles {
handle.await.unwrap();
}
}
// Test WorkerFactory
#[test]
fn test_create_regular_worker() {
let worker: Box<dyn Worker> = Box::new(
@@ -1240,7 +1187,6 @@ mod tests {
#[test]
fn test_create_prefill_worker() {
// With bootstrap port
let worker1: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://prefill:8080")
.worker_type(WorkerType::Prefill {
@@ -1256,7 +1202,6 @@ mod tests {
}
);
// Without bootstrap port
let worker2: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://prefill:8080")
.worker_type(WorkerType::Prefill {
@@ -1283,7 +1228,6 @@ mod tests {
assert_eq!(worker.worker_type(), WorkerType::Decode);
}
// Test WorkerLoadGuard
#[test]
fn test_load_guard_single_worker() {
use crate::core::BasicWorkerBuilder;
@@ -1297,7 +1241,6 @@ mod tests {
assert_eq!(worker.load(), 1);
}
// Guard dropped, load decremented
assert_eq!(worker.load(), 0);
}
@@ -1325,13 +1268,11 @@ mod tests {
{
let _guard = WorkerLoadGuard::new_multi(worker_refs);
// All loads incremented
assert_eq!(workers[0].load(), 1);
assert_eq!(workers[1].load(), 1);
assert_eq!(workers[2].load(), 1);
}
// All loads decremented
assert_eq!(workers[0].load(), 0);
assert_eq!(workers[1].load(), 0);
assert_eq!(workers[2].load(), 0);
@@ -1347,29 +1288,21 @@ mod tests {
);
assert_eq!(worker.load(), 0);
// Clone for use inside catch_unwind
let worker_clone = Arc::clone(&worker);
// Use AssertUnwindSafe wrapper for the test
// This is safe because we're only testing the load counter behavior,
// not the grpc_client which is None for HTTP workers
use std::panic::AssertUnwindSafe;
// This will panic, but the guard should still clean up
let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
let _guard = WorkerLoadGuard::new(worker_clone.as_ref());
assert_eq!(worker_clone.load(), 1);
panic!("Test panic");
}));
// Verify panic occurred
assert!(result.is_err());
// Load should be decremented even after panic
assert_eq!(worker.load(), 0);
}
// Test helper functions
#[test]
fn test_urls_to_workers() {
let urls = vec!["http://w1:8080".to_string(), "http://w2:8080".to_string()];
@@ -1400,23 +1333,17 @@ mod tests {
assert_eq!(urls, vec!["http://w1:8080", "http://w2:8080"]);
}
// Test synchronous health check wrapper
#[test]
fn test_check_health_sync_wrapper() {
// We can't easily test the actual HTTP call without mocking,
// but we can verify the sync wrapper works
use crate::core::BasicWorkerBuilder;
let worker = BasicWorkerBuilder::new("http://test:8080")
.worker_type(WorkerType::Regular)
.build();
// This will fail because there's no server at this URL,
// but it tests that the sync wrapper doesn't panic
let result = worker.check_health();
assert!(result.is_err());
}
// Performance test for load counter
#[test]
fn test_load_counter_performance() {
use crate::core::BasicWorkerBuilder;
@@ -1436,12 +1363,9 @@ mod tests {
let ops_per_sec = iterations as f64 / duration.as_secs_f64();
println!("Load counter operations per second: {:.0}", ops_per_sec);
// Should be well over 1M ops/sec
assert!(ops_per_sec > 1_000_000.0);
}
// ===== Tests for DPAwareWorker =====
#[test]
fn test_dp_aware_worker_creation() {
let dp_worker = DPAwareWorkerBuilder::new("http://worker1:8080", 2, 4)
@@ -1562,8 +1486,6 @@ mod tests {
assert_eq!(dp_worker.processed_requests(), 1);
}
// ===== Tests for WorkerFactory async methods =====
#[tokio::test]
async fn test_factory_create_dp_aware() {
let worker = WorkerFactory::create_dp_aware(
@@ -1610,26 +1532,21 @@ mod tests {
.worker_type(WorkerType::Regular)
.build();
// Initial state should be available
assert!(worker.is_available());
assert_eq!(worker.circuit_breaker().state(), CircuitState::Closed);
// Record some failures
worker.record_outcome(false);
worker.record_outcome(false);
// Still available (default threshold is 5)
assert!(worker.is_available());
// Record more failures to open circuit
worker.record_outcome(false);
worker.record_outcome(false);
worker.record_outcome(false);
// Circuit should be open, worker not available
assert!(!worker.is_available());
assert!(worker.is_healthy()); // Still healthy
assert!(!worker.circuit_breaker().can_execute()); // But circuit is open
assert!(worker.is_healthy());
assert!(!worker.circuit_breaker().can_execute());
}
#[test]
@@ -1647,20 +1564,16 @@ mod tests {
.circuit_breaker_config(config)
.build();
// Should open after 2 failures
worker.record_outcome(false);
assert!(worker.is_available());
worker.record_outcome(false);
assert!(!worker.is_available());
// Wait for timeout
thread::sleep(Duration::from_millis(150));
// Should be half-open
assert!(worker.is_available());
assert_eq!(worker.circuit_breaker().state(), CircuitState::HalfOpen);
// Success should close it
worker.record_outcome(true);
assert_eq!(worker.circuit_breaker().state(), CircuitState::Closed);
}
@@ -1671,24 +1584,18 @@ mod tests {
.worker_type(WorkerType::Regular)
.build();
// Should have circuit breaker
assert!(dp_worker.is_available());
// Record failures
for _ in 0..5 {
dp_worker.record_outcome(false);
}
// Should not be available
assert!(!dp_worker.is_available());
assert_eq!(dp_worker.circuit_breaker().state(), CircuitState::Open);
}
// ===== Integration tests =====
#[tokio::test]
async fn test_mixed_worker_types() {
// Create a mix of worker types
let regular: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://regular:8080")
.worker_type(WorkerType::Regular)
@@ -1739,22 +1646,19 @@ mod tests {
dp_aware_decode,
];
// Test that they all implement Worker trait properly
for worker in &workers {
assert!(worker.is_healthy());
assert_eq!(worker.load(), 0);
assert_eq!(worker.processed_requests(), 0);
}
// Test specific behaviors
assert!(!workers[0].is_dp_aware()); // regular
assert!(!workers[1].is_dp_aware()); // prefill
assert!(!workers[2].is_dp_aware()); // decode
assert!(workers[3].is_dp_aware()); // dp_aware_regular
assert!(workers[4].is_dp_aware()); // dp_aware_prefill
assert!(workers[5].is_dp_aware()); // dp_aware_decode
assert!(!workers[0].is_dp_aware());
assert!(!workers[1].is_dp_aware());
assert!(!workers[2].is_dp_aware());
assert!(workers[3].is_dp_aware());
assert!(workers[4].is_dp_aware());
assert!(workers[5].is_dp_aware());
// Test worker types
assert_eq!(workers[0].worker_type(), WorkerType::Regular);
assert_eq!(
workers[1].worker_type(),