fix: scope KV router cancellation to router lifetime - #10730
Conversation
|
👋 Hi AmeenP! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
Bot
added
external-contribution
AmeenP
changed the title
5d4f5a7 to
88a89ec
Compare
AmeenP
changed the title
Signed-off-by: AmeenP <ameenp360@gmail.com>
88a89ec to
ca361d1
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 88a89ec8cb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| kv_router_config.validate()?; | ||
| let component = endpoint.component(); | ||
| let cancellation_token = component.drt().primary_token(); | ||
| let cancellation_token = component.drt().child_token(); |
There was a problem hiding this comment.
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 👍 / 👎.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
WalkthroughAcross the KV router subsystem, cancellation tokens are no longer derived internally from ChangesKV Router Cancellation Token Externalization
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
lib/llm/src/discovery/worker_monitor.rs (1)
555-576:⚠️ Potential issue | 🟠 Major | ⚡ Quick winCancel the per-attempt token when monitor startup fails.
The bridge task is spawned before
list_and_watchcan fail; on the error path onlystartedis reset, so retries can leave bridge tasks parked until runtime or monitor shutdown. Cancel this child token before returning the error.Proposed fix
Err(e) => { tracing::error!("KvWorkerMonitor: failed to create discovery stream: {}", e); + cancellation_token.cancel(); // Reset started flag so retry can work self.started.store(false, Ordering::SeqCst); return Err(e); }🤖 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/discovery/worker_monitor.rs` around lines 555 - 576, The bridge task spawned at the beginning of the code block uses a child cancellation token that remains active even when list_and_watch fails. In the error handling block where the started flag is reset and the error is returned, add a call to cancel the cancellation_token (the child token created at the start) before returning the Err(e). This ensures the bridge task is properly cleaned up on retry attempts rather than remaining parked until runtime or monitor shutdown.lib/llm/src/kv_router/scheduler.rs (1)
353-364:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winCancel the scheduler token used by the test.
Line 363 now passes an independent token, so Line 395 no longer stops the scheduler or metrics task. Keep the token in a variable and cancel that instead.
Proposed fix
+ let cancellation_token = CancellationToken::new(); let scheduler = KvScheduler::start( component.clone(), 64, cfg_rx, DefaultWorkerSelector::new(Some(config.clone()), "decode"), &config, None, None::<Arc<NoopOverlapScoresRefresh>>, None, "decode", - CancellationToken::new(), + cancellation_token.clone(), ) .await .unwrap(); ... - component.drt().primary_token().cancel(); + cancellation_token.cancel();Also applies to: 395-395
🤖 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/scheduler.rs` around lines 353 - 364, The test is creating a CancellationToken inline in the KvScheduler::start call at line 353-364, but then at line 395 needs to cancel the scheduler. Since the token is not stored in a variable, line 395 cannot reference it to perform the cancellation. Extract the CancellationToken::new() call into a variable before the KvScheduler::start invocation, pass that variable instead of the inline token creation, and then use that same variable at line 395 to cancel the scheduler token.lib/llm/src/kv_router.rs (1)
228-287:⚠️ Potential issue | 🟠 Major | ⚡ Quick winCancel the router token on
KvRouter::newfailure paths.This token is only cancelled in
Drop, butDropnever runs ifKvRouter::newfails afterIndexer::new,KvScheduler::start,start_subscriber, orensure_served_indexer_servicehas started background work. Add a construction guard that cancels the token on early return, then disarm it immediately before returningOk(Self { ... }).Sketch of the guard pattern
+struct CancelOnDrop(Option<tokio_util::sync::CancellationToken>); + +impl CancelOnDrop { + fn new(token: tokio_util::sync::CancellationToken) -> Self { + Self(Some(token)) + } + + fn disarm(mut self) { + self.0.take(); + } +} + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + if let Some(token) = self.0.take() { + token.cancel(); + } + } +} + ... let component = endpoint.component(); let cancellation_token = component.drt().child_token(); + let cancellation_guard = CancelOnDrop::new(cancellation_token.clone()); ... - Ok(Self { + let router = Self { indexer, scheduler, workers_with_configs, block_size, kv_router_config, prefill_load_estimator, cancellation_token, client, is_eagle, _served_indexer_handle: served_indexer_handle, shared_cache, - }) + }; + cancellation_guard.disarm(); + Ok(router)Also applies to: 296-327
🤖 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.rs` around lines 228 - 287, Implement a cancellation guard pattern to ensure the cancellation_token is properly cancelled if KvRouter::new fails after background work has started. Create a guard variable that wraps the cancellation_token immediately after obtaining it from component.drt().child_token(). The guard should automatically cancel the token in its Drop implementation. Before the final Ok(Self { ... }) return statement, call a method to disarm the guard to prevent cancellation on successful initialization. This ensures that if any of the subsequent operations (Indexer::new, KvScheduler::start, start_subscriber, or ensure_served_indexer_service) fail and the function returns early, the background tasks will be properly cancelled rather than left running orphaned.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@lib/llm/src/kv_router/indexer/recovery/worker_query.rs`:
- Around line 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.
---
Outside diff comments:
In `@lib/llm/src/discovery/worker_monitor.rs`:
- Around line 555-576: The bridge task spawned at the beginning of the code
block uses a child cancellation token that remains active even when
list_and_watch fails. In the error handling block where the started flag is
reset and the error is returned, add a call to cancel the cancellation_token
(the child token created at the start) before returning the Err(e). This ensures
the bridge task is properly cleaned up on retry attempts rather than remaining
parked until runtime or monitor shutdown.
In `@lib/llm/src/kv_router.rs`:
- Around line 228-287: Implement a cancellation guard pattern to ensure the
cancellation_token is properly cancelled if KvRouter::new fails after background
work has started. Create a guard variable that wraps the cancellation_token
immediately after obtaining it from component.drt().child_token(). The guard
should automatically cancel the token in its Drop implementation. Before the
final Ok(Self { ... }) return statement, call a method to disarm the guard to
prevent cancellation on successful initialization. This ensures that if any of
the subsequent operations (Indexer::new, KvScheduler::start, start_subscriber,
or ensure_served_indexer_service) fail and the function returns early, the
background tasks will be properly cancelled rather than left running orphaned.
In `@lib/llm/src/kv_router/scheduler.rs`:
- Around line 353-364: The test is creating a CancellationToken inline in the
KvScheduler::start call at line 353-364, but then at line 395 needs to cancel
the scheduler. Since the token is not stored in a variable, line 395 cannot
reference it to perform the cancellation. Extract the CancellationToken::new()
call into a variable before the KvScheduler::start invocation, pass that
variable instead of the inline token creation, and then use that same variable
at line 395 to cancel the scheduler token.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 58736094-7527-4f74-b4b6-0debfcc4c86d
📒 Files selected for processing (8)
lib/llm/src/discovery/worker_monitor.rslib/llm/src/kv_router.rslib/llm/src/kv_router/indexer/mod.rslib/llm/src/kv_router/indexer/recovery/jetstream.rslib/llm/src/kv_router/indexer/recovery/subscriber.rslib/llm/src/kv_router/indexer/recovery/worker_query.rslib/llm/src/kv_router/indexer/side.rslib/llm/src/kv_router/scheduler.rs
| 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 { |
There was a problem hiding this comment.
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.
|
/ok to test ca361d1 |
|
Closing as superseded: #11390 (merged 2026-07-09) landed the core of this design — router-owned |
2 participants
Summary
This scopes KV-router and worker-monitor background work to the lifetime of the owner that created it, instead of using process/runtime-level cancellation for WorkerSet-scoped work.
Refs #10729.
What changed
KvRouter::newusingcomponent.drt().child_token().WorkerQueryClientrecovery/discovery.component.drt().primary_token()inside router-owned indexer/subscriber/recovery paths.KvWorkerMonitoran owned cancellation token that is canceled when the last monitor handle is dropped, while still bridging runtime shutdown into that owned token.Why
On LoRA-enabled KV-router deployments, frontend RSS and thread count can grow until OOM. The suspected root cause is that WorkerSet/router-scoped background tasks survive WorkerSet/router removal because they are tied to runtime/process cancellation instead of owner lifetime.
LoRA deployments make this especially visible because model-card/adapter churn repeatedly exercises WorkerSet/router teardown and rebuild.
Validation
Passed locally on current
origin/mainafter rebase:Follow-up validation
Summary by CodeRabbit