[router] refactor router and worker management 2/n (#10666)

This commit is contained in:
Simo Lin
2025-09-19 15:37:57 -04:00
committed by GitHub
parent dab4663b4e
commit 00eb5eb721
11 changed files with 483 additions and 157 deletions

View File

@@ -101,6 +101,14 @@ impl TestContext {
// Create app context
let app_context = common::create_test_context(config.clone());
// Initialize workers in the registry before creating router
if !worker_urls.is_empty() {
use sglang_router_rs::routers::WorkerInitializer;
WorkerInitializer::initialize_workers(&config, &app_context.worker_registry)
.await
.expect("Failed to initialize workers");
}
// Create router
let router = RouterFactory::create_router(&app_context).await.unwrap();
let router = Arc::from(router);

View File

@@ -39,9 +39,20 @@ impl TestContext {
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
}
config.mode = RoutingMode::Regular { worker_urls };
config.mode = RoutingMode::Regular {
worker_urls: worker_urls.clone(),
};
let app_context = common::create_test_context(config.clone());
// Initialize workers in the registry before creating router
if !worker_urls.is_empty() {
use sglang_router_rs::routers::WorkerInitializer;
WorkerInitializer::initialize_workers(&config, &app_context.worker_registry)
.await
.expect("Failed to initialize workers");
}
let app_context = common::create_test_context(config);
let router = RouterFactory::create_router(&app_context).await.unwrap();
let router = Arc::from(router);

View File

@@ -40,9 +40,20 @@ impl TestContext {
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
}
config.mode = RoutingMode::Regular { worker_urls };
config.mode = RoutingMode::Regular {
worker_urls: worker_urls.clone(),
};
let app_context = common::create_test_context(config.clone());
// Initialize workers in the registry before creating router
if !worker_urls.is_empty() {
use sglang_router_rs::routers::WorkerInitializer;
WorkerInitializer::initialize_workers(&config, &app_context.worker_registry)
.await
.expect("Failed to initialize workers");
}
let app_context = common::create_test_context(config);
let router = RouterFactory::create_router(&app_context).await.unwrap();
let router = Arc::from(router);

View File

@@ -207,19 +207,21 @@ mod test_pd_routing {
history_backend: sglang_router_rs::config::HistoryBackend::Memory,
};
// Router creation will fail due to health checks, but config should be valid
let app_context =
sglang_router_rs::server::AppContext::new(config, reqwest::Client::new(), 64, None)
.expect("Failed to create AppContext");
let app_context = std::sync::Arc::new(app_context);
let result = RouterFactory::create_router(&app_context).await;
assert!(result.is_err());
let error_msg = result.unwrap_err();
// Error should be about health/timeout, not configuration
assert!(
error_msg.contains("healthy") || error_msg.contains("timeout"),
"Unexpected error: {}",
error_msg
result.is_ok(),
"Router creation should succeed with empty worker"
);
// Verify that no workers are registered since we didn't initialize them
let stats = app_context.worker_registry.stats();
assert_eq!(
stats.total_workers, 0,
"No workers should be registered without initialization"
);
}
}