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

Tiny extract select_worker_min_load (#14648)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-08 21:22:10 -08:00
committed by GitHub
parent af60cad05d
commit 817daba062
1 changed files with 64 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/policies/cache_aware.rs
+54 -32
View File
@@ -210,47 +210,25 @@ impl CacheAwarePolicy {
);
}
}
}
impl LoadBalancingPolicy for CacheAwarePolicy {
fn select_worker(
fn select_worker_min_load(
&self,
workers: &[Arc<dyn Worker>],
request_text: Option<&str>,
request_text: &Option<&str>,
healthy_indices: &[usize],
model_id: &str,
// TODO may skip passing this arg (and compute inside function) if this is not bottleneck
max_load: usize,
min_load: usize,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
// Determine the model for this set of workers (router pre-filters by model)
// All workers should be from the same model
let first_model = workers[healthy_indices[0]].model_id();
let model_id = if first_model.is_empty() || first_model == "unknown" {
"default"
} else {
first_model
};
// Get current load statistics - compute min/max in single pass without allocation
let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| {
let load = w.load();
(min.min(load), max.max(load))
});
let min_load = if min_load == usize::MAX { 0 } else { min_load };
// Check if load is imbalanced
let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold
&& (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold);
if is_imbalanced {
// Log load balancing trigger
// TODO may use `&str`
let worker_loads: Vec<(String, usize)> = workers
.iter()
.map(|w| (w.url().to_string(), w.load()))
.collect();
// TODO may change text
debug!(
"Load balancing triggered | max: {} | min: {} | workers: {:?}",
max_load, min_load, worker_loads
@@ -287,7 +265,51 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
RouterMetrics::record_processed_request(workers[min_load_idx].url());
RouterMetrics::record_policy_decision(self.name(), workers[min_load_idx].url());
return Some(min_load_idx);
Some(min_load_idx)
}
}
impl LoadBalancingPolicy for CacheAwarePolicy {
fn select_worker(
&self,
workers: &[Arc<dyn Worker>],
request_text: Option<&str>,
) -> Option<usize> {
let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() {
return None;
}
// Determine the model for this set of workers (router pre-filters by model)
// All workers should be from the same model
let first_model = workers[healthy_indices[0]].model_id();
let model_id = if first_model.is_empty() || first_model == "unknown" {
"default"
} else {
first_model
};
// Get current load statistics - compute min/max in single pass without allocation
let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| {
let load = w.load();
(min.min(load), max.max(load))
});
let min_load = if min_load == usize::MAX { 0 } else { min_load };
// Check if load is imbalanced
let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold
&& (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold);
if is_imbalanced {
return self.select_worker_min_load(
workers,
&request_text,
&healthy_indices,
model_id,
max_load,
min_load,
);
}
// Use cache-aware routing when balanced