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

[router] refactor worker to builder pattern 5/n (#10653)

Browse Source
This commit is contained in:
Simo Lin
2025-09-19 05:43:23 -04:00
committed by GitHub
parent 3fa3c22ae2
commit 873d858b28
4 changed files with 197 additions and 277 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
sgl-router/tests/test_pd_routing.rs
View File

@@ -4,7 +4,7 @@ mod test_pd_routing {
use sglang_router_rs::config::{
CircuitBreakerConfig, ConnectionMode, PolicyConfig, RetryConfig, RouterConfig, RoutingMode,
};
use sglang_router_rs::core::{WorkerFactory, WorkerType};
use sglang_router_rs::core::{BasicWorkerBuilder, Worker, WorkerType};
use sglang_router_rs::routers::http::pd_types::get_hostname;
use sglang_router_rs::routers::http::pd_types::PDSelectionPolicy;
use sglang_router_rs::routers::RouterFactory;
@@ -46,11 +46,16 @@ mod test_pd_routing {
#[test]
fn test_worker_types() {
use sglang_router_rs::core::{WorkerFactory, WorkerType};
use sglang_router_rs::core::{BasicWorkerBuilder, Worker, WorkerType};
// Test worker creation for prefill servers
let prefill_worker =
WorkerFactory::create_prefill("http://prefill:8080".to_string(), Some(9000));
let prefill_worker: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://prefill:8080")
.worker_type(WorkerType::Prefill {
bootstrap_port: Some(9000),
})
.build(),
);
assert_eq!(prefill_worker.url(), "http://prefill:8080");
match prefill_worker.worker_type() {
WorkerType::Prefill { bootstrap_port } => {
@@ -60,7 +65,11 @@ mod test_pd_routing {
}
// Test worker creation for decode servers
let decode_worker = WorkerFactory::create_decode("http://decode:8080".to_string());
let decode_worker: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://decode:8080")
.worker_type(WorkerType::Decode)
.build(),
);
assert_eq!(decode_worker.url(), "http://decode:8080");
match decode_worker.worker_type() {
WorkerType::Decode => (),
@@ -68,7 +77,11 @@ mod test_pd_routing {
}
// Test regular worker creation
let regular_worker = WorkerFactory::create_regular("http://regular:8080".to_string());
let regular_worker: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://regular:8080")
.worker_type(WorkerType::Regular)
.build(),
);
assert_eq!(regular_worker.url(), "http://regular:8080");
match regular_worker.worker_type() {
WorkerType::Regular => (),
@@ -277,8 +290,13 @@ mod test_pd_routing {
});
// Create a prefill worker to simulate injection
let prefill_worker =
WorkerFactory::create_prefill("http://prefill1:8080".to_string(), Some(9000));
let prefill_worker: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://prefill1:8080")
.worker_type(WorkerType::Prefill {
bootstrap_port: Some(9000),
})
.build(),
);
// Extract bootstrap port from worker type
let bootstrap_port = match prefill_worker.worker_type() {
@@ -660,7 +678,7 @@ mod test_pd_routing {
#[test]
fn test_bootstrap_injection_with_benchmark_requests() {
use sglang_router_rs::core::{WorkerFactory, WorkerType};
use sglang_router_rs::core::{BasicWorkerBuilder, Worker, WorkerType};
// Test bootstrap injection with actual benchmark request patterns
let mut benchmark_request = json!({
@@ -675,8 +693,13 @@ mod test_pd_routing {
});
// Create a prefill worker to simulate injection
let prefill_worker =
WorkerFactory::create_prefill("http://prefill:8080".to_string(), Some(9000));
let prefill_worker: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://prefill:8080")
.worker_type(WorkerType::Prefill {
bootstrap_port: Some(9000),
})
.build(),
);
// Extract bootstrap port from worker type
let bootstrap_port = match prefill_worker.worker_type() {
@@ -806,8 +829,13 @@ mod test_pd_routing {
});
// Create a prefill worker to simulate injection
let prefill_worker =
WorkerFactory::create_prefill("http://prefill:8080".to_string(), Some(9000));
let prefill_worker: Box<dyn Worker> = Box::new(
BasicWorkerBuilder::new("http://prefill:8080")
.worker_type(WorkerType::Prefill {
bootstrap_port: Some(9000),
})
.build(),
);
// Extract bootstrap port from worker type
let bootstrap_port = match prefill_worker.worker_type() {