diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 4a67d9dcc..c1d7ee3df 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -34,7 +34,8 @@ 2. Load Balancing (Shortest Queue) ------------------------------------------- This strategy tracks pending request counts per worker and routes new requests - to the least busy worker when the system is detected to be imbalanced. + to the least busy worker when the system is detected to be imbalanced. Ties + are randomly broken. Configuration Parameters: ------------------------ @@ -63,7 +64,7 @@ use std::sync::Arc; use async_trait::async_trait; use dashmap::DashMap; -use rand::Rng; +use rand::{seq::IteratorRandom, Rng}; use smg_mesh::{tree_ops::TreeOperation, OptionalMeshSyncManager}; use tracing::{debug, warn}; @@ -325,11 +326,21 @@ impl CacheAwarePolicy { ); } - // Use shortest queue when imbalanced - let min_load_idx = healthy_indices + // Use shortest queue when imbalanced. Tie break randomly. + // Snapshot load() (live atomic count of load). Without snapshot + // there could be no workers found matching min_load because of + // load update. + let loads: Vec<(usize, usize)> = healthy_indices .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied()?; + .map(|&idx| (idx, workers[idx].load())) + .collect(); + let min_load = loads.iter().map(|&(_, load)| load).min()?; + let min_load_idx = loads + .iter() + .copied() + .filter(|&(_, load)| load == min_load) + .map(|(idx, _)| idx) + .choose(&mut rand::rng())?; // Even in imbalanced mode, update the tree to maintain cache state if let Some(text) = request_text { @@ -438,11 +449,21 @@ impl LoadBalancingPolicy for CacheAwarePolicy { .position(|w| w.url() == tenant_url) .filter(|&idx| workers[idx].is_healthy()) } else { - // Low cache match: use worker with minimum load - healthy_indices + // Low cache match: use worker with minimum load. Tie break randomly. + // Snapshot load() (live atomic count of load). Without snapshot + // there could be no workers found matching min_load because of + // load update. + let loads: Vec<(usize, usize)> = healthy_indices + .iter() + .map(|&idx| (idx, workers[idx].load())) + .collect(); + let min_load = loads.iter().map(|&(_, load)| load).min()?; + loads .iter() - .min_by_key(|&&idx| workers[idx].load()) .copied() + .filter(|&(_, load)| load == min_load) + .map(|(idx, _)| idx) + .choose(&mut rand::rng()) }; if let Some(idx) = selected_idx { @@ -642,6 +663,113 @@ mod tests { } } + // In imbalanced mode the overloaded worker must be avoided AND the remaining + // tied (min-load) workers must be spread across via random tie-breaking. + #[tokio::test] + async fn test_cache_aware_imbalanced_random_tie_break() { + let policy = CacheAwarePolicy::with_config(CacheAwareConfig { + cache_threshold: 0.5, + balance_abs_threshold: 5, + balance_rel_threshold: 2.0, + eviction_interval_secs: 0, + max_tree_size: 10000, + }); + + let num_workers = 5; + let mut workers: Vec> = Vec::new(); + for j in 0..num_workers { + workers.push(Arc::new( + BasicWorkerBuilder::new(&format!("http://w{}:8000", j)) + .worker_type(WorkerType::Regular) + .build(), + )); + } + + // Overload worker 0: max=50, min=0 => (50-0) > 5 AND 50 > 0*2.0 => imbalanced. + for _ in 0..50 { + workers[0].increment_load(); + } + policy.init_workers(&workers); + + // Reuse the SAME prompt: the imbalanced branch bypasses cache affinity, + // so even a guaranteed cache hit must not pin all traffic to one worker. + let mut selection_counts = vec![0; num_workers]; + for _ in 0..100 { + let info = SelectWorkerInfo { + request_text: Some("same_prompt"), + ..Default::default() + }; + let idx = policy.select_worker(&workers, &info).await.unwrap(); + selection_counts[idx] += 1; + } + + // The overloaded worker is never selected in imbalanced mode. + assert_eq!( + selection_counts[0], 0, + "overloaded worker was selected in imbalanced mode: {:?}", + selection_counts + ); + // Every tied (load-0) worker is selected at least once. + for idx in 1..num_workers { + assert!( + selection_counts[idx] > 0, + "tied worker {} was never selected in imbalanced mode: {:?}", + idx, + selection_counts + ); + } + } + + // Verify random tie breaking for cache misses and low load situations. + // Important for when concurrency is lower than the number of workers. + // Without random tie breaking there will be multiple workers with 0 load + // and all requests will be sent to the same replica which doesn't utilize + // the available memory on all workers for the KV cache. + #[tokio::test] + async fn test_cache_aware_random_tie_break_cold_start() { + let config = CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }; + let policy = CacheAwarePolicy::with_config(config); + + // Create 5 workers + let num_workers = 5; + let mut workers: Vec> = Vec::new(); + for j in 0..num_workers { + workers.push(Arc::new( + BasicWorkerBuilder::new(&format!("http://w{}:8000", j)) + .worker_type(WorkerType::Regular) + .build(), + )); + } + policy.init_workers(&workers); + + // Send 100 requests with unique prompts to simulate 100 cache misses. + let mut selection_counts = vec![0; num_workers]; + for i in 0..100 { + let prompt = format!("{}_unique_request", i); + let info = SelectWorkerInfo { + request_text: Some(&prompt), + ..Default::default() + }; + let idx = policy.select_worker(&workers, &info).await.unwrap(); + selection_counts[idx] += 1; + } + + // Verify all workers were selected at least once when prompts are + // unique. + for (idx, &count) in selection_counts.iter().enumerate() { + assert!( + count > 0, + "Worker {} was never selected; tie-breaking failed to spread across \ + all tied workers: {:?}", + idx, + selection_counts + ); + } + } + #[tokio::test] async fn test_cache_aware_worker_removal() { let config = CacheAwareConfig { @@ -941,6 +1069,7 @@ mod tests { async fn test_pd_pool_isolation_two_policies() { let config = CacheAwareConfig { eviction_interval_secs: 0, + cache_threshold: 0.0, ..Default::default() }; let prefill_policy = CacheAwarePolicy::with_config(config.clone()); @@ -1005,6 +1134,7 @@ mod tests { async fn test_pd_pool_isolation_shared_policy_regression() { let config = CacheAwareConfig { eviction_interval_secs: 0, + cache_threshold: 0.0, ..Default::default() }; let policy = CacheAwarePolicy::with_config(config);