publish kv events over js

Signed-off-by: PeaBrane <yanrpei@gmail.com>

Author: PeaBrane

lib/bindings/python/rust/llm/kv.rs

@@ -870,12 +870,14 @@ impl KvPushRouter {
             // Get component from endpoint
             let component = endpoint.inner.component();
 
-            // Create KvRouter
+            // Create KvRouter with a unique consumer UUID
+            let consumer_uuid = uuid::Uuid::new_v4().to_string();
             let kv_router = llm_rs::kv_router::KvRouter::new(
                 component.clone(),
                 block_size as u32,
                 None, // default selector
                 Some(kv_router_config.inner()),
+                consumer_uuid,
             )
             .await
             .map_err(to_pyerr)?;

lib/llm/src/discovery/model_manager.rs

@@ -218,10 +218,11 @@ impl ModelManager {
             .drt()
             .etcd_client()
             .ok_or_else(|| anyhow::anyhow!("KV routing requires etcd (dynamic mode)"))?;
+        let router_uuid = uuid::Uuid::new_v4();
         let router_key = format!(
             "kv_routers/{}/{}",
             Slug::from_string(model_name),
-            uuid::Uuid::new_v4()
+            router_uuid
         );
         etcd_client
             .kv_create(
@@ -237,6 +238,7 @@ impl ModelManager {
             kv_cache_block_size,
             Some(selector),
             kv_router_config,
+            router_uuid.to_string(),
         )
         .await?;
         let new_kv_chooser = Arc::new(chooser);

lib/llm/src/kv_router.rs

@@ -15,6 +15,7 @@ use dynamo_runtime::{
     },
     prelude::*,
     protocols::annotated::Annotated,
+    transports::nats::NatsQueue,
 };
 use futures::stream::{self, StreamExt};
 use serde::{Deserialize, Serialize};
@@ -47,7 +48,9 @@ use crate::{
     protocols::common::llm_backend::LLMEngineOutput,
 };
 
-use dynamo_runtime::traits::events::EventSubscriber;
+use dynamo_runtime::traits::events::EventPublisher;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
 
 // [gluo TODO] shouldn't need to be public
 // this should be discovered from the component
@@ -56,7 +59,7 @@ use dynamo_runtime::traits::events::EventSubscriber;
 pub const KV_METRICS_ENDPOINT: &str = "load_metrics";
 
 // for metric publishing (push-based)
-pub const KV_EVENT_SUBJECT: &str = "kv_events";
+pub const KV_EVENT_SUBJECT: &str = "kv-events";
 pub const KV_HIT_RATE_SUBJECT: &str = "kv-hit-rate";
 pub const KV_METRICS_SUBJECT: &str = "kv_metrics";
 
@@ -165,11 +168,80 @@ pub struct KvRouter {
 }
 
 impl KvRouter {
+    /// Start a background task to consume events from NatsQueue and forward them to the indexer
+    async fn start_event_consumer(
+        component: Component,
+        consumer_uuid: String,
+        kv_events_tx: mpsc::Sender<RouterEvent>,
+        cancellation_token: CancellationToken,
+    ) -> Result<()> {
+        let stream_name = format!("{}_{}", component.subject(), KV_EVENT_SUBJECT)
+            .replace(['/', '\\', '.', '_'], "-");
+        let nats_server =
+            std::env::var("NATS_SERVER").unwrap_or_else(|_| "nats://localhost:4222".to_string());
+        let mut nats_queue = NatsQueue::new_with_consumer(
+            stream_name,
+            nats_server,
+            std::time::Duration::from_secs(300), // Very long timeout (5 minutes)
+            consumer_uuid,
+        );
+
+        nats_queue.connect().await?;
+
+        tokio::spawn(async move {
+            loop {
+                tokio::select! {
+                    _ = cancellation_token.cancelled() => {
+                        tracing::info!("KvRouter event consumer received cancellation signal");
+                        break;
+                    }
+                    result = nats_queue.dequeue_task(Some(std::time::Duration::from_secs(300))) => {
+                        match result {
+                            Ok(Some(bytes)) => {
+                                let event: RouterEvent = match serde_json::from_slice(&bytes) {
+                                    Ok(event) => event,
+                                    Err(e) => {
+                                        tracing::warn!("Failed to deserialize RouterEvent: {:?}", e);
+                                        continue;
+                                    }
+                                };
+
+                                // Forward the RouterEvent to the indexer
+                                if let Err(e) = kv_events_tx.send(event).await {
+                                    tracing::warn!(
+                                        "failed to send kv event to indexer; shutting down: {:?}",
+                                        e
+                                    );
+                                    break;
+                                }
+                            },
+                            Ok(None) => {
+                                tracing::trace!("KvRouter dequeue timeout, continuing");
+                            },
+                            Err(e) => {
+                                tracing::error!("Failed to dequeue task: {:?}", e);
+                                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                            }
+                        }
+                    }
+                }
+            }
+
+            // Clean up the queue and remove the durable consumer
+            if let Err(e) = nats_queue.shutdown().await {
+                tracing::warn!("Failed to shutdown NatsQueue: {}", e);
+            }
+        });
+
+        Ok(())
+    }
+
     pub async fn new(
         component: Component,
         block_size: u32,
         selector: Option<Box<dyn WorkerSelector + Send + Sync>>,
         kv_router_config: Option<KvRouterConfig>,
+        consumer_uuid: String,
     ) -> Result<Self> {
         let kv_router_config = kv_router_config.unwrap_or_default();
 
@@ -230,31 +302,15 @@ impl KvRouter {
         )
         .await?;
 
-        // [gluo TODO] try subscribe_with_type::<RouterEvent>,
-        // error checking below will be different.
+        // Start event consumer if using KvIndexer
         if let Indexer::KvIndexer(ref kv_indexer) = indexer {
-            let mut kv_events_rx = component.subscribe(KV_EVENT_SUBJECT).await?;
-            let kv_events_tx = kv_indexer.event_sender();
-
-            tokio::spawn(async move {
-                while let Some(event) = kv_events_rx.next().await {
-                    let event: RouterEvent = match serde_json::from_slice(&event.payload) {
-                        Ok(event) => event,
-                        Err(e) => {
-                            tracing::warn!("Failed to deserialize RouterEvent: {:?}", e);
-                            // Choosing warn and continue to process other events from other workers
-                            // A bad event likely signals a problem with a worker, but potentially other workers are still healthy
-                            continue;
-                        }
-                    };
-                    if let Err(e) = kv_events_tx.send(event).await {
-                        tracing::warn!(
-                            "failed to send kv event to indexer; shutting down: {:?}",
-                            e
-                        );
-                    }
-                }
-            });
+            Self::start_event_consumer(
+                component.clone(),
+                consumer_uuid,
+                kv_indexer.event_sender(),
+                cancellation_token.clone(),
+            )
+            .await?;
         }
 
         tracing::info!("KV Routing initialized");

lib/llm/src/kv_router/publisher.rs

@@ -29,6 +29,7 @@ use dynamo_runtime::{
         network::Ingress,
     },
     protocols::annotated::Annotated,
+    transports::nats::NatsQueue,
 };
 use futures::stream;
 use std::sync::Arc;
@@ -132,16 +133,26 @@ impl KvEventPublisher {
             )?);
         }
 
-        component
-            .drt()
-            .runtime()
-            .secondary()
-            .spawn(start_event_processor(
-                component,
-                worker_id,
-                cancellation_token.clone(),
-                rx,
-            ));
+        let stream_name = format!("{}.{}", component.subject(), KV_EVENT_SUBJECT)
+            .replace(['/', '\\', '.', '_'], "-");
+        let nats_server =
+            std::env::var("NATS_SERVER").unwrap_or_else(|_| "nats://localhost:4222".to_string());
+        // Create NatsQueue without consumer since we're only publishing
+        let mut nats_queue = NatsQueue::new_without_consumer(
+            stream_name,
+            nats_server,
+            std::time::Duration::from_secs(60), // Default timeout
+        );
+
+        // Connect the NatsQueue before passing it to the event processor
+        let cancellation_token_clone = cancellation_token.clone();
+        component.drt().runtime().secondary().spawn(async move {
+            if let Err(e) = nats_queue.connect().await {
+                tracing::error!("Failed to connect NatsQueue: {}", e);
+                return;
+            }
+            start_event_processor(nats_queue, worker_id, cancellation_token_clone, rx).await
+        });
 
         Ok(Self {
             kv_block_size,
@@ -197,7 +208,7 @@ async fn start_event_processor<P: EventPublisher + Send + Sync + 'static>(
                 // Encapsulate in a router event and publish.
                 tracing::trace!("Event processor for worker_id {} processing event: {:?}", worker_id, event.data);
                 let router_event = RouterEvent::new(worker_id, event);
-                if let Err(e) = publisher.publish(KV_EVENT_SUBJECT, &router_event).await {
+                if let Err(e) = publisher.publish("queue", &router_event).await {
                     tracing::error!("Failed to publish event: {}", e);
                 }
             }
@@ -845,7 +856,7 @@ mod tests_startup_helpers {
         let published = published.lock().unwrap();
         assert_eq!(published.len(), 1);
         let (subject, _) = &published[0];
-        assert_eq!(subject, &KV_EVENT_SUBJECT.to_string());
+        assert_eq!(subject, "queue");
     }
 
     //--------------------------------------------------------------------

lib/runtime/src/transports/nats.rs

@@ -450,11 +450,31 @@ pub struct NatsQueue {
 }
 
 impl NatsQueue {
-    /// Create a new NatsQueue with the given configuration
+    /// Create a new NatsQueue with the default "worker-group" consumer
     pub fn new(stream_name: String, nats_server: String, dequeue_timeout: time::Duration) -> Self {
         // Sanitize stream name to remove path separators (like in Python version)
+        // rupei: are we sure NATs stream name accepts '_'?
         let sanitized_stream_name = stream_name.replace(['/', '\\'], "_");
+        let subject = format!("{}.*", sanitized_stream_name);
+
+        Self {
+            stream_name: sanitized_stream_name,
+            nats_server,
+            dequeue_timeout,
+            client: None,
+            subject,
+            subscriber: None,
+            consumer_name: Some("worker-group".to_string()),
+        }
+    }
 
+    /// Create a new NatsQueue without a consumer (publisher-only mode)
+    pub fn new_without_consumer(
+        stream_name: String,
+        nats_server: String,
+        dequeue_timeout: time::Duration,
+    ) -> Self {
+        let sanitized_stream_name = stream_name.replace(['/', '\\'], "_");
         let subject = format!("{}.*", sanitized_stream_name);
 
         Self {
@@ -511,20 +531,18 @@ impl NatsQueue {
                 client.jetstream().create_stream(stream_config).await?;
             }
 
-            // Create persistent subscriber
-            let consumer_config = jetstream::consumer::pull::Config {
-                durable_name: Some(
-                    self.consumer_name
-                        .clone()
-                        .unwrap_or_else(|| "worker-group".to_string()),
-                ),
-                ..Default::default()
-            };
+            // Create persistent subscriber only if consumer_name is set
+            if let Some(ref consumer_name) = self.consumer_name {
+                let consumer_config = jetstream::consumer::pull::Config {
+                    durable_name: Some(consumer_name.clone()),
+                    ..Default::default()
+                };
 
-            let stream = client.jetstream().get_stream(&self.stream_name).await?;
-            let subscriber = stream.create_consumer(consumer_config).await?;
+                let stream = client.jetstream().get_stream(&self.stream_name).await?;
+                let subscriber = stream.create_consumer(consumer_config).await?;
+                self.subscriber = Some(subscriber);
+            }
 
-            self.subscriber = Some(subscriber);
             self.client = Some(client);
         }
 
@@ -767,6 +785,14 @@ impl EventPublisher for NatsQueue {
         event_name: impl AsRef<str> + Send + Sync,
         bytes: Vec<u8>,
     ) -> Result<()> {
+        // Check if event_name is "queue", otherwise warn
+        if event_name.as_ref() != "queue" {
+            tracing::warn!(
+                "Expected event_name to be 'queue', but got '{}'",
+                event_name.as_ref()
+            );
+        }
+
         let subject = format!("{}.{}", self.subject(), event_name.as_ref());
 
         // Note: enqueue_task requires &mut self, but EventPublisher requires &self
@@ -1019,9 +1045,14 @@ mod tests {
             consumer2_name,
         );
 
-        // Connect both queues (first one creates the stream, second one reuses it)
+        // Create a third queue without consumer (publisher-only)
+        let mut queue3 =
+            NatsQueue::new_without_consumer(stream_name.clone(), nats_server.clone(), timeout);
+
+        // Connect all queues (first one creates the stream, others reuse it)
         queue1.connect().await.expect("Failed to connect queue1");
         queue2.connect().await.expect("Failed to connect queue2");
+        queue3.connect().await.expect("Failed to connect queue3");
 
         // Send 4 messages using the EventPublisher trait
         let message_strings = vec![