test(router): cover round-robin unset dp-rank flow

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<Vec<mpsc::UnboundedSender<DirectRequest>>>,
     senders_ready: Notify,
     engine_args: MockEngineArgs,
+    unset_dp_rank_counter: AtomicU32,
     /// Bootstrap server for prefill workers in disaggregated mode
     bootstrap_server: Arc<OnceCell<Arc<BootstrapServer>>>,
     /// 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<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
     ) -> Result<ManyOut<LLMEngineOutput>, 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 {

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,

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",

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,