[router] refactor worker to builder pattern 3/n (#10647)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user