From baaf12c7f4ec53133ec0eba25cc528ba1db49e9f Mon Sep 17 00:00:00 2001 From: Mingxing Zhang Date: Sun, 5 Apr 2026 17:31:08 +0800 Subject: [PATCH 1/3] cache-aware: log load snapshot and selection reason --- sgl-model-gateway/src/policies/cache_aware.rs | 57 ++++++++++++------- .../src/routers/http/pd_router.rs | 36 ++++++------ 2 files changed, 55 insertions(+), 38 deletions(-) diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index dfc3ef6462c5..c935e7d7a6d3 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -65,7 +65,7 @@ use async_trait::async_trait; use dashmap::DashMap; use rand::Rng; use smg_mesh::{tree_ops::TreeOperation, OptionalMeshSyncManager}; -use tracing::{debug, warn}; +use tracing::{debug, info, warn}; use super::{ get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask, @@ -292,15 +292,8 @@ impl CacheAwarePolicy { max_load: usize, min_load: usize, ) -> Option { - // Log load balancing trigger (only compute worker loads if debug enabled) - if tracing::enabled!(tracing::Level::DEBUG) { - let worker_loads: Vec<(&str, usize)> = - workers.iter().map(|w| (w.url(), w.load())).collect(); - debug!( - "Load balancing triggered | max: {} | min: {} | workers: {:?}", - max_load, min_load, worker_loads - ); - } + let worker_loads: Vec<(&str, usize)> = + workers.iter().map(|w| (w.url(), w.load())).collect(); // Use shortest queue when imbalanced let min_load_idx = healthy_indices @@ -341,6 +334,14 @@ impl CacheAwarePolicy { // Increment processed counter workers[min_load_idx].increment_processed(); + info!( + selected_worker = workers[min_load_idx].url(), + policy_reason = "load", + max_load, + min_load, + worker_loads = ?worker_loads, + "Cache-aware policy selected worker" + ); Some(min_load_idx) } @@ -404,22 +405,30 @@ impl LoadBalancingPolicy for CacheAwarePolicy { }; // Select worker without String allocation - let selected_idx = if match_rate > self.config.cache_threshold { + let (selected_idx, selection_reason) = if match_rate > self.config.cache_threshold { // Cache hit path: find worker by URL (compare &str directly, no allocation) let tenant_url: &str = &result.tenant; - workers - .iter() - .position(|w| w.url() == tenant_url) - .filter(|&idx| workers[idx].is_healthy()) + ( + workers + .iter() + .position(|w| w.url() == tenant_url) + .filter(|&idx| workers[idx].is_healthy()), + "cache_aware", + ) } else { // Low cache match: use worker with minimum load - healthy_indices - .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied() + ( + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied(), + "cache_aware", + ) }; if let Some(idx) = selected_idx { + let worker_loads: Vec<(&str, usize)> = + workers.iter().map(|w| (w.url(), w.load())).collect(); // Update the tree with this request (use worker URL directly, no allocation) tree.insert(text, workers[idx].url()); @@ -438,6 +447,16 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Increment processed counter workers[idx].increment_processed(); + info!( + selected_worker = workers[idx].url(), + policy_reason = selection_reason, + match_rate, + cache_threshold = self.config.cache_threshold, + max_load, + min_load, + worker_loads = ?worker_loads, + "Cache-aware policy selected worker" + ); return Some(idx); } diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index a939b6427a95..0112844f3366 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -419,7 +419,6 @@ impl PDRouter { &self, res: reqwest::Response, context: &PDRequestContext<'_>, - prefill: Arc, decode: Arc, ) -> Response { let status = res.status(); @@ -454,7 +453,6 @@ impl PDRouter { context.return_logprob, Some(decode_url), Some(response_headers), - prefill, decode, ) } else { @@ -539,10 +537,10 @@ impl PDRouter { decode: Arc, _start_time: Instant, ) -> 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(), headers)); + // Track prefill load for both streaming and non-streaming requests. + // For streaming, prefill should be released immediately after prefill completes. + let mut prefill_guard = Some(WorkerLoadGuard::new(prefill.clone(), headers)); + // Decode load is held until response completion (or stream end/abort). let _decode_guard = (!context.is_stream).then(|| WorkerLoadGuard::new(decode.clone(), headers)); @@ -596,7 +594,7 @@ impl PDRouter { ); return self - .handle_decode_error_response(res, &context, prefill, decode) + .handle_decode_error_response(res, &context, decode) .await; } @@ -625,6 +623,10 @@ impl PDRouter { }; if context.is_stream { + // Prefill has completed at this point. Release its load now so + // long decode streams do not keep prefill marked as busy. + prefill_guard.take(); + // Streaming response let prefill_logprobs = if context.return_logprob { prefill_body @@ -646,7 +648,6 @@ impl PDRouter { context.return_logprob, None, Some(response_headers), - prefill, decode, ) } else { @@ -839,7 +840,6 @@ impl PDRouter { return_logprob: bool, decode_url: Option, headers: Option, - prefill: Arc, decode: Arc, ) -> Response { use crate::core::AttachedBody; @@ -882,10 +882,9 @@ impl PDRouter { let stream = UnboundedReceiverStream::new(rx); let body = Body::from_stream(stream); - let guards = vec![ - WorkerLoadGuard::new(prefill, headers.as_ref()), - WorkerLoadGuard::new(decode, headers.as_ref()), - ]; + // For streaming requests, only decode load should be tied to stream lifecycle. + // Prefill load is released as soon as prefill response completes. + let guards = vec![WorkerLoadGuard::new(decode, headers.as_ref())]; let mut response = Response::new(body); *response.status_mut() = status; @@ -1550,20 +1549,19 @@ mod tests { false, None, None, - prefill_ref.clone(), decode_ref.clone(), ); - // Guards are now attached to response body, so load should be 1 - assert_eq!(prefill_ref.load(), 1); + // Only decode load is attached to the streaming response lifecycle. + assert_eq!(prefill_ref.load(), 0); assert_eq!(decode_ref.load(), 1); tx.send(bytes::Bytes::from("test data")).unwrap(); sleep(Duration::from_millis(10)).await; - // Load still 1 while response body exists - assert_eq!(prefill_ref.load(), 1); + // Decode load stays 1 while response body exists. + assert_eq!(prefill_ref.load(), 0); assert_eq!(decode_ref.load(), 1); drop(tx); @@ -1572,7 +1570,7 @@ mod tests { drop(response); } - // Guards dropped when response dropped + // Decode guard dropped when response dropped. assert_eq!(prefill_ref.load(), 0); assert_eq!(decode_ref.load(), 0); } From 0ec6cd784530aae7989fc16bebd236587a69cef6 Mon Sep 17 00:00:00 2001 From: Mingxing Zhang Date: Sun, 5 Apr 2026 17:39:34 +0800 Subject: [PATCH 2/3] pd-router: use request headers for streaming load guards --- sgl-model-gateway/src/routers/http/pd_router.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 0112844f3366..2fb940687ab2 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -452,6 +452,7 @@ impl PDRouter { None, context.return_logprob, Some(decode_url), + context.headers.as_ref(), Some(response_headers), decode, ) @@ -537,6 +538,8 @@ impl PDRouter { decode: Arc, _start_time: Instant, ) -> Response { + let request_headers = headers.cloned(); + // Track prefill load for both streaming and non-streaming requests. // For streaming, prefill should be released immediately after prefill completes. let mut prefill_guard = Some(WorkerLoadGuard::new(prefill.clone(), headers)); @@ -647,6 +650,7 @@ impl PDRouter { prefill_logprobs, context.return_logprob, None, + request_headers.as_ref(), Some(response_headers), decode, ) @@ -839,7 +843,8 @@ impl PDRouter { prefill_logprobs: Option, return_logprob: bool, decode_url: Option, - headers: Option, + request_headers: Option<&HeaderMap>, + response_headers: Option, decode: Arc, ) -> Response { use crate::core::AttachedBody; @@ -884,12 +889,12 @@ impl PDRouter { // For streaming requests, only decode load should be tied to stream lifecycle. // Prefill load is released as soon as prefill response completes. - let guards = vec![WorkerLoadGuard::new(decode, headers.as_ref())]; + let guards = vec![WorkerLoadGuard::new(decode, request_headers)]; let mut response = Response::new(body); *response.status_mut() = status; - let mut response_headers = headers.unwrap_or_default(); + let mut response_headers = response_headers.unwrap_or_default(); response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); *response.headers_mut() = response_headers; @@ -1549,6 +1554,7 @@ mod tests { false, None, None, + None, decode_ref.clone(), ); From 6d38fc22d3d5fc4941e80f829171bb9bd7e069d2 Mon Sep 17 00:00:00 2001 From: Mingxing Zhang Date: Sun, 5 Apr 2026 17:39:40 +0800 Subject: [PATCH 3/3] pd-router: document why prefill guard is not stream-attached --- sgl-model-gateway/src/routers/http/pd_router.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 2fb940687ab2..4a62140bf1fc 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -628,6 +628,13 @@ impl PDRouter { if context.is_stream { // Prefill has completed at this point. Release its load now so // long decode streams do not keep prefill marked as busy. + // + // Why not keep prefill guard on the streaming response body? + // - Prefill work is already finished before we return the stream. + // - Holding prefill guard until stream end inflates prefill load and + // can mislead imbalance detection in cache-aware/pd scheduling. + // - Decode guard still remains attached to stream lifecycle to reflect + // actual ongoing decode work (including long streams / client aborts). prefill_guard.take(); // Streaming response @@ -888,7 +895,9 @@ impl PDRouter { let body = Body::from_stream(stream); // For streaming requests, only decode load should be tied to stream lifecycle. - // Prefill load is released as soon as prefill response completes. + // Prefill load is released in execute_dual_dispatch_internal immediately after + // prefill completion. Keeping prefill guard here would intentionally delay + // release until stream end, which is incorrect for PD prefill capacity tracking. let guards = vec![WorkerLoadGuard::new(decode, request_headers)]; let mut response = Response::new(body);