From a8361894bc6a26713744ba53f403bc920b720039 Mon Sep 17 00:00:00 2001 From: Mingxing Zhang Date: Sun, 5 Apr 2026 16:55:25 +0800 Subject: [PATCH] smg: align cache_aware load fallback semantics with power_of_two --- sgl-model-gateway/src/core/worker_manager.rs | 12 ++-- sgl-model-gateway/src/policies/cache_aware.rs | 69 +++++++++++++++---- sgl-model-gateway/src/policies/registry.rs | 45 ++++++++++++ 3 files changed, 108 insertions(+), 18 deletions(-) diff --git a/sgl-model-gateway/src/core/worker_manager.rs b/sgl-model-gateway/src/core/worker_manager.rs index 4a21ae3fb802..b7ea80042b07 100644 --- a/sgl-model-gateway/src/core/worker_manager.rs +++ b/sgl-model-gateway/src/core/worker_manager.rs @@ -347,10 +347,10 @@ impl LoadMonitor { loop { interval_timer.tick().await; - let power_of_two_policies = policy_registry.get_all_power_of_two_policies(); + let external_load_policies = policy_registry.get_all_external_load_policies(); - if power_of_two_policies.is_empty() { - debug!("No PowerOfTwo policies found, skipping load fetch"); + if external_load_policies.is_empty() { + debug!("No load-aware policies found, skipping load fetch"); continue; } @@ -363,11 +363,11 @@ impl LoadMonitor { if !loads.is_empty() { debug!( - "Fetched loads from {} workers, updating {} PowerOfTwo policies", + "Fetched loads from {} workers, updating {} load-aware policies", loads.len(), - power_of_two_policies.len() + external_load_policies.len() ); - for policy in &power_of_two_policies { + for policy in &external_load_policies { policy.update_loads(&loads); } let _ = tx.send(loads); diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index dfc3ef6462c5..563090487759 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -59,7 +59,10 @@ during the next eviction cycle. */ -use std::sync::Arc; +use std::{ + collections::HashMap, + sync::{Arc, RwLock}, +}; use async_trait::async_trait; use dashmap::DashMap; @@ -84,6 +87,8 @@ use crate::core::{Worker, UNKNOWN_MODEL_ID}; pub struct CacheAwarePolicy { config: CacheAwareConfig, trees: Arc>>, + /// Cached load information from external monitoring (/get_load num_tokens sum). + cached_loads: RwLock>, mesh_sync: OptionalMeshSyncManager, _eviction_task: Option, } @@ -124,6 +129,7 @@ impl CacheAwarePolicy { Self { config, trees, + cached_loads: RwLock::new(HashMap::new()), mesh_sync: None, _eviction_task: eviction_task, } @@ -288,6 +294,7 @@ impl CacheAwarePolicy { workers: &[Arc], request_text: &Option<&str>, healthy_indices: &[usize], + effective_loads: &[(usize, isize)], model_id: &str, max_load: usize, min_load: usize, @@ -303,10 +310,11 @@ impl CacheAwarePolicy { } // Use shortest queue when imbalanced - let min_load_idx = healthy_indices + let min_load_idx = effective_loads .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied()?; + .min_by_key(|(_, load)| *load) + .map(|(idx, _)| *idx) + .or_else(|| healthy_indices.first().copied())?; // Even in imbalanced mode, update the tree to maintain cache state if let Some(text) = request_text { @@ -364,12 +372,42 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // All workers should be from the same model let model_id = normalize_model_key(workers[healthy_indices[0]].model_id()); - // 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)) + // Resolve effective loads using the same fallback semantics as PowerOfTwoPolicy: + // 1) Prefer cached external token loads only when ALL compared workers have entries. + // 2) If any worker is missing external load data, fallback to local request counters for ALL workers. + let load_guard = self.cached_loads.read().ok(); + let token_loads: Option)>> = load_guard.as_ref().map(|m| { + healthy_indices + .iter() + .map(|&idx| (idx, m.get(workers[idx].url()).copied())) + .collect() }); - let min_load = if min_load == usize::MAX { 0 } else { min_load }; + let effective_loads: Vec<(usize, isize)> = match token_loads { + Some(loads) if loads.iter().all(|(_, l)| l.is_some()) => loads + .into_iter() + .map(|(idx, l)| (idx, l.unwrap())) + .collect(), + _ => healthy_indices + .iter() + .map(|&idx| (idx, workers[idx].load() as isize)) + .collect(), + }; + + let (min_load_i, max_load_i) = effective_loads + .iter() + .fold((isize::MAX, isize::MIN), |(min, max), (_, load)| { + (min.min(*load), max.max(*load)) + }); + let min_load = if min_load_i == isize::MAX { + 0 + } else { + min_load_i.max(0) as usize + }; + let max_load = if max_load_i == isize::MIN { + 0 + } else { + max_load_i.max(0) as usize + }; // Check if load is imbalanced let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold @@ -380,6 +418,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { workers, &request_text, &healthy_indices, + &effective_loads, model_id, max_load, min_load, @@ -413,10 +452,10 @@ impl LoadBalancingPolicy for CacheAwarePolicy { .filter(|&idx| workers[idx].is_healthy()) } else { // Low cache match: use worker with minimum load - healthy_indices + effective_loads .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied() + .min_by_key(|(_, load)| *load) + .map(|(idx, _)| *idx) }; if let Some(idx) = selected_idx { @@ -492,6 +531,12 @@ impl LoadBalancingPolicy for CacheAwarePolicy { "cache_aware" } + fn update_loads(&self, loads: &HashMap) { + if let Ok(mut cached) = self.cached_loads.write() { + *cached = loads.clone(); + } + } + fn needs_request_text(&self) -> bool { true // Cache-aware policy needs request text for cache affinity } diff --git a/sgl-model-gateway/src/policies/registry.rs b/sgl-model-gateway/src/policies/registry.rs index 2e0e4b9eb0e7..988eac2995fa 100644 --- a/sgl-model-gateway/src/policies/registry.rs +++ b/sgl-model-gateway/src/policies/registry.rs @@ -294,6 +294,51 @@ impl PolicyRegistry { power_of_two_policies } + /// Get all policies that consume externally sampled worker loads. + /// + /// Currently includes: + /// - power_of_two + /// - cache_aware + pub fn get_all_external_load_policies(&self) -> Vec> { + let mut policies = Vec::new(); + + let should_include = |name: &str| name == "power_of_two" || name == "cache_aware"; + + if should_include(self.default_policy.name()) { + policies.push(Arc::clone(&self.default_policy)); + } + + let prefill_policy_opt = self.prefill_policy.get(); + let decode_policy_opt = self.decode_policy.get(); + + if let Some(policy) = prefill_policy_opt { + if should_include(policy.name()) && !Arc::ptr_eq(policy, &self.default_policy) { + policies.push(Arc::clone(policy)); + } + } + + if let Some(policy) = decode_policy_opt { + if should_include(policy.name()) + && !Arc::ptr_eq(policy, &self.default_policy) + && !prefill_policy_opt.is_some_and(|p| Arc::ptr_eq(p, policy)) + { + policies.push(Arc::clone(policy)); + } + } + + for entry in self.model_policies.iter() { + let policy = entry.value(); + if should_include(policy.name()) { + let already_added = policies.iter().any(|p| Arc::ptr_eq(p, policy)); + if !already_added { + policies.push(Arc::clone(policy)); + } + } + } + + policies + } + /// Initialize cache-aware policy with workers if applicable /// This should be called after workers are registered for a model pub fn init_cache_aware_policy(&self, model_id: &str, workers: &[Arc]) {