Unverified Commit 5dccf697 authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[router] create worker removal step and clean up worker manager (#11921)

parent eec9e471
...@@ -66,7 +66,7 @@ pub fn create_test_context(config: RouterConfig) -> Arc<AppContext> { ...@@ -66,7 +66,7 @@ pub fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
let worker_job_queue = Arc::new(OnceLock::new()); let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new()); let workflow_engine = Arc::new(OnceLock::new());
Arc::new(AppContext::new( let app_context = Arc::new(AppContext::new(
config, config,
client, client,
rate_limiter, rate_limiter,
...@@ -81,7 +81,32 @@ pub fn create_test_context(config: RouterConfig) -> Arc<AppContext> { ...@@ -81,7 +81,32 @@ pub fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
load_monitor, load_monitor,
worker_job_queue, worker_job_queue,
workflow_engine, workflow_engine,
)) ));
// Initialize JobQueue after AppContext is created
let weak_context = Arc::downgrade(&app_context);
let job_queue = sglang_router_rs::core::JobQueue::new(
sglang_router_rs::core::JobQueueConfig::default(),
weak_context,
);
app_context
.worker_job_queue
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use sglang_router_rs::core::workflow::{
create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine.register_workflow(create_worker_registration_workflow());
engine.register_workflow(create_worker_removal_workflow());
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
app_context
} }
// Tokenizer download configuration // Tokenizer download configuration
......
...@@ -7,7 +7,6 @@ use reqwest::Client; ...@@ -7,7 +7,6 @@ use reqwest::Client;
use serde_json::json; use serde_json::json;
use sglang_router_rs::{ use sglang_router_rs::{
config::{RouterConfig, RoutingMode}, config::{RouterConfig, RoutingMode},
core::WorkerManager,
routers::{RouterFactory, RouterTrait}, routers::{RouterFactory, RouterTrait},
}; };
...@@ -51,13 +50,6 @@ impl TestContext { ...@@ -51,13 +50,6 @@ impl TestContext {
let app_context = common::create_test_context(config.clone()); let app_context = common::create_test_context(config.clone());
// Initialize workers in the registry before creating router
if !worker_urls.is_empty() {
WorkerManager::initialize_workers(&config, &app_context.worker_registry, None)
.await
.expect("Failed to initialize workers");
}
let router = RouterFactory::create_router(&app_context).await.unwrap(); let router = RouterFactory::create_router(&app_context).await.unwrap();
let router = Arc::from(router); let router = Arc::from(router);
......
...@@ -8,7 +8,6 @@ use reqwest::Client; ...@@ -8,7 +8,6 @@ use reqwest::Client;
use serde_json::json; use serde_json::json;
use sglang_router_rs::{ use sglang_router_rs::{
config::{RouterConfig, RoutingMode}, config::{RouterConfig, RoutingMode},
core::WorkerManager,
routers::{RouterFactory, RouterTrait}, routers::{RouterFactory, RouterTrait},
}; };
...@@ -52,13 +51,6 @@ impl TestContext { ...@@ -52,13 +51,6 @@ impl TestContext {
let app_context = common::create_test_context(config.clone()); let app_context = common::create_test_context(config.clone());
// Initialize workers in the registry before creating router
if !worker_urls.is_empty() {
WorkerManager::initialize_workers(&config, &app_context.worker_registry, None)
.await
.expect("Failed to initialize workers");
}
let router = RouterFactory::create_router(&app_context).await.unwrap(); let router = RouterFactory::create_router(&app_context).await.unwrap();
let router = Arc::from(router); let router = Arc::from(router);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment