From e99ee0c695f55a70d05ee2165807358af9713363 Mon Sep 17 00:00:00 2001 From: Praneth Paruchuri Date: Thu, 11 Dec 2025 00:45:38 +0530 Subject: [PATCH] [model-gateway] Fix incompatible metric comparison in` PowerOfTwo` policy (#14823) --- .../src/policies/power_of_two.rs | 195 ++++++++++++++++-- 1 file changed, 177 insertions(+), 18 deletions(-) diff --git a/sgl-model-gateway/src/policies/power_of_two.rs b/sgl-model-gateway/src/policies/power_of_two.rs index 18a2cdca2..9f7856eb0 100644 --- a/sgl-model-gateway/src/policies/power_of_two.rs +++ b/sgl-model-gateway/src/policies/power_of_two.rs @@ -27,18 +27,6 @@ impl PowerOfTwoPolicy { cached_loads: RwLock::new(HashMap::new()), } } - - fn get_worker_load(&self, worker: &dyn Worker) -> isize { - // First check cached loads (from external monitoring) - if let Ok(loads) = self.cached_loads.read() { - if let Some(&load) = loads.get(worker.url()) { - return load; - } - } - - // Fall back to local load counter - worker.load() as isize - } } impl LoadBalancingPolicy for PowerOfTwoPolicy { @@ -66,12 +54,35 @@ impl LoadBalancingPolicy for PowerOfTwoPolicy { let worker_idx1 = healthy_indices[idx1]; let worker_idx2 = healthy_indices[idx2]; + let worker1 = &workers[worker_idx1]; + let worker2 = &workers[worker_idx2]; - // Compare loads and select the less loaded one - let load1 = self.get_worker_load(workers[worker_idx1].as_ref()); - let load2 = self.get_worker_load(workers[worker_idx2].as_ref()); + // Access cached loads safely + let loads_guard = self.cached_loads.read().ok(); - // Log selection for debugging + // Try to get high-fidelity token loads for BOTH workers + let load1_tokens = loads_guard + .as_ref() + .and_then(|m| m.get(worker1.url()).copied()); + let load2_tokens = loads_guard + .as_ref() + .and_then(|m| m.get(worker2.url()).copied()); + + // If either worker is missing token data (e.g. monitor failure), + // we must degrade BOTH to request counts to ensure fairness. + let (load1, load2) = match (load1_tokens, load2_tokens) { + (Some(t1), Some(t2)) => { + // Both have token data. Compare Tokens. + (t1, t2) + } + _ => { + // If One or both are missing token data. + // Fallback to local request counts for BOTH. + (worker1.load() as isize, worker2.load() as isize) + } + }; + + // Select worker with lower load let selected_idx = if load1 <= load2 { worker_idx1 } else { @@ -80,9 +91,9 @@ impl LoadBalancingPolicy for PowerOfTwoPolicy { debug!( "Power-of-two selection: {}={} vs {}={} -> selected {}", - workers[worker_idx1].url(), + worker1.url(), load1, - workers[worker_idx2].url(), + worker2.url(), load2, workers[selected_idx].url() ); @@ -207,4 +218,152 @@ mod tests { // With single worker, should always select it assert_eq!(policy.select_worker(&workers, None), Some(0)); } + + #[test] + fn test_reproduce_incompatible_metric_bug() { + use std::{collections::HashMap, sync::Arc}; + + use crate::core::{BasicWorkerBuilder, WorkerType}; + + // 1. Setup the policy + let policy = PowerOfTwoPolicy::new(); + + // 2. Create Worker A: Idle (0 reqs), but has high token usage in cache + let worker_a = BasicWorkerBuilder::new("http://worker_a:8000") + .worker_type(WorkerType::Regular) + .build(); + + // 3. Create Worker B: Busy (5 reqs), but missing from cache + let worker_b = BasicWorkerBuilder::new("http://worker_b:8000") + .worker_type(WorkerType::Regular) + .build(); + + // Manually increment load on Worker B to simulate active requests + for _ in 0..5 { + worker_b.increment_load(); + } + + let workers: Vec> = vec![Arc::new(worker_a), Arc::new(worker_b)]; + + // 4. Simulate LoadMonitor update: + // Only Worker A gets a token report. Worker B is missing (e.g. monitor failure). + let mut loads = HashMap::new(); + loads.insert("http://worker_a:8000".to_string(), 50_000); // 50k tokens load + policy.update_loads(&loads); + + // 5. Run selection + let selected_idx = policy + .select_worker(&workers, None) + .expect("Should select a worker"); + + // 6. Verify the Fix + // Logic: + // - Worker A has token load (50k) but Worker B has NO token load. + // - Policy should fallback to request counts for BOTH. + // - A has 0 requests, B has 5 requests. + // - 0 <= 5, so A should be selected. + + if selected_idx == 0 { + println!("Bug Fixed: System correctly fell back to request counts and selected idle Worker A."); + } else { + println!( + "Bug PERSISTS: Selected Worker B (Load: 5 reqs) over Worker A (Load: 50k tokens)" + ); + } + + // Assert that the CORRECT worker (A, index 0) is selected + assert_eq!( + selected_idx, 0, + "The policy failed to handle incompatible metrics. Should select idle Worker A." + ); + } + #[test] + fn test_power_of_two_edge_cases() { + use std::{collections::HashMap, sync::Arc}; + + use crate::core::{BasicWorkerBuilder, WorkerType}; + + let policy = PowerOfTwoPolicy::new(); + + // Helper to create a worker with specific request load + let create_worker = |url: &str, reqs: usize| { + let w = BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .build(); + for _ in 0..reqs { + w.increment_load(); + } + Arc::new(w) + }; + + // Scenario 1: Happy Path (Both have Token Data) + // Worker A: 10 requests, but only 1,000 tokens (Light usage) -> Should be CHOSEN + // Worker B: 2 requests, but 100,000 tokens (Heavy usage) -> Should be AVOIDED + // This proves we use high-fidelity metrics when available, ignoring request counts. + let w_a = create_worker("http://a:8000", 10); + let w_b = create_worker("http://b:8000", 2); + let workers_1: Vec> = vec![w_a.clone(), w_b.clone()]; + + let mut loads_1 = HashMap::new(); + loads_1.insert("http://a:8000".to_string(), 1_000); + 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(); + assert_eq!( + idx_1, 0, + "Happy Path Failed: Should select Worker A (fewer tokens) despite higher request count" + ); + + // Scenario 2: Partial Failure (Worker A has tokens, Worker B is missing) + // Worker A: 10 requests, 1,000 tokens (Cached) + // Worker B: 2 requests, MISSING cache + // Logic: Fallback to requests -> Compare 10 (A) vs 2 (B) -> Select B + let w_c = create_worker("http://c:8000", 10); + let w_d = create_worker("http://d:8000", 2); + let workers_2: Vec> = vec![w_c.clone(), w_d.clone()]; + + let mut loads_2 = HashMap::new(); + loads_2.insert("http://c:8000".to_string(), 1_000); + // http://d:8000 is MISSING + policy.update_loads(&loads_2); + + let idx_2 = policy.select_worker(&workers_2, None).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) + // Worker A: 2 requests, MISSING cache + // Worker B: 10 requests, 1,000 tokens (Cached) + // Logic: Fallback to requests -> Compare 2 (A) vs 10 (B) -> Select A + let w_e = create_worker("http://e:8000", 2); + let w_f = create_worker("http://f:8000", 10); + let workers_3: Vec> = vec![w_e.clone(), w_f.clone()]; + + let mut loads_3 = HashMap::new(); + // http://e:8000 is MISSING + 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(); + 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) + // Worker A: 5 requests + // Worker B: 3 requests + // Logic: Requests vs Requests -> Select B + let w_g = create_worker("http://g:8000", 5); + let w_h = create_worker("http://h:8000", 3); + let workers_4: Vec> = vec![w_g.clone(), w_h.clone()]; + + let loads_4 = HashMap::new(); + policy.update_loads(&loads_4); + + let idx_4 = policy.select_worker(&workers_4, None).unwrap(); + assert_eq!( + idx_4, 1, + "Total Fail Failed: Should select Worker B based on request count" + ); + + println!("All edge case tests passed successfully."); + } }