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

[router] refactor worker to builder pattern 3/n (#10647)

Browse Source
This commit is contained in:
Simo Lin
2025-09-19 01:52:57 -04:00
committed by GitHub
parent 56b991b12d
commit ac2a723bb3
11 changed files with 267 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
sgl-router/src/policies/cache_aware.rs
View File

@@ -451,7 +451,7 @@ impl Drop for CacheAwarePolicy {
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use crate::core::{BasicWorkerBuilder, WorkerType};
#[test]
fn test_cache_aware_with_balanced_load() {
@@ -462,14 +462,16 @@ mod tests {
};
let policy = CacheAwarePolicy::with_config(config);
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Initialize the policy with workers
@@ -497,8 +499,12 @@ mod tests {
max_tree_size: 10000,
});
let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular);
let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular);
let worker1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build();
let worker2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build();
// Create significant load imbalance
for _ in 0..20 {
@@ -524,14 +530,16 @@ mod tests {
};
let policy = CacheAwarePolicy::with_config(config);
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
policy.init_workers(&workers);

29
sgl-router/src/policies/mod.rs
View File

@@ -121,23 +121,26 @@ pub(crate) fn get_healthy_worker_indices(workers: &[Arc<dyn Worker>]) -> Vec<usi
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use crate::core::{BasicWorkerBuilder, WorkerType};
#[test]
fn test_get_healthy_worker_indices() {
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// All healthy initially

41
sgl-router/src/policies/power_of_two.rs
View File

@@ -119,14 +119,20 @@ impl Default for PowerOfTwoPolicy {
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use crate::core::{BasicWorkerBuilder, WorkerType};
#[test]
fn test_power_of_two_selection() {
let policy = PowerOfTwoPolicy::new();
let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular);
let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular);
let worker3 = BasicWorker::new("http://w3:8000".to_string(), WorkerType::Regular);
let worker1 = BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build();
let worker2 = BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build();
let worker3 = BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.build();
// Set different loads
for _ in 0..10 {
@@ -157,14 +163,16 @@ mod tests {
fn test_power_of_two_with_cached_loads() {
let policy = PowerOfTwoPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Update cached loads
@@ -190,10 +198,11 @@ mod tests {
#[test]
fn test_power_of_two_single_worker() {
let policy = PowerOfTwoPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
))];
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
)];
// With single worker, should always select it
assert_eq!(policy.select_worker(&workers, None), Some(0));

56
sgl-router/src/policies/random.rs
View File

@@ -51,25 +51,28 @@ impl LoadBalancingPolicy for RandomPolicy {
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use crate::core::{BasicWorkerBuilder, WorkerType};
use std::collections::HashMap;
#[test]
fn test_random_selection() {
let policy = RandomPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Test multiple selections to ensure randomness
@@ -89,14 +92,16 @@ mod tests {
fn test_random_with_unhealthy_workers() {
let policy = RandomPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Mark first worker as unhealthy
@@ -111,10 +116,11 @@ mod tests {
#[test]
fn test_random_no_healthy_workers() {
let policy = RandomPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
))];
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
)];
workers[0].set_healthy(false);
assert_eq!(policy.select_worker(&workers, None), None);

74
sgl-router/src/policies/round_robin.rs
View File

@@ -60,24 +60,27 @@ impl LoadBalancingPolicy for RoundRobinPolicy {
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BasicWorker, WorkerType};
use crate::core::{BasicWorkerBuilder, WorkerType};
#[test]
fn test_round_robin_selection() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Should select workers in order: 0, 1, 2, 0, 1, 2, ...
@@ -92,18 +95,21 @@ mod tests {
fn test_round_robin_with_unhealthy_workers() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w3:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w3:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Mark middle worker as unhealthy
@@ -120,14 +126,16 @@ mod tests {
fn test_round_robin_reset() {
let policy = RoundRobinPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![
Arc::new(BasicWorker::new(
"http://w1:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(BasicWorker::new(
"http://w2:8000".to_string(),
WorkerType::Regular,
)),
Arc::new(
BasicWorkerBuilder::new("http://w1:8000")
.worker_type(WorkerType::Regular)
.build(),
),
Arc::new(
BasicWorkerBuilder::new("http://w2:8000")
.worker_type(WorkerType::Regular)
.build(),
),
];
// Advance the counter