Skip to content

fix: scope KV router cancellation to router lifetime #10730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
AmeenP wants to merge 1 commit into from
Closed

fix: scope KV router cancellation to router lifetime #10730

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 25 additions & 1 deletion lib/llm/src/discovery/worker_monitor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use std::sync::RwLock;
use std::sync::atomic::{AtomicBool, Ordering};

use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;

use dashmap::DashMap;
use dynamo_kv_router::protocols::ActiveLoad;
Expand Down Expand Up @@ -398,6 +399,18 @@ pub struct KvWorkerMonitor {
thresholds: Arc<RwLock<LoadThresholdConfig>>,
/// Guard to ensure start_monitoring() only runs once across clones
started: Arc<AtomicBool>,
/// Cancels the background monitoring task when the last monitor clone is dropped.
cancellation: Arc<MonitorCancellation>,
}

struct MonitorCancellation {
token: CancellationToken,
}

impl Drop for MonitorCancellation {
fn drop(&mut self) {
self.token.cancel();
}
}

impl KvWorkerMonitor {
Expand All @@ -422,6 +435,9 @@ impl KvWorkerMonitor {
worker_load_states: Arc::new(DashMap::new()),
thresholds: Arc::new(RwLock::new(config)),
started: Arc::new(AtomicBool::new(false)),
cancellation: Arc::new(MonitorCancellation {
token: CancellationToken::new(),
}),
}
}

Expand Down Expand Up @@ -536,7 +552,15 @@ impl WorkerLoadMonitor for KvWorkerMonitor {
let endpoint = &self.client.endpoint;
let component = endpoint.component();

let cancellation_token = component.drt().child_token();
let cancellation_token = self.cancellation.token.child_token();
let runtime_cancellation = component.drt().child_token();
let cancellation_bridge = cancellation_token.clone();
tokio::spawn(async move {
tokio::select! {
_ = runtime_cancellation.cancelled() => cancellation_bridge.cancel(),
_ = cancellation_bridge.cancelled() => {}
}
});

// Watch for runtime config updates from model deployment cards via discovery interface
let discovery = component.drt().discovery();
Expand Down
13 changes: 10 additions & 3 deletions lib/llm/src/kv_router.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,12 +225,13 @@ 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();
let cancellation_token = component.drt().child_token();

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep router token off endpoint shutdown

Runtime::child_token() is tied to the endpoint-shutdown token (lib/runtime/src/runtime.rs:293-295), and shutdown cancels that in Phase 1 before waiting for graceful endpoints (lib/runtime/src/runtime.rs:322-338). Because this router token is now passed into the KV indexer/subscriber paths, single-thread KV routers (router_event_threads=1) can stop their KvIndexer during the graceful-drain window and make in-flight requests hit IndexerOffline; the previous primary_token() stayed alive until Phase 3. Use a router-owned child of the primary token, while still canceling it in Drop.

Useful? React with 👍 / 👎.

let min_initial_workers = min_initial_workers_from_env()?;

let indexer = Indexer::new(
component,
&kv_router_config,
cancellation_token.clone(),
block_size,
model_name.as_deref(),
)
Expand Down Expand Up @@ -269,15 +270,21 @@ where
overlap_scores_refresh,
Some(overloaded_worker_provider),
worker_type,
cancellation_token.clone(),
)
.await?;

// Start KV event subscription if needed — skip when using a remote indexer.
if kv_router_config.use_remote_indexer {
tracing::info!("Skipping KV event subscription (using remote indexer)");
} else if kv_router_config.should_subscribe_to_kv_events() {
indexer::start_subscriber(component.clone(), &kv_router_config, indexer.clone())
.await?;
indexer::start_subscriber(
component.clone(),
&kv_router_config,
indexer.clone(),
cancellation_token.clone(),
)
.await?;
} else {
tracing::info!(
"Skipping KV event subscription (use_kv_events={}, overlap_score_credit={})",
Expand Down
21 changes: 15 additions & 6 deletions lib/llm/src/kv_router/indexer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 lookup;
mod recording;
Expand Down Expand Up @@ -69,6 +70,7 @@ impl Indexer {
pub async fn new(
component: &Component,
kv_router_config: &KvRouterConfig,
cancellation_token: CancellationToken,
block_size: u32,
model_name: Option<&str>,
) -> Result<Self> {
Expand Down Expand Up @@ -97,7 +99,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,
cancellation_token,
block_size,
);
return Ok(Self::Remote {
primary: Arc::new(remote),
approx,
Expand Down Expand Up @@ -128,7 +135,6 @@ impl Indexer {
});
}

let cancellation_token = component.drt().primary_token();
return Ok(Self::KvIndexer {
primary: KvIndexer::new_with_pruning(
cancellation_token,
Expand All @@ -142,7 +148,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,
cancellation_token.clone(),
block_size,
);

if kv_router_config.router_event_threads > 1 {
let kv_indexer_metrics = KvIndexerMetrics::from_component(component);
Expand All @@ -163,8 +174,6 @@ 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,
Expand Down
2 changes: 1 addition & 1 deletion lib/llm/src/kv_router/indexer/recovery/jetstream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
17 changes: 13 additions & 4 deletions lib/llm/src/kv_router/indexer/recovery/subscriber.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,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.
///
Expand All @@ -27,9 +28,8 @@ async fn start_kv_router_background_event_plane(
component: Component,
indexer: Indexer,
transport_kind: EventTransportKind,
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
Expand All @@ -45,7 +45,8 @@ async fn start_kv_router_background_event_plane(

// WorkerQueryClient handles its own discovery loop for lifecycle + initial recovery.
// No blocking wait — recovery happens asynchronously as endpoints are discovered.
let worker_query_client = WorkerQueryClient::spawn(component.clone(), indexer).await?;
let worker_query_client =
WorkerQueryClient::spawn(component.clone(), indexer, cancellation_token.clone()).await?;
let kv_event_subject = format!(
"namespace.{}.component.{}.{}",
component.namespace().name(),
Expand Down Expand Up @@ -116,6 +117,7 @@ pub async fn start_subscriber(
component: Component,
kv_router_config: &KvRouterConfig,
indexer: Indexer,
cancellation_token: CancellationToken,
) -> Result<()> {
let transport_kind = component.drt().default_event_transport_kind();

Expand All @@ -140,6 +142,7 @@ pub async fn start_subscriber(
consumer_id,
indexer,
kv_router_config,
cancellation_token,
)
.await
} else {
Expand All @@ -156,6 +159,12 @@ pub async fn start_subscriber(
tracing::info!("Using NATS Core subscription (local_indexer mode)");
}

start_kv_router_background_event_plane(component, indexer, transport_kind).await
start_kv_router_background_event_plane(
component,
indexer,
transport_kind,
cancellation_token,
)
.await
}
}
7 changes: 5 additions & 2 deletions lib/llm/src/kv_router/indexer/recovery/worker_query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,12 +87,15 @@ impl WorkerQueryClient {
/// The background loop watches `ComponentEndpoints` discovery for query endpoints,
/// recovers each `(worker_id, dp_rank)` as it appears, and sends worker removal
/// events when all dp_ranks for a worker disappear.
pub async fn spawn(component: Component, indexer: Indexer) -> Result<Arc<Self>> {
pub async fn spawn(
component: Component,
indexer: Indexer,
cancel_token: tokio_util::sync::CancellationToken,
) -> Result<Arc<Self>> {
let transport = Arc::new(RuntimeWorkerQueryTransport::new(&component).await?);
let client = Self::new(component.clone(), indexer, transport);

let client_bg = client.clone();
let cancel_token = component.drt().primary_token();
tokio::spawn(async move {
if let Err(e) = client_bg.run_discovery_loop(cancel_token).await {
Comment on lines +90 to 100

@coderabbitai coderabbitai Bot Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Store the cancellation token and apply it to recovery tasks too.

The new token only stops the discovery loop. Recovery tasks spawned by spawn_recovery_task still hold Arc<WorkerQueryClient> and can continue through jitter, semaphore waits, retries, and worker queries after the router token is cancelled.

Direction for the fix 
 pub struct WorkerQueryClient {
     component: Component,
     transport: Arc<dyn WorkerQueryTransport>,
     /// Indexer for applying recovered events and worker removals.
     indexer: Indexer,
+    cancel_token: tokio_util::sync::CancellationToken,
     worker_states: DashMap<WorkerId, Arc<Mutex<WorkerState>>>,
     query_endpoints: WorkerQueryEndpointDirectory,
     recovery_semaphore: Arc<Semaphore>,
 }

Then pass the token through new, clone it in spawn_recovery_task, and wrap jitter sleeps, semaphore acquisition, retry backoffs, and query_worker(...) waits in tokio::select! branches that return when cancel_token.cancelled() fires.

Also applies to: 358-384, 531-590

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/kv_router/indexer/recovery/worker_query.rs` around lines 90 -
100, The cancellation token is only applied to the discovery loop in the spawn
function, but recovery tasks can continue running after cancellation. Store the
cancel_token as a field in the struct by adding it to the `new` method signature
and storing it in Self. In the spawn function at lines 90-100, pass the
cancel_token to the `new` call. In the spawn_recovery_task method at lines
358-384, clone the stored cancel_token and pass it to the task spawning logic.
In the recovery task implementation at lines 531-590, wrap all awaitable
operations (jitter sleeps, semaphore acquisition, retry backoffs, and
query_worker calls) in tokio::select! branches that also listen for the
cancel_token.cancelled() signal to allow immediate cancellation when the token
fires.

tracing::error!("WorkerQueryClient discovery loop failed: {e}");
Expand Down
5 changes: 3 additions & 2 deletions lib/llm/src/kv_router/indexer/side.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -28,6 +29,7 @@ impl SideIndexer {
pub(super) fn new_predict_on_route(
component: &Component,
kv_router_config: &KvRouterConfig,
cancellation_token: CancellationToken,
block_size: u32,
) -> Option<Self> {
let ttl_secs = kv_router_config.router_predicted_ttl_secs?;
Expand All @@ -51,7 +53,6 @@ impl SideIndexer {
)));
}

let cancellation_token = component.drt().primary_token();
Some(Self::KvIndexer(KvIndexer::new_with_pruning(
cancellation_token,
block_size,
Expand Down
7 changes: 5 additions & 2 deletions lib/llm/src/kv_router/scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,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
Expand Down Expand Up @@ -66,6 +67,7 @@ where
overlap_scores_refresh: Option<Arc<RF>>,
overloaded_worker_provider: Option<OverloadedWorkerProvider>,
worker_type: &'static str,
cancellation_token: CancellationToken,
) -> Result<Self, KvSchedulerError> {
let initial_workers: HashMap<WorkerId, ModelRuntimeConfig> =
workers_with_configs.borrow().clone();
Expand Down Expand Up @@ -107,13 +109,13 @@ where
overloaded_worker_provider,
kv_router_config.router_queue_recheck_interval(),
kv_router_config.router_track_prefill_tokens,
component.drt().child_token(),
cancellation_token.clone(),
worker_type,
watch_worker_configs,
));

let metrics_scheduler = Arc::clone(&inner);
let metrics_cancel_token = component.drt().child_token();
let metrics_cancel_token = cancellation_token;
let mut queue_updates = inner.subscribe_queue_updates();
tokio::spawn(async move {
let mut recheck_interval = tokio::time::interval(Duration::from_secs(60));
Expand Down Expand Up @@ -358,6 +360,7 @@ mod tests {
None::<Arc<NoopOverlapScoresRefresh>>,
None,
"decode",
CancellationToken::new(),
)
.await
.unwrap();
Expand Down
Loading