should be functional

Author: PeaBrane

components/metrics/src/lib.rs

@@ -84,7 +84,7 @@ use std::net::SocketAddr;
 use std::time::Duration as StdDuration;
 
 use dynamo_llm::kv_router::protocols::{ForwardPassMetrics, LoadMetrics};
-use dynamo_llm::kv_router::scheduler::Endpoint;
+use dynamo_llm::kv_router::scoring::Endpoint;
 use dynamo_llm::kv_router::scoring::ProcessedEndpoints;
 
 use dynamo_runtime::{

lib/bindings/python/Cargo.lock

@@ -15,7 +15,6 @@ use dynamo_runtime::{
     protocols::annotated::Annotated,
 };
 use futures::stream::{self, StreamExt};
-use tokio::sync::Mutex;
 
 pub mod approx;
 pub mod indexer;
@@ -138,10 +137,6 @@ pub struct KvRouter {
     scheduler: KvScheduler,
 
     block_size: u32,
-
-    // To ensure blocking reads / writes
-    // TODO: benchmark tradeoffs
-    find_best_match_mutex: Mutex<()>,
 }
 
 impl KvRouter {
@@ -178,13 +173,8 @@ impl KvRouter {
             ))
         };
 
-        let scheduler = KvScheduler::start(
-            component.namespace().clone(),
-            block_size,
-            instances_rx,
-            selector,
-        )
-        .await?;
+        let scheduler =
+            KvScheduler::start(component.clone(), block_size, instances_rx, selector).await?;
 
         // [gluo TODO] try subscribe_with_type::<RouterEvent>,
         // error checking below will be different.
@@ -218,7 +208,6 @@ impl KvRouter {
             indexer,
             scheduler,
             block_size,
-            find_best_match_mutex: Mutex::new(()), // Add this
         })
     }
 
@@ -230,28 +219,19 @@ impl KvRouter {
         context_id: &str,
         tokens: &[u32],
     ) -> anyhow::Result<(i64, u32)> {
-        // Acquire mutex to serialize access
-        // TODO: may as well make all the subroutines synchronous if benchmarking favors this
-        let _guard = self.find_best_match_mutex.lock().await;
-
         let isl_tokens = tokens.len();
 
         let block_hashes = compute_block_hash_for_seq(tokens, self.block_size);
-        let seq_hashes = compute_seq_hash_for_block(&block_hashes);
 
         let overlap_scores = self.indexer.find_matches(block_hashes.clone()).await?;
 
         let best_worker_id = self
             .scheduler
-            .schedule(
-                context_id.to_string(),
-                isl_tokens,
-                seq_hashes.clone(),
-                overlap_scores.clone(),
-            )
+            .schedule(context_id.to_string(), isl_tokens, overlap_scores.clone())
             .await?;
 
         if let Indexer::ApproxKvIndexer(ref indexer) = self.indexer {
+            let seq_hashes = compute_seq_hash_for_block(&block_hashes);
             indexer
                 .process_routing_decision(best_worker_id, block_hashes, seq_hashes)
                 .await
@@ -267,15 +247,10 @@ impl KvRouter {
     }
 
     /// Free all blocks associated with a request
-    pub async fn mark_prefill_completed(&self, request_id: &String) {
+    pub async fn mark_prefill_completed(&self, request_id: &str) {
         self.scheduler.mark_prefill_completed(request_id).await
     }
 
-    /// Free all blocks associated with a request
-    pub async fn free(&self, request_id: &String) {
-        self.scheduler.free(request_id).await
-    }
-
     /// Get the block size this router was configured with
     pub fn block_size(&self) -> u32 {
         self.block_size
@@ -362,8 +337,6 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
                     while let Some(item) = response_stream.next().await {
                         yield item;
                     }
-
-                    chooser.free(&context_id).await;
                 });
                 Ok(ResponseStream::new(wrapped_stream, stream_context))
             }

lib/llm/src/kv_router.rs

@@ -18,7 +18,7 @@ use std::sync::Once;
 pub use crate::kv_router::protocols::{ForwardPassMetrics, LoadMetrics, PredictiveLoadMetrics};
 use crate::kv_router::KV_METRICS_ENDPOINT;
 
-use crate::kv_router::scheduler::Endpoint;
+use crate::kv_router::scoring::Endpoint;
 use crate::kv_router::ProcessedEndpoints;
 use dynamo_runtime::component::Component;
 use dynamo_runtime::{service::EndpointInfo, utils::Duration, Result};

lib/llm/src/kv_router/metrics_aggregator.rs

@@ -2,25 +2,41 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use anyhow::Result;
-use dynamo_runtime::component::{Component, Namespace};
+use dynamo_runtime::component::Component;
 use dynamo_runtime::traits::events::{EventPublisher, EventSubscriber};
+use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use tokio::sync::RwLock;
+// Remove the Mutex import since we're using DashMap
 
 use super::protocols::{PrefillEvent, PrefillEventData};
 use crate::kv_router::PREFILL_SUBJECT;
 use dashmap::DashMap;
+use std::collections::HashMap;
+use std::hash::Hash;
+
+pub fn get_snapshot<K, V>(state: &DashMap<K, V>) -> HashMap<K, V>
+where
+    K: Clone + Hash + Eq,
+    V: Copy,
+{
+    state
+        .iter()
+        .map(|entry| (entry.key().clone(), *entry.value()))
+        .collect()
+}
 
 /// A counter that tracks pending prefill tokens for each request.
 ///
 /// This struct maintains a local hashmap of request_id to token count,
 /// a running sum of all tokens, and subscribes to prefill events over NATS
 /// to keep the counts synchronized across components.
+#[derive(Clone)]
 pub struct PrefillCounter {
     state: Arc<RwLock<PrefillCounterState>>,
-    namespace: Namespace,
+    component: Component,
 }
 
 struct PrefillCounterState {
@@ -78,14 +94,14 @@ impl PrefillCounter {
 
         let counter = Self {
             state: state.clone(),
-            namespace: component.namespace().clone(),
+            component: component.clone(),
         };
 
         let state_clone = state.clone();
-        let namespace_clone = counter.namespace.clone();
+        let component_clone = component.clone();
 
         tokio::spawn(async move {
-            if let Err(e) = Self::subscribe_to_events(state_clone, namespace_clone).await {
+            if let Err(e) = Self::subscribe_to_events(state_clone, component_clone).await {
                 tracing::error!("Error in prefill events subscription: {}", e);
             }
         });
@@ -97,9 +113,9 @@ impl PrefillCounter {
     /// TODO: somehow try to block events that are sent by itself
     async fn subscribe_to_events(
         state: Arc<RwLock<PrefillCounterState>>,
-        namespace: Namespace,
+        component: Component,
     ) -> Result<()> {
-        let mut subscriber = namespace
+        let mut subscriber = component
             .subscribe_with_type::<PrefillEvent>(PREFILL_SUBJECT)
             .await?;
 
@@ -135,7 +151,7 @@ impl PrefillCounter {
                         .tokens_map
                         .insert(event.request_id.clone(), new_tokens);
                 }
-                PrefillEventData::CompletePrefill(_) => {
+                PrefillEventData::CompletePrefill => {
                     let state_read = state.read().await;
                     if !state_read.contains_key(&event.request_id) {
                         continue;
@@ -164,7 +180,7 @@ impl PrefillCounter {
                 PrefillEventData::NewPrefill(tokens)
             },
         };
-        self.namespace.publish(PREFILL_SUBJECT, &event).await?;
+        self.component.publish(PREFILL_SUBJECT, &event).await?;
 
         Ok(old_value)
     }
@@ -173,12 +189,12 @@ impl PrefillCounter {
         let state = self.state.write().await;
         let removed_tokens = state.remove(request_id);
 
-        if let Some(tokens) = removed_tokens {
+        if removed_tokens.is_some() {
             let event = PrefillEvent {
                 request_id: request_id.to_string(),
-                data: PrefillEventData::CompletePrefill(tokens),
+                data: PrefillEventData::CompletePrefill,
             };
-            self.namespace.publish(PREFILL_SUBJECT, &event).await?;
+            self.component.publish(PREFILL_SUBJECT, &event).await?;
         }
 
         Ok(removed_tokens)
@@ -203,6 +219,92 @@ impl PrefillCounter {
         let state = self.state.read().await;
         state.tokens_map.is_empty()
     }
+
+    /// Returns a snapshot of the current state as a HashMap
+    pub async fn snapshot(&self) -> HashMap<String, usize> {
+        let state = self.state.read().await;
+        get_snapshot(&state.tokens_map)
+    }
+}
+
+/// A collection of PrefillCounters for multiple workers
+pub struct PrefillCountersMultiWorker {
+    pub counters: DashMap<i64, PrefillCounter>,
+    pub request_to_workers: DashMap<String, i64>,
+    component: Component,
+}
+
+impl PrefillCountersMultiWorker {
+    pub fn new(component: Component) -> Self {
+        Self {
+            counters: DashMap::new(),
+            request_to_workers: DashMap::new(),
+            component,
+        }
+    }
+
+    pub async fn add_prefill(
+        &self,
+        worker_id: i64,
+        request_id: String,
+        new_tokens: usize,
+    ) -> Result<()> {
+        if let Some(existing_worker_id) = self.request_to_workers.get(&request_id) {
+            tracing::warn!(
+                "Request {} already exists for worker {}, but trying to add to worker {}",
+                request_id,
+                *existing_worker_id,
+                worker_id
+            );
+        }
+        self.request_to_workers
+            .insert(request_id.clone(), worker_id);
+
+        if let Some(counter) = self.counters.get(&worker_id) {
+            counter.insert(request_id, new_tokens).await?;
+        } else {
+            tracing::warn!(
+                "Worker {} does not exist, creating new PrefillCounter",
+                worker_id
+            );
+            let new_counter = PrefillCounter::new(self.component.clone());
+            new_counter.insert(request_id, new_tokens).await?;
+            self.counters.insert(worker_id, new_counter);
+        }
+
+        Ok(())
+    }
+
+    pub async fn remove_prefill(&self, request_id: &str) -> Result<Option<usize>> {
+        let Some((_request_id, worker_id)) = self.request_to_workers.remove(request_id) else {
+            tracing::warn!("Request {} not found", request_id);
+            return Ok(None);
+        };
+
+        if let Some(counter) = self.counters.get(&worker_id) {
+            counter.remove(request_id).await
+        } else {
+            tracing::warn!(
+                "Worker {} not found in counters for request {}",
+                worker_id,
+                request_id
+            );
+            Ok(None)
+        }
+    }
+
+    /// Get the running sums for all workers as a HashMap<i64, usize>
+    pub async fn running_sums(&self) -> HashMap<i64, usize> {
+        let futures = FuturesUnordered::new();
+
+        for entry in self.counters.iter() {
+            let worker_id = *entry.key();
+            let counter = entry.value().clone();
+            futures.push(async move { (worker_id, counter.running_sum().await) });
+        }
+
+        futures.collect::<HashMap<_, _>>().await
+    }
 }
 
 #[cfg(test)]
@@ -222,22 +324,17 @@ mod integration_tests {
         let runtime = Runtime::from_current()?;
         let distributed = DistributedRuntime::from_settings(runtime.clone()).await?;
 
-        // Create namespace and components for two counters
+        // Create namespace and a single component
         let namespace = distributed.namespace("test_prefill_counter")?;
-        let component1 = namespace
-            .component("counter1")?
-            .service_builder()
-            .create()
-            .await?;
-        let component2 = namespace
-            .component("counter2")?
+        let component = namespace
+            .component("shared_counter")?
             .service_builder()
             .create()
             .await?;
 
-        // Create two PrefillCounter instances
-        let counter1 = PrefillCounter::new(component1);
-        let counter2 = PrefillCounter::new(component2);
+        // Create two PrefillCounter instances using the same component (cloned)
+        let counter1 = PrefillCounter::new(component.clone());
+        let counter2 = PrefillCounter::new(component.clone());
 
         // Give some time for subscribers to initialize
         tokio::time::sleep(Duration::from_millis(2000)).await;