feat: [router] serve worker KV query over dynamo endpoint instead of nats (#5451)

Author: biswapanda

examples/backends/vllm/launch/agg_router_replicas.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+set -e
+trap 'echo Cleaning up...; kill 0' EXIT
+
+# Set deterministic hash for KV event IDs
+export PYTHONHASHSEED=0
+
+# Common configuration
+MODEL="Qwen/Qwen3-0.6B"
+BLOCK_SIZE=64
+
+# run two routers (different HTTP + system ports)
+# Note: use --router-reset-states only on one router to avoid wiping shared state twice.
+DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT_R1:-8091} \
+python -m dynamo.frontend \
+    --router-mode kv \
+    --router-reset-states \
+    --http-port ${DYN_HTTP_PORT_R1:-8000} &
+
+DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT_R2:-8092} \
+python -m dynamo.frontend \
+    --router-mode kv \
+    --http-port ${DYN_HTTP_PORT_R2:-8001} &
+
+# run workers (enable local indexer so routers can query on restart)
+DYN_LOCAL_INDEXER=true \
+DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT1:-8081} \
+CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager \
+    --connector none \
+    --enable-local-indexer true \
+    --kv-events-config '{"publisher":"zmq","topic":"kv-events","endpoint":"tcp://*:20080","enable_kv_cache_events":true}' &
+
+DYN_LOCAL_INDEXER=true \
+DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT2:-8082} \
+VLLM_NIXL_SIDE_CHANNEL_PORT=20097 \
+CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager \
+    --connector none \
+    --enable-local-indexer true \
+    --kv-events-config '{"publisher":"zmq","topic":"kv-events","endpoint":"tcp://*:20081","enable_kv_cache_events":true}'

lib/llm/src/kv_router.rs

@@ -79,7 +79,7 @@ pub const RADIX_STATE_BUCKET: &str = "radix-bucket";
 pub const RADIX_STATE_FILE: &str = "radix-state";
 
 // for worker-local kvindexer query
-pub const WORKER_KV_INDEXER_QUERY_SUBJECT: &str = "worker_kv_indexer_query";
+pub const WORKER_KV_INDEXER_QUERY_ENDPOINT: &str = "worker_kv_indexer_query";
 pub const WORKER_KV_INDEXER_BUFFER_SIZE: usize = 1024; // store 1024 most recent events in worker buffer
 
 // for router discovery registration

lib/llm/src/kv_router/indexer.rs

@@ -35,6 +35,7 @@ use async_trait::async_trait;
 use dynamo_runtime::{
     component::Component,
     metrics::{MetricsHierarchy, prometheus_names::kvrouter},
+    protocols::maybe_error::MaybeError,
 };
 use prometheus::{IntCounterVec, Opts};
 use serde::{Deserialize, Serialize};
@@ -142,6 +143,21 @@ pub enum WorkerKvQueryResponse {
     },
     /// Invalid range: end_id < start_id
     InvalidRange { start_id: u64, end_id: u64 },
+    /// Query failed on worker (serialized error)
+    Error(String),
+}
+
+impl MaybeError for WorkerKvQueryResponse {
+    fn from_err(err: Box<dyn std::error::Error + Send + Sync>) -> Self {
+        WorkerKvQueryResponse::Error(err.to_string())
+    }
+
+    fn err(&self) -> Option<anyhow::Error> {
+        match self {
+            WorkerKvQueryResponse::Error(msg) => Some(anyhow::Error::msg(msg.clone())),
+            _ => None,
+        }
+    }
 }
 
 /// A block in the Radix Tree.

lib/llm/src/kv_router/publisher.rs

@@ -16,20 +16,17 @@ use tokio_util::sync::CancellationToken;
 use zeromq::{Socket, SocketRecv, SubSocket};
 
 use dynamo_runtime::metrics::{MetricsHierarchy, prometheus_names::kvstats};
-use dynamo_runtime::traits::{
-    DistributedRuntimeProvider, events::EventPublisher, events::EventSubscriber,
-};
+use dynamo_runtime::traits::{DistributedRuntimeProvider, events::EventPublisher};
 use dynamo_runtime::{
     component::{Component, Namespace},
     transports::nats::{NatsQueue, Slug},
 };
-use futures::StreamExt;
 
 use crate::kv_router::{
     KV_EVENT_SUBJECT, KV_METRICS_SUBJECT, WORKER_KV_INDEXER_BUFFER_SIZE,
-    WORKER_KV_INDEXER_QUERY_SUBJECT,
-    indexer::{KvIndexerMetrics, LocalKvIndexer, RouterEvent, WorkerKvQueryRequest},
+    indexer::{KvIndexerMetrics, LocalKvIndexer, RouterEvent},
     protocols::*,
+    worker_query::start_worker_kv_query_endpoint,
 };
 use dynamo_runtime::config::environment_names::nats as env_nats;
 
@@ -173,11 +170,10 @@ impl KvEventPublisher {
                 .drt()
                 .runtime()
                 .secondary()
-                .spawn(start_worker_kv_query_service(
+                .spawn(start_worker_kv_query_endpoint(
                     component,
                     worker_id,
                     local_indexer,
-                    cancellation_token.clone(),
                 ))
         });
 
@@ -311,80 +307,6 @@ async fn start_event_processor<P: EventPublisher + Send + Sync + 'static>(
     }
 }
 
-// Processor for Router -> LocalKvIndexer query service
-async fn start_worker_kv_query_service(
-    component: Component,
-    worker_id: u64,
-    local_indexer: Arc<LocalKvIndexer>,
-    cancellation_token: CancellationToken,
-) {
-    // Create NATS subscriber on a subject specific to worker's id
-    let subject = format!("{}.{}", WORKER_KV_INDEXER_QUERY_SUBJECT, worker_id);
-    let mut subscriber = match component.subscribe(&subject).await {
-        Ok(sub) => sub,
-        Err(e) => {
-            tracing::error!(
-                "Query service failed to subscribe for worker {worker_id} on subject {subject}: {e}"
-            );
-            return;
-        }
-    };
-    tracing::info!("Query service listening on NATS for worker {worker_id} on subject {subject}");
-
-    // Receive query request from router, retrieve event(s) from LocalKvIndexer, return response
-    loop {
-        tokio::select! {
-            _ = cancellation_token.cancelled() => {
-                tracing::info!("Query service received cancellation signal for worker {worker_id}");
-                break;
-            }
-
-            msg = subscriber.next() => {
-                let Some(msg) = msg else {
-                    tracing::warn!("Query service NATS stream ended for worker {worker_id}");
-                    break;
-                };
-
-                // deserialize from msg (async_nats::Message)
-                let request: WorkerKvQueryRequest = match serde_json::from_slice(&msg.payload) {
-                    Ok(request) => request,
-                    Err(e) => {
-                        tracing::error!("Failed to deserialize WorkerKvQueryRequest for worker {worker_id}: {e}");
-                        continue;
-                    }
-                };
-
-                tracing::debug!("Received query request for worker {worker_id}: {request:?}");
-
-                // Query events based on optional start/end ids
-                let response = local_indexer
-                    .get_events_in_id_range(request.start_event_id, request.end_event_id)
-                    .await;
-
-                // Send reply back (if reply subject exists)
-                if let Some(reply_subject) = msg.reply {
-                    let payload = match serde_json::to_vec(&response) {
-                        Ok(p) => p,
-                        Err(e) => {
-                            tracing::error!("Failed to serialize response for worker {worker_id}: {e}");
-                            continue;
-                        }
-                    };
-
-                    // Publish through DRT/NATS directly instead of namespace (adds a prefix)
-                    if let Err(e) = component
-                        .drt()
-                        .kv_router_nats_publish(reply_subject.to_string(), payload.into())
-                        .await
-                    {
-                        tracing::error!("Failed to send reply for worker {worker_id}: {e}");
-                    }
-                }
-            }
-        }
-    }
-}
-
 /// Calculate exponential backoff duration based on consecutive error count
 fn calculate_backoff_ms(consecutive_errors: u32) -> u64 {
     std::cmp::min(
@@ -1864,6 +1786,9 @@ mod tests_startup_helpers {
         let missed_events = match response {
             crate::kv_router::indexer::WorkerKvQueryResponse::Events(e) => e,
             crate::kv_router::indexer::WorkerKvQueryResponse::TreeDump(e) => e,
+            crate::kv_router::indexer::WorkerKvQueryResponse::Error(message) => {
+                panic!("Unexpected error response: {message}")
+            }
             other => panic!("Unexpected response: {:?}", other),
         };
         assert_eq!(

lib/llm/src/kv_router/subscriber.rs

@@ -232,6 +232,9 @@ pub async fn recover_from_worker(
         WorkerKvQueryResponse::InvalidRange { start_id, end_id } => {
             anyhow::bail!("Invalid range: end_id ({end_id}) < start_id ({start_id})");
         }
+        WorkerKvQueryResponse::Error(message) => {
+            anyhow::bail!("Worker {worker_id} query failed: {message}");
+        }
     };
 
     let events_count = events.len();