Improve cache-aware logging and fix prefill load handling for streaming responses

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<usize> {
-        // 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);
             }

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

@@ -419,7 +419,6 @@ impl PDRouter {
         &self,
         res: reqwest::Response,
         context: &PDRequestContext<'_>,
-        prefill: Arc<dyn Worker>,
         decode: Arc<dyn Worker>,
     ) -> Response {
         let status = res.status();
@@ -453,8 +452,8 @@ impl PDRouter {
                 None,
                 context.return_logprob,
                 Some(decode_url),
+                context.headers.as_ref(),
                 Some(response_headers),
-                prefill,
                 decode,
             )
         } else {
@@ -539,10 +538,12 @@ impl PDRouter {
         decode: Arc<dyn Worker>,
         _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));
+        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));
+        // 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 +597,7 @@ impl PDRouter {
                     );
 
                     return self
-                        .handle_decode_error_response(res, &context, prefill, decode)
+                        .handle_decode_error_response(res, &context, decode)
                         .await;
                 }
 
@@ -625,6 +626,17 @@ 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
                     let prefill_logprobs = if context.return_logprob {
                         prefill_body
@@ -645,8 +657,8 @@ impl PDRouter {
                         prefill_logprobs,
                         context.return_logprob,
                         None,
+                        request_headers.as_ref(),
                         Some(response_headers),
-                        prefill,
                         decode,
                     )
                 } else {
@@ -838,8 +850,8 @@ impl PDRouter {
         prefill_logprobs: Option<Value>,
         return_logprob: bool,
         decode_url: Option<String>,
-        headers: Option<HeaderMap>,
-        prefill: Arc<dyn Worker>,
+        request_headers: Option<&HeaderMap>,
+        response_headers: Option<HeaderMap>,
         decode: Arc<dyn Worker>,
     ) -> Response {
         use crate::core::AttachedBody;
@@ -882,15 +894,16 @@ 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 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);
         *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;
 
@@ -1550,20 +1563,20 @@ mod tests {
                 false,
                 None,
                 None,
-                prefill_ref.clone(),
+                None,
                 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 +1585,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);
     }