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

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

View File

@@ -1,4 +1,4 @@
use sglang_router_rs::core::{BasicWorker, Worker, WorkerType};
use sglang_router_rs::core::{BasicWorkerBuilder, Worker, WorkerType};
use sglang_router_rs::policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy};
use std::collections::HashMap;
use std::sync::Arc;
@@ -16,13 +16,17 @@ fn test_backward_compatibility_with_empty_model_id() {
let policy = CacheAwarePolicy::with_config(config);
// Create workers with empty model_id (simulating existing routers)
let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular);
let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
.worker_type(WorkerType::Regular)
.build();
// No model_id label - should default to "unknown"
let mut labels2 = HashMap::new();
labels2.insert("model_id".to_string(), "unknown".to_string());
let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular)
.with_labels(labels2);
let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
.worker_type(WorkerType::Regular)
.labels(labels2)
.build();
// Add workers - should both go to "default" tree
policy.add_worker(&worker1);
@@ -53,23 +57,31 @@ fn test_mixed_model_ids() {
let policy = CacheAwarePolicy::with_config(config);
// Create workers with different model_id scenarios
let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular);
let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
.worker_type(WorkerType::Regular)
.build();
// No model_id label - defaults to "unknown" which goes to "default" tree
let mut labels2 = HashMap::new();
labels2.insert("model_id".to_string(), "llama-3".to_string());
let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular)
.with_labels(labels2);
let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
.worker_type(WorkerType::Regular)
.labels(labels2)
.build();
let mut labels3 = HashMap::new();
labels3.insert("model_id".to_string(), "unknown".to_string());
let worker3 = BasicWorker::new("http://worker3:8080".to_string(), WorkerType::Regular)
.with_labels(labels3);
let worker3 = BasicWorkerBuilder::new("http://worker3:8080")
.worker_type(WorkerType::Regular)
.labels(labels3)
.build();
let mut labels4 = HashMap::new();
labels4.insert("model_id".to_string(), "llama-3".to_string());
let worker4 = BasicWorker::new("http://worker4:8080".to_string(), WorkerType::Regular)
.with_labels(labels4);
let worker4 = BasicWorkerBuilder::new("http://worker4:8080")
.worker_type(WorkerType::Regular)
.labels(labels4)
.build();
// Add all workers
policy.add_worker(&worker1);
@@ -108,10 +120,14 @@ fn test_remove_worker_by_url_backward_compat() {
// Create workers with different model_ids
let mut labels1 = HashMap::new();
labels1.insert("model_id".to_string(), "llama-3".to_string());
let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular)
.with_labels(labels1);
let worker1 = BasicWorkerBuilder::new("http://worker1:8080")
.worker_type(WorkerType::Regular)
.labels(labels1)
.build();
let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular);
let worker2 = BasicWorkerBuilder::new("http://worker2:8080")
.worker_type(WorkerType::Regular)
.build();
// No model_id label - defaults to "unknown"
// Add workers