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

186
sgl-router/src/service_discovery.rs
View File

@@ -1,4 +1,6 @@
use crate::routers::RouterTrait;
use crate::core::WorkerManager;
use crate::protocols::worker_spec::WorkerConfigRequest;
use crate::server::AppContext;
use futures::{StreamExt, TryStreamExt};
use k8s_openapi::api::core::v1::Pod;
@@ -175,7 +177,7 @@ impl PodInfo {
pub async fn start_service_discovery(
config: ServiceDiscoveryConfig,
router: Arc<dyn RouterTrait>,
app_context: Arc<AppContext>,
) -> Result<task::JoinHandle<()>, kube::Error> {
// Don't initialize anything if service discovery is disabled
if !config.enabled {
@@ -277,13 +279,13 @@ pub async fn start_service_discovery(
// Clone again for the next closure
let tracked_pods_clone2 = Arc::clone(&tracked_pods_clone);
let router_clone = Arc::clone(&router);
let app_context_clone = Arc::clone(&app_context);
let config_clone2 = Arc::clone(&config_arc);
match filtered_stream
.try_for_each(move |pod| {
let tracked_pods_inner = Arc::clone(&tracked_pods_clone2);
let router_inner = Arc::clone(&router_clone);
let app_context_inner = Arc::clone(&app_context_clone);
let config_inner = Arc::clone(&config_clone2);
async move {
@@ -294,16 +296,15 @@ pub async fn start_service_discovery(
handle_pod_deletion(
&pod_info,
tracked_pods_inner,
router_inner,
app_context_inner,
port,
config_inner.pd_mode,
)
.await;
} else {
handle_pod_event(
&pod_info,
tracked_pods_inner,
router_inner,
app_context_inner,
port,
config_inner.pd_mode,
)
@@ -347,7 +348,7 @@ pub async fn start_service_discovery(
async fn handle_pod_event(
pod_info: &PodInfo,
tracked_pods: Arc<Mutex<HashSet<PodInfo>>>,
router: Arc<dyn RouterTrait>,
app_context: Arc<AppContext>,
port: u16,
pd_mode: bool,
) {
@@ -380,40 +381,44 @@ async fn handle_pod_event(
pod_info.name, pod_info.pod_type, worker_url
);
// Handle PD mode with specific pod types
let result = if pd_mode && pod_info.pod_type.is_some() {
// Need to import PDRouter type
use crate::routers::http::pd_router::PDRouter;
// Try to downcast to PDRouter
if let Some(pd_router) = router.as_any().downcast_ref::<PDRouter>() {
match &pod_info.pod_type {
Some(PodType::Prefill) => pd_router
.add_prefill_server(
worker_url.clone(),
pd_router.api_key.clone(),
pod_info.bootstrap_port,
)
.await
.map_err(|e| e.to_string()),
Some(PodType::Decode) => pd_router
.add_decode_server(worker_url.clone(), pd_router.api_key.clone())
.await
.map_err(|e| e.to_string()),
Some(PodType::Regular) | None => {
// Fall back to regular add_worker for regular pods
router.add_worker(&worker_url, &pd_router.api_key).await
}
}
} else {
Err("PD mode enabled but router is not a PDRouter".to_string())
// Build worker config based on pod type and routing mode
let worker_type = if pd_mode {
match &pod_info.pod_type {
Some(PodType::Prefill) => Some("prefill".to_string()),
Some(PodType::Decode) => Some("decode".to_string()),
Some(PodType::Regular) | None => None,
}
} else {
// Regular mode or no pod type specified
// In pod, no need api key
router.add_worker(&worker_url, &None).await
None
};
// Only set bootstrap_port for prefill workers in PD mode
let bootstrap_port = if pd_mode {
match &pod_info.pod_type {
Some(PodType::Prefill) => pod_info.bootstrap_port,
_ => None,
}
} else {
None
};
let config = WorkerConfigRequest {
url: worker_url.clone(),
model_id: None,
worker_type,
priority: None,
cost: None,
labels: HashMap::new(),
bootstrap_port,
tokenizer_path: None,
reasoning_parser: None,
tool_parser: None,
chat_template: None,
api_key: None,
};
let result = WorkerManager::add_worker_from_config(&config, &app_context).await;
match result {
Ok(_) => {
debug!("Worker added: {}", worker_url);
@@ -433,9 +438,8 @@ async fn handle_pod_event(
async fn handle_pod_deletion(
pod_info: &PodInfo,
tracked_pods: Arc<Mutex<HashSet<PodInfo>>>,
router: Arc<dyn RouterTrait>,
app_context: Arc<AppContext>,
port: u16,
pd_mode: bool,
) {
let worker_url = pod_info.worker_url(port);
@@ -456,35 +460,8 @@ async fn handle_pod_deletion(
pod_info.name, pod_info.pod_type, worker_url
);
// Handle PD mode removal
if pd_mode && pod_info.pod_type.is_some() {
use crate::routers::http::pd_router::PDRouter;
// Try to downcast to PDRouter for PD-specific removal
if let Some(pd_router) = router.as_any().downcast_ref::<PDRouter>() {
match &pod_info.pod_type {
Some(PodType::Prefill) => {
if let Err(e) = pd_router.remove_prefill_server(&worker_url).await {
error!("Failed to remove prefill server {}: {}", worker_url, e);
}
}
Some(PodType::Decode) => {
if let Err(e) = pd_router.remove_decode_server(&worker_url).await {
error!("Failed to remove decode server {}: {}", worker_url, e);
}
}
Some(PodType::Regular) | None => {
// Fall back to regular remove_worker
router.remove_worker(&worker_url);
}
}
} else {
// PD mode but not a PDRouter, use generic removal
router.remove_worker(&worker_url);
}
} else {
// Regular mode removal
router.remove_worker(&worker_url);
if let Err(e) = WorkerManager::remove_worker(&worker_url, &app_context) {
error!("Failed to remove worker {}: {}", worker_url, e);
}
} else {
// This case might occur if a pod is deleted before it was ever marked healthy and added.
@@ -582,12 +559,10 @@ mod tests {
}
}
// Helper to create a Router instance for testing event handlers
async fn create_test_router() -> Arc<dyn RouterTrait> {
// Helper to create an AppContext instance for testing event handlers
async fn create_test_app_context() -> Arc<AppContext> {
use crate::config::RouterConfig;
use crate::middleware::TokenBucket;
use crate::routers::http::router::Router;
use crate::server::AppContext;
// Create a minimal RouterConfig for testing with very short timeout
let router_config = RouterConfig {
@@ -596,7 +571,7 @@ mod tests {
}; // Very short timeout for tests
// Create AppContext with minimal components
let app_context = Arc::new(AppContext {
Arc::new(AppContext {
client: reqwest::Client::new(),
router_config: router_config.clone(),
rate_limiter: Arc::new(TokenBucket::new(1000, 1000)),
@@ -609,10 +584,7 @@ mod tests {
tool_parser_registry: None, // HTTP mode doesn't need tool parser
router_manager: None, // Test doesn't need router manager
response_storage: Arc::new(crate::data_connector::MemoryResponseStorage::new()),
});
let router = Router::new(&app_context).await.unwrap();
Arc::new(router) as Arc<dyn RouterTrait>
})
}
// Helper to create a PD config for testing
@@ -914,7 +886,7 @@ mod tests {
#[tokio::test]
async fn test_handle_pod_event_add_unhealthy_pod() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "pod1".into(),
@@ -929,21 +901,18 @@ mod tests {
handle_pod_event(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
false, // pd_mode = false
)
.await;
assert!(!tracked_pods.lock().unwrap().contains(&pod_info));
assert!(!router
.get_worker_urls()
.contains(&pod_info.worker_url(port)));
}
#[tokio::test]
async fn test_handle_pod_deletion_non_existing_pod() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "pod1".into(),
@@ -958,19 +927,17 @@ mod tests {
handle_pod_deletion(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
false, // pd_mode = false
)
.await;
assert!(tracked_pods.lock().unwrap().is_empty());
assert!(router.get_worker_urls().is_empty());
}
#[tokio::test]
async fn test_handle_pd_pod_event_prefill_pod() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "prefill-pod".into(),
@@ -983,23 +950,23 @@ mod tests {
let port = 8080u16;
// This test validates the structure but won't actually add workers since
// we're using a regular router instead of PD router
// the test worker URL won't be reachable
handle_pod_event(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
false, // pd_mode = false, so it should fallback to regular handling
true, // pd_mode = true for PD pod
)
.await;
// Pod should not be tracked since router.add_worker will fail for non-running server
// Pod should not be tracked since add_worker_from_config will fail for non-running server
assert!(!tracked_pods.lock().unwrap().contains(&pod_info));
}
#[tokio::test]
async fn test_handle_pd_pod_event_decode_pod() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "decode-pod".into(),
@@ -1014,19 +981,19 @@ mod tests {
handle_pod_event(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
false, // pd_mode = false, so it should fallback to regular handling
true, // pd_mode = true for PD pod
)
.await;
// Pod should not be tracked since router.add_worker will fail for non-running server
// Pod should not be tracked since add_worker_from_config will fail for non-running server
assert!(!tracked_pods.lock().unwrap().contains(&pod_info));
}
#[tokio::test]
async fn test_handle_pd_pod_deletion_tracked_pod() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "test-pod".into(),
@@ -1048,9 +1015,8 @@ mod tests {
handle_pod_deletion(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
false, // pd_mode = false
)
.await;
@@ -1060,7 +1026,7 @@ mod tests {
#[tokio::test]
async fn test_handle_pd_pod_deletion_untracked_pod() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "untracked-pod".into(),
@@ -1077,9 +1043,8 @@ mod tests {
handle_pod_deletion(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
true, // pd_mode = true
)
.await;
@@ -1089,7 +1054,7 @@ mod tests {
#[tokio::test]
async fn test_unified_handler_regular_mode() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "regular-pod".into(),
@@ -1105,19 +1070,19 @@ mod tests {
handle_pod_event(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
false, // pd_mode = false
)
.await;
// Pod should not be tracked since router.add_worker will fail for non-running server
// Pod should not be tracked since add_worker_from_url will fail for non-running server
assert!(!tracked_pods.lock().unwrap().contains(&pod_info));
}
#[tokio::test]
async fn test_unified_handler_pd_mode_with_prefill() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "prefill-pod".into(),
@@ -1133,19 +1098,19 @@ mod tests {
handle_pod_event(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
true, // pd_mode = true
)
.await;
// Pod should not be tracked since router.add_pd_worker will fail for regular router
// Pod should not be tracked since add_worker_from_config will fail for non-running server
assert!(!tracked_pods.lock().unwrap().contains(&pod_info));
}
#[tokio::test]
async fn test_unified_handler_deletion_with_pd_mode() {
let router = create_test_router().await;
let app_context = create_test_app_context().await;
let tracked_pods = Arc::new(Mutex::new(HashSet::new()));
let pod_info = PodInfo {
name: "decode-pod".into(),
@@ -1168,9 +1133,8 @@ mod tests {
handle_pod_deletion(
&pod_info,
Arc::clone(&tracked_pods),
Arc::clone(&router),
Arc::clone(&app_context),
port,
true, // pd_mode = true
)
.await;