fix(router): use independent P/D policies in gRPC PD selection

model_gateway/src/policies/round_robin.rs

@@ -165,4 +165,51 @@ mod tests {
         policy.reset();
         assert_eq!(policy.select_worker(&workers, &info), Some(0));
     }
+
+    #[test]
+    fn test_independent_policies_cover_all_workers_across_two_pools() {
+        fn make_workers(prefix: &str, n: usize) -> Vec<Arc<dyn Worker>> {
+            (0..n)
+                .map(|i| {
+                    Arc::new(
+                        BasicWorkerBuilder::new(format!("http://{prefix}{i}:8000"))
+                            .worker_type(WorkerType::Regular)
+                            .health_config(no_health_check())
+                            .build(),
+                    ) as Arc<dyn Worker>
+                })
+                .collect()
+        }
+
+        let prefill_workers = make_workers("p", 4);
+        let decode_workers = make_workers("d", 4);
+        let info = SelectWorkerInfo::default();
+
+        let shared = RoundRobinPolicy::new();
+        let mut shared_prefill = [0usize; 4];
+        let mut shared_decode = [0usize; 4];
+        for _ in 0..40 {
+            let p = shared.select_worker(&prefill_workers, &info).unwrap();
+            let d = shared.select_worker(&decode_workers, &info).unwrap();
+            shared_prefill[p] += 1;
+            shared_decode[d] += 1;
+        }
+        assert_eq!(shared_prefill, [20, 0, 20, 0]);
+        assert_eq!(shared_decode, [0, 20, 0, 20]);
+
+        let prefill_policy = RoundRobinPolicy::new();
+        let decode_policy = RoundRobinPolicy::new();
+        let mut indep_prefill = [0usize; 4];
+        let mut indep_decode = [0usize; 4];
+        for _ in 0..40 {
+            let p = prefill_policy
+                .select_worker(&prefill_workers, &info)
+                .unwrap();
+            let d = decode_policy.select_worker(&decode_workers, &info).unwrap();
+            indep_prefill[p] += 1;
+            indep_decode[d] += 1;
+        }
+        assert_eq!(indep_prefill, [10, 10, 10, 10]);
+        assert_eq!(indep_decode, [10, 10, 10, 10]);
+    }
 }

model_gateway/src/routers/grpc/common/stages/worker_selection.rs

@@ -370,8 +370,9 @@ impl WorkerSelectionStage {
             return None;
         }
 
-        // Select using policies
-        let policy = self.policy_registry.get_policy_or_default(model_id);
+        // Independent P/D policies so stateful ones (e.g. round_robin) don't share a counter.
+        let prefill_policy = self.policy_registry.get_prefill_policy();
+        let decode_policy = self.policy_registry.get_decode_policy();
 
         // Get cached hash ring for consistent hashing (O(log n) lookup)
         let hash_ring = self.worker_registry.get_hash_ring(model_id);
@@ -385,29 +386,28 @@ impl WorkerSelectionStage {
             hash_ring,
             leg: WorkerLeg::Prefill,
         };
-        let prefill_idx = self
-            .policy_registry
-            .select_worker(&policy, &available_prefill, &info)?;
+        let prefill_idx =
+            self.policy_registry
+                .select_worker(&prefill_policy, &available_prefill, &info)?;
         info.leg = WorkerLeg::Decode;
-        let decode_idx = self
-            .policy_registry
-            .select_worker(&policy, &available_decode, &info)?;
+        let decode_idx =
+            self.policy_registry
+                .select_worker(&decode_policy, &available_decode, &info)?;
 
         let model = model_id;
-        let policy_name = policy.name();
 
         // Record worker selection metrics for both prefill and decode
         Metrics::record_worker_selection(
             metrics_labels::WORKER_PREFILL,
             metrics_labels::CONNECTION_GRPC,
             model,
-            policy_name,
+            prefill_policy.name(),
         );
         Metrics::record_worker_selection(
             metrics_labels::WORKER_DECODE,
             metrics_labels::CONNECTION_GRPC,
             model,
-            policy_name,
+            decode_policy.name(),
         );
 
         Some((
@@ -660,3 +660,119 @@ fn hex_encode(bytes: &[u8]) -> String {
     }
     out
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use openai_protocol::worker::HealthCheckConfig;
+
+    use super::*;
+    use crate::{
+        config::types::PolicyConfig,
+        policies::PolicyFactory,
+        worker::{BasicWorkerBuilder, ConnectionMode, ModelCard},
+    };
+
+    fn no_health_check() -> HealthCheckConfig {
+        HealthCheckConfig {
+            disable_health_check: true,
+            ..Default::default()
+        }
+    }
+
+    fn register_pd_workers(
+        registry: &WorkerRegistry,
+        model_id: &str,
+        n: usize,
+    ) -> (Vec<String>, Vec<String>) {
+        let mut prefill_urls = Vec::with_capacity(n);
+        let mut decode_urls = Vec::with_capacity(n);
+
+        for i in 0..n {
+            let url = format!("grpc://127.0.0.1:{}", 8000 + i);
+            prefill_urls.push(url.clone());
+            registry
+                .register(Arc::new(
+                    BasicWorkerBuilder::new(url)
+                        .model(ModelCard::new(model_id))
+                        .worker_type(WorkerType::Prefill)
+                        .connection_mode(ConnectionMode::Grpc)
+                        .health_config(no_health_check())
+                        .build(),
+                ))
+                .unwrap();
+        }
+
+        for i in 0..n {
+            let url = format!("grpc://127.0.0.1:{}", 8100 + i);
+            decode_urls.push(url.clone());
+            registry
+                .register(Arc::new(
+                    BasicWorkerBuilder::new(url)
+                        .model(ModelCard::new(model_id))
+                        .worker_type(WorkerType::Decode)
+                        .connection_mode(ConnectionMode::Grpc)
+                        .health_config(no_health_check())
+                        .build(),
+                ))
+                .unwrap();
+        }
+
+        (prefill_urls, decode_urls)
+    }
+
+    #[test]
+    fn select_pd_pair_round_robin_covers_all_workers_with_independent_policies() {
+        let model_id = "test-model";
+        let worker_registry = Arc::new(WorkerRegistry::new());
+        let (prefill_urls, decode_urls) = register_pd_workers(&worker_registry, model_id, 4);
+
+        let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin));
+        // Mirror production PD startup: create two independent RoundRobin instances.
+        policy_registry
+            .set_prefill_policy(PolicyFactory::create_from_config(&PolicyConfig::RoundRobin));
+        policy_registry
+            .set_decode_policy(PolicyFactory::create_from_config(&PolicyConfig::RoundRobin));
+
+        let prefill_policy = policy_registry.get_prefill_policy();
+        let decode_policy = policy_registry.get_decode_policy();
+        assert_eq!(prefill_policy.name(), "round_robin");
+        assert_eq!(decode_policy.name(), "round_robin");
+        assert!(
+            !Arc::ptr_eq(&prefill_policy, &decode_policy),
+            "prefill/decode must not share one RoundRobinPolicy counter"
+        );
+
+        let stage = WorkerSelectionStage::new(
+            worker_registry,
+            policy_registry,
+            WorkerSelectionMode::PrefillDecode,
+        );
+
+        let mut prefill_hits: HashMap<String, usize> = HashMap::new();
+        let mut decode_hits: HashMap<String, usize> = HashMap::new();
+        for _ in 0..40 {
+            let (prefill, decode, _) = stage
+                .select_pd_pair(model_id, None, None, None)
+                .expect("select_pd_pair should return a pair");
+            *prefill_hits.entry(prefill.url().to_string()).or_default() += 1;
+            *decode_hits.entry(decode.url().to_string()).or_default() += 1;
+        }
+
+        for url in &prefill_urls {
+            assert_eq!(
+                prefill_hits.get(url).copied().unwrap_or(0),
+                10,
+                "prefill worker {url} should receive 10 of 40 selections"
+            );
+        }
+        for url in &decode_urls {
+            assert_eq!(
+                decode_hits.get(url).copied().unwrap_or(0),
+                10,
+                "decode worker {url} should receive 10 of 40 selections"
+            );
+        }
+    }
+}