chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Tiny refactor select_workers API for future passing more information (#15596)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-25 10:31:10 +08:00
committed by GitHub
parent 38dd4fbb66
commit 1ba897f330
10 changed files with 1463 additions and 1188 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2448
sgl-model-gateway/src/policies/bucket.rs
View File

File diff suppressed because it is too large Load Diff

64
sgl-model-gateway/src/policies/cache_aware.rs
View File

@@ -74,7 +74,7 @@ use tracing::debug;
use super::{
get_healthy_worker_indices, normalize_model_key, tree::Tree, CacheAwareConfig,
LoadBalancingPolicy,
LoadBalancingPolicy, SelectWorkerInfo,
};
use crate::core::Worker;
@@ -284,11 +284,8 @@ impl CacheAwarePolicy {
}
impl LoadBalancingPolicy for CacheAwarePolicy {
fn select_worker(
&self,
workers: &[Arc<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize> {
fn select_worker(&self, workers: &[Arc<dyn Worker>], info: &SelectWorkerInfo) -> Option<usize> {
let request_text = info.request_text;
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
@@ -458,14 +455,35 @@ mod tests {
policy.init_workers(&workers);
// First request should be distributed
let idx1 = policy.select_worker(&workers, Some("hello world")).unwrap();
let idx1 = policy
.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("hello world"),
},
)
.unwrap();
// Same request should go to same worker (cache hit)
let idx2 = policy.select_worker(&workers, Some("hello world")).unwrap();
let idx2 = policy
.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("hello world"),
},
)
.unwrap();
assert_eq!(idx1, idx2);
// Similar request should also go to same worker
let idx3 = policy.select_worker(&workers, Some("hello")).unwrap();
let idx3 = policy
.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("hello"),
},
)
.unwrap();
assert_eq!(idx1, idx3);
}
@@ -496,8 +514,11 @@ mod tests {
policy.init_workers(&workers);
// Should select worker2 (lower load) despite cache affinity
let info = SelectWorkerInfo {
request_text: Some("test"),
};
for _ in 0..5 {
let idx = policy.select_worker(&workers, Some("test")).unwrap();
let idx = policy.select_worker(&workers, &info).unwrap();
assert_eq!(idx, 1); // Should always pick worker2
}
}
@@ -525,15 +546,32 @@ mod tests {
policy.init_workers(&workers);
// Route some requests
policy.select_worker(&workers, Some("test1"));
policy.select_worker(&workers, Some("test2"));
policy.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test1"),
},
);
policy.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test2"),
},
);
// Remove a worker
policy.remove_worker_by_url("http://w1:8000");
workers[0].set_healthy(false);
// All requests should now go to worker2
let idx = policy.select_worker(&workers, Some("test1")).unwrap();
let idx = policy
.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test1"),
},
)
.unwrap();
assert_eq!(idx, 1);
}
}

17
sgl-model-gateway/src/policies/mod.rs
View File

@@ -33,11 +33,11 @@ pub trait LoadBalancingPolicy: Send + Sync + Debug {
///
/// This is used for regular routing mode where requests go to a single worker.
/// Now uses Arc<dyn Worker> for better performance and to avoid unnecessary cloning.
fn select_worker(
&self,
workers: &[Arc<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize>;
///
/// # Arguments
/// * `workers` - Available workers to select from
/// * `info` - Additional information for routing decisions
fn select_worker(&self, workers: &[Arc<dyn Worker>], info: &SelectWorkerInfo) -> Option<usize>;
/// Update policy state after request completion
///
@@ -135,6 +135,13 @@ pub(crate) fn normalize_model_key(model_id: &str) -> &str {
}
}
/// Information passed to policy for worker selection
#[derive(Debug, Default, Clone)]
pub struct SelectWorkerInfo<'a> {
/// Request text for cache-aware routing
pub request_text: Option<&'a str>,
}
#[cfg(test)]
mod tests {
use super::*;

33
sgl-model-gateway/src/policies/power_of_two.rs
View File

@@ -8,7 +8,7 @@ use std::{
use rand::Rng;
use tracing::debug;
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
use crate::core::Worker;
/// Power-of-two choices policy
@@ -33,7 +33,7 @@ impl LoadBalancingPolicy for PowerOfTwoPolicy {
fn select_worker(
&self,
workers: &[Arc<dyn Worker>],
_request_text: Option<&str>,
_info: &SelectWorkerInfo,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
@@ -157,8 +157,9 @@ mod tests {
// Run multiple selections
let mut selected_counts = [0; 3];
let info = SelectWorkerInfo::default();
for _ in 0..100 {
if let Some(idx) = policy.select_worker(&workers, None) {
if let Some(idx) = policy.select_worker(&workers, &info) {
selected_counts[idx] += 1;
}
}
@@ -192,8 +193,9 @@ mod tests {
// Should prefer worker2 with lower cached load
let mut w2_selected = 0;
let info = SelectWorkerInfo::default();
for _ in 0..50 {
if let Some(idx) = policy.select_worker(&workers, None) {
if let Some(idx) = policy.select_worker(&workers, &info) {
if idx == 1 {
w2_selected += 1;
}
@@ -214,7 +216,10 @@ mod tests {
)];
// With single worker, should always select it
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(
policy.select_worker(&workers, &SelectWorkerInfo::default()),
Some(0)
);
}
#[test]
@@ -251,7 +256,7 @@ mod tests {
// 5. Run selection
let selected_idx = policy
.select_worker(&workers, None)
.select_worker(&workers, &SelectWorkerInfo::default())
.expect("Should select a worker");
// 6. Verify the Fix
@@ -307,7 +312,9 @@ mod tests {
loads_1.insert("http://b:8000".to_string(), 100_000);
policy.update_loads(&loads_1);
let idx_1 = policy.select_worker(&workers_1, None).unwrap();
let idx_1 = policy
.select_worker(&workers_1, &SelectWorkerInfo::default())
.unwrap();
assert_eq!(
idx_1, 0,
"Happy Path Failed: Should select Worker A (fewer tokens) despite higher request count"
@@ -326,7 +333,9 @@ mod tests {
// http://d:8000 is MISSING
policy.update_loads(&loads_2);
let idx_2 = policy.select_worker(&workers_2, None).unwrap();
let idx_2 = policy
.select_worker(&workers_2, &SelectWorkerInfo::default())
.unwrap();
assert_eq!(idx_2, 1, "Partial Fail 1 Failed: Should fallback to requests and select Worker B (fewer requests)");
// Scenario 3: Partial Failure (Worker A is missing, Worker B has tokens)
@@ -342,7 +351,9 @@ mod tests {
loads_3.insert("http://f:8000".to_string(), 1_000);
policy.update_loads(&loads_3);
let idx_3 = policy.select_worker(&workers_3, None).unwrap();
let idx_3 = policy
.select_worker(&workers_3, &SelectWorkerInfo::default())
.unwrap();
assert_eq!(idx_3, 0, "Partial Fail 2 Failed: Should fallback to requests and select Worker A (fewer requests)");
// Scenario 4: Total Failure (Both missing)
@@ -356,7 +367,9 @@ mod tests {
let loads_4 = HashMap::new();
policy.update_loads(&loads_4);
let idx_4 = policy.select_worker(&workers_4, None).unwrap();
let idx_4 = policy
.select_worker(&workers_4, &SelectWorkerInfo::default())
.unwrap();
assert_eq!(
idx_4, 1,
"Total Fail Failed: Should select Worker B based on request count"

16
sgl-model-gateway/src/policies/random.rs
View File

@@ -4,7 +4,7 @@ use std::sync::Arc;
use rand::Rng;
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
use crate::core::Worker;
/// Random selection policy
@@ -23,7 +23,7 @@ impl LoadBalancingPolicy for RandomPolicy {
fn select_worker(
&self,
workers: &[Arc<dyn Worker>],
_request_text: Option<&str>,
_info: &SelectWorkerInfo,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
@@ -76,7 +76,7 @@ mod tests {
let mut counts = HashMap::new();
for _ in 0..100 {
if let Some(idx) = policy.select_worker(&workers, None) {
if let Some(idx) = policy.select_worker(&workers, &SelectWorkerInfo::default()) {
*counts.entry(idx).or_insert(0) += 1;
}
}
@@ -107,7 +107,10 @@ mod tests {
// Should always select the healthy worker (index 1)
for _ in 0..10 {
assert_eq!(policy.select_worker(&workers, None), Some(1));
assert_eq!(
policy.select_worker(&workers, &SelectWorkerInfo::default()),
Some(1)
);
}
}
@@ -121,6 +124,9 @@ mod tests {
)];
workers[0].set_healthy(false);
assert_eq!(policy.select_worker(&workers, None), None);
assert_eq!(
policy.select_worker(&workers, &SelectWorkerInfo::default()),
None
);
}
}

31
sgl-model-gateway/src/policies/round_robin.rs
View File

@@ -5,7 +5,7 @@ use std::sync::{
Arc,
};
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
use crate::core::Worker;
/// Round-robin selection policy
@@ -28,7 +28,7 @@ impl LoadBalancingPolicy for RoundRobinPolicy {
fn select_worker(
&self,
workers: &[Arc<dyn Worker>],
_request_text: Option<&str>,
_info: &SelectWorkerInfo,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
@@ -83,11 +83,12 @@ mod tests {
];
// Should select workers in order: 0, 1, 2, 0, 1, 2, ...
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
assert_eq!(policy.select_worker(&workers, None), Some(2));
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
let info = SelectWorkerInfo::default();
assert_eq!(policy.select_worker(&workers, &info), Some(0));
assert_eq!(policy.select_worker(&workers, &info), Some(1));
assert_eq!(policy.select_worker(&workers, &info), Some(2));
assert_eq!(policy.select_worker(&workers, &info), Some(0));
assert_eq!(policy.select_worker(&workers, &info), Some(1));
}
#[test]
@@ -115,10 +116,11 @@ mod tests {
workers[1].set_healthy(false);
// Should skip unhealthy worker: 0, 2, 0, 2, ...
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(2));
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(2));
let info = SelectWorkerInfo::default();
assert_eq!(policy.select_worker(&workers, &info), Some(0));
assert_eq!(policy.select_worker(&workers, &info), Some(2));
assert_eq!(policy.select_worker(&workers, &info), Some(0));
assert_eq!(policy.select_worker(&workers, &info), Some(2));
}
#[test]
@@ -138,11 +140,12 @@ mod tests {
];
// Advance the counter
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, None), Some(1));
let info = SelectWorkerInfo::default();
assert_eq!(policy.select_worker(&workers, &info), Some(0));
assert_eq!(policy.select_worker(&workers, &info), Some(1));
// Reset should start from beginning
policy.reset();
assert_eq!(policy.select_worker(&workers, None), Some(0));
assert_eq!(policy.select_worker(&workers, &info), Some(0));
}
}

9
sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs
View File

@@ -10,7 +10,7 @@ use super::PipelineStage;
use crate::{
core::{ConnectionMode, Worker, WorkerRegistry, WorkerType},
observability::metrics::{metrics_labels, Metrics},
policies::PolicyRegistry,
policies::{PolicyRegistry, SelectWorkerInfo},
routers::{
error,
grpc::context::{RequestContext, WorkerSelection},
@@ -146,7 +146,7 @@ impl WorkerSelectionStage {
};
// Select worker using the policy
let idx = policy.select_worker(&available, text)?;
let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?;
let selected = available[idx].clone();
// Record worker selection metric
@@ -203,8 +203,9 @@ impl WorkerSelectionStage {
None => self.policy_registry.get_default_policy(),
};
let prefill_idx = policy.select_worker(&available_prefill, text)?;
let decode_idx = policy.select_worker(&available_decode, text)?;
let info = SelectWorkerInfo { request_text: text };
let prefill_idx = policy.select_worker(&available_prefill, &info)?;
let decode_idx = policy.select_worker(&available_decode, &info)?;
let model = model_id.unwrap_or("default");
let policy_name = policy.name();

4
sgl-model-gateway/src/routers/http/pd_router.rs
View File

@@ -26,7 +26,7 @@ use crate::{
metrics::{bool_to_static_str, metrics_labels, Metrics},
otel_trace::inject_trace_context_http,
},
policies::{LoadBalancingPolicy, PolicyRegistry},
policies::{LoadBalancingPolicy, PolicyRegistry, SelectWorkerInfo},
protocols::{
chat::{ChatCompletionRequest, ChatMessage, MessageContent},
common::{InputIds, StringOrArray},
@@ -784,7 +784,7 @@ impl PDRouter {
}
let selected_idx = policy
.select_worker(&available_workers, request_text)
.select_worker(&available_workers, &SelectWorkerInfo { request_text })
.ok_or_else(|| {
format!(
"Policy {} failed to select a {} worker",

4
sgl-model-gateway/src/routers/http/router.rs
View File

@@ -27,7 +27,7 @@ use crate::{
metrics::{bool_to_static_str, metrics_labels, Metrics},
otel_trace::inject_trace_context_http,
},
policies::PolicyRegistry,
policies::{PolicyRegistry, SelectWorkerInfo},
protocols::{
chat::ChatCompletionRequest,
classify::ClassifyRequest,
@@ -168,7 +168,7 @@ impl Router {
None => self.policy_registry.get_default_policy(),
};
let idx = policy.select_worker(&available, text)?;
let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?;
// Record worker selection metric (Layer 3)
Metrics::record_worker_selection(

25
sgl-model-gateway/tests/cache_aware_backward_compat_test.rs
View File

@@ -2,7 +2,7 @@ use std::{collections::HashMap, sync::Arc};
use sgl_model_gateway::{
core::{BasicWorkerBuilder, Worker, WorkerType},
policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy},
policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, SelectWorkerInfo},
};
#[test]
@@ -40,7 +40,12 @@ 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, Some("test request"));
let selected = policy.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test request"),
},
);
assert!(selected.is_some(), "Should select a worker");
// Remove workers - should work without errors
@@ -97,12 +102,15 @@ fn test_mixed_model_ids() {
let default_workers: Vec<Arc<dyn Worker>> =
vec![Arc::new(worker1.clone()), Arc::new(worker3.clone())];
let selected = policy.select_worker(&default_workers, Some("test request"));
let info = SelectWorkerInfo {
request_text: Some("test request"),
};
let selected = policy.select_worker(&default_workers, &info);
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, Some("test request"));
let selected = policy.select_worker(&llama_workers, &info);
assert!(selected.is_some(), "Should select from llama-3 workers");
let all_workers: Vec<Arc<dyn Worker>> = vec![
@@ -111,7 +119,7 @@ fn test_mixed_model_ids() {
Arc::new(worker3.clone()),
Arc::new(worker4.clone()),
];
let selected = policy.select_worker(&all_workers, Some("test request"));
let selected = policy.select_worker(&all_workers, &info);
assert!(selected.is_some(), "Should select from all workers");
}
@@ -144,6 +152,11 @@ 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, Some("test"));
let selected = policy.select_worker(
&workers,
&SelectWorkerInfo {
request_text: Some("test"),
},
);
assert_eq!(selected, Some(0), "Should only have worker2 left");
}