push based publishing

Signed-off-by: PeaBrane <[email protected]>

Author: PeaBrane

lib/llm/src/kv_router/scheduler.rs

@@ -14,12 +14,11 @@ use std::time::Duration;
 use tokio::sync::{RwLock, watch};
 
 use super::KV_HIT_RATE_SUBJECT;
-use super::KV_METRICS_SUBJECT;
 use super::KvRouterConfig;
 use super::RouterConfigOverride;
 use super::WorkerSelector;
 use super::indexer::OverlapScores;
-use super::protocols::{ActiveLoad, DpRank, WorkerId, WorkerSelectionResult, WorkerWithDpRank};
+use super::protocols::{DpRank, WorkerId, WorkerSelectionResult, WorkerWithDpRank};
 use super::sequence::{ActiveSequencesMultiWorker, SequenceError};
 
 use crate::tokens::SequenceHash;
@@ -288,47 +287,6 @@ impl KvScheduler {
             tracing::trace!("background endpoint subscriber shutting down");
         });
 
-        // Background task to publish active load metrics every 5 seconds
-        let slots_metrics = slots.clone();
-        let ns_metrics = component.namespace().clone();
-        let metrics_cancel_token = component.drt().child_token();
-        tokio::spawn(async move {
-            let mut interval = tokio::time::interval(Duration::from_secs(5));
-            tracing::trace!("active load metrics publishing task started");
-
-            loop {
-                tokio::select! {
-                    _ = metrics_cancel_token.cancelled() => {
-                        tracing::trace!("active load metrics publishing task shutting down");
-                        break;
-                    }
-                    _ = interval.tick() => {
-                        // Query active tokens and blocks from all workers
-                        // Note: active_tokens is always tracked, but active_blocks may be empty
-                        // if --no-track-active-blocks is set
-                        let active_tokens = slots_metrics.active_tokens().await;
-                        let active_blocks = slots_metrics.active_blocks().await;
-
-                        // Publish ActiveLoad for each worker/dp_rank (iterate over tokens since always tracked)
-                        for (worker, tokens) in active_tokens.iter() {
-                            let blocks = active_blocks.get(worker).copied();
-                            let active_load = ActiveLoad {
-                                worker_id: worker.worker_id,
-                                dp_rank: worker.dp_rank,
-                                active_decode_blocks: blocks.map(|b| b as u64),
-                                active_prefill_tokens: Some(*tokens as u64),
-                            };
-
-                            if let Err(e) = ns_metrics.publish(KV_METRICS_SUBJECT, &active_load).await {
-                                tracing::warn!("Failed to publish active load metrics: {:?}", e);
-                            }
-                        }
-                    }
-                }
-            }
-            tracing::trace!("active load metrics publishing task exiting");
-        });
-
         Ok(KvScheduler { request_tx, slots })
     }
 

lib/llm/src/kv_router/sequence.rs

@@ -38,8 +38,10 @@ use std::time::Duration;
 use tokio::time::Instant;
 use uuid::Uuid;
 
-use super::protocols::{ActiveSequenceEvent, ActiveSequenceEventData, WorkerWithDpRank};
-use crate::kv_router::ACTIVE_SEQUENCES_SUBJECT;
+use super::protocols::{
+    ActiveLoad, ActiveSequenceEvent, ActiveSequenceEventData, WorkerWithDpRank,
+};
+use crate::kv_router::{ACTIVE_SEQUENCES_SUBJECT, KV_METRICS_SUBJECT};
 use crate::local_model::runtime_config::ModelRuntimeConfig;
 use dynamo_runtime::CancellationToken;
 
@@ -701,6 +703,9 @@ impl ActiveSequencesMultiWorker {
             self.request_to_worker.remove(expired_id);
         }
 
+        // Publish ActiveLoad metrics for this worker
+        self.publish_active_load_for_worker(worker).await;
+
         Ok(())
     }
 
@@ -744,6 +749,9 @@ impl ActiveSequencesMultiWorker {
 
         self.request_to_worker.remove(request_id);
 
+        // Publish ActiveLoad metrics for this worker
+        self.publish_active_load_for_worker(worker).await;
+
         Ok(())
     }
 
@@ -790,9 +798,66 @@ impl ActiveSequencesMultiWorker {
             })
             .map_err(|_| SequenceError::WorkerChannelClosed)?;
 
+        // Publish ActiveLoad metrics for this worker
+        self.publish_active_load_for_worker(worker).await;
+
         Ok(())
     }
 
+    /// Helper method to query a single worker for active blocks/tokens and publish ActiveLoad
+    async fn publish_active_load_for_worker(&self, worker: WorkerWithDpRank) {
+        let Some(sender) = self.senders.get(&worker) else {
+            tracing::warn!("Worker {worker:?} not found when publishing ActiveLoad");
+            return;
+        };
+
+        // Query active blocks
+        let (blocks_tx, blocks_rx) = tokio::sync::oneshot::channel();
+        if sender
+            .send(UpdateSequences::ActiveBlocks { resp_tx: blocks_tx })
+            .is_err()
+        {
+            tracing::warn!("Failed to send ActiveBlocks query to worker {worker:?}");
+            return;
+        }
+
+        // Query active tokens
+        let (tokens_tx, tokens_rx) = tokio::sync::oneshot::channel();
+        if sender
+            .send(UpdateSequences::ActiveTokens { resp_tx: tokens_tx })
+            .is_err()
+        {
+            tracing::warn!("Failed to send ActiveTokens query to worker {worker:?}");
+            return;
+        }
+
+        // Await both responses
+        let (active_blocks, active_tokens) = match tokio::join!(blocks_rx, tokens_rx) {
+            (Ok(blocks), Ok(tokens)) => (blocks, tokens),
+            _ => {
+                tracing::warn!("Failed to receive active blocks/tokens from worker {worker:?}");
+                return;
+            }
+        };
+
+        // Publish ActiveLoad
+        let active_load = ActiveLoad {
+            worker_id: worker.worker_id,
+            dp_rank: worker.dp_rank,
+            active_decode_blocks: Some(active_blocks as u64),
+            active_prefill_tokens: Some(active_tokens as u64),
+        };
+
+        if let Err(e) = self
+            .component
+            .namespace()
+            .publish(KV_METRICS_SUBJECT, &active_load)
+            .await
+        {
+            tracing::warn!("Failed to publish ActiveLoad for worker {worker:?}: {e:?}");
+        }
+    }
+
     /// Get the number of workers
     pub fn num_workers(&self) -> usize {
         self.senders.len()