clone mdc rx, os env args, test header

components/src/dynamo/vllm/args.py

@@ -207,8 +207,10 @@ def parse_args() -> Config:
     )
     parser.add_argument(
         "--enable-local-indexer",
-        action="store_true",
-        help="Enable worker-local KV indexer for tracking this worker's own KV cache state.",
+        type=str,
+        choices=["true", "false"],
+        default=os.environ.get("DYN_LOCAL_INDEXER", "false"),
+        help="Enable worker-local KV indexer for tracking this worker's own KV cache state (can also be toggled with env var DYN_LOCAL_INDEXER).",
     )
     parser.add_argument(
         "--use-vllm-tokenizer",
@@ -220,6 +222,7 @@ def parse_args() -> Config:
 
     parser = AsyncEngineArgs.add_cli_args(parser)
     args = parser.parse_args()
+    args.enable_local_indexer = str(args.enable_local_indexer).lower() == "true"
     engine_args = AsyncEngineArgs.from_cli_args(args)
 
     # Workaround for vLLM GIL contention bug with NIXL connector when using UniProcExecutor.

lib/llm/src/kv_router.rs

@@ -313,16 +313,6 @@ impl KvRouter {
                 card.runtime_config
             });
 
-        // Watch for local indexer states via discovery interface (separate stream needed
-        // because streams are consumed by watch_and_extract_field)
-        let discovery_stream_local_indexer = discovery
-            .list_and_watch(discovery_key, Some(cancellation_token.clone()))
-            .await?;
-        let local_indexer_rx = watch_and_extract_field(
-            discovery_stream_local_indexer,
-            |card: ModelDeploymentCard| card.runtime_config.enable_local_indexer,
-        );
-
         let indexer = if kv_router_config.overlap_score_weight == 0.0 {
             // When overlap_score_weight is zero, we don't need to track prefixes
             Indexer::None
@@ -353,7 +343,7 @@ impl KvRouter {
             component.clone(),
             block_size,
             instance_ids_rx,
-            runtime_configs_rx,
+            runtime_configs_rx.clone(),
             selector,
             kv_router_config.router_replica_sync,
             consumer_id.clone(),
@@ -363,7 +353,7 @@ impl KvRouter {
         // Initialize worker query client using namespace abstraction
         // (created before background task so we can use it for startup recovery)
         let worker_query_client =
-            worker_query::WorkerQueryClient::new(component.clone(), local_indexer_rx);
+            worker_query::WorkerQueryClient::new(component.clone(), runtime_configs_rx.clone());
         tracing::info!("Worker query client initialized");
 
         // Start KV event subscriber background process (only when use_kv_events is enabled)

lib/llm/src/kv_router/publisher.rs

@@ -2133,7 +2133,7 @@ mod test_integration_publisher_with_kvindexer {
     use dynamo_runtime::protocols::annotated::Annotated;
 
     /// Integration test: KvPushRouter end-to-end routing with mock engines.
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     #[ignore] // Requires NATS/etcd. Run with: cargo test --package dynamo-llm --lib --features integration test_distributed_kvindexer_e2e -- --ignored --nocapture
     async fn test_distributed_kvindexer_e2e() -> anyhow::Result<()> {
         const BLOCK_SIZE: u32 = 4;
@@ -2339,7 +2339,7 @@ mod test_integration_publisher_with_kvindexer {
         Ok(())
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     #[ignore]
     async fn test_distributed_kvindexer_e2e_startup() -> anyhow::Result<()> {
         const BLOCK_SIZE: u32 = 4;

lib/llm/src/kv_router/worker_query.rs

@@ -12,6 +12,7 @@ use tokio::sync::watch;
 use crate::kv_router::WORKER_KV_INDEXER_QUERY_SUBJECT;
 use crate::kv_router::indexer::{WorkerKvQueryRequest, WorkerKvQueryResponse};
 use crate::kv_router::protocols::WorkerId;
+use crate::local_model::runtime_config::ModelRuntimeConfig;
 
 /// Router-side client for querying worker local KV indexers
 ///
@@ -21,27 +22,27 @@ use crate::kv_router::protocols::WorkerId;
 pub struct WorkerQueryClient {
     component: Component,
     /// Watch receiver for enable_local_indexer state per worker
-    local_indexer_rx: watch::Receiver<HashMap<WorkerId, bool>>,
+    model_runtime_config_rx: watch::Receiver<HashMap<WorkerId, ModelRuntimeConfig>>,
 }
 
 impl WorkerQueryClient {
     /// Create a new WorkerQueryClient with a watch receiver for local indexer states
     pub fn new(
         component: Component,
-        local_indexer_rx: watch::Receiver<HashMap<WorkerId, bool>>,
+        model_runtime_config_rx: watch::Receiver<HashMap<WorkerId, ModelRuntimeConfig>>,
     ) -> Self {
         Self {
             component,
-            local_indexer_rx,
+            model_runtime_config_rx,
         }
     }
 
     /// Check if a worker has local indexer enabled
     pub fn has_local_indexer(&self, worker_id: WorkerId) -> bool {
-        self.local_indexer_rx
+        self.model_runtime_config_rx
             .borrow()
             .get(&worker_id)
-            .copied()
+            .map(|config| config.enable_local_indexer)
             .unwrap_or(false)
     }