From 08687cc855ac6d91f468bf65862ab945c1990f86 Mon Sep 17 00:00:00 2001 From: davilu-nvidia <134493379+davilu-nvidia@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:41:28 +0800 Subject: [PATCH 1/4] fix(kv-router): scope KvRouter cancellation to a child token The KvRouter captured the primary (root) CancellationToken and cancels it in its Drop impl to tear down its own background tasks (indexer, scheduler, overlap refresher). Because the primary token is the root of the runtime's cancellation tree, cancelling it propagates to the entire DistributedRuntime -- including the etcd lease keep-alive task. As a result, dropping a KvRouter during a rebuild (e.g. when all workers for a model exit and a new worker starts, triggering the router/indexer to be recreated) tore down the etcd lease and brought the whole runtime down with an etcd lease error. Use a child token so Drop only cancels this router's background tasks and does not affect the runtime-wide lease or other components. Signed-off-by: davilu-nvidia <134493379+davilu-nvidia@users.noreply.github.com> --- lib/llm/src/kv_router.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index a40873b6fbd7..8d94d8c3428e 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -254,7 +254,14 @@ 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(); + // Use a child token, not the primary (root) token. The `Drop` impl + // below cancels this token to tear down only this router's background + // tasks (indexer, scheduler, overlap refresher, ...). Cancelling the + // primary token would propagate to the entire DistributedRuntime -- + // including the etcd lease keep-alive -- so dropping a KvRouter during a + // rebuild (e.g. all workers for a model exit and a new one starts) would + // tear down the lease and bring down the whole runtime. + let cancellation_token = component.drt().child_token(); let min_initial_workers = min_initial_workers_from_env()?; let indexer = Indexer::new( From 73949e24e787b1c5e8822eb9d22d67ceed1322fd Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 8 Jul 2026 09:30:58 -0700 Subject: [PATCH 2/4] fix(kv-router): own background task cancellation Signed-off-by: PeaBrane --- lib/llm/src/kv_router.rs | 16 ++++----- lib/llm/src/kv_router/indexer/mod.rs | 25 ++++++++----- .../kv_router/indexer/recovery/jetstream.rs | 2 +- .../kv_router/indexer/recovery/subscriber.rs | 8 +++-- .../indexer/recovery/worker_query.rs | 35 ++++++++++++++----- lib/llm/src/kv_router/indexer/side.rs | 5 +-- lib/llm/src/kv_router/scheduler.rs | 11 ++++-- lib/llm/src/kv_router/sequence.rs | 13 ++++--- 8 files changed, 78 insertions(+), 37 deletions(-) diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 8d94d8c3428e..91aa3070b0dc 100644 --- a/lib/llm/src/kv_router.rs +++ b/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>, - cancellation_token: tokio_util::sync::CancellationToken, + cancellation_token: CancellationToken, client: Client, is_eagle: bool, _served_indexer_handle: Option, @@ -254,14 +255,9 @@ where let kv_router_config = kv_router_config.unwrap_or_default(); kv_router_config.validate()?; let component = endpoint.component(); - // Use a child token, not the primary (root) token. The `Drop` impl - // below cancels this token to tear down only this router's background - // tasks (indexer, scheduler, overlap refresher, ...). Cancelling the - // primary token would propagate to the entire DistributedRuntime -- - // including the etcd lease keep-alive -- so dropping a KvRouter during a - // rebuild (e.g. all workers for a model exit and a new one starts) would - // tear down the lease and bring down the whole runtime. + // 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( @@ -269,6 +265,7 @@ where &kv_router_config, block_size, model_name.as_deref(), + cancellation_token.child_token(), ) .await?; @@ -307,6 +304,7 @@ where Some(overloaded_worker_provider), model_name.as_deref(), worker_type, + cancellation_token.child_token(), ) .await?; @@ -321,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 { @@ -349,6 +348,7 @@ where }; tracing::info!("KV Routing initialized"); + let cancellation_token = cancellation_guard.disarm(); Ok(Self { indexer, scheduler, diff --git a/lib/llm/src/kv_router/indexer/mod.rs b/lib/llm/src/kv_router/indexer/mod.rs index a9d098cd8af5..6020df8c22b4 100644 --- a/lib/llm/src/kv_router/indexer/mod.rs +++ b/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 { 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, diff --git a/lib/llm/src/kv_router/indexer/recovery/jetstream.rs b/lib/llm/src/kv_router/indexer/recovery/jetstream.rs index bb1da97dd494..e7686deac999 100644 --- a/lib/llm/src/kv_router/indexer/recovery/jetstream.rs +++ b/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 diff --git a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs index e424bf5d659a..e48d5095e790 100644 --- a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs +++ b/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 } diff --git a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs index bfaf67e01d04..24bc6f00d1f4 100644 --- a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs +++ b/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,9 +67,10 @@ pub struct WorkerQueryClient { worker_states: DashMap>>, query_endpoints: Arc, recovery_semaphore: Arc, + 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, + recovery_cancels: DashMap, } impl WorkerQueryClient { @@ -76,6 +78,7 @@ impl WorkerQueryClient { component: Component, indexer: Indexer, transport: Arc, + cancellation_token: CancellationToken, ) -> Arc { 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> { 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) } diff --git a/lib/llm/src/kv_router/indexer/side.rs b/lib/llm/src/kv_router/indexer/side.rs index abaaa162c2d5..91eea37bf5f5 100644 --- a/lib/llm/src/kv_router/indexer/side.rs +++ b/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 { 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, diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 6b289e589e86..843bdbf254d3 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/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 where @@ -60,6 +61,7 @@ where overloaded_worker_provider: Option, model_name: Option<&str>, worker_type: &'static str, + cancellation_token: CancellationToken, ) -> Result { let initial_workers: HashMap = 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(); } } diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index 1baf9e61fc6f..914faf668fb7 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -16,11 +16,11 @@ pub use dynamo_kv_router::sequence::{ActiveSequences, RequestId}; use anyhow::Result; use dynamo_runtime::component::Component; -use dynamo_runtime::traits::DistributedRuntimeProvider; use dynamo_runtime::transports::event_plane::{EventPublisher, EventSubscriber}; use std::collections::HashMap; use std::sync::Arc; use std::task::{Context, Poll}; +use tokio_util::sync::CancellationToken; use super::metrics::{RouterWorkerStatusMetrics, WORKER_LOAD_METRICS}; use crate::kv_router::{ACTIVE_SEQUENCES_SUBJECT, KV_METRICS_SUBJECT}; @@ -133,6 +133,7 @@ pub async fn create_multi_worker_sequences( replica_sync: bool, router_id: u64, worker_type: &'static str, + cancellation_token: CancellationToken, ) -> Result> { let event_publisher = EventPublisher::for_component(&component, ACTIVE_SEQUENCES_SUBJECT).await?; @@ -172,12 +173,10 @@ pub async fn create_multi_worker_sequences( .await? .typed::(); let subscriber = RuntimeSequenceSubscriber { inner: subscriber }; - let cancel_token = component.drt().runtime().child_token(); - arc.start_replica_sync(subscriber, cancel_token); + arc.start_replica_sync(subscriber, cancellation_token.child_token()); } - let expiry_cancel = component.drt().runtime().child_token(); - arc.start_periodic_force_expiry_across_all_workers(expiry_cancel); + arc.start_periodic_force_expiry_across_all_workers(cancellation_token.child_token()); Ok(arc) } @@ -224,6 +223,7 @@ mod tests { true, 1, crate::discovery::WORKER_TYPE_DECODE, + CancellationToken::new(), ) .await?; let seq_manager_2 = create_multi_worker_sequences( @@ -233,6 +233,7 @@ mod tests { true, 2, crate::discovery::WORKER_TYPE_DECODE, + CancellationToken::new(), ) .await?; @@ -378,6 +379,7 @@ mod tests { true, 1, crate::discovery::WORKER_TYPE_DECODE, + CancellationToken::new(), ) .await?; let seq_manager_2 = create_multi_worker_sequences( @@ -387,6 +389,7 @@ mod tests { true, 2, crate::discovery::WORKER_TYPE_DECODE, + CancellationToken::new(), ) .await?; From e569307d6495666759bd5b85236753cd0ddfbb6a Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 8 Jul 2026 15:57:16 -0700 Subject: [PATCH 3/4] test(router): shut down ephemeral remote runtimes Signed-off-by: PeaBrane --- tests/router/common.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/router/common.py b/tests/router/common.py index 302384ade0b2..b7c6ece41b22 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -708,7 +708,7 @@ async def wait_for_served_indexer( raise TimeoutError("Timed out waiting for served indexer endpoints to register") - async def test_sync(): + async def run_test(runtimes): endpoint_path = ( f"{engine_workers.namespace}.{engine_workers.component_name}.generate" ) @@ -744,6 +744,7 @@ async def make_router( return runtime, endpoint, kv_router except Exception as error: last_error = error + runtime.shutdown() if not (serve_indexer or use_remote_indexer): raise del endpoint @@ -761,6 +762,7 @@ async def make_router( runtime_a, endpoint_a, router_a = await make_router( serve_indexer=True, use_remote_indexer=False ) + runtimes.append(runtime_a) serving_runtimes.append(runtime_a) serving_endpoints.append(endpoint_a) serving_routers.append(router_a) @@ -769,6 +771,7 @@ async def make_router( runtime_b, endpoint_b, router_b = await make_router( serve_indexer=True, use_remote_indexer=False ) + runtimes.append(runtime_b) serving_runtimes.append(runtime_b) serving_endpoints.append(endpoint_b) serving_routers.append(router_b) @@ -779,11 +782,12 @@ async def make_router( expected_record_instances=0 if use_kv_events else 1, ) - _, consumer_endpoint, consumer_router = await make_router( + consumer_runtime, consumer_endpoint, consumer_router = await make_router( serve_indexer=False, use_remote_indexer=True, router_predicted_ttl_secs=router_predicted_ttl_secs, ) + runtimes.append(consumer_runtime) worker_ids = sorted( await poll_for_worker_instances( @@ -903,6 +907,14 @@ async def make_router( consumer_endpoint, expected_num_instances, max_wait_time=120 ) + async def test_sync(): + runtimes = [] + try: + await run_test(runtimes) + finally: + for runtime in runtimes: + runtime.shutdown() + asyncio.run(test_sync()) From 56f828f73f547883d3bfc56ca6b124eeed2b22ad Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 8 Jul 2026 16:14:20 -0700 Subject: [PATCH 4/4] test(router): shut down owned runtimes Signed-off-by: PeaBrane --- tests/router/common.py | 31 ++++++++++----- tests/router/e2e_harness.py | 25 +++++++----- tests/router/helper.py | 14 +++++++ tests/router/test_mocker_output_replay_e2e.py | 18 +++++---- tests/router/test_router_e2e_with_mockers.py | 17 +++++---- tests/router/test_slot_tracker_e2e.py | 38 ++++++++++++++++--- 6 files changed, 102 insertions(+), 41 deletions(-) diff --git a/tests/router/common.py b/tests/router/common.py index b7c6ece41b22..97cf00c3294d 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -22,6 +22,7 @@ _nats_server, assert_event_dumps_equal, get_runtime, + managed_runtime, poll_for_worker_instances, send_inflight_requests, send_request_via_python_kv_router, @@ -1869,7 +1870,7 @@ def _test_router_indexers_sync( raise ValueError("nats_server is required when test_nats_interruption=True") # Use async to manage the test flow - async def test_sync(): + async def run_test(runtime_stack): # Create KvRouterConfig with lower snapshot threshold for testing kv_router_config = KvRouterConfig( router_snapshot_threshold=20, @@ -1881,8 +1882,8 @@ async def test_sync(): # If standalone indexer mode, launch workers one-by-one and register. # We need to create a temporary endpoint just to discover worker IDs. if standalone_indexer_url: - tmp_runtime = get_runtime( - store_backend, request_plane, event_plane=event_plane + tmp_runtime = runtime_stack.enter_context( + managed_runtime(store_backend, request_plane, event_plane=event_plane) ) tmp_endpoint = tmp_runtime.endpoint( f"{engine_workers.namespace}.{engine_workers.component_name}.generate" @@ -1979,7 +1980,9 @@ async def send_requests_to_router(router, num_requests, router_name, endpoint): # Create first runtime and endpoint for router 1 logger.info("Creating first KV router with its own runtime") - runtime1 = get_runtime(store_backend, request_plane, event_plane=event_plane) + runtime1 = runtime_stack.enter_context( + managed_runtime(store_backend, request_plane, event_plane=event_plane) + ) endpoint1 = runtime1.endpoint( f"{engine_workers.namespace}.{engine_workers.component_name}.generate" ) @@ -2091,7 +2094,9 @@ async def send_requests_to_router(router, num_requests, router_name, endpoint): # Create second runtime and endpoint for router 2 logger.info("Creating second KV router with its own runtime") - runtime2 = get_runtime(store_backend, request_plane, event_plane=event_plane) + runtime2 = runtime_stack.enter_context( + managed_runtime(store_backend, request_plane, event_plane=event_plane) + ) endpoint2 = runtime2.endpoint( f"{engine_workers.namespace}.{engine_workers.component_name}.generate" ) @@ -2284,7 +2289,10 @@ async def send_requests_to_router(router, num_requests, router_name, endpoint): "Skipping NATS consumers verification (local indexer uses NATS Core, not JetStream)" ) - # Run the async test + async def test_sync(): + with contextlib.ExitStack() as runtime_stack: + await run_test(runtime_stack) + asyncio.run(test_sync()) logger.info("Indexers sync test completed successfully") @@ -2673,7 +2681,7 @@ def _test_router_decisions_disagg_round_robin_prefill_dp_rank( frontend_port, ) - async def test_sync(): + async def run_test(runtime): frontend_url = f"http://localhost:{frontend_port}" chat_url = f"{frontend_url}/v1/chat/completions" await wait_for_frontend_ready( @@ -2687,9 +2695,6 @@ async def test_sync(): request_plane=request_plane, ) - runtime = get_runtime( - store_backend=store_backend, request_plane=request_plane - ) prefill_endpoint = runtime.endpoint( f"{prefill_workers.namespace}.prefill.generate" ) @@ -2767,6 +2772,12 @@ def stored_blocks_by_dp_rank(events_json: str) -> dict[int, int]: final_counts = stored_blocks_by_dp_rank(await observer_router.dump_events()) return prefill_worker_id, baseline_counts, final_counts + async def test_sync(): + with managed_runtime( + store_backend=store_backend, request_plane=request_plane + ) as runtime: + return await run_test(runtime) + prefill_worker_id, baseline_counts, final_counts = asyncio.run(test_sync()) delta_counts = { diff --git a/tests/router/e2e_harness.py b/tests/router/e2e_harness.py index d00cc48bdfa9..d489feae066e 100644 --- a/tests/router/e2e_harness.py +++ b/tests/router/e2e_harness.py @@ -13,7 +13,7 @@ _test_router_decisions_disagg, _test_router_indexers_sync, ) -from tests.router.helper import generate_random_suffix, get_runtime +from tests.router.helper import generate_random_suffix, managed_runtime from tests.utils.constants import DefaultPort from tests.utils.port_utils import allocate_ports, deallocate_ports from tests.utils.test_output import resolve_test_output_path @@ -157,11 +157,6 @@ def __exit__(self, exc_type, exc_val, exc_tb): time.sleep(self.cleanup_delay_seconds) -def get_engine_endpoint(engine_workers, request_plane: str, component_name: str): - runtime = get_runtime(request_plane=request_plane) - return runtime.endpoint(f"{engine_workers.namespace}.{component_name}.generate") - - def _create_engine_process( *, engine_process_cls, @@ -264,8 +259,13 @@ def run_router_decisions_test( default_process_kwargs=default_process_kwargs, engine_process_kwargs=engine_process_kwargs, ) - with process as engine_workers: - endpoint = get_engine_endpoint(engine_workers, request_plane, component_name) + with ( + process as engine_workers, + managed_runtime(request_plane=request_plane) as runtime, + ): + endpoint = runtime.endpoint( + f"{engine_workers.namespace}.{component_name}.generate" + ) scenario_kwargs = dict(test_kwargs or {}) for argument, attribute in ( ("standalone_indexer_url", "standalone_indexer_url"), @@ -306,8 +306,13 @@ def run_cache_salt_isolation_test( default_process_kwargs={"num_workers": 2, "single_gpu": True}, engine_process_kwargs=None, ) - with process as engine_workers: - endpoint = get_engine_endpoint(engine_workers, request_plane, component_name) + with ( + process as engine_workers, + managed_runtime(request_plane=request_plane) as runtime, + ): + endpoint = runtime.endpoint( + f"{engine_workers.namespace}.{component_name}.generate" + ) _test_router_cache_salt_isolation( engine_workers, endpoint, diff --git a/tests/router/helper.py b/tests/router/helper.py index 3f67848e46c1..da920d131379 100644 --- a/tests/router/helper.py +++ b/tests/router/helper.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +import contextlib import json import logging import os @@ -585,6 +586,19 @@ def get_runtime( ) +@contextlib.contextmanager +def managed_runtime( + store_backend: str = "etcd", + request_plane: str = "tcp", + event_plane: Optional[str] = None, +): + runtime = get_runtime(store_backend, request_plane, event_plane) + try: + yield runtime + finally: + runtime.shutdown() + + async def check_nats_consumers(namespace: str, expected_count: Optional[int] = None): """Check NATS consumers for the KV events stream. diff --git a/tests/router/test_mocker_output_replay_e2e.py b/tests/router/test_mocker_output_replay_e2e.py index 10f905b80c0d..92cda7c19300 100644 --- a/tests/router/test_mocker_output_replay_e2e.py +++ b/tests/router/test_mocker_output_replay_e2e.py @@ -12,7 +12,7 @@ from dynamo.llm import KvRouter, KvRouterConfig from tests.router.common import _create_kv_router_with_timeout -from tests.router.helper import get_runtime, wait_for_workers_ready +from tests.router.helper import managed_runtime, wait_for_workers_ready from tests.router.mocker_process import MockerProcess from tests.utils.constants import ROUTER_MODEL_NAME @@ -210,14 +210,16 @@ def test_mocker_output_replay_generate_from_request_multi_turn( "response_replay_trace_path": replay_trace_path, } - with MockerProcess( - request, - mocker_args=mocker_args, - num_mockers=NUM_MOCKERS, - request_plane=request_plane, - ) as mockers: + with ( + MockerProcess( + request, + mocker_args=mocker_args, + num_mockers=NUM_MOCKERS, + request_plane=request_plane, + ) as mockers, + managed_runtime(request_plane=request_plane) as runtime, + ): logger.info("Started mocker replay test endpoint: %s", mockers.endpoint) - runtime = get_runtime(request_plane=request_plane) endpoint = runtime.endpoint( f"{mockers.namespace}.{mockers.component_name}.generate" ) diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index 658cdf6a787b..fdc9dcbb42dd 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -44,6 +44,7 @@ from tests.router.helper import ( generate_random_suffix, get_runtime, + managed_runtime, poll_for_worker_instances, topology_env, ) @@ -600,18 +601,20 @@ def test_kv_router_bindings( "durable_kv_events": durable_kv_events, } - with MockerProcess( - request, - mocker_args=mocker_args, - num_mockers=NUM_MOCKERS, - request_plane=request_plane, - ) as mockers: + with ( + MockerProcess( + request, + mocker_args=mocker_args, + num_mockers=NUM_MOCKERS, + request_plane=request_plane, + ) as mockers, + managed_runtime(request_plane=request_plane) as runtime, + ): # Start mocker instances logger.info(f"Starting {NUM_MOCKERS} mocker instances") logger.info(f"All mockers using endpoint: {mockers.endpoint}") # Get runtime and create endpoint - runtime = get_runtime(request_plane=request_plane) endpoint = runtime.endpoint( f"{mockers.namespace}.{mockers.component_name}.generate" ) diff --git a/tests/router/test_slot_tracker_e2e.py b/tests/router/test_slot_tracker_e2e.py index 7932229366ea..3aebc7f0d8e3 100644 --- a/tests/router/test_slot_tracker_e2e.py +++ b/tests/router/test_slot_tracker_e2e.py @@ -19,6 +19,7 @@ """ import asyncio +import contextlib import gc import itertools from collections.abc import AsyncIterator @@ -28,7 +29,7 @@ from dynamo.llm import KvRouter, KvRouterConfig from tests.router.common import _create_kv_router_with_timeout -from tests.router.helper import generate_random_suffix, get_runtime +from tests.router.helper import generate_random_suffix, managed_runtime from tests.router.mocker_process import MockerProcess, launch_disagg_workers from tests.utils.constants import ROUTER_MODEL_NAME @@ -75,8 +76,11 @@ def _create_router( router_config: KvRouterConfig, discovery_backend: str, request_plane: str, + runtime_stack: contextlib.ExitStack, ): - runtime = get_runtime(discovery_backend, request_plane) + runtime = runtime_stack.enter_context( + managed_runtime(discovery_backend, request_plane) + ) endpoint = runtime.endpoint( f"{engine_workers.namespace}.{engine_workers.component_name}.generate" ) @@ -341,12 +345,13 @@ def test_unexpected_drop_and_normal_completion( discovery_backend, request_plane, ) -> None: - async def run() -> None: + async def run_test(runtime_stack: contextlib.ExitStack) -> None: runtime, endpoint, router = _create_router( aggregated_mocker, _router_config(), discovery_backend, request_plane, + runtime_stack, ) _ = runtime, endpoint baseline = await _snapshot(router) @@ -397,6 +402,10 @@ async def run() -> None: del dropped_stream, completed_stream gc.collect() + async def run() -> None: + with contextlib.ExitStack() as runtime_stack: + await run_test(runtime_stack) + asyncio.run(run()) @@ -405,12 +414,13 @@ def test_router_observed_first_token_marks_prefill_complete( discovery_backend, request_plane, ) -> None: - async def run() -> None: + async def run_test(runtime_stack: contextlib.ExitStack) -> None: runtime, endpoint, router = _create_router( aggregated_mocker, _router_config(), discovery_backend, request_plane, + runtime_stack, ) _ = runtime, endpoint baseline = await _snapshot(router) @@ -444,6 +454,10 @@ async def run() -> None: description="decode cleanup after first-token drop", ) + async def run() -> None: + with contextlib.ExitStack() as runtime_stack: + await run_test(runtime_stack) + asyncio.run(run()) @@ -452,19 +466,21 @@ def test_bidirectional_replica_lifecycle( discovery_backend, request_plane, ) -> None: - async def run() -> None: + async def run_test(runtime_stack: contextlib.ExitStack) -> None: config = _router_config(router_replica_sync=True) runtime_a, endpoint_a, router_a = _create_router( aggregated_mocker, config, discovery_backend, request_plane, + runtime_stack, ) runtime_b, endpoint_b, router_b = _create_router( aggregated_mocker, config, discovery_backend, request_plane, + runtime_stack, ) _ = runtime_a, endpoint_a, runtime_b, endpoint_b baseline_a = await _snapshot(router_a) @@ -475,6 +491,10 @@ async def run() -> None: await _assert_replicated_lifecycle(router_a, router_b, baseline_a, baseline_b) await _assert_replicated_lifecycle(router_b, router_a, baseline_b, baseline_a) + async def run() -> None: + with contextlib.ExitStack() as runtime_stack: + await run_test(runtime_stack) + asyncio.run(run()) @@ -483,19 +503,21 @@ def test_disaggregated_role_attribution( discovery_backend, request_plane, ) -> None: - async def run() -> None: + async def run_test(runtime_stack: contextlib.ExitStack) -> None: prefill_workers, decode_workers = disagg_mockers prefill_runtime, prefill_endpoint, prefill_router = _create_router( prefill_workers, _router_config(router_track_active_blocks=False), discovery_backend, request_plane, + runtime_stack, ) decode_runtime, decode_endpoint, decode_router = _create_router( decode_workers, _router_config(router_track_prefill_tokens=False), discovery_backend, request_plane, + runtime_stack, ) _ = ( prefill_runtime, @@ -549,4 +571,8 @@ async def run() -> None: del prefill_stream, decode_stream gc.collect() + async def run() -> None: + with contextlib.ExitStack() as runtime_stack: + await run_test(runtime_stack) + asyncio.run(run())