fix(kv-router): scope KvRouter cancellation to a child token

lib/llm/src/kv_router.rs

@@ -21,6 +21,7 @@ use dynamo_kv_router::{
     },
 };
 use dynamo_runtime::{
+    CancellationToken,
     component::{Client, Endpoint},
     discovery::DiscoveryQuery,
     error::{DynamoError, ErrorType},
@@ -219,7 +220,7 @@ where
     block_size: u32,
     kv_router_config: KvRouterConfig,
     prefill_load_estimator: Option<Arc<dyn PrefillLoadEstimator>>,
-    cancellation_token: tokio_util::sync::CancellationToken,
+    cancellation_token: CancellationToken,
     client: Client,
     is_eagle: bool,
     _served_indexer_handle: Option<ServedIndexerHandle>,
@@ -254,14 +255,17 @@ where
         let kv_router_config = kv_router_config.unwrap_or_default();
         kv_router_config.validate()?;
         let component = endpoint.component();
-        let cancellation_token = component.drt().primary_token();
+        // Router-owned tasks derive from this token so a rebuild cannot cancel the runtime.
+        let cancellation_token = component.drt().child_token();
+        let cancellation_guard = cancellation_token.clone().drop_guard();
         let min_initial_workers = min_initial_workers_from_env()?;
 
         let indexer = Indexer::new(
             component,
             &kv_router_config,
             block_size,
             model_name.as_deref(),
+            cancellation_token.child_token(),
         )
         .await?;
 
@@ -300,6 +304,7 @@ where
             Some(overloaded_worker_provider),
             model_name.as_deref(),
             worker_type,
+            cancellation_token.child_token(),
         )
         .await?;
 
@@ -314,6 +319,7 @@ where
                 workers_with_configs.clone(),
                 model_name.clone().unwrap_or_else(|| "unknown".to_string()),
                 worker_type,
+                cancellation_token.child_token(),
             )
             .await?;
         } else {
@@ -342,6 +348,7 @@ where
         };
 
         tracing::info!("KV Routing initialized");
+        let cancellation_token = cancellation_guard.disarm();
         Ok(Self {
             indexer,
             scheduler,

lib/llm/src/kv_router/indexer/mod.rs

@@ -20,8 +20,9 @@ use dynamo_kv_router::{
 pub(crate) use dynamo_kv_router::indexer::TieredMatchDetails;
 #[allow(unused_imports)]
 pub(crate) use dynamo_kv_router::indexer::WireTieredMatchDetails;
-use dynamo_runtime::{component::Component, traits::DistributedRuntimeProvider};
+use dynamo_runtime::component::Component;
 use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 
 mod embedding_cache;
 mod lookup;
@@ -77,6 +78,7 @@ impl Indexer {
         kv_router_config: &KvRouterConfig,
         block_size: u32,
         model_name: Option<&str>,
+        cancellation_token: CancellationToken,
     ) -> Result<Self> {
         if kv_router_config.overlap_score_credit == 0.0 {
             return Ok(Self::None);
@@ -103,7 +105,12 @@ impl Indexer {
             );
             let remote =
                 RemoteIndexer::new(component, model_name, kv_router_config.use_kv_events).await?;
-            let approx = SideIndexer::new_predict_on_route(component, kv_router_config, block_size);
+            let approx = SideIndexer::new_predict_on_route(
+                component,
+                kv_router_config,
+                block_size,
+                cancellation_token.child_token(),
+            );
             return Ok(Self::Remote {
                 primary: Arc::new(remote),
                 approx,
@@ -135,10 +142,9 @@ impl Indexer {
                 });
             }
 
-            let cancellation_token = component.drt().primary_token();
             return Ok(Self::KvIndexer {
                 primary: KvIndexer::new_with_pruning(
-                    cancellation_token,
+                    cancellation_token.child_token(),
                     block_size,
                     kv_indexer_metrics.clone(),
                     prune_config,
@@ -153,7 +159,12 @@ impl Indexer {
             });
         }
 
-        let approx = SideIndexer::new_predict_on_route(component, kv_router_config, block_size);
+        let approx = SideIndexer::new_predict_on_route(
+            component,
+            kv_router_config,
+            block_size,
+            cancellation_token.child_token(),
+        );
 
         if kv_router_config.router_event_threads > 1 {
             let kv_indexer_metrics = KvIndexerMetrics::from_component(component);
@@ -175,11 +186,9 @@ impl Indexer {
         }
 
         let kv_indexer_metrics = KvIndexerMetrics::from_component(component);
-        let cancellation_token = component.drt().primary_token();
-
         Ok(Self::KvIndexer {
             primary: KvIndexer::new_with_pruning(
-                cancellation_token,
+                cancellation_token.child_token(),
                 block_size,
                 kv_indexer_metrics.clone(),
                 None,

lib/llm/src/kv_router/indexer/recovery/jetstream.rs

@@ -227,8 +227,8 @@ pub(crate) async fn start_kv_router_background(
     consumer_id: String,
     indexer: Indexer,
     kv_router_config: &KvRouterConfig,
+    cancellation_token: CancellationToken,
 ) -> Result<()> {
-    let cancellation_token = component.drt().primary_token();
     let router_snapshot_threshold = kv_router_config.router_snapshot_threshold;
     let router_reset_states = kv_router_config.router_reset_states;
     // Set up NATS connections

lib/llm/src/kv_router/indexer/recovery/subscriber.rs

@@ -13,6 +13,7 @@ use dynamo_runtime::{
     component::Component, discovery::EventTransportKind, prelude::*,
     transports::event_plane::EventSubscriber,
 };
+use tokio_util::sync::CancellationToken;
 
 /// Start a simplified background task for event consumption using the event plane.
 ///
@@ -31,9 +32,8 @@ async fn start_kv_router_background_event_plane(
     workers_with_configs: RuntimeConfigWatch,
     model: String,
     worker_type: &'static str,
+    cancellation_token: CancellationToken,
 ) -> Result<()> {
-    let cancellation_token = component.drt().primary_token();
-
     // Subscribe to KV events BEFORE spawning the discovery/recovery loop.
     // This ensures no events are lost between the initial dump fetch and the
     // subscription becoming active — the tree state at fetch time is guaranteed
@@ -55,6 +55,7 @@ async fn start_kv_router_background_event_plane(
         workers_with_configs,
         model,
         worker_type,
+        cancellation_token.child_token(),
     )
     .await?;
     let kv_event_subject = format!(
@@ -130,6 +131,7 @@ pub async fn start_subscriber(
     workers_with_configs: RuntimeConfigWatch,
     model: String,
     worker_type: &'static str,
+    cancellation_token: CancellationToken,
 ) -> Result<()> {
     let transport_kind = component.drt().default_event_transport_kind();
 
@@ -154,6 +156,7 @@ pub async fn start_subscriber(
             consumer_id,
             indexer,
             kv_router_config,
+            cancellation_token,
         )
         .await
     } else {
@@ -177,6 +180,7 @@ pub async fn start_subscriber(
             workers_with_configs,
             model,
             worker_type,
+            cancellation_token,
         )
         .await
     }

lib/llm/src/kv_router/indexer/recovery/worker_query.rs

@@ -12,6 +12,7 @@ use dynamo_runtime::traits::DistributedRuntimeProvider;
 use futures::StreamExt;
 use rand::Rng;
 use tokio::sync::{Mutex, Semaphore};
+use tokio_util::sync::CancellationToken;
 
 use super::worker_query_directory::{DiscoveredQueryEndpoint, WorkerQueryEndpointDirectory};
 #[cfg(test)]
@@ -66,16 +67,18 @@ pub struct WorkerQueryClient {
     worker_states: DashMap<WorkerId, Arc<Mutex<WorkerState>>>,
     query_endpoints: Arc<WorkerQueryEndpointDirectory>,
     recovery_semaphore: Arc<Semaphore>,
+    cancellation_token: CancellationToken,
     /// Per-rank cancellation for in-flight recovery tasks; cancelled on rank
     /// removal so retry backoff stops polling workers that no longer exist.
-    recovery_cancels: DashMap<RecoveryKey, tokio_util::sync::CancellationToken>,
+    recovery_cancels: DashMap<RecoveryKey, CancellationToken>,
 }
 
 impl WorkerQueryClient {
     fn new(
         component: Component,
         indexer: Indexer,
         transport: Arc<dyn WorkerQueryTransport>,
+        cancellation_token: CancellationToken,
     ) -> Arc<Self> {
         Arc::new(Self {
             component,
@@ -84,6 +87,7 @@ impl WorkerQueryClient {
             worker_states: DashMap::new(),
             query_endpoints: Arc::new(WorkerQueryEndpointDirectory::default()),
             recovery_semaphore: Arc::new(Semaphore::new(RECOVERY_CONCURRENCY_LIMIT)),
+            cancellation_token,
             recovery_cancels: DashMap::new(),
         })
     }
@@ -102,14 +106,18 @@ impl WorkerQueryClient {
         workers_with_configs: RuntimeConfigWatch,
         model: String,
         worker_type: &'static str,
+        cancellation_token: CancellationToken,
     ) -> Result<Arc<Self>> {
         let transport = Arc::new(RuntimeWorkerQueryTransport::new(&component).await?);
-        let client = Self::new(component.clone(), indexer, transport);
+        let client = Self::new(
+            component.clone(),
+            indexer,
+            transport,
+            cancellation_token.clone(),
+        );
 
-        let discovery_cancel = component.drt().primary_token();
-        // TODO: Parent recovery tasks with a router-scoped token once the subscriber
-        // lifecycle owns one instead of relying on the runtime-wide token.
-        let health_cancel = discovery_cancel.child_token();
+        let discovery_cancel = cancellation_token.child_token();
+        let health_cancel = cancellation_token.child_token();
         spawn_kv_event_source_health_monitor(
             component.clone(),
             workers_with_configs,
@@ -428,7 +436,13 @@ impl WorkerQueryClient {
                         // the recovery identity, or use a generation that survives
                         // removal.
                         (worker_state.epoch == epoch && worker_state.ranks.contains_key(&key.1))
-                            .then(|| client.recovery_cancels.entry(key).or_default().clone())
+                            .then(|| {
+                                client
+                                    .recovery_cancels
+                                    .entry(key)
+                                    .or_insert_with(|| client.cancellation_token.child_token())
+                                    .clone()
+                            })
                     }
                     None => None,
                 };
@@ -856,7 +870,12 @@ mod tests {
         let component = make_test_component(name).await;
         let (kv_indexer, indexer) = make_test_indexer();
         let transport = Arc::new(MockWorkerQueryTransport::default());
-        let client = WorkerQueryClient::new(component, indexer, transport.clone());
+        let client = WorkerQueryClient::new(
+            component,
+            indexer,
+            transport.clone(),
+            CancellationToken::new(),
+        );
         (client, transport, kv_indexer)
     }
 

lib/llm/src/kv_router/indexer/side.rs

@@ -14,7 +14,8 @@ use dynamo_kv_router::{
     },
     protocols::{DpRank, OverlapScores, WorkerId, WorkerWithDpRank},
 };
-use dynamo_runtime::{component::Component, traits::DistributedRuntimeProvider};
+use dynamo_runtime::component::Component;
+use tokio_util::sync::CancellationToken;
 
 use super::lookup::HashInput;
 
@@ -29,6 +30,7 @@ impl SideIndexer {
         component: &Component,
         kv_router_config: &KvRouterConfig,
         block_size: u32,
+        cancellation_token: CancellationToken,
     ) -> Option<Self> {
         let ttl_secs = kv_router_config.router_predicted_ttl_secs?;
         let prune_config = Some(PruneConfig {
@@ -51,7 +53,6 @@ impl SideIndexer {
             )));
         }
 
-        let cancellation_token = component.drt().primary_token();
         Some(Self::KvIndexer(KvIndexer::new_with_pruning(
             cancellation_token,
             block_size,

lib/llm/src/kv_router/scheduler.rs

@@ -30,6 +30,7 @@ use dynamo_tokens::SequenceHash;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::time::Duration;
+use tokio_util::sync::CancellationToken;
 
 pub struct KvScheduler<Sel = DefaultWorkerSelector, RF = NoopOverlapScoresRefresh>
 where
@@ -60,6 +61,7 @@ where
         overloaded_worker_provider: Option<OverloadedWorkerProvider>,
         model_name: Option<&str>,
         worker_type: &'static str,
+        cancellation_token: CancellationToken,
     ) -> Result<Self, KvSchedulerError> {
         let initial_workers: HashMap<WorkerId, ModelRuntimeConfig> =
             workers_with_configs.borrow().clone();
@@ -72,6 +74,7 @@ where
             kv_router_config.router_replica_sync,
             router_id,
             worker_type,
+            cancellation_token.child_token(),
         )
         .await
         .map_err(|e| KvSchedulerError::InitFailed(e.to_string()))?;
@@ -107,14 +110,14 @@ where
             overloaded_worker_provider,
             kv_router_config.router_queue_recheck_interval(),
             kv_router_config.router_track_prefill_tokens,
-            component.drt().child_token(),
+            cancellation_token.child_token(),
             worker_type,
             watch_worker_configs,
         ));
 
         let metrics_scheduler = Arc::clone(&inner);
         let background_metrics = queue_metrics.clone();
-        let metrics_cancel_token = component.drt().child_token();
+        let metrics_cancel_token = cancellation_token.child_token();
         let mut queue_updates = inner.subscribe_queue_updates();
         tokio::spawn(async move {
             let mut recheck_interval = tokio::time::interval(Duration::from_secs(60));
@@ -471,6 +474,7 @@ mod tests {
             router_track_active_blocks: false,
             ..Default::default()
         };
+        let cancellation_token = CancellationToken::new();
 
         let scheduler = KvScheduler::start(
             component.clone(),
@@ -483,6 +487,7 @@ mod tests {
             None,
             Some("test-model"),
             "decode",
+            cancellation_token.clone(),
         )
         .await
         .unwrap();
@@ -514,6 +519,6 @@ mod tests {
         .await
         .unwrap();
 
-        component.drt().primary_token().cancel();
+        cancellation_token.cancel();
     }
 }