minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Change routing policy API to be async to support more policies (#17048)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-17 17:12:51 +08:00
committed by GitHub
parent d2c863878c
commit 9c2530642c
16 changed files with 350 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs
+27 -23
View File
@@ -5,8 +5,8 @@ use smg::{
policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, SelectWorkerInfo},
};
#[test]
fn test_backward_compatibility_with_empty_model_id() {
#[tokio::test]
async fn test_backward_compatibility_with_empty_model_id() {
let config = CacheAwareConfig {
cache_threshold: 0.5,
balance_abs_threshold: 2,
@@ -40,13 +40,15 @@ fn test_backward_compatibility_with_empty_model_id() {
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker1.clone()), Arc::new(worker2.clone())];
// Select worker - should work without errors
let selected = policy.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test request"),
..Default::default()
},
);
let selected = policy
.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test request"),
..Default::default()
},
)
.await;
assert!(selected.is_some(), "Should select a worker");
// Remove workers - should work without errors
@@ -54,8 +56,8 @@ fn test_backward_compatibility_with_empty_model_id() {
policy.remove_worker(&worker2);
}
#[test]
fn test_mixed_model_ids() {
#[tokio::test]
async fn test_mixed_model_ids() {
let config = CacheAwareConfig {
cache_threshold: 0.5,
balance_abs_threshold: 2,
@@ -107,12 +109,12 @@ fn test_mixed_model_ids() {
request_text: Some("test request"),
..Default::default()
};
let selected = policy.select_worker(&default_workers, &info);
let selected = policy.select_worker(&default_workers, &info).await;
assert!(selected.is_some(), "Should select from default workers");
let llama_workers: Vec<Arc<dyn Worker>> =
vec![Arc::new(worker2.clone()), Arc::new(worker4.clone())];
let selected = policy.select_worker(&llama_workers, &info);
let selected = policy.select_worker(&llama_workers, &info).await;
assert!(selected.is_some(), "Should select from llama-3 workers");
let all_workers: Vec<Arc<dyn Worker>> = vec![
@@ -121,12 +123,12 @@ fn test_mixed_model_ids() {
Arc::new(worker3.clone()),
Arc::new(worker4.clone()),
];
let selected = policy.select_worker(&all_workers, &info);
let selected = policy.select_worker(&all_workers, &info).await;
assert!(selected.is_some(), "Should select from all workers");
}
#[test]
fn test_remove_worker_by_url_backward_compat() {
#[tokio::test]
async fn test_remove_worker_by_url_backward_compat() {
let config = CacheAwareConfig::default();
let policy = CacheAwarePolicy::with_config(config);
@@ -154,12 +156,14 @@ fn test_remove_worker_by_url_backward_compat() {
policy.remove_worker_by_url("http://worker1:8080");
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker2.clone())];
let selected = policy.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test"),
..Default::default()
},
);
let selected = policy
.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test"),
..Default::default()
},
)
.await;
assert_eq!(selected, Some(0), "Should only have worker2 left");
}