test_router_e2e_with_mockers.py 57.2 KB
Newer Older
1
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
# SPDX-License-Identifier: Apache-2.0
3

4
5
6
7
8
9
# NOTE: These tests run reliably in serial but have encountered intermittent failures
# under pytest-xdist parallel execution (-n auto). Each test spawns its own
# DistributedRuntime with isolated etcd/NATS and unique namespaces, but the Rust
# runtime may use process-global state (e.g. lazy_static / OnceLock singletons for
# endpoint tables) that races under concurrent xdist workers. Do not add
# @pytest.mark.parallel until DRT endpoint registration is confirmed thread-safe.
10
#
11
import asyncio
12
13
import logging
import os
14
15
import sys
from contextlib import contextmanager
16
from pathlib import Path
17
from typing import Any, Dict, Iterator, Optional
18

19
import aiohttp
20
21
import pytest

22
from tests.router.common import (
23
    _test_busy_threshold_endpoint,
24
    _test_disagg_direct_mode,
25
    _test_python_router_bindings,
26
    _test_remote_indexer_decisions,
27
28
    _test_router_basic,
    _test_router_decisions,
29
    _test_router_decisions_disagg,
30
    _test_router_decisions_disagg_round_robin_prefill_dp_rank,
31
    _test_router_indexers_sync,
32
33
34
35
    _test_router_overload_503,
    _test_router_query_instance_id,
    _test_router_two_routers,
)
36
37
38
39
40
41
from tests.router.helper import (
    generate_random_suffix,
    get_kv_indexer_command,
    get_runtime,
    wait_for_indexer_workers_active,
)
Alec's avatar
Alec committed
42
from tests.utils.constants import ROUTER_MODEL_NAME
43
from tests.utils.managed_process import ManagedProcess
44
45
46
47
48
from tests.utils.port_utils import (
    allocate_contiguous_ports,
    allocate_ports,
    deallocate_ports,
)
49

50
51
52
53
logger = logging.getLogger(__name__)

MODEL_NAME = ROUTER_MODEL_NAME

54
55
56
57
pytestmark = [
    pytest.mark.pre_merge,
    pytest.mark.gpu_0,
    pytest.mark.integration,
58
    pytest.mark.model(MODEL_NAME),
59
]
60
61
NUM_MOCKERS = 2
SPEEDUP_RATIO = 10.0
62
63
64
BASE_PORT = 9100  # Base port for general test allocations (frontend, system, etc.)
BASE_PORT_BOOTSTRAP = 10100  # Base port for disagg bootstrap rendezvous
BASE_PORT_ZMQ = 11100  # Base port for ZMQ KV event publishing
65
NUM_REQUESTS = 100
66
BLOCK_SIZE = 16
67
PLANNER_PROFILE_DATA_DIR = (
68
69
    Path(__file__).resolve().parents[2]
    / "components/src/dynamo/planner/tests/data/profiling_results/H200_TP1P_TP1D"
70
)
71
72
73
74
75
76
77
ROUTER_AIC_CONFIG = {
    "aic_backend": "vllm",
    "aic_system": "h200_sxm",
    "aic_backend_version": "0.12.0",
    "aic_tp_size": 1,
    "aic_model_path": "Qwen/Qwen3-32B",
}
78
79
80
ROUND_ROBIN_MOCKER_SKIP_REASON = (
    "Flaky on CI: tcp nondurable round-robin mocker router path timed out"
)
81
82
83
84
85
86
87


def _require_router_aic() -> dict[str, Any]:
    pytest.importorskip(
        "aiconfigurator", reason="router AIC test requires aiconfigurator"
    )
    return ROUTER_AIC_CONFIG.copy()
88
89


90
def get_unique_ports(
91
92
93
94
    request,
    num_ports: int = 1,
    store_backend: str = "etcd",
    request_plane: str = "nats",
95
    registration_order: str = "prefill_first",
96
) -> list[int]:
97
    """Allocate random free ports for xdist-safe router tests.
98

99
100
101
    This replaces the previous "test-name offset" scheme with the shared flock-backed
    allocator from `tests.utils.port_utils`, which avoids collisions across pytest-xdist
    worker processes.
102

103
104
105
106
    Notes:
    - The extra parameters are kept for call-site compatibility (they no longer affect
      the chosen ports).
    - Ports are released at the end of the test via a pytest finalizer.
107
    """
108
109
110
    _ = (store_backend, request_plane, registration_order)
    ports = allocate_ports(num_ports, BASE_PORT)
    request.addfinalizer(lambda: deallocate_ports(ports))
111
112
113
    return ports


114
115
116
117
118
119
120
121
122
123
124
125
# Shared test payload for all tests
TEST_PAYLOAD: Dict[str, Any] = {
    "model": MODEL_NAME,
    "messages": [
        {
            "role": "user",
            "content": "In a quiet meadow tucked between rolling hills, a plump gray rabbit nibbled on clover beneath the shade of a gnarled oak tree. Its ears twitched at the faint rustle of leaves, but it remained calm, confident in the safety of its burrow just a few hops away. The late afternoon sun warmed its fur, and tiny dust motes danced in the golden light as bees hummed lazily nearby. Though the rabbit lived a simple life, every day was an adventure of scents, shadows, and snacks—an endless search for the tastiest patch of greens and the softest spot to nap.",
        }
    ],
    "stream": True,
    "max_tokens": 10,
}
126
127
128
129
130
131
132
133
134
135
136
SOAK_TEST_PAYLOAD: Dict[str, Any] = {
    "model": MODEL_NAME,
    "messages": [
        {
            "role": "user",
            "content": "one two three four five six seven eight nine ten",
        }
    ],
    "stream": False,
    "max_tokens": 1,
}
137

138

139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def _build_mocker_command(
    endpoint: str,
    store_backend: str,
    num_workers: int,
    mocker_args: Dict[str, Any],
    worker_type: Optional[str] = None,
) -> list[str]:
    """Build the mocker CLI command with all arguments.

    Args:
        endpoint: The dynamo endpoint string
        store_backend: Storage backend ("etcd" or "file")
        num_workers: Number of workers to spawn (uses --num-workers flag)
        mocker_args: Dictionary of mocker arguments
        worker_type: Optional worker type ("prefill" or "decode") for disagg mode

    Returns:
        List of command arguments for subprocess
    """
    command = [
159
        sys.executable,
160
161
162
163
164
165
        "-m",
        "dynamo.mocker",
        "--model-path",
        MODEL_NAME,
        "--endpoint",
        endpoint,
166
        "--discovery-backend",
167
168
169
170
171
172
173
        store_backend,
        "--num-workers",
        str(num_workers),
    ]

    # Add worker type flag for disaggregated mode
    if worker_type == "prefill":
174
        command.extend(["--disaggregation-mode", "prefill"])
175
    elif worker_type == "decode":
176
        command.extend(["--disaggregation-mode", "decode"])
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202

    # Add individual CLI arguments from mocker_args
    if "speedup_ratio" in mocker_args:
        command.extend(["--speedup-ratio", str(mocker_args["speedup_ratio"])])
    if "block_size" in mocker_args:
        command.extend(["--block-size", str(mocker_args["block_size"])])
    if "num_gpu_blocks" in mocker_args:
        command.extend(
            ["--num-gpu-blocks-override", str(mocker_args["num_gpu_blocks"])]
        )
    if "max_num_seqs" in mocker_args:
        command.extend(["--max-num-seqs", str(mocker_args["max_num_seqs"])])
    if "max_num_batched_tokens" in mocker_args:
        command.extend(
            ["--max-num-batched-tokens", str(mocker_args["max_num_batched_tokens"])]
        )
    if "enable_prefix_caching" in mocker_args:
        if mocker_args["enable_prefix_caching"]:
            command.append("--enable-prefix-caching")
        else:
            command.append("--no-enable-prefix-caching")
    if "enable_chunked_prefill" in mocker_args:
        if mocker_args["enable_chunked_prefill"]:
            command.append("--enable-chunked-prefill")
        else:
            command.append("--no-enable-chunked-prefill")
203
204
    if "preemption_mode" in mocker_args:
        command.extend(["--preemption-mode", str(mocker_args["preemption_mode"])])
205
206
    if "dp_size" in mocker_args:
        command.extend(["--data-parallel-size", str(mocker_args["dp_size"])])
207
208
209
210
211
212
213
214
215
216
217
218
219
220
    if "planner_profile_data" in mocker_args:
        command.extend(
            ["--planner-profile-data", str(mocker_args["planner_profile_data"])]
        )
    if mocker_args.get("aic_perf_model") is True:
        command.append("--aic-perf-model")
    if "aic_system" in mocker_args:
        command.extend(["--aic-system", str(mocker_args["aic_system"])])
    if "aic_backend_version" in mocker_args:
        command.extend(
            ["--aic-backend-version", str(mocker_args["aic_backend_version"])]
        )
    if "aic_tp_size" in mocker_args:
        command.extend(["--aic-tp-size", str(mocker_args["aic_tp_size"])])
221
222
223
    # Use --durable-kv-events to enable JetStream mode (local indexer disabled)
    if mocker_args.get("durable_kv_events") is True:
        command.append("--durable-kv-events")
224
225
    if "bootstrap_ports" in mocker_args:
        command.extend(["--bootstrap-ports", mocker_args["bootstrap_ports"]])
226
227
    if "zmq_kv_events_ports" in mocker_args:
        command.extend(["--zmq-kv-events-ports", mocker_args["zmq_kv_events_ports"]])
228
229
    if "zmq_replay_ports" in mocker_args:
        command.extend(["--zmq-replay-ports", mocker_args["zmq_replay_ports"]])
230
231
232
233

    return command


234
class MockerProcess:
235
236
237
    """Manages mocker engine instances with shared tokio runtime via --num-workers.

    When standalone_indexer=True, launches mockers one-by-one (each as --num-workers 1)
238
    and runs a standalone HTTP KV indexer binary alongside them. Call launch_workers_with_indexer()
239
240
    in async context to start mockers and register their ZMQ ports with the indexer.
    """
241

242
243
244
245
246
    def __init__(
        self,
        request,
        mocker_args: Optional[Dict[str, Any]] = None,
        num_mockers: int = 1,
247
        store_backend: str = "etcd",
248
        request_plane: str = "nats",
249
        zmq_kv_events: bool = False,
250
        standalone_indexer: bool = False,
251
        model_name: str = "mocker",
252
        zmq_replay: bool = False,
253
    ):
254
255
        namespace_suffix = generate_random_suffix()
        self.namespace = f"test-namespace-{namespace_suffix}"
256
        self.component_name = "mocker"
257
        self.model_name = model_name
258
        self.endpoint = f"dyn://{self.namespace}.{self.component_name}.generate"
259
        self.num_workers = num_mockers
260
        self._zmq_kv_events_ports: list[int] = []
261
        self._zmq_replay_ports: list[int] = []
262
263
264
265
266
267
        self._standalone_indexer = standalone_indexer
        self._standalone_indexer_port: Optional[int] = None
        self._standalone_indexer_b_port: Optional[int] = None
        self._indexer_process: Optional[ManagedProcess] = None
        self._indexer_b_process: Optional[ManagedProcess] = None
        self._mocker_processes: list[ManagedProcess] = []
268
269
270
271
272
273
274
        self._request = request
        self._store_backend = store_backend
        self._request_plane = request_plane
        self._mocker_args_orig: Dict[str, Any] = (mocker_args or {}).copy()
        self.worker_id_to_zmq_ports: dict[int, dict[int, str]] = {}

        mocker_args = self._mocker_args_orig.copy()
275
276
        # Store dp_size for DP-aware test functions
        self.dp_size = mocker_args.get("dp_size")
277
278
        # Alias for consistency with vLLM/SGLang workers
        self.data_parallel_size = self.dp_size
279

280
281
        # Allocate contiguous ZMQ port blocks for KV event publishing because
        # the mocker binds base_port + dp_rank for each DP rank.
282
283
        if zmq_kv_events:
            dp_size = mocker_args.get("dp_size", 1)
284
285
            self._zmq_kv_events_ports = allocate_contiguous_ports(
                num_mockers, dp_size, BASE_PORT_ZMQ
286
287
            )
            bases = [self._zmq_kv_events_ports[i * dp_size] for i in range(num_mockers)]
288
289
            if not standalone_indexer:
                mocker_args["zmq_kv_events_ports"] = ",".join(str(p) for p in bases)
290
291
292
293
294
            logger.info(
                f"Allocated ZMQ KV event ports {self._zmq_kv_events_ports} "
                f"(bases: {bases}) for {num_mockers} workers"
            )

295
        # Allocate contiguous ZMQ replay port blocks with the same layout.
296
297
        if zmq_replay and zmq_kv_events:
            dp_size = mocker_args.get("dp_size", 1)
298
299
            self._zmq_replay_ports = allocate_contiguous_ports(
                num_mockers, dp_size, BASE_PORT_ZMQ + 1000
300
301
302
303
            )
            replay_bases = [
                self._zmq_replay_ports[i * dp_size] for i in range(num_mockers)
            ]
304
305
            if not standalone_indexer:
                mocker_args["zmq_replay_ports"] = ",".join(str(p) for p in replay_bases)
306
307
308
309
310
            logger.info(
                f"Allocated ZMQ replay ports {self._zmq_replay_ports} "
                f"(bases: {replay_bases}) for {num_mockers} workers"
            )

311
312
313
314
315
316
        if standalone_indexer:
            # Allocate ports for standalone indexer A and B (P2P recovery peer)
            indexer_ports = allocate_ports(2, BASE_PORT)
            self._standalone_indexer_port = indexer_ports[0]
            self._standalone_indexer_b_port = indexer_ports[1]
            request.addfinalizer(lambda: deallocate_ports(indexer_ports))
317
            # Don't build a single mocker command — we'll launch per-worker in launch_workers_with_indexer
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
            self._process = None
        else:
            command = _build_mocker_command(
                endpoint=self.endpoint,
                store_backend=store_backend,
                num_workers=num_mockers,
                mocker_args=mocker_args,
            )

            env = os.environ.copy()
            env["DYN_REQUEST_PLANE"] = request_plane

            self._process = ManagedProcess(
                command=command,
                env=env,
                timeout=60,
                display_output=True,
                health_check_ports=[],
                health_check_urls=[],
                log_dir=request.node.name,
                terminate_all_matching_process_names=False,
339
                display_name="dynamo-mocker",
340
341
342
343
            )
        logger.info(
            f"Created mocker process with {num_mockers} worker(s), endpoint: {self.endpoint}"
            f"{', standalone_indexer=True' if standalone_indexer else ''}"
344
345
        )

346
347
348
349
350
    @property
    def standalone_indexer_url(self) -> Optional[str]:
        if self._standalone_indexer_port is not None:
            return f"http://localhost:{self._standalone_indexer_port}"
        return None
351

352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
    @property
    def standalone_indexer_b_url(self) -> Optional[str]:
        if self._standalone_indexer_b_port is not None:
            return f"http://localhost:{self._standalone_indexer_b_port}"
        return None

    def __enter__(self):
        if self._standalone_indexer:
            # Launch the standalone indexer binary
            block_size = self._mocker_args_orig.get("block_size", BLOCK_SIZE)
            indexer_cmd = [
                *get_kv_indexer_command(),
                "--block-size",
                str(block_size),
                "--port",
                str(self._standalone_indexer_port),
            ]
            self._indexer_process = ManagedProcess(
                command=indexer_cmd,
                timeout=120,
                display_output=True,
                health_check_ports=[self._standalone_indexer_port],
                health_check_urls=[],
                log_dir=self._request.node.name,
                terminate_all_matching_process_names=False,
                display_name="dynamo-kv-indexer",
            )
            logger.info(
                f"Starting standalone indexer on port {self._standalone_indexer_port}"
            )
            self._indexer_process.__enter__()
383
            # Don't start mocker processes yet — launch_workers_with_indexer will do it
384
385
386
387
388
        else:
            logger.info(f"Starting mocker process with {self.num_workers} worker(s)")
            self._process.__enter__()
        return self

389
390
    async def launch_workers_with_indexer(self, endpoint):
        """Launch workers one-by-one and register each with the standalone indexer.
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453

        For each mocker:
        1. Launch a mocker process with --num-workers 1
        2. Poll endpoint.client().instance_ids() until a new worker_id appears
        3. POST /register to the indexer with the worker_id and its ZMQ addresses

        Args:
            endpoint: The dynamo endpoint object to discover worker IDs.
        """
        client = await endpoint.client()
        known_ids: set[int] = set()
        dp_size = self._mocker_args_orig.get("dp_size", 1)

        for i in range(self.num_workers):
            # Build per-mocker args with its own ZMQ base port
            mocker_args = self._mocker_args_orig.copy()
            base_port = self._zmq_kv_events_ports[i * dp_size]
            mocker_args["zmq_kv_events_ports"] = str(base_port)
            if self._zmq_replay_ports:
                replay_base = self._zmq_replay_ports[i * dp_size]
                mocker_args["zmq_replay_ports"] = str(replay_base)

            command = _build_mocker_command(
                endpoint=self.endpoint,
                store_backend=self._store_backend,
                num_workers=1,
                mocker_args=mocker_args,
            )

            env = os.environ.copy()
            env["DYN_REQUEST_PLANE"] = self._request_plane

            proc = ManagedProcess(
                command=command,
                env=env,
                timeout=60,
                display_output=True,
                health_check_ports=[],
                health_check_urls=[],
                log_dir=self._request.node.name,
                terminate_all_matching_process_names=False,
                display_name=f"mocker-{i}",
            )
            proc.__enter__()
            self._mocker_processes.append(proc)

            # Poll for the new worker_id
            new_worker_id = None
            for _ in range(120):
                ids = set(client.instance_ids())
                new = ids - known_ids
                if new:
                    new_worker_id = new.pop()
                    known_ids.add(new_worker_id)
                    break
                await asyncio.sleep(0.5)

            if new_worker_id is None:
                raise RuntimeError(
                    f"Timed out waiting for mocker {i} to register "
                    f"(known_ids={known_ids})"
                )

454
455
456
457
            # Register each dp_rank endpoint with the standalone indexer.
            # The mocker binds on base_port + dp_rank (contiguous), so we must
            # use the same formula here rather than indexing into the allocated
            # port list, which may contain gaps when intervening ports are busy.
458
459
            zmq_addresses = {}
            register_url = f"{self.standalone_indexer_url}/register"
460
461
462
            replay_base = (
                self._zmq_replay_ports[i * dp_size] if self._zmq_replay_ports else None
            )
463
464
            async with aiohttp.ClientSession() as session:
                for dp_rank in range(dp_size):
465
                    port = base_port + dp_rank
466
467
468
469
470
471
472
473
474
475
476
477
                    endpoint = f"tcp://127.0.0.1:{port}"
                    zmq_addresses[dp_rank] = endpoint

                    payload = {
                        "instance_id": new_worker_id,
                        "endpoint": endpoint,
                        "dp_rank": dp_rank,
                        "model_name": self.model_name,
                        "block_size": self._mocker_args_orig.get(
                            "block_size", BLOCK_SIZE
                        ),
                    }
478
479
480
481
                    if replay_base is not None:
                        payload[
                            "replay_endpoint"
                        ] = f"tcp://127.0.0.1:{replay_base + dp_rank}"
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
                    async with session.post(register_url, json=payload) as resp:
                        if resp.status != 201:
                            body = await resp.text()
                            raise RuntimeError(
                                f"Failed to register instance {new_worker_id} "
                                f"dp_rank {dp_rank}: {resp.status} {body}"
                            )

            self.worker_id_to_zmq_ports[new_worker_id] = zmq_addresses

            logger.info(
                f"Mocker {i}: worker_id={new_worker_id}, "
                f"zmq_addresses={zmq_addresses}"
            )

        await wait_for_indexer_workers_active(
            self.standalone_indexer_url, self.worker_id_to_zmq_ports
        )
        logger.info(
            f"All {self.num_workers} mockers launched and registered with indexer"
        )

    def launch_indexer(self):
        """Launch a second standalone indexer (Indexer B) with --peers pointing to Indexer A.

        Workers are passed via --workers so ZMQ sockets connect before recovery
        runs, ensuring the subscription handshake completes during the recovery
        delay and no events are lost to the ZMQ slow-joiner problem.
        """
        if not self._standalone_indexer or self._standalone_indexer_b_port is None:
            raise RuntimeError("launch_indexer requires standalone_indexer=True")
        if not self.worker_id_to_zmq_ports:
            raise RuntimeError("launch_indexer requires workers to be registered first")

        block_size = self._mocker_args_orig.get("block_size", BLOCK_SIZE)

        # Build --workers arg: "worker_id:dp_rank=zmq_addr,..."
        worker_entries = []
        for worker_id, zmq_addresses in self.worker_id_to_zmq_ports.items():
            for dp_rank, zmq_endpoint in zmq_addresses.items():
                worker_entries.append(f"{worker_id}:{dp_rank}={zmq_endpoint}")
        workers_arg = ",".join(worker_entries)

        indexer_b_cmd = [
            *get_kv_indexer_command(),
            "--block-size",
            str(block_size),
            "--port",
            str(self._standalone_indexer_b_port),
            "--peers",
            f"http://localhost:{self._standalone_indexer_port}",
            "--workers",
            workers_arg,
            "--model-name",
            self.model_name,
        ]
        self._indexer_b_process = ManagedProcess(
            command=indexer_b_cmd,
            timeout=120,
541
            display_output=True,
542
            health_check_ports=[self._standalone_indexer_b_port],
543
            health_check_urls=[],
544
            log_dir=self._request.node.name,
545
            terminate_all_matching_process_names=False,
546
            display_name="dynamo-kv-indexer-b",
547
548
        )
        logger.info(
549
550
            f"Starting standalone indexer B on port {self._standalone_indexer_b_port} "
            f"with peer http://localhost:{self._standalone_indexer_port}"
551
        )
552
        self._indexer_b_process.__enter__()
553

554
    def __exit__(self, exc_type, exc_val, exc_tb):
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
        logger.info("Stopping mocker process(es)")
        # Stop individual mocker processes (standalone_indexer mode)
        for proc in self._mocker_processes:
            try:
                proc.__exit__(exc_type, exc_val, exc_tb)
            except Exception as e:
                logger.warning(f"Error stopping mocker process: {e}")
        self._mocker_processes.clear()
        # Stop standalone indexer B (P2P recovery peer)
        if self._indexer_b_process is not None:
            try:
                self._indexer_b_process.__exit__(exc_type, exc_val, exc_tb)
            except Exception as e:
                logger.warning(f"Error stopping indexer B process: {e}")
            self._indexer_b_process = None
        # Stop standalone indexer A
        if self._indexer_process is not None:
            try:
                self._indexer_process.__exit__(exc_type, exc_val, exc_tb)
            except Exception as e:
                logger.warning(f"Error stopping indexer process: {e}")
            self._indexer_process = None
        # Stop single mocker process (non-standalone mode)
578
579
        if self._process is not None:
            self._process.__exit__(exc_type, exc_val, exc_tb)
580
581
582
583
        if self._zmq_kv_events_ports:
            deallocate_ports(self._zmq_kv_events_ports)
            logger.info(f"Deallocated ZMQ KV event ports {self._zmq_kv_events_ports}")
            self._zmq_kv_events_ports = []
584
585
586
587
        if self._zmq_replay_ports:
            deallocate_ports(self._zmq_replay_ports)
            logger.info(f"Deallocated ZMQ replay ports {self._zmq_replay_ports}")
            self._zmq_replay_ports = []
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607


class DisaggMockerProcess:
    """Manages prefill or decode mocker instances for disaggregated serving.

    Uses --num-workers for shared tokio runtime. For disaggregated serving:
    - Prefill workers: worker_type="prefill", endpoint is namespace.prefill.generate
    - Decode workers: worker_type="decode", endpoint is namespace.backend.generate

    Both prefill and decode workers should share the same namespace for proper discovery.
    """

    def __init__(
        self,
        request,
        namespace: str,
        worker_type: str,
        mocker_args: Optional[Dict[str, Any]] = None,
        num_mockers: int = 1,
        store_backend: str = "etcd",
608
        request_plane: str = "nats",
609
        enable_bootstrap: bool = False,
610
611
612
613
    ):
        if worker_type not in ("prefill", "decode"):
            raise ValueError(
                f"worker_type must be 'prefill' or 'decode', got {worker_type}"
614
            )
615
616
617
618

        self.namespace = namespace
        self.worker_type = worker_type
        self.num_workers = num_mockers
619
        self._bootstrap_ports: list[int] = []
620
621
622
623
624
625
626
627
628

        # Set component name and endpoint based on worker type
        if worker_type == "prefill":
            self.component_name = "prefill"
            self.endpoint = f"dyn://{self.namespace}.prefill.generate"
        else:
            self.component_name = "backend"
            self.endpoint = f"dyn://{self.namespace}.backend.generate"

629
630
631
632
        mocker_args = (mocker_args or {}).copy()

        # Allocate bootstrap ports for prefill workers if enabled (one per worker)
        if enable_bootstrap and worker_type == "prefill":
633
            self._bootstrap_ports = allocate_ports(num_mockers, BASE_PORT_BOOTSTRAP)
634
635
636
637
638
639
            mocker_args["bootstrap_ports"] = ",".join(
                str(p) for p in self._bootstrap_ports
            )
            logger.info(
                f"Allocated bootstrap ports {self._bootstrap_ports} for {num_mockers} prefill workers"
            )
640
641
642
643
644
645
646
647
648

        command = _build_mocker_command(
            endpoint=self.endpoint,
            store_backend=store_backend,
            num_workers=num_mockers,
            mocker_args=mocker_args,
            worker_type=worker_type,
        )

649
650
651
        env = os.environ.copy()
        env["DYN_REQUEST_PLANE"] = request_plane

652
653
        self._process = ManagedProcess(
            command=command,
654
            env=env,
655
656
657
658
659
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
660
            terminate_all_matching_process_names=False,
661
            display_name=f"dynamo-mocker-{worker_type}",
662
663
664
665
666
        )
        logger.info(
            f"Created {worker_type} mocker process with {num_mockers} worker(s), "
            f"endpoint: {self.endpoint}"
        )
667

668
669
670
671
672
    @property
    def bootstrap_ports(self) -> list[int]:
        """Return the allocated bootstrap ports, if any."""
        return self._bootstrap_ports

673
    def __enter__(self):
674
675
676
677
        logger.info(
            f"Starting {self.worker_type} mocker process with {self.num_workers} worker(s)"
        )
        self._process.__enter__()
678
        return self
679

680
    def __exit__(self, exc_type, exc_val, exc_tb):
681
682
        logger.info(f"Stopping {self.worker_type} mocker process")
        self._process.__exit__(exc_type, exc_val, exc_tb)
683
684
685
686
687
        # Deallocate bootstrap ports if we allocated any
        if self._bootstrap_ports:
            deallocate_ports(self._bootstrap_ports)
            logger.info(f"Deallocated bootstrap ports {self._bootstrap_ports}")
            self._bootstrap_ports = []
688
689


690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
@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


761
@pytest.mark.timeout(180)  # planner-profile mocker setup can exceed 120s on CI CPUs
762
@pytest.mark.parametrize(
763
    "router_mode,durable_kv_events,mocker_args_override",
764
    [
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
        pytest.param("kv", False, {}, id="kv-nondurable"),
        pytest.param(
            "kv",
            False,
            {"planner_profile_data": PLANNER_PROFILE_DATA_DIR},
            id="kv-planner",
        ),
        pytest.param(
            "kv",
            False,
            {"aic_perf_model": True, "aic_system": "h200_sxm"},
            id="kv-aic",
        ),
        pytest.param("kv", True, {}, id="kv-durable"),
        pytest.param("round-robin", False, {}, id="roundrobin"),
        pytest.param("random", False, {}, id="random"),
        pytest.param("power-of-two", False, {}, id="power-of-two"),
782
783
784
    ],
    indirect=["durable_kv_events"],
)
785
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
786
@pytest.mark.skip(reason=ROUND_ROBIN_MOCKER_SKIP_REASON)
787
def test_mocker_router(
788
789
790
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
791
    router_mode,
792
    request_plane,
793
    durable_kv_events,
794
    mocker_args_override,
795
):
796
    """Test router with multiple mocker engine instances across all router modes.
797

798
799
    Covers kv, round-robin, and random routing. Tests both NATS and TCP request planes.
    """
800
    # runtime_services starts etcd and optionally nats based on request_plane
801
802
803
    logger.info(
        f"Starting mocker router test: router_mode={router_mode}, request_plane={request_plane}"
    )
804

805
806
807
808
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
809
        "durable_kv_events": durable_kv_events,
810
    }
811
    mocker_args.update(mocker_args_override)
812

813
814
815
816
817
818
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
819
        # Start mocker instances with the new CLI interface
820
821
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
822

823
        # Get unique port for this test
824
825
826
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
827
828

        # Run basic router test (starts router internally and waits for workers to be ready)
829
830
831
832
        _test_router_basic(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
833
            frontend_port=frontend_port,
834
835
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
836
            request_plane=request_plane,
837
            router_mode=router_mode,
838
            min_initial_workers=mockers.num_workers,
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
        )


@pytest.mark.timeout(180)
@pytest.mark.parametrize("router_mode", ["kv", "round-robin", "random"])
@pytest.mark.parametrize(
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
)
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
def test_mocker_router_soak(
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    router_mode,
    durable_kv_events,
    request_plane,
):
    mocker_args = {
        "speedup_ratio": 1000.0,
        "block_size": BLOCK_SIZE,
        "durable_kv_events": durable_kv_events,
    }

    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
    ) as mockers:
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]

        _test_router_basic(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=SOAK_TEST_PAYLOAD,
            num_requests=1024,
            request_plane=request_plane,
            router_mode=router_mode,
            min_initial_workers=mockers.num_workers,
882
883
884
        )


885
@pytest.mark.parametrize("store_backend", ["etcd", "file"])
886
@pytest.mark.parametrize(
887
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
888
)  # Use NATS Core (local indexer)
889
@pytest.mark.timeout(180)  # bumped for xdist contention (was 60s; ~19.86s serial avg)
890
891
def test_mocker_two_kv_router(
    request,
892
    runtime_services_dynamic_ports,
893
894
895
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
896
    durable_kv_events,
897
):
898
899
900
    """
    Test with two KV routers and multiple mocker engine instances.
    Alternates requests between the two routers to test load distribution.
901
    Tests with both etcd and file storage backends.
902
903
904
    """

    # runtime_services starts etcd and nats
905
906
907
    logger.info(
        f"Starting mocker two KV router test with {store_backend} storage backend"
    )
908

909
910
911
912
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
913
        "durable_kv_events": durable_kv_events,
914
    }
915

916
917
918
919
920
921
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
    ) as mockers:
922
        # Start mocker instances with the new CLI interface
923
924
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
925

926
927
928
929
930
        # Get unique ports for this test (2 ports for two routers)
        router_ports = get_unique_ports(
            request, num_ports=2, store_backend=store_backend
        )

931
932
933
934
935
        # Run two-router test (starts KV routers internally and manages their lifecycle)
        _test_router_two_routers(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
936
            router_ports=router_ports,
937
938
939
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            store_backend=store_backend,
940
            skip_consumer_verification=not durable_kv_events,  # Skip JetStream checks in NATS Core mode
941
942
943
        )


944
@pytest.mark.parametrize(
945
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
946
)  # Use NATS Core (local indexer)
947
@pytest.mark.timeout(45)  # ~3x average (~13.10s), rounded up (when enabled)
Alec's avatar
Alec committed
948
def test_mocker_kv_router_overload_503(
949
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
950
):
951
    """Test that KV router returns 503 when mocker workers are overloaded."""
952
    logger.info("Starting mocker KV router overload test for 503 status")
953
    # Create mocker args dictionary with limited resources - use local indexer (NATS Core mode)
954
    mocker_args = {
955
        "speedup_ratio": 0.01,
956
957
        "block_size": 4,  # Smaller block size
        "num_gpu_blocks": 64,  # Limited GPU blocks to exhaust quickly
958
        "durable_kv_events": durable_kv_events,
959
    }
960

961
    with MockerProcess(request, mocker_args=mocker_args, num_mockers=1) as mockers:
962
        # Start single mocker instance with limited resources
963
964
        logger.info("Starting single mocker instance with limited resources")
        logger.info(f"Mocker using endpoint: {mockers.endpoint}")
965

966
967
968
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

969
970
971
972
973
974
975
        # Run overload 503 test
        _test_router_overload_503(
            engine_workers=mockers,
            block_size=4,  # Match the mocker's block size
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
976
            blocks_threshold=0.2,
977
        )
978

979

980
@pytest.mark.timeout(90)  # bumped for xdist contention (was 22s; ~7.10s serial avg)
981
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
982
@pytest.mark.parametrize(
983
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
984
)  # Use NATS Core (local indexer)
985
def test_kv_router_bindings(
986
987
988
989
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
990
    durable_kv_events,
991
):
992
993
    """Test KvRouter Python bindings with mocker engines."""
    logger.info("Starting KvRouter bindings test")
994
995
996
997
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
998
        "durable_kv_events": durable_kv_events,
999
    }
1000

1001
1002
1003
1004
1005
1006
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
1007
        # Start mocker instances
1008
1009
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
1010

1011
        # Get runtime and create endpoint
1012
        runtime = get_runtime(request_plane=request_plane)
1013
1014
1015
        endpoint = runtime.endpoint(
            f"{mockers.namespace}.{mockers.component_name}.generate"
        )
1016

1017
1018
1019
        # Run Python router bindings test
        _test_python_router_bindings(
            engine_workers=mockers,
1020
1021
            endpoint=endpoint,
            block_size=BLOCK_SIZE,
1022
1023
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
1024
        )
1025

1026

1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
@pytest.mark.parametrize(
    "store_backend,durable_kv_events,request_plane",
    [
        ("etcd", True, "nats"),  # JetStream mode - uses JetStream
        ("etcd", False, "tcp"),  # NATS core mode (with gap detection) - no JetStream
        ("file", True, "nats"),  # File backend - uses JetStream
    ],
    ids=[
        "jetstream",
        "nats_core",
        "file",
    ],
    indirect=["request_plane", "durable_kv_events"],
)
@pytest.mark.timeout(300)
def test_indexers_sync(
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
    durable_kv_events,
    request_plane,
):
    """
    Test that two KV routers have synchronized indexer states after processing requests.
    This test verifies that both routers converge to the same internal state.

    Tests with three configurations:
    - jetstream: etcd backend, JetStream for KV events, NATS request plane
    - nats_core: etcd backend, NATS Core with gap detection, TCP request plane
    - file: file backend, JetStream for KV events, NATS request plane
    """
    logger.info(
        f"Starting indexers sync test: store_backend={store_backend}, "
        f"durable_kv_events={durable_kv_events}, request_plane={request_plane}"
    )

    # Use the dynamic-port fixture to avoid hardcoded localhost:4222/2379 in parallel runs.
    nats_process, _etcd_process = runtime_services_dynamic_ports

    # Create mocker args dictionary
    # Use 2 DP ranks to test per-dp_rank event ID tracking and recovery
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "durable_kv_events": durable_kv_events,
        "dp_size": 2,
    }

    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
        request_plane=request_plane,
        zmq_kv_events=True,
        zmq_replay=True,
        standalone_indexer=True,
        model_name=MODEL_NAME,
    ) as mockers:
        # Start mocker instances (2 workers x 2 DP ranks = 4 independent event streams)
        logger.info(f"Starting {NUM_MOCKERS} mocker instances with dp_size=2")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")

        # Use the common test implementation (creates its own runtimes for each router)
        # Note: Consumer verification is done inside _test_router_indexers_sync while routers are alive
        # When using durable_kv_events=True, use JetStream mode for the router
        _test_router_indexers_sync(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
            store_backend=store_backend,
            request_plane=request_plane,
            test_nats_interruption=not durable_kv_events,
            nats_server=nats_process if not durable_kv_events else None,
            durable_kv_events=durable_kv_events,
            standalone_indexer_url=mockers.standalone_indexer_url,
            standalone_indexer_b_url=mockers.standalone_indexer_b_url,
            test_zmq_replay=True,
        )

        logger.info("Indexers sync test completed successfully")


1113
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
1114
@pytest.mark.parametrize(
1115
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
1116
)  # Use NATS Core (local indexer)
Alec's avatar
Alec committed
1117
def test_query_instance_id_returns_worker_and_tokens(
1118
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
1119
):
1120
    """Test query_instance_id annotation with mocker engines."""
1121
    logger.info("Starting KV router query_instance_id annotation test")
1122
1123
1124
1125
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1126
        "durable_kv_events": durable_kv_events,
1127
    }
1128

1129
1130
1131
    with MockerProcess(
        request, mocker_args=mocker_args, num_mockers=NUM_MOCKERS
    ) as mockers:
1132
        # Start mocker instances
1133
1134
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
1135

1136
1137
1138
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

1139
1140
1141
1142
1143
1144
1145
1146
        # Run query_instance_id annotation test
        _test_router_query_instance_id(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
        )
1147

1148

1149
@pytest.mark.timeout(300)  # bumped for xdist contention (was 29s; ~9.55s serial avg)
Yan Ru Pei's avatar
Yan Ru Pei committed
1150
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
1151
@pytest.mark.parametrize(
1152
    "durable_kv_events,use_kv_events,zmq_kv_events,use_remote_indexer",
1153
    [
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
        (True, True, False, False),  # JetStream mode with KV events
        (False, True, False, False),  # NATS Core mode with local indexer (default)
        (False, True, False, True),  # NATS Core mode with a served remote indexer
        (False, False, False, False),  # Approximate mode (--no-kv-events)
        (
            False,
            False,
            False,
            True,
        ),  # Approximate mode with a singleton served remote indexer
        (False, True, True, False),  # ZMQ mode: mocker → ZMQ PUB → relay → NATS
    ],
    ids=[
        "jetstream",
        "nats_core",
        "nats_core_remote",
        "no_kv_events",
        "no_kv_events_remote",
        "zmq",
1173
    ],
1174
    indirect=["durable_kv_events"],
1175
)
1176
def test_router_decisions(
1177
1178
1179
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
1180
    durable_kv_events,
1181
    use_kv_events,
1182
    request_plane,
1183
    zmq_kv_events,
1184
    use_remote_indexer,
1185
1186
1187
):
    """Validate KV cache prefix reuse and dp_rank routing by sending progressive requests with overlapping prefixes.

1188
    Parameterized to test:
1189
1190
    - JetStream mode: KV events via NATS JetStream (durable)
    - NATS Core mode (default): KV events via NATS Core with local indexer on workers
1191
    - NATS Core mode with a served remote indexer
1192
1193
    - Approximate mode (--no-kv-events): No KV events, router predicts cache state
      based on routing decisions with TTL-based expiration and pruning
1194
    - Approximate mode with a singleton served remote indexer
1195
    """
1196
    # runtime_services_dynamic_ports handles NATS and etcd startup
1197
    logger.info(
1198
1199
1200
1201
        "Starting test router decisions: durable_kv_events=%s, use_kv_events=%s, use_remote_indexer=%s",
        durable_kv_events,
        use_kv_events,
        use_remote_indexer,
1202
    )
1203

Yan Ru Pei's avatar
Yan Ru Pei committed
1204
    # Create mocker args dictionary with dp_size=4
1205
    # durable_kv_events=True enables JetStream mode; False (default) uses NATS Core with local indexer
Yan Ru Pei's avatar
Yan Ru Pei committed
1206
1207
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
1208
        "block_size": 8,
Yan Ru Pei's avatar
Yan Ru Pei committed
1209
        "dp_size": 4,
1210
        "durable_kv_events": durable_kv_events and use_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
1211
    }
1212

1213
1214
1215
1216
1217
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
1218
        zmq_kv_events=zmq_kv_events,
1219
        standalone_indexer=zmq_kv_events,
1220
        model_name=MODEL_NAME,
1221
    ) as mockers:
1222
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
1223

1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
        if use_remote_indexer:
            _test_remote_indexer_decisions(
                mockers,
                MODEL_NAME,
                block_size=8,
                use_kv_events=use_kv_events,
                test_dp_rank=True,
                request_plane=request_plane,
            )
            return

1235
        runtime = get_runtime(request_plane=request_plane)
1236
        endpoint = runtime.endpoint(f"{mockers.namespace}.mocker.generate")
1237

1238
        _test_router_decisions(
1239
1240
1241
1242
1243
1244
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=use_kv_events,
1245
            durable_kv_events=durable_kv_events,
1246
            standalone_indexer_url=mockers.standalone_indexer_url,
1247
1248
        )

1249

1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
@pytest.mark.timeout(300)
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
def test_router_decisions_router_aic(
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
):
    """Validate agg KV-router decisions with router-side AIC enabled on the NATS Core path."""
    logger.info("Starting agg router decisions test with router-side AIC enabled")

    router_aic_config = _require_router_aic()
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": 8,
        "dp_size": 4,
        "durable_kv_events": False,
    }

    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
        model_name=MODEL_NAME,
    ) as mockers:
        runtime = get_runtime(request_plane=request_plane)
        endpoint = runtime.endpoint(f"{mockers.namespace}.mocker.generate")

        _test_router_decisions(
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=True,
            durable_kv_events=False,
            router_aic_config=router_aic_config,
        )


1291
@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"])
1292
1293
1294
@pytest.mark.parametrize(
    "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"]
)
1295
@pytest.mark.timeout(180)  # bumped for xdist contention (was 59s; ~19.51s serial avg)
1296
def test_router_decisions_disagg(
1297
1298
1299
1300
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    registration_order,
1301
    enable_disagg_bootstrap,
1302
1303
1304
1305
1306
):
    """Validate KV cache prefix reuse in disaggregated prefill-decode setup.

    Tests that progressive requests with overlapping prefixes are routed to the
    same prefill worker due to KV cache reuse.
1307

1308
1309
1310
    Parameterized to test:
    - registration_order: prefill_first vs decode_first
    - enable_disagg_bootstrap: without vs with bootstrap rendezvous
1311
    """
1312
    # runtime_services_dynamic_ports handles NATS and etcd startup
1313
1314
    logger.info(
        f"Starting disaggregated router prefix reuse test "
1315
        f"(registration_order={registration_order}, bootstrap={enable_disagg_bootstrap})"
1316
    )
1317
1318
1319
1320
1321

    # Generate shared namespace for prefill and decode workers
    namespace_suffix = generate_random_suffix()
    shared_namespace = f"test-namespace-{namespace_suffix}"

1322
    # Create mocker args - use NATS Core with local indexer (default mode)
1323
1324
1325
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1326
        # durable_kv_events defaults to False (NATS Core mode)
1327
    }
1328

1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
    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,
1349
            request_plane="nats",
1350
            enable_bootstrap=enable_disagg_bootstrap,
1351
        )
1352

1353

1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
@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,
    )
1373

1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
    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,
    }
1385

1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
    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",
        )

    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)
1412
1413


1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
@pytest.mark.timeout(180)
def test_router_decisions_disagg_router_aic(
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
):
    """Validate disagg KV-router decisions with router-side AIC enabled on the default startup path."""
    logger.info("Starting disaggregated router prefix reuse test with router-side AIC")

    router_aic_config = _require_router_aic()
    namespace_suffix = generate_random_suffix()
    shared_namespace = f"test-namespace-{namespace_suffix}"
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
    }

    with DisaggMockerProcess(
        request,
        namespace=shared_namespace,
        worker_type="prefill",
        mocker_args=mocker_args,
        num_mockers=4,
        request_plane="nats",
        enable_bootstrap=False,
    ) as prefill_workers:
        logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

        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}")

            frontend_port = get_unique_ports(
                request, num_ports=1, registration_order="prefill_first"
            )[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",
                router_aic_config=router_aic_config,
            )


1468
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
1469
@pytest.mark.parametrize(
1470
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
1471
)  # Use NATS Core (local indexer)
1472
@pytest.mark.timeout(120)  # bumped for xdist contention (was 39s; ~12.84s serial avg)
1473
def test_busy_threshold_endpoint(
1474
1475
1476
1477
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
1478
    durable_kv_events,
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
):
    """Test that the /busy_threshold endpoint can be hit and responds correctly.

    TODO: This doesn't actually test any e2e rejection for now. A proper test would:
    1. Set a very low threshold
    2. Send enough requests to exceed the threshold
    3. Verify that subsequent requests are rejected with 503

    For now, this test only verifies the endpoint is accessible and returns valid responses.
    """
1489
    # runtime_services_dynamic_ports handles NATS and etcd startup
1490
1491
1492
    logger.info(
        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
    )
1493

1494
1495
1496
1497
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1498
        "durable_kv_events": durable_kv_events,
1499
    }
1500

1501
1502
1503
1504
1505
1506
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
1507
1508
1509
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")

1510
1511
1512
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
1513
1514
1515
1516
1517
1518
1519

        _test_busy_threshold_endpoint(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
1520
            request_plane=request_plane,
1521
        )
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579


@pytest.mark.timeout(180)
def test_disagg_direct_mode_epp_headers(
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
):
    """E2E: disaggregated serving with Direct routing mode (simulating GAIE EPP).

    This test verifies the EPP-driven routing path used in the GAIE deploy recipe:
      - Frontend runs with --router-mode direct (no autonomous worker selection)
      - Worker IDs are supplied via x-worker-instance-id / x-prefill-instance-id headers

    Validates:
      1. Requests with explicit headers succeed and report correct worker IDs
      2. Requests without headers are rejected (Direct mode enforces header routing)
    """
    logger.info("Starting disaggregated Direct-mode EPP headers E2E test")

    namespace_suffix = generate_random_suffix()
    shared_namespace = f"test-namespace-{namespace_suffix}"

    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
    }

    with DisaggMockerProcess(
        request,
        namespace=shared_namespace,
        worker_type="prefill",
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane="nats",
    ) as prefill_workers:
        logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

        with DisaggMockerProcess(
            request,
            namespace=shared_namespace,
            worker_type="decode",
            mocker_args=mocker_args,
            num_mockers=2,
            request_plane="nats",
        ) as decode_workers:
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")

            frontend_port = get_unique_ports(request, num_ports=1)[0]

            _test_disagg_direct_mode(
                prefill_workers=prefill_workers,
                decode_workers=decode_workers,
                request=request,
                frontend_port=frontend_port,
                test_payload=TEST_PAYLOAD,
                request_plane="nats",
            )