diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index a40873b6fbd7..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,7 +255,9 @@ 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( @@ -262,6 +265,7 @@ where &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, 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?; diff --git a/tests/router/common.py b/tests/router/common.py index 302384ade0b2..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, @@ -708,7 +709,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 +745,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 +763,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 +772,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 +783,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 +908,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()) @@ -1857,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, @@ -1869,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" @@ -1967,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" ) @@ -2079,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" ) @@ -2272,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") @@ -2661,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( @@ -2675,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" ) @@ -2755,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())