From 1834401e7c4376d319b843215128f1ba2c922efb Mon Sep 17 00:00:00 2001
From: Praneth Paruchuri <pranethparuchuri@gmail.com>
Date: Fri, 12 Dec 2025 10:45:12 +0530
Subject: [PATCH] [model-gateway] optimize worker selection (#14894)

---
 sgl-model-gateway/Cargo.toml                  |  4 ++
 .../benches/router_registry_bench.rs          | 61 +++++++++++++++++++
 sgl-model-gateway/src/core/worker_registry.rs | 19 ++++++
 .../src/routers/router_manager.rs             |  9 +--
 4 files changed, 86 insertions(+), 7 deletions(-)
 create mode 100644 sgl-model-gateway/benches/router_registry_bench.rs

diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml
index 5e3a24b82..fa0e32880 100644
--- a/sgl-model-gateway/Cargo.toml
+++ b/sgl-model-gateway/Cargo.toml
@@ -146,6 +146,10 @@ name = "request_processing"
 harness = false
 path = "benches/request_processing.rs"
 
+[[bench]]
+name = "router_registry_bench"
+harness = false
+
 [[bench]]
 name = "tokenizer_benchmark"
 harness = false
diff --git a/sgl-model-gateway/benches/router_registry_bench.rs b/sgl-model-gateway/benches/router_registry_bench.rs
new file mode 100644
index 000000000..d6fa4958a
--- /dev/null
+++ b/sgl-model-gateway/benches/router_registry_bench.rs
@@ -0,0 +1,61 @@
+use std::{collections::HashMap, sync::Arc};
+
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+use sgl_model_gateway::core::{
+    BasicWorkerBuilder, CircuitBreakerConfig, WorkerRegistry, WorkerType,
+};
+
+// Helper to populate registry
+fn setup_registry(count: usize) -> Arc<WorkerRegistry> {
+    let registry = Arc::new(WorkerRegistry::new());
+
+    for i in 0..count {
+        let mut labels = HashMap::new();
+        labels.insert("model_id".to_string(), "benchmark-model".to_string());
+
+        let worker_type = if i % 2 == 0 {
+            WorkerType::Regular
+        } else {
+            WorkerType::Decode
+        };
+
+        let worker = BasicWorkerBuilder::new(format!("http://worker-{}:8000", i))
+            .worker_type(worker_type)
+            .labels(labels)
+            .circuit_breaker_config(CircuitBreakerConfig::default())
+            .build();
+
+        registry.register(Arc::from(worker));
+    }
+    registry
+}
+
+fn bench_optimizations(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Registry Optimizations");
+
+    // We test with 5000 workers to simulate high load
+    let size = 5000;
+    let registry = setup_registry(size);
+
+    //  The OLD method (Slow: Allocates vector + Clones ARCs)
+    group.bench_function(BenchmarkId::new("Old: get_all()", size), |b| {
+        b.iter(|| {
+            black_box(registry.get_all());
+        });
+    });
+
+    //  The NEW method (Fast: O(1) Lookup, Zero Allocation)
+    group.bench_function(
+        BenchmarkId::new("New: get_worker_distribution()", size),
+        |b| {
+            b.iter(|| {
+                black_box(registry.get_worker_distribution());
+            });
+        },
+    );
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_optimizations);
+criterion_main!(benches);
diff --git a/sgl-model-gateway/src/core/worker_registry.rs b/sgl-model-gateway/src/core/worker_registry.rs
index 314e12fc6..f96b88230 100644
--- a/sgl-model-gateway/src/core/worker_registry.rs
+++ b/sgl-model-gateway/src/core/worker_registry.rs
@@ -382,6 +382,25 @@ impl WorkerRegistry {
         }
     }
 
+    /// Get counts of regular and PD workers efficiently (O(1))
+    /// This avoids the overhead of get_all() which allocates memory and iterates all workers
+    pub fn get_worker_distribution(&self) -> (usize, usize) {
+        // Use the existing type_workers index for O(1) lookup
+        let regular_count = self
+            .type_workers
+            .get(&WorkerType::Regular)
+            .map(|v| v.len())
+            .unwrap_or(0);
+
+        // Get total workers count efficiently from DashMap
+        let total_workers = self.workers.len();
+
+        // PD workers are any workers that are not Regular
+        let pd_count = total_workers.saturating_sub(regular_count);
+
+        (regular_count, pd_count)
+    }
+
     /// Start a health checker for all workers in the registry
     /// This should be called once after the registry is populated with workers
     pub fn start_health_checker(&self, check_interval_secs: u64) -> crate::core::HealthChecker {
diff --git a/sgl-model-gateway/src/routers/router_manager.rs b/sgl-model-gateway/src/routers/router_manager.rs
index ceafc5625..fb46c2602 100644
--- a/sgl-model-gateway/src/routers/router_manager.rs
+++ b/sgl-model-gateway/src/routers/router_manager.rs
@@ -269,13 +269,8 @@ impl RouterManager {
         let mut best_router = None;
         let mut best_score = 0.0;
 
-        // Cache worker list to avoid duplicate get_all() calls
-        let all_workers = self.worker_registry.get_all();
-        let num_regular_workers = all_workers
-            .iter()
-            .filter(|w| matches!(w.worker_type(), WorkerType::Regular))
-            .count();
-        let num_pd_workers = all_workers.len() - num_regular_workers;
+        //  Uses O(1) lookups instead of allocating a full vector of workers via get_all()
+        let (num_regular_workers, num_pd_workers) = self.worker_registry.get_worker_distribution();
 
         for router in candidate_routers {
             let mut score = 1.0;