Support tracking worker routing key loads in gateway (#16765)

sgl-model-gateway/src/core/worker.rs

@@ -39,6 +39,75 @@ static WORKER_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
         .expect("Failed to create worker HTTP client")
 });
 
+pub struct WorkerRoutingKeyLoad {
+    url: String,
+    active_routing_keys: dashmap::DashMap<String, usize>,
+}
+
+impl WorkerRoutingKeyLoad {
+    pub fn new(url: impl Into<String>) -> Self {
+        Self {
+            url: url.into(),
+            active_routing_keys: dashmap::DashMap::new(),
+        }
+    }
+
+    pub fn value(&self) -> usize {
+        self.active_routing_keys.len()
+    }
+
+    pub fn increment(&self, routing_key: &str) {
+        *self
+            .active_routing_keys
+            .entry(routing_key.to_string())
+            .or_insert(0) += 1;
+        self.update_metrics();
+    }
+
+    pub fn decrement(&self, routing_key: &str) {
+        use dashmap::mapref::entry::Entry;
+
+        match self.active_routing_keys.entry(routing_key.to_string()) {
+            Entry::Occupied(mut entry) => {
+                let counter = entry.get_mut();
+                if *counter > 0 {
+                    *counter -= 1;
+                    if *counter == 0 {
+                        entry.remove();
+                    }
+                } else {
+                    tracing::warn!(
+                        worker_url = %self.url,
+                        routing_key = %routing_key,
+                        "Attempted to decrement routing key counter that is already at 0"
+                    );
+                }
+            }
+            Entry::Vacant(_) => {
+                tracing::warn!(
+                    worker_url = %self.url,
+                    routing_key = %routing_key,
+                    "Attempted to decrement non-existent routing key"
+                );
+            }
+        }
+        self.update_metrics();
+    }
+
+    fn update_metrics(&self) {
+        Metrics::set_worker_routing_keys_active(&self.url, self.value());
+    }
+}
+
+impl fmt::Debug for WorkerRoutingKeyLoad {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("WorkerRoutingKeyLoad")
+            .field("url", &self.url)
+            .field("active_routing_keys", &self.value())
+            .finish()
+    }
+}
+
 /// Core worker abstraction that represents a backend service
 #[async_trait]
 pub trait Worker: Send + Sync + fmt::Debug {
@@ -111,6 +180,9 @@ pub trait Worker: Send + Sync + fmt::Debug {
     /// Reset the load counter to 0 (for sync/recovery)
     fn reset_load(&self) {}
 
+    /// Get the worker routing key load tracker
+    fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad;
+
     /// Get the number of processed requests
     fn processed_requests(&self) -> usize;
 
@@ -545,6 +617,7 @@ impl WorkerMetadata {
 pub struct BasicWorker {
     pub metadata: WorkerMetadata,
     pub load_counter: Arc<AtomicUsize>,
+    pub worker_routing_key_load: Arc<WorkerRoutingKeyLoad>,
     pub processed_counter: Arc<AtomicUsize>,
     pub healthy: Arc<AtomicBool>,
     pub consecutive_failures: Arc<AtomicUsize>,
@@ -698,6 +771,10 @@ impl Worker for BasicWorker {
         self.update_running_requests_metrics();
     }
 
+    fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad {
+        &self.worker_routing_key_load
+    }
+
     fn processed_requests(&self) -> usize {
         self.processed_counter.load(Ordering::Relaxed)
     }
@@ -942,6 +1019,10 @@ impl Worker for DPAwareWorker {
         self.base_worker.reset_load();
     }
 
+    fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad {
+        self.base_worker.worker_routing_key_load()
+    }
+
     fn processed_requests(&self) -> usize {
         self.base_worker.processed_requests()
     }
@@ -1017,18 +1098,34 @@ impl Worker for DPAwareWorker {
 /// immediately but the stream continues in the background.
 pub struct WorkerLoadGuard {
     worker: Arc<dyn Worker>,
+    routing_key: Option<String>,
 }
 
 impl WorkerLoadGuard {
-    pub fn new(worker: Arc<dyn Worker>) -> Self {
+    pub fn new(worker: Arc<dyn Worker>, headers: Option<&http::HeaderMap>) -> Self {
+        use crate::routers::header_utils::extract_routing_key;
+
         worker.increment_load();
-        Self { worker }
+
+        let routing_key = extract_routing_key(headers).map(String::from);
+
+        if let Some(ref key) = routing_key {
+            worker.worker_routing_key_load().increment(key);
+        }
+
+        Self {
+            worker,
+            routing_key,
+        }
     }
 }
 
 impl Drop for WorkerLoadGuard {
     fn drop(&mut self) {
         self.worker.decrement_load();
+        if let Some(ref key) = self.routing_key {
+            self.worker.worker_routing_key_load().decrement(key);
+        }
     }
 }
 
@@ -1886,4 +1983,153 @@ mod tests {
         // Not found
         assert!(metadata.find_model("unknown-model").is_none());
     }
+
+    #[test]
+    fn test_worker_routing_key_load_increment_decrement() {
+        let load = WorkerRoutingKeyLoad::new("http://test:8000");
+        assert_eq!(load.value(), 0);
+
+        load.increment("key1");
+        assert_eq!(load.value(), 1);
+
+        load.increment("key2");
+        assert_eq!(load.value(), 2);
+
+        load.increment("key1");
+        assert_eq!(load.value(), 2);
+
+        load.decrement("key1");
+        assert_eq!(load.value(), 2);
+
+        load.decrement("key1");
+        assert_eq!(load.value(), 1);
+
+        load.decrement("key2");
+        assert_eq!(load.value(), 0);
+    }
+
+    #[test]
+    fn test_worker_routing_key_load_cleanup_on_zero() {
+        let load = WorkerRoutingKeyLoad::new("http://test:8000");
+
+        load.increment("key1");
+        load.increment("key2");
+        load.increment("key3");
+        assert_eq!(load.active_routing_keys.len(), 3);
+
+        load.decrement("key1");
+        assert_eq!(load.active_routing_keys.len(), 2);
+
+        load.decrement("key2");
+        assert_eq!(load.active_routing_keys.len(), 1);
+
+        load.decrement("key3");
+        assert_eq!(load.active_routing_keys.len(), 0);
+    }
+
+    #[test]
+    fn test_worker_routing_key_load_multiple_requests_same_key() {
+        let load = WorkerRoutingKeyLoad::new("http://test:8000");
+
+        load.increment("key-1");
+        load.increment("key-1");
+        load.increment("key-1");
+        assert_eq!(load.value(), 1);
+
+        load.decrement("key-1");
+        assert_eq!(load.value(), 1);
+
+        load.decrement("key-1");
+        assert_eq!(load.value(), 1);
+
+        load.decrement("key-1");
+        assert_eq!(load.value(), 0);
+        assert_eq!(load.active_routing_keys.len(), 0);
+    }
+
+    #[test]
+    fn test_worker_routing_key_load_decrement_nonexistent() {
+        let load = WorkerRoutingKeyLoad::new("http://test:8000");
+        load.decrement("nonexistent");
+        assert_eq!(load.value(), 0);
+    }
+
+    #[test]
+    fn test_worker_load_guard_with_routing_key() {
+        use crate::core::BasicWorkerBuilder;
+
+        let worker: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("http://test:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        );
+
+        assert_eq!(worker.load(), 0);
+        assert_eq!(worker.worker_routing_key_load().value(), 0);
+
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-routing-key", "key-123".parse().unwrap());
+
+        {
+            let _guard = WorkerLoadGuard::new(worker.clone(), Some(&headers));
+            assert_eq!(worker.load(), 1);
+            assert_eq!(worker.worker_routing_key_load().value(), 1);
+        }
+
+        assert_eq!(worker.load(), 0);
+        assert_eq!(worker.worker_routing_key_load().value(), 0);
+    }
+
+    #[test]
+    fn test_worker_load_guard_without_routing_key() {
+        use crate::core::BasicWorkerBuilder;
+
+        let worker: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("http://test:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        );
+
+        assert_eq!(worker.load(), 0);
+        assert_eq!(worker.worker_routing_key_load().value(), 0);
+
+        {
+            let _guard = WorkerLoadGuard::new(worker.clone(), None);
+            assert_eq!(worker.load(), 1);
+            assert_eq!(worker.worker_routing_key_load().value(), 0);
+        }
+
+        assert_eq!(worker.load(), 0);
+        assert_eq!(worker.worker_routing_key_load().value(), 0);
+    }
+
+    #[test]
+    fn test_worker_load_guard_multiple_same_routing_key() {
+        use crate::core::BasicWorkerBuilder;
+
+        let worker: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("http://test:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        );
+
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-routing-key", "key-123".parse().unwrap());
+
+        let guard1 = WorkerLoadGuard::new(worker.clone(), Some(&headers));
+        assert_eq!(worker.load(), 1);
+        assert_eq!(worker.worker_routing_key_load().value(), 1);
+
+        let guard2 = WorkerLoadGuard::new(worker.clone(), Some(&headers));
+        assert_eq!(worker.load(), 2);
+        assert_eq!(worker.worker_routing_key_load().value(), 1);
+
+        drop(guard1);
+        assert_eq!(worker.load(), 1);
+        assert_eq!(worker.worker_routing_key_load().value(), 1);
+
+        drop(guard2);
+        assert_eq!(worker.load(), 0);
+        assert_eq!(worker.worker_routing_key_load().value(), 0);
+    }
 }

sgl-model-gateway/src/core/worker_builder.rs

@@ -6,7 +6,7 @@ use super::{
     model_type::ModelType,
     worker::{
         BasicWorker, ConnectionMode, DPAwareWorker, HealthConfig, RuntimeType, WorkerMetadata,
-        WorkerType,
+        WorkerRoutingKeyLoad, WorkerType,
     },
 };
 use crate::{observability::metrics::Metrics, routers::grpc::client::GrpcClient};
@@ -193,6 +193,7 @@ impl BasicWorkerBuilder {
         BasicWorker {
             metadata,
             load_counter: Arc::new(AtomicUsize::new(0)),
+            worker_routing_key_load: Arc::new(WorkerRoutingKeyLoad::new(&self.url)),
             processed_counter: Arc::new(AtomicUsize::new(0)),
             healthy: Arc::new(AtomicBool::new(healthy)),
             consecutive_failures: Arc::new(AtomicUsize::new(0)),

sgl-model-gateway/src/observability/metrics.rs

@@ -916,6 +916,16 @@ impl Metrics {
         .set(count as f64);
     }
 
+    /// Set active routing keys per worker
+    pub fn set_worker_routing_keys_active(worker: &str, count: usize) {
+        let worker_interned = intern_string(worker);
+        gauge!(
+            "smg_worker_routing_keys_active",
+            "worker" => worker_interned
+        )
+        .set(count as f64);
+    }
+
     /// Set worker health status
     pub fn set_worker_health(worker_url: &str, healthy: bool) {
         let worker_interned = intern_string(worker_url);

sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -72,7 +72,7 @@ impl PipelineStage for RequestExecutionStage {
             )
         })?;
 
-        ctx.state.load_guards = Some(LoadGuards::from(workers));
+        ctx.state.load_guards = Some(LoadGuards::new(workers, ctx.input.headers.as_ref()));
 
         // Extract dispatch metadata for tracing span
         let request_id = ctx

sgl-model-gateway/src/routers/grpc/context.rs

@@ -167,15 +167,15 @@ pub(crate) enum LoadGuards {
     },
 }
 
-impl From<&WorkerSelection> for LoadGuards {
-    fn from(selection: &WorkerSelection) -> Self {
+impl LoadGuards {
+    pub fn new(selection: &WorkerSelection, headers: Option<&HeaderMap>) -> Self {
         match selection {
             WorkerSelection::Single { worker } => LoadGuards::Single {
-                _guard: WorkerLoadGuard::new(worker.clone()),
+                _guard: WorkerLoadGuard::new(worker.clone(), headers),
             },
             WorkerSelection::Dual { prefill, decode } => LoadGuards::Dual {
-                _prefill: WorkerLoadGuard::new(prefill.clone()),
-                _decode: WorkerLoadGuard::new(decode.clone()),
+                _prefill: WorkerLoadGuard::new(prefill.clone(), headers),
+                _decode: WorkerLoadGuard::new(decode.clone(), headers),
             },
         }
     }

sgl-model-gateway/src/routers/http/pd_router.rs

@@ -539,8 +539,10 @@ impl PDRouter {
     ) -> Response {
         // For non-streaming: use guard for automatic load management
         // For streaming: load will be managed in create_streaming_response
-        let _prefill_guard = (!context.is_stream).then(|| WorkerLoadGuard::new(prefill.clone()));
-        let _decode_guard = (!context.is_stream).then(|| WorkerLoadGuard::new(decode.clone()));
+        let _prefill_guard =
+            (!context.is_stream).then(|| WorkerLoadGuard::new(prefill.clone(), headers));
+        let _decode_guard =
+            (!context.is_stream).then(|| WorkerLoadGuard::new(decode.clone(), headers));
 
         let mut headers_with_trace = headers.cloned().unwrap_or_default();
         inject_trace_context_http(&mut headers_with_trace);
@@ -878,13 +880,17 @@ impl PDRouter {
         let mut response = Response::new(body);
         *response.status_mut() = status;
 
+        // Attach load guards to response body for proper RAII lifecycle
+        // Guards are dropped when response body is consumed or client disconnects
+        let guards = vec![
+            WorkerLoadGuard::new(prefill, headers.as_ref()),
+            WorkerLoadGuard::new(decode, headers.as_ref()),
+        ];
+
         let mut headers = headers.unwrap_or_default();
         headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream"));
         *response.headers_mut() = headers;
 
-        // Attach load guards to response body for proper RAII lifecycle
-        // Guards are dropped when response body is consumed or client disconnects
-        let guards = vec![WorkerLoadGuard::new(prefill), WorkerLoadGuard::new(decode)];
         AttachedBody::wrap_response(response, guards)
     }
 
@@ -1460,8 +1466,8 @@ mod tests {
             true,
         ));
 
-        let _prefill_guard = WorkerLoadGuard::new(prefill_worker.clone());
-        let _decode_guard = WorkerLoadGuard::new(decode_worker.clone());
+        let _prefill_guard = WorkerLoadGuard::new(prefill_worker.clone(), None);
+        let _decode_guard = WorkerLoadGuard::new(decode_worker.clone(), None);
 
         assert_eq!(prefill_worker.load(), 1);
         assert_eq!(decode_worker.load(), 1);

sgl-model-gateway/src/routers/http/router.rs

@@ -294,7 +294,7 @@ impl Router {
         };
 
         let load_guard =
-            (policy.name() == "cache_aware").then(|| WorkerLoadGuard::new(worker.clone()));
+            (policy.name() == "cache_aware").then(|| WorkerLoadGuard::new(worker.clone(), headers));
 
         // Note: Using borrowed reference avoids heap allocation
         events::RequestSentEvent { url: worker.url() }.emit();

sgl-model-gateway/tests/load_guard_raii_test.rs

@@ -37,7 +37,7 @@ async fn test_guard_dropped_when_response_body_consumed() {
     let response = Response::new(body);
 
     // Attach guard
-    let guard = WorkerLoadGuard::new(worker.clone());
+    let guard = WorkerLoadGuard::new(worker.clone(), None);
     assert_eq!(worker.load(), 1);
 
     let guarded_response = AttachedBody::wrap_response(response, guard);
@@ -62,7 +62,7 @@ async fn test_guard_dropped_when_response_dropped_without_consumption() {
         let body = Body::from("Hello, World!");
         let response = Response::new(body);
 
-        let guard = WorkerLoadGuard::new(worker.clone());
+        let guard = WorkerLoadGuard::new(worker.clone(), None);
         assert_eq!(worker.load(), 1);
 
         let _guarded_response = AttachedBody::wrap_response(response, guard);
@@ -86,7 +86,7 @@ async fn test_streaming_guard_dropped_when_stream_ends() {
     let (tx, rx) = mpsc::unbounded_channel::<Bytes>();
 
     let response = create_sse_response(rx);
-    let guard = WorkerLoadGuard::new(worker.clone());
+    let guard = WorkerLoadGuard::new(worker.clone(), None);
     assert_eq!(worker.load(), 1);
 
     let guarded_response = AttachedBody::wrap_response(response, guard);
@@ -133,7 +133,7 @@ async fn test_streaming_guard_dropped_on_client_disconnect() {
     let (tx, rx) = mpsc::unbounded_channel::<Bytes>();
 
     let response = create_sse_response(rx);
-    let guard = WorkerLoadGuard::new(worker.clone());
+    let guard = WorkerLoadGuard::new(worker.clone(), None);
     assert_eq!(worker.load(), 1);
 
     let guarded_response = AttachedBody::wrap_response(response, guard);
@@ -171,8 +171,8 @@ async fn test_multiple_guards_all_dropped() {
         let response = Response::new(body);
 
         // Create guards for both workers (simulates dual prefill/decode)
-        let guard1 = WorkerLoadGuard::new(worker1.clone());
-        let guard2 = WorkerLoadGuard::new(worker2.clone());
+        let guard1 = WorkerLoadGuard::new(worker1.clone(), None);
+        let guard2 = WorkerLoadGuard::new(worker2.clone(), None);
         assert_eq!(worker1.load(), 1);
         assert_eq!(worker2.load(), 1);
 
@@ -197,7 +197,7 @@ async fn test_guard_with_empty_body() {
         let body = Body::empty();
         let response = Response::new(body);
 
-        let guard = WorkerLoadGuard::new(worker.clone());
+        let guard = WorkerLoadGuard::new(worker.clone(), None);
         assert_eq!(worker.load(), 1);
 
         let guarded_response = AttachedBody::wrap_response(response, guard);