[model-gateway] Fix incompatible metric comparison in PowerOfTwo policy (#14823)

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<Arc<dyn Worker>> = 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<Arc<dyn Worker>> = 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<Arc<dyn Worker>> = 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<Arc<dyn Worker>> = 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<Arc<dyn Worker>> = 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.");
+    }
 }