From b4518e74710d853d69c2b8140e99fcbf66889800 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 8 Apr 2026 07:20:11 +0000 Subject: [PATCH 1/7] try --- lib/llm/src/kv_router/prefill_router/execution.rs | 13 ++++++------- lib/llm/src/kv_router/prefill_router/mod.rs | 2 +- lib/llm/src/kv_router/prefill_router/types.rs | 2 +- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/lib/llm/src/kv_router/prefill_router/execution.rs b/lib/llm/src/kv_router/prefill_router/execution.rs index 4f189230795e..823c79517a13 100644 --- a/lib/llm/src/kv_router/prefill_router/execution.rs +++ b/lib/llm/src/kv_router/prefill_router/execution.rs @@ -43,11 +43,10 @@ impl PrefillRouter { let dp_rank = req .routing .as_ref() - .and_then(|r| r.prefill_dp_rank.or(r.dp_rank)) - .unwrap_or(0); + .and_then(|r| r.prefill_dp_rank.or(r.dp_rank)); tracing::debug!( worker_id = id, - dp_rank = dp_rank, + dp_rank = ?dp_rank, "Using pre-selected prefill worker for bootstrap" ); (id, dp_rank) @@ -99,7 +98,7 @@ impl PrefillRouter { tracing::debug!( worker_id = worker_id, - dp_rank = dp_rank, + dp_rank = ?dp_rank, bootstrap_host = %host, bootstrap_port = port, bootstrap_room = bootstrap_room, @@ -266,7 +265,7 @@ impl PrefillRouter { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, - ) -> Result<(u64, u32)> { + )-> Result<(u64, Option)> { let prefill_router = self .prefill_router .get() @@ -288,7 +287,7 @@ impl PrefillRouter { allowed_worker_ids, ) .await?; - Ok((worker.worker_id, worker.dp_rank)) + Ok((worker.worker_id, Some(worker.dp_rank))) } InnerPrefillRouter::SimpleRouter(r) => { let worker_id = if update_states { @@ -297,7 +296,7 @@ impl PrefillRouter { r.peek_next_worker() } .ok_or_else(|| anyhow::anyhow!("No workers available for prefill"))?; - Ok((worker_id, 0)) + Ok((worker_id, None)) } } } diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index 4df7b37b41c4..2a30470125ec 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -141,7 +141,7 @@ impl let routing = prefill_req.routing_mut(); routing.prefill_worker_id = Some(worker_id); - routing.dp_rank = Some(dp_rank); + routing.dp_rank = dp_rank; prefill_req.bootstrap_info = Some(bootstrap_info.clone()); let prefill_context = diff --git a/lib/llm/src/kv_router/prefill_router/types.rs b/lib/llm/src/kv_router/prefill_router/types.rs index 1acf1de53d86..2fed2c2776cb 100644 --- a/lib/llm/src/kv_router/prefill_router/types.rs +++ b/lib/llm/src/kv_router/prefill_router/types.rs @@ -36,7 +36,7 @@ pub(super) enum PrefillOutcome { pub(super) enum PrefillResolveDecision { Resolved { worker_id: u64, - dp_rank: u32, + dp_rank: Option, bootstrap_info: BootstrapInfo, }, Unavailable, From c213ec6b52103b08087ae6c279b9eac6cd769e15 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 8 Apr 2026 07:25:41 +0000 Subject: [PATCH 2/7] add tmp print --- .../src/dynamo/sglang/request_handlers/llm/prefill_handler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py b/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py index a0b16c1d4a85..d4acea27e377 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py @@ -141,6 +141,8 @@ async def generate( if dp_rank is not None and dp_rank == _DP_RANK_UNSET: dp_rank = None + print("dprank", dp_rank) + trace_header = build_trace_headers(context) if self.enable_trace else None results = await self.engine.async_generate( From b103baecbf4de1aeb8a4597a931b563d92394c51 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 8 Apr 2026 08:24:51 +0000 Subject: [PATCH 3/7] fix: align C binding return type with PrefillRouter and remove debug print PrefillRouter::query_prefill_worker returns Option for dp_rank. The C FFI wrapper was declaring u32, causing E0308 in clippy. Map None to u32::MAX (NO_DP_RANK sentinel) so the Python side sees _DP_RANK_UNSET. --- .../src/dynamo/sglang/request_handlers/llm/prefill_handler.py | 2 -- lib/bindings/c/src/lib.rs | 4 +++- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py b/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py index d4acea27e377..a0b16c1d4a85 100644 --- a/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py +++ b/components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py @@ -141,8 +141,6 @@ async def generate( if dp_rank is not None and dp_rank == _DP_RANK_UNSET: dp_rank = None - print("dprank", dp_rank) - trace_header = build_trace_headers(context) if self.enable_trace else None results = await self.engine.async_generate( diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index fe1ced1db16d..c423e57e3ef3 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -455,7 +455,7 @@ impl RouterHandles { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, - ) -> Result<(u64, u32), QueryRouterResult> { + ) -> Result<(u64, Option), QueryRouterResult> { if let Some(ref ids) = allowed_worker_ids { self.prefill_router.register_workers(ids); } @@ -1214,6 +1214,8 @@ pub unsafe extern "C" fn route_prefill_request( .query_prefill_worker(&tokens, None, false, None, 0.0, allowed_worker_ids) .await?; + let prefill_dp_rank = prefill_dp_rank.unwrap_or(u32::MAX); + tracing::info!( prefill_worker_id = prefill_worker_id, prefill_dp_rank = prefill_dp_rank, From 514d6b11e0c0af43295dbf507eb8c727638fdeb6 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Wed, 8 Apr 2026 08:29:50 +0000 Subject: [PATCH 4/7] lint --- lib/llm/src/kv_router/prefill_router/execution.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/llm/src/kv_router/prefill_router/execution.rs b/lib/llm/src/kv_router/prefill_router/execution.rs index 823c79517a13..72d20441532d 100644 --- a/lib/llm/src/kv_router/prefill_router/execution.rs +++ b/lib/llm/src/kv_router/prefill_router/execution.rs @@ -265,7 +265,7 @@ impl PrefillRouter { lora_name: Option, priority_jump: f64, allowed_worker_ids: Option>, - )-> Result<(u64, Option)> { + ) -> Result<(u64, Option)> { let prefill_router = self .prefill_router .get() From 0d6d9721e26d34947547dcd312ec13777b955b15 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 8 Apr 2026 08:15:27 -0700 Subject: [PATCH 5/7] teach mockers to round-robin unset dp rank Signed-off-by: PeaBrane --- lib/llm/src/mocker.rs | 18 ++- tests/router/common.py | 139 +++++++++++++++++++ tests/router/router_process.py | 5 +- tests/router/test_router_e2e_with_mockers.py | 99 ++++++++++++- 4 files changed, 252 insertions(+), 9 deletions(-) diff --git a/lib/llm/src/mocker.rs b/lib/llm/src/mocker.rs index 3581b59879d7..2c34a415b0b2 100644 --- a/lib/llm/src/mocker.rs +++ b/lib/llm/src/mocker.rs @@ -8,6 +8,7 @@ use std::collections::VecDeque; use std::sync::Arc; +use std::sync::atomic::{AtomicU32, Ordering}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use crate::backend::ExecutionContext; @@ -297,6 +298,7 @@ pub struct MockEngine { request_senders: OnceCell>>, senders_ready: Notify, engine_args: MockEngineArgs, + unset_dp_rank_counter: AtomicU32, /// Bootstrap server for prefill workers in disaggregated mode bootstrap_server: Arc>>, /// Keep schedulers alive so their CancelGuards don't fire prematurely. @@ -311,11 +313,20 @@ impl MockEngine { request_senders: OnceCell::new(), senders_ready: Notify::new(), engine_args, + unset_dp_rank_counter: AtomicU32::new(0), bootstrap_server: Arc::new(OnceCell::new()), _schedulers: OnceCell::new(), } } + fn resolve_dp_rank(&self, request: &PreprocessedRequest) -> u32 { + if let Some(dp_rank) = request.routing.as_ref().and_then(|routing| routing.dp_rank) { + return dp_rank; + } + + self.unset_dp_rank_counter.fetch_add(1, Ordering::Relaxed) % self.engine_args.dp_size + } + pub async fn start(&self, component: Component) -> Result<()> { // Use primary_token() instead of child_token() so the mocker continues running // during graceful shutdown (Phase 1/2) and only stops in Phase 3. @@ -583,12 +594,7 @@ impl AsyncEngine, ManyOut, Error> ) -> Result, Error> { let (request, ctx) = input.into_parts(); - // Extract dp_rank from routing hints (defaults to 0 if not set) - let dp_rank = request - .routing - .as_ref() - .and_then(|r| r.dp_rank) - .unwrap_or(0); + let dp_rank = self.resolve_dp_rank(&request); // Validate dp_rank if dp_rank >= self.engine_args.dp_size { diff --git a/tests/router/common.py b/tests/router/common.py index cc0d83c6020c..5836b47d730e 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -1779,6 +1779,145 @@ async def send_progressive_requests(): ) +def _test_router_decisions_disagg_round_robin_prefill_dp_rank( + prefill_workers, + decode_workers, + block_size: int, + request, + frontend_port: int, + test_payload: dict, + expected_prefill_dp_ranks: int, + store_backend: str = "etcd", + request_plane: str = "nats", +): + """Verify disaggregated round-robin requests store prefill KV blocks across DP ranks.""" + + with FrontendRouterProcess( + request, + block_size, + frontend_port, + decode_workers.namespace, + store_backend, + enforce_disagg=True, + request_plane=request_plane, + router_mode="round-robin", + min_initial_workers=decode_workers.num_workers, + ): + logger.info( + "Starting round-robin frontend on port %s for disagg prefill dp-rank test", + frontend_port, + ) + + async def test_sync(): + frontend_url = f"http://localhost:{frontend_port}" + chat_url = f"{frontend_url}/v1/chat/completions" + await wait_for_frontend_ready( + frontend_url=frontend_url, + expected_num_workers=decode_workers.num_workers, + timeout=120, + ) + + runtime = get_runtime( + store_backend=store_backend, request_plane=request_plane + ) + prefill_endpoint = runtime.endpoint( + f"{prefill_workers.namespace}.prefill.generate" + ) + + with min_initial_workers_env(prefill_workers.num_workers): + observer_router = KvRouter( + endpoint=prefill_endpoint, + block_size=block_size, + kv_router_config=KvRouterConfig( + router_snapshot_threshold=20, + use_kv_events=True, + durable_kv_events=False, + router_event_threads=4, + router_track_prefill_tokens=True, + router_prefill_load_model="none", + ), + ) + + client = await prefill_endpoint.client() + worker_ids: list[int] = [] + deadline = asyncio.get_running_loop().time() + 60 + while asyncio.get_running_loop().time() < deadline: + worker_ids = sorted(set(client.instance_ids())) + if len(worker_ids) >= prefill_workers.num_workers: + break + await asyncio.sleep(1.0) + + assert len(worker_ids) == prefill_workers.num_workers, ( + f"Timed out waiting for prefill workers. " + f"Found {worker_ids}, expected {prefill_workers.num_workers}" + ) + prefill_worker_id = worker_ids[0] + + def stored_blocks_by_dp_rank(events_json: str) -> dict[int, int]: + counts = {dp_rank: 0 for dp_rank in range(expected_prefill_dp_ranks)} + for event in json.loads(events_json): + if event.get("worker_id") != prefill_worker_id: + continue + stored = event.get("event", {}).get("data", {}).get("stored") + if stored is None: + continue + dp_rank = event.get("event", {}).get("dp_rank", 0) + counts[dp_rank] = counts.get(dp_rank, 0) + len( + stored.get("blocks", []) + ) + return counts + + await asyncio.sleep(2.0) + baseline_counts = stored_blocks_by_dp_rank( + await observer_router.dump_events() + ) + + async with aiohttp.ClientSession() as session: + for request_idx in range(expected_prefill_dp_ranks * 2): + prompt_tokens = " ".join( + f"prefill-{request_idx}-token-{token_idx}" + for token_idx in range(block_size * 3) + ) + payload = { + **test_payload, + "stream": False, + "max_tokens": 1, + "messages": [ + { + "role": "user", + "content": prompt_tokens, + } + ], + } + async with session.post(chat_url, json=payload) as response: + assert response.status == 200, ( + f"Request {request_idx + 1} failed with status " + f"{response.status}: {await response.text()}" + ) + await response.text() + await asyncio.sleep(0.5) + + await asyncio.sleep(2.0) + final_counts = stored_blocks_by_dp_rank(await observer_router.dump_events()) + return prefill_worker_id, baseline_counts, final_counts + + prefill_worker_id, baseline_counts, final_counts = asyncio.run(test_sync()) + + delta_counts = { + dp_rank: final_counts.get(dp_rank, 0) - baseline_counts.get(dp_rank, 0) + for dp_rank in range(expected_prefill_dp_ranks) + } + active_dp_ranks = sorted( + dp_rank for dp_rank, block_count in delta_counts.items() if block_count > 0 + ) + + assert active_dp_ranks == list(range(expected_prefill_dp_ranks)), ( + f"Expected round-robin prefill requests for worker {prefill_worker_id} " + f"to store KV blocks on dp_ranks {list(range(expected_prefill_dp_ranks))}, " + f"but saw deltas {delta_counts}" + ) + + def _test_router_decisions( engine_workers, endpoint, diff --git a/tests/router/router_process.py b/tests/router/router_process.py index c52826fbd2ba..2523d4016dce 100644 --- a/tests/router/router_process.py +++ b/tests/router/router_process.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import os +import sys from tests.utils.managed_process import ManagedProcess @@ -34,7 +35,7 @@ def __init__( use_remote_indexer: bool = False, ): command = [ - "python3", + sys.executable, "-m", "dynamo.frontend", "--router-mode", @@ -141,7 +142,7 @@ def __init__( request_plane: str = "nats", ): command = [ - "python3", + sys.executable, "-m", "dynamo.frontend", "--router-mode", diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index 5f5e421a426d..d46a96aa51db 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -11,6 +11,7 @@ import asyncio import logging import os +import sys from pathlib import Path from typing import Any, Dict, Optional @@ -25,6 +26,7 @@ _test_router_basic, _test_router_decisions, _test_router_decisions_disagg, + _test_router_decisions_disagg_round_robin_prefill_dp_rank, _test_router_indexers_sync, _test_router_overload_503, _test_router_query_instance_id, @@ -139,7 +141,7 @@ def _build_mocker_command( List of command arguments for subprocess """ command = [ - "python", + sys.executable, "-m", "dynamo.mocker", "--model-path", @@ -1278,6 +1280,101 @@ def test_router_decisions_disagg( ) +@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"]) +@pytest.mark.parametrize( + "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"] +) +@pytest.mark.timeout(180) +def test_router_decisions_disagg_round_robin_prefill_dp_rank( + request, + runtime_services_dynamic_ports, + predownload_tokenizers, + registration_order, + enable_disagg_bootstrap, +): + """Verify round-robin disagg prefill requests spread KV stores across DP ranks.""" + logger.info( + "Starting disaggregated round-robin prefill dp-rank test " + "(registration_order=%s, bootstrap=%s)", + registration_order, + enable_disagg_bootstrap, + ) + + namespace_suffix = generate_random_suffix() + shared_namespace = f"test-namespace-{namespace_suffix}" + prefill_mocker_args = { + "speedup_ratio": SPEEDUP_RATIO, + "block_size": BLOCK_SIZE, + "dp_size": 4, + } + decode_mocker_args = { + "speedup_ratio": SPEEDUP_RATIO, + "block_size": BLOCK_SIZE, + } + + def run_case(prefill_workers, decode_workers): + frontend_port = get_unique_ports( + request, num_ports=1, registration_order=registration_order + )[0] + _test_router_decisions_disagg_round_robin_prefill_dp_rank( + prefill_workers=prefill_workers, + decode_workers=decode_workers, + block_size=BLOCK_SIZE, + request=request, + frontend_port=frontend_port, + test_payload=TEST_PAYLOAD, + expected_prefill_dp_ranks=prefill_mocker_args["dp_size"], + request_plane="nats", + ) + + if registration_order == "prefill_first": + with DisaggMockerProcess( + request, + namespace=shared_namespace, + worker_type="prefill", + mocker_args=prefill_mocker_args, + num_mockers=1, + request_plane="nats", + enable_bootstrap=enable_disagg_bootstrap, + ) as prefill_workers: + logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}") + + with DisaggMockerProcess( + request, + namespace=shared_namespace, + worker_type="decode", + mocker_args=decode_mocker_args, + num_mockers=1, + request_plane="nats", + ) as decode_workers: + logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") + run_case(prefill_workers, decode_workers) + else: + with DisaggMockerProcess( + request, + namespace=shared_namespace, + worker_type="decode", + mocker_args=decode_mocker_args, + num_mockers=1, + request_plane="nats", + ) as decode_workers: + logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") + + with DisaggMockerProcess( + request, + namespace=shared_namespace, + worker_type="prefill", + mocker_args=prefill_mocker_args, + num_mockers=1, + request_plane="nats", + enable_bootstrap=enable_disagg_bootstrap, + ) as prefill_workers: + logger.info( + f"Prefill workers using endpoint: {prefill_workers.endpoint}" + ) + run_case(prefill_workers, decode_workers) + + @pytest.mark.timeout(180) def test_router_decisions_disagg_router_aic( request, From ea945e73a899f571d046e2af535a8f437cac49c5 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 8 Apr 2026 09:34:22 -0700 Subject: [PATCH 6/7] test: stop copy-pasting disagg launches Signed-off-by: PeaBrane --- tests/router/test_router_e2e_with_mockers.py | 235 +++++++++---------- 1 file changed, 105 insertions(+), 130 deletions(-) diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index d46a96aa51db..5980570ae66b 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -12,8 +12,9 @@ import logging import os import sys +from contextlib import contextmanager from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, Iterator, Optional import aiohttp import pytest @@ -670,6 +671,77 @@ def __exit__(self, exc_type, exc_val, exc_tb): self._bootstrap_ports = [] +@contextmanager +def _launch_disagg_workers( + request, + namespace: str, + registration_order: str, + *, + prefill_mocker_args: Dict[str, Any], + decode_mocker_args: Dict[str, Any], + num_prefill_mockers: int, + num_decode_mockers: int, + enable_disagg_bootstrap: bool, + request_plane: str = "nats", +) -> Iterator[tuple[DisaggMockerProcess, DisaggMockerProcess]]: + if registration_order not in ("prefill_first", "decode_first"): + raise ValueError(f"Unexpected registration order: {registration_order}") + + if registration_order == "prefill_first": + logger.info("Starting %s prefill mocker instances (first)", num_prefill_mockers) + with DisaggMockerProcess( + request, + namespace=namespace, + worker_type="prefill", + mocker_args=prefill_mocker_args, + num_mockers=num_prefill_mockers, + request_plane=request_plane, + enable_bootstrap=enable_disagg_bootstrap, + ) as prefill_workers: + logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}") + + logger.info( + "Starting %s decode mocker instances (second)", num_decode_mockers + ) + with DisaggMockerProcess( + request, + namespace=namespace, + worker_type="decode", + mocker_args=decode_mocker_args, + num_mockers=num_decode_mockers, + request_plane=request_plane, + ) as decode_workers: + logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") + yield prefill_workers, decode_workers + return + + logger.info("Starting %s decode mocker instances (first)", num_decode_mockers) + with DisaggMockerProcess( + request, + namespace=namespace, + worker_type="decode", + mocker_args=decode_mocker_args, + num_mockers=num_decode_mockers, + request_plane=request_plane, + ) as decode_workers: + logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") + + logger.info( + "Starting %s prefill mocker instances (second)", num_prefill_mockers + ) + with DisaggMockerProcess( + request, + namespace=namespace, + worker_type="prefill", + mocker_args=prefill_mocker_args, + num_mockers=num_prefill_mockers, + request_plane=request_plane, + enable_bootstrap=enable_disagg_bootstrap, + ) as prefill_workers: + logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}") + yield prefill_workers, decode_workers + + @pytest.mark.timeout(180) # planner-profile mocker setup can exceed 120s on CI CPUs @pytest.mark.parametrize( "router_mode,durable_kv_events,mocker_args_override", @@ -1194,90 +1266,28 @@ def test_router_decisions_disagg( # durable_kv_events defaults to False (NATS Core mode) } - if registration_order == "prefill_first": - # Start prefill workers first - logger.info("Starting 4 prefill mocker instances (first)") - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="prefill", - mocker_args=mocker_args, - num_mockers=4, - request_plane="nats", - enable_bootstrap=enable_disagg_bootstrap, - ) as prefill_workers: - logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}") - - # Then start decode workers - logger.info("Starting 4 decode mocker instances (second)") - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="decode", - mocker_args=mocker_args, - num_mockers=4, - request_plane="nats", - ) as decode_workers: - logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") - - # Get unique port for this test - frontend_port = get_unique_ports( - request, num_ports=1, registration_order=registration_order - )[0] - - # Run disagg routing test - _test_router_decisions_disagg( - prefill_workers=prefill_workers, - decode_workers=decode_workers, - block_size=BLOCK_SIZE, - request=request, - frontend_port=frontend_port, - test_payload=TEST_PAYLOAD, - request_plane="nats", - ) - else: - # Start decode workers first - logger.info("Starting 4 decode mocker instances (first)") - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="decode", - mocker_args=mocker_args, - num_mockers=4, + with _launch_disagg_workers( + request, + shared_namespace, + registration_order, + prefill_mocker_args=mocker_args, + decode_mocker_args=mocker_args, + num_prefill_mockers=4, + num_decode_mockers=4, + enable_disagg_bootstrap=enable_disagg_bootstrap, + ) as (prefill_workers, decode_workers): + frontend_port = get_unique_ports( + request, num_ports=1, registration_order=registration_order + )[0] + _test_router_decisions_disagg( + prefill_workers=prefill_workers, + decode_workers=decode_workers, + block_size=BLOCK_SIZE, + request=request, + frontend_port=frontend_port, + test_payload=TEST_PAYLOAD, request_plane="nats", - ) as decode_workers: - logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") - - # Then start prefill workers - logger.info("Starting 4 prefill mocker instances (second)") - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="prefill", - mocker_args=mocker_args, - num_mockers=4, - request_plane="nats", - enable_bootstrap=enable_disagg_bootstrap, - ) as prefill_workers: - logger.info( - f"Prefill workers using endpoint: {prefill_workers.endpoint}" - ) - - # Get unique port for this test - frontend_port = get_unique_ports( - request, num_ports=1, registration_order=registration_order - )[0] - - # Run disagg routing test - _test_router_decisions_disagg( - prefill_workers=prefill_workers, - decode_workers=decode_workers, - block_size=BLOCK_SIZE, - request=request, - frontend_port=frontend_port, - test_payload=TEST_PAYLOAD, - request_plane="nats", - ) + ) @pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"]) @@ -1327,52 +1337,17 @@ def run_case(prefill_workers, decode_workers): request_plane="nats", ) - if registration_order == "prefill_first": - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="prefill", - mocker_args=prefill_mocker_args, - num_mockers=1, - request_plane="nats", - enable_bootstrap=enable_disagg_bootstrap, - ) as prefill_workers: - logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}") - - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="decode", - mocker_args=decode_mocker_args, - num_mockers=1, - request_plane="nats", - ) as decode_workers: - logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") - run_case(prefill_workers, decode_workers) - else: - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="decode", - mocker_args=decode_mocker_args, - num_mockers=1, - request_plane="nats", - ) as decode_workers: - logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}") - - with DisaggMockerProcess( - request, - namespace=shared_namespace, - worker_type="prefill", - mocker_args=prefill_mocker_args, - num_mockers=1, - request_plane="nats", - enable_bootstrap=enable_disagg_bootstrap, - ) as prefill_workers: - logger.info( - f"Prefill workers using endpoint: {prefill_workers.endpoint}" - ) - run_case(prefill_workers, decode_workers) + with _launch_disagg_workers( + request, + shared_namespace, + registration_order, + prefill_mocker_args=prefill_mocker_args, + decode_mocker_args=decode_mocker_args, + num_prefill_mockers=1, + num_decode_mockers=1, + enable_disagg_bootstrap=enable_disagg_bootstrap, + ) as (prefill_workers, decode_workers): + run_case(prefill_workers, decode_workers) @pytest.mark.timeout(180) From 911a3e5e3891914860fd20a782732b4022c81a75 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 8 Apr 2026 10:30:15 -0700 Subject: [PATCH 7/7] test: name spawned router children Signed-off-by: PeaBrane --- tests/router/router_process.py | 2 ++ tests/router/test_router_e2e_with_mockers.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/router/router_process.py b/tests/router/router_process.py index 2523d4016dce..bfd6255829f7 100644 --- a/tests/router/router_process.py +++ b/tests/router/router_process.py @@ -114,6 +114,7 @@ def __init__( ], log_dir=request.node.name, terminate_all_matching_process_names=False, + display_name=f"dynamo-frontend-{router_mode}", ) self.port = frontend_port self.router_mode = router_mode @@ -170,6 +171,7 @@ def __init__( ], log_dir=request.node.name, terminate_all_matching_process_names=False, + display_name="dynamo-frontend-direct", ) self.port = frontend_port diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index 5980570ae66b..3e756910e6ce 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -322,6 +322,7 @@ def __init__( health_check_urls=[], log_dir=request.node.name, terminate_all_matching_process_names=False, + display_name="dynamo-mocker", ) logger.info( f"Created mocker process with {num_mockers} worker(s), endpoint: {self.endpoint}" @@ -643,6 +644,7 @@ def __init__( health_check_urls=[], log_dir=request.node.name, terminate_all_matching_process_names=False, + display_name=f"dynamo-mocker-{worker_type}", ) logger.info( f"Created {worker_type} mocker process with {num_mockers} worker(s), "