From 1f0e3d7fd84a55ba4494d492be9e4f78d9de8baf Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 13 Jan 2026 08:07:17 +0800 Subject: [PATCH] Support tracking worker routing key loads in gateway (#16765) --- sgl-model-gateway/src/core/worker.rs | 250 +++++++++++++++++- sgl-model-gateway/src/core/worker_builder.rs | 3 +- .../src/observability/metrics.rs | 10 + .../grpc/common/stages/request_execution.rs | 2 +- sgl-model-gateway/src/routers/grpc/context.rs | 10 +- .../src/routers/http/pd_router.rs | 20 +- sgl-model-gateway/src/routers/http/router.rs | 2 +- .../tests/load_guard_raii_test.rs | 14 +- 8 files changed, 287 insertions(+), 24 deletions(-) diff --git a/sgl-model-gateway/src/core/worker.rs b/sgl-model-gateway/src/core/worker.rs index 9e4e078dd..86310df80 100644 --- a/sgl-model-gateway/src/core/worker.rs +++ b/sgl-model-gateway/src/core/worker.rs @@ -39,6 +39,75 @@ static WORKER_CLIENT: LazyLock = LazyLock::new(|| { .expect("Failed to create worker HTTP client") }); +pub struct WorkerRoutingKeyLoad { + url: String, + active_routing_keys: dashmap::DashMap, +} + +impl WorkerRoutingKeyLoad { + pub fn new(url: impl Into) -> 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, + pub worker_routing_key_load: Arc, pub processed_counter: Arc, pub healthy: Arc, pub consecutive_failures: Arc, @@ -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, + routing_key: Option, } impl WorkerLoadGuard { - pub fn new(worker: Arc) -> Self { + pub fn new(worker: Arc, 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 = 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 = 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 = 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); + } } diff --git a/sgl-model-gateway/src/core/worker_builder.rs b/sgl-model-gateway/src/core/worker_builder.rs index cb6260829..f62cdd57c 100644 --- a/sgl-model-gateway/src/core/worker_builder.rs +++ b/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)), diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index f5814abcb..7554975fc 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/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); diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs b/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs index 48262dca8..08c8b40f6 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/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 diff --git a/sgl-model-gateway/src/routers/grpc/context.rs b/sgl-model-gateway/src/routers/grpc/context.rs index 7a6a8eab8..0ff12a4c6 100644 --- a/sgl-model-gateway/src/routers/grpc/context.rs +++ b/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), }, } } diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 78a40a30d..12b5fd3a9 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/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); diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index f76b0372e..8cfa5b01f 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/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(); diff --git a/sgl-model-gateway/tests/load_guard_raii_test.rs b/sgl-model-gateway/tests/load_guard_raii_test.rs index 3f7634307..4a1dbade7 100644 --- a/sgl-model-gateway/tests/load_guard_raii_test.rs +++ b/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::(); 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::(); 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);