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

[router] refactor router and worker management 1/n (#10664)

Browse Source
This commit is contained in:
Simo Lin
2025-09-19 09:19:57 -04:00
committed by GitHub
parent 68cdc1893d
commit 36efd5be8a
2 changed files with 119 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
sgl-router/src/routers/grpc/pd_router.rs
View File

@@ -2,11 +2,11 @@
use crate::config::types::RetryConfig;
use crate::core::{
BasicWorkerBuilder, CircuitBreakerConfig, HealthChecker, HealthConfig, Worker, WorkerType,
BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, WorkerRegistry, WorkerType,
};
use crate::grpc::SglangSchedulerClient;
use crate::metrics::RouterMetrics;
use crate::policies::LoadBalancingPolicy;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterTrait, WorkerManagement};
use crate::tokenizer::traits::Tokenizer;
@@ -19,21 +19,17 @@ use axum::{
response::{IntoResponse, Response},
};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::sync::Arc;
use std::time::Duration;
use tracing::{info, warn};
/// gRPC PD (Prefill-Decode) router implementation for SGLang
#[allow(dead_code)] // Fields will be used once implementation is complete
pub struct GrpcPDRouter {
/// Prefill worker connections
prefill_workers: Arc<RwLock<Vec<Arc<dyn Worker>>>>,
/// Decode worker connections
decode_workers: Arc<RwLock<Vec<Arc<dyn Worker>>>>,
/// gRPC clients for prefill workers
prefill_grpc_clients: Arc<RwLock<HashMap<String, SglangSchedulerClient>>>,
/// gRPC clients for decode workers
decode_grpc_clients: Arc<RwLock<HashMap<String, SglangSchedulerClient>>>,
/// Centralized worker registry
worker_registry: Arc<WorkerRegistry>,
/// Centralized policy registry
policy_registry: Arc<PolicyRegistry>,
/// Load balancing policy for prefill
prefill_policy: Arc<dyn LoadBalancingPolicy>,
/// Load balancing policy for decode
@@ -44,9 +40,6 @@ pub struct GrpcPDRouter {
reasoning_parser_factory: ParserFactory,
/// Tool parser registry for function/tool calls
tool_parser_registry: &'static ParserRegistry,
/// Worker health checkers
_prefill_health_checker: Option<HealthChecker>,
_decode_health_checker: Option<HealthChecker>,
/// Configuration
timeout_secs: u64,
interval_secs: u64,
@@ -65,6 +58,10 @@ impl GrpcPDRouter {
decode_policy: Arc<dyn LoadBalancingPolicy>,
ctx: &Arc<crate::server::AppContext>,
) -> Result<Self, String> {
// Get registries from context
let worker_registry = ctx.worker_registry.clone();
let policy_registry = ctx.policy_registry.clone();
// Update metrics
RouterMetrics::set_active_workers(prefill_urls.len() + decode_urls.len());
@@ -126,10 +123,9 @@ impl GrpcPDRouter {
return Err("Failed to connect to any gRPC workers".to_string());
}
// Create Prefill Worker trait objects with gRPC connection mode
let prefill_workers: Vec<Arc<dyn Worker>> = prefill_urls
.iter()
.map(|(url, bootstrap_port)| {
// Create Prefill Worker trait objects with gRPC connection mode and register them
for (url, bootstrap_port) in &prefill_urls {
if let Some(client) = prefill_grpc_clients.remove(url) {
let worker = BasicWorkerBuilder::new(url.clone())
.worker_type(WorkerType::Prefill {
bootstrap_port: *bootstrap_port,
@@ -145,15 +141,17 @@ impl GrpcPDRouter {
failure_threshold: ctx.router_config.health_check.failure_threshold,
success_threshold: ctx.router_config.health_check.success_threshold,
})
.grpc_client(client)
.build();
Arc::new(worker) as Arc<dyn Worker>
})
.collect();
// Create Decode Worker trait objects with gRPC connection mode
let decode_workers: Vec<Arc<dyn Worker>> = decode_urls
.iter()
.map(|url| {
// Register worker in the centralized registry
worker_registry.register(Arc::new(worker));
}
}
// Create Decode Worker trait objects with gRPC connection mode and register them
for url in &decode_urls {
if let Some(client) = decode_grpc_clients.remove(url) {
let worker = BasicWorkerBuilder::new(url.clone())
.worker_type(WorkerType::Decode)
.connection_mode(crate::core::ConnectionMode::Grpc { port: None })
@@ -165,12 +163,23 @@ impl GrpcPDRouter {
failure_threshold: ctx.router_config.health_check.failure_threshold,
success_threshold: ctx.router_config.health_check.success_threshold,
})
.grpc_client(client)
.build();
Arc::new(worker) as Arc<dyn Worker>
})
.collect();
// Initialize policies with workers if needed
// Register worker in the centralized registry
worker_registry.register(Arc::new(worker));
}
}
// Initialize policies with workers if needed - filter for gRPC workers only
let prefill_workers = worker_registry.get_workers_filtered(
None, // any model
Some(WorkerType::Prefill {
bootstrap_port: None,
}),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false, // include unhealthy workers during initialization
);
if let Some(cache_aware) = prefill_policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
@@ -178,6 +187,12 @@ impl GrpcPDRouter {
cache_aware.init_workers(&prefill_workers);
}
let decode_workers = worker_registry.get_workers_filtered(
None, // any model
Some(WorkerType::Decode),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false, // include unhealthy workers during initialization
);
if let Some(cache_aware) = decode_policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
@@ -185,30 +200,16 @@ impl GrpcPDRouter {
cache_aware.init_workers(&decode_workers);
}
let prefill_workers = Arc::new(RwLock::new(prefill_workers));
let decode_workers = Arc::new(RwLock::new(decode_workers));
let prefill_health_checker = crate::core::start_health_checker(
Arc::clone(&prefill_workers),
ctx.router_config.worker_startup_check_interval_secs,
);
let decode_health_checker = crate::core::start_health_checker(
Arc::clone(&decode_workers),
ctx.router_config.worker_startup_check_interval_secs,
);
// No need for local health checkers - WorkerRegistry handles health checking
Ok(GrpcPDRouter {
prefill_workers,
decode_workers,
prefill_grpc_clients: Arc::new(RwLock::new(prefill_grpc_clients)),
decode_grpc_clients: Arc::new(RwLock::new(decode_grpc_clients)),
worker_registry,
policy_registry,
prefill_policy,
decode_policy,
tokenizer,
reasoning_parser_factory,
tool_parser_registry,
_prefill_health_checker: Some(prefill_health_checker),
_decode_health_checker: Some(decode_health_checker),
timeout_secs: ctx.router_config.worker_startup_timeout_secs,
interval_secs: ctx.router_config.worker_startup_check_interval_secs,
dp_aware: ctx.router_config.dp_aware,
@@ -221,15 +222,23 @@ impl GrpcPDRouter {
impl std::fmt::Debug for GrpcPDRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let prefill_workers = self.worker_registry.get_workers_filtered(
None,
Some(WorkerType::Prefill {
bootstrap_port: None,
}),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false,
);
let decode_workers = self.worker_registry.get_workers_filtered(
None,
Some(WorkerType::Decode),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false,
);
f.debug_struct("GrpcPDRouter")
.field(
"prefill_workers_count",
&self.prefill_workers.read().unwrap().len(),
)
.field(
"decode_workers_count",
&self.decode_workers.read().unwrap().len(),
)
.field("prefill_workers_count", &prefill_workers.len())
.field("decode_workers_count", &decode_workers.len())
.field("timeout_secs", &self.timeout_secs)
.field("interval_secs", &self.interval_secs)
.field("dp_aware", &self.dp_aware)
@@ -351,6 +360,28 @@ impl WorkerManagement for GrpcPDRouter {
fn remove_worker(&self, _worker_url: &str) {}
fn get_worker_urls(&self) -> Vec<String> {
vec![]
let mut urls = Vec::new();
// Get gRPC prefill worker URLs only
let prefill_workers = self.worker_registry.get_workers_filtered(
None,
Some(WorkerType::Prefill {
bootstrap_port: None,
}),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false,
);
urls.extend(prefill_workers.iter().map(|w| w.url().to_string()));
// Get gRPC decode worker URLs only
let decode_workers = self.worker_registry.get_workers_filtered(
None,
Some(WorkerType::Decode),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false,
);
urls.extend(decode_workers.iter().map(|w| w.url().to_string()));
urls
}
}

58
sgl-router/src/routers/grpc/router.rs
View File

@@ -2,11 +2,11 @@
use crate::config::types::RetryConfig;
use crate::core::{
BasicWorkerBuilder, CircuitBreakerConfig, HealthChecker, HealthConfig, Worker, WorkerType,
BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, WorkerRegistry, WorkerType,
};
use crate::grpc::SglangSchedulerClient;
use crate::metrics::RouterMetrics;
use crate::policies::LoadBalancingPolicy;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterTrait, WorkerManagement};
use crate::tokenizer::traits::Tokenizer;
@@ -19,17 +19,17 @@ use axum::{
response::{IntoResponse, Response},
};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::sync::Arc;
use std::time::Duration;
use tracing::{info, warn};
/// gRPC router implementation for SGLang
#[allow(dead_code)] // Fields will be used once implementation is complete
pub struct GrpcRouter {
/// Worker connections
workers: Arc<RwLock<Vec<Arc<dyn Worker>>>>,
/// gRPC clients for each worker
grpc_clients: Arc<RwLock<HashMap<String, SglangSchedulerClient>>>,
/// Centralized worker registry
worker_registry: Arc<WorkerRegistry>,
/// Centralized policy registry
policy_registry: Arc<PolicyRegistry>,
/// Load balancing policy
policy: Arc<dyn LoadBalancingPolicy>,
/// Tokenizer for handling text encoding/decoding
@@ -38,8 +38,6 @@ pub struct GrpcRouter {
reasoning_parser_factory: ParserFactory,
/// Tool parser registry for function/tool calls
tool_parser_registry: &'static ParserRegistry,
/// Worker health checker
_health_checker: Option<HealthChecker>,
/// Configuration
timeout_secs: u64,
interval_secs: u64,
@@ -102,10 +100,11 @@ impl GrpcRouter {
return Err("Failed to connect to any gRPC workers".to_string());
}
// Create Worker trait objects with gRPC connection mode
let mut workers: Vec<Arc<dyn Worker>> = Vec::new();
// Get registries from context
let worker_registry = ctx.worker_registry.clone();
let policy_registry = ctx.policy_registry.clone();
// Move clients from the HashMap to the workers
// Create Worker trait objects with gRPC connection mode and register them
for url in &worker_urls {
if let Some(client) = grpc_clients.remove(url) {
let worker = BasicWorkerBuilder::new(url.clone())
@@ -122,12 +121,21 @@ impl GrpcRouter {
.grpc_client(client)
.build();
workers.push(Arc::new(worker) as Arc<dyn Worker>);
// Register worker in the centralized registry
worker_registry.register(Arc::new(worker));
} else {
warn!("No gRPC client for worker {}, skipping", url);
}
}
// Get only gRPC workers from registry for policy initialization
let workers = worker_registry.get_workers_filtered(
None, // any model
Some(WorkerType::Regular),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false, // include unhealthy workers during initialization
);
// Initialize policy with workers if needed
if let Some(cache_aware) = policy
.as_any()
@@ -136,20 +144,15 @@ impl GrpcRouter {
cache_aware.init_workers(&workers);
}
let workers = Arc::new(RwLock::new(workers));
let health_checker = crate::core::start_health_checker(
Arc::clone(&workers),
ctx.router_config.worker_startup_check_interval_secs,
);
// No need for local health checkers - WorkerRegistry handles health checking
Ok(GrpcRouter {
workers,
grpc_clients: Arc::new(RwLock::new(grpc_clients)),
worker_registry,
policy_registry,
policy,
tokenizer,
reasoning_parser_factory,
tool_parser_registry,
_health_checker: Some(health_checker),
timeout_secs: ctx.router_config.worker_startup_timeout_secs,
interval_secs: ctx.router_config.worker_startup_check_interval_secs,
dp_aware: ctx.router_config.dp_aware,
@@ -162,8 +165,9 @@ impl GrpcRouter {
impl std::fmt::Debug for GrpcRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let stats = self.worker_registry.stats();
f.debug_struct("GrpcRouter")
.field("workers_count", &self.workers.read().unwrap().len())
.field("workers_count", &stats.total_workers)
.field("timeout_secs", &self.timeout_secs)
.field("interval_secs", &self.interval_secs)
.field("dp_aware", &self.dp_aware)
@@ -285,9 +289,13 @@ impl WorkerManagement for GrpcRouter {
fn remove_worker(&self, _worker_url: &str) {}
fn get_worker_urls(&self) -> Vec<String> {
self.workers
.read()
.unwrap()
self.worker_registry
.get_workers_filtered(
None, // any model
Some(WorkerType::Regular),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false, // include all workers
)
.iter()
.map(|w| w.url().to_string())
.collect()