test_router_e2e_with_mockers.py 44 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

# Parallelization: Hermetic tests (xdist-safe via dynamic ports + per-test namespaces).
# Tested on: Linux container.
# Combined pre_merge wall time (this file):
# - Serialized: 304.01s.
# - Parallel (-n auto): 34.55s (269.46s saved, 8.80x).
9
10
11
12
13
14
#
# NOTE: TCP request plane is NOT tested here. These tests use --num-workers > 1 which spawns
# multiple workers in a single process sharing one TCP server. The shared TCP server uses
# endpoint_path (e.g., "generate") as the routing key, causing handler collisions when multiple
# workers register the same endpoint. This is a test-only limitation; production deployments
# with separate processes per worker work correctly with TCP.
15
import asyncio
16
17
import logging
import os
18
from typing import Any, Dict, Optional
19

20
import aiohttp
21
22
import pytest

23
from tests.router.common import (  # utilities
24
    _test_busy_threshold_endpoint,
25
26
27
    _test_python_router_bindings,
    _test_router_basic,
    _test_router_decisions,
28
    _test_router_decisions_disagg,
29
30
31
32
33
34
35
    _test_router_indexers_sync,
    _test_router_overload_503,
    _test_router_query_instance_id,
    _test_router_two_routers,
    generate_random_suffix,
    get_runtime,
)
Alec's avatar
Alec committed
36
from tests.utils.constants import ROUTER_MODEL_NAME
37
from tests.utils.managed_process import ManagedProcess
38
from tests.utils.port_utils import allocate_ports, deallocate_ports
39

40
41
42
43
logger = logging.getLogger(__name__)

MODEL_NAME = ROUTER_MODEL_NAME

44
45
46
47
pytestmark = [
    pytest.mark.pre_merge,
    pytest.mark.gpu_0,
    pytest.mark.integration,
48
    pytest.mark.model(MODEL_NAME),
49
]
50
51
NUM_MOCKERS = 2
SPEEDUP_RATIO = 10.0
52
53
54
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
55
NUM_REQUESTS = 100
56
BLOCK_SIZE = 16
57
58


59
def get_unique_ports(
60
61
62
63
    request,
    num_ports: int = 1,
    store_backend: str = "etcd",
    request_plane: str = "nats",
64
    registration_order: str = "prefill_first",
65
) -> list[int]:
66
    """Allocate random free ports for xdist-safe router tests.
67

68
69
70
    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.
71

72
73
74
75
    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.
76
    """
77
78
79
    _ = (store_backend, request_plane, registration_order)
    ports = allocate_ports(num_ports, BASE_PORT)
    request.addfinalizer(lambda: deallocate_ports(ports))
80
81
82
    return ports


83
84
85
86
87
88
89
90
91
92
93
94
95
# 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,
}

96

97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
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 = [
        "python",
        "-m",
        "dynamo.mocker",
        "--model-path",
        MODEL_NAME,
        "--endpoint",
        endpoint,
124
        "--discovery-backend",
125
126
127
128
129
130
131
        store_backend,
        "--num-workers",
        str(num_workers),
    ]

    # Add worker type flag for disaggregated mode
    if worker_type == "prefill":
132
        command.extend(["--disaggregation-mode", "prefill"])
133
    elif worker_type == "decode":
134
        command.extend(["--disaggregation-mode", "decode"])
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160

    # 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")
161
162
    if "preemption_mode" in mocker_args:
        command.extend(["--preemption-mode", str(mocker_args["preemption_mode"])])
163
164
    if "dp_size" in mocker_args:
        command.extend(["--data-parallel-size", str(mocker_args["dp_size"])])
165
166
167
    # 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")
168
169
    if "bootstrap_ports" in mocker_args:
        command.extend(["--bootstrap-ports", mocker_args["bootstrap_ports"]])
170
171
    if "zmq_kv_events_ports" in mocker_args:
        command.extend(["--zmq-kv-events-ports", mocker_args["zmq_kv_events_ports"]])
172
173
174
175

    return command


176
class MockerProcess:
177
178
179
180
181
182
    """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)
    and runs a standalone HTTP KV indexer binary alongside them. Call launch_mockers_with_indexer()
    in async context to start mockers and register their ZMQ ports with the indexer.
    """
183

184
185
186
187
188
    def __init__(
        self,
        request,
        mocker_args: Optional[Dict[str, Any]] = None,
        num_mockers: int = 1,
189
        store_backend: str = "etcd",
190
        request_plane: str = "nats",
191
        zmq_kv_events: bool = False,
192
        standalone_indexer: bool = False,
193
        model_name: str = "mocker",
194
    ):
195
196
        namespace_suffix = generate_random_suffix()
        self.namespace = f"test-namespace-{namespace_suffix}"
197
        self.component_name = "mocker"
198
        self.model_name = model_name
199
        self.endpoint = f"dyn://{self.namespace}.{self.component_name}.generate"
200
        self.num_workers = num_mockers
201
        self._zmq_kv_events_ports: list[int] = []
202
203
        self._standalone_indexer = standalone_indexer
        self._standalone_indexer_port: Optional[int] = None
204
        self._standalone_indexer_b_port: Optional[int] = None
205
        self._indexer_process: Optional[ManagedProcess] = None
206
        self._indexer_b_process: Optional[ManagedProcess] = None
207
208
209
210
211
212
213
214
        self._mocker_processes: list[ManagedProcess] = []
        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()
215
216
        # Store dp_size for DP-aware test functions
        self.dp_size = mocker_args.get("dp_size")
217
218
        # Alias for consistency with vLLM/SGLang workers
        self.data_parallel_size = self.dp_size
219

220
221
222
223
224
225
226
227
228
229
        # Allocate ZMQ base ports for KV event publishing.
        # Each worker's DP ranks bind on base_port + dp_rank, so we need bases
        # spaced dp_size apart. Allocate num_mockers * dp_size ports total,
        # then pick every dp_size'th port as a base.
        if zmq_kv_events:
            dp_size = mocker_args.get("dp_size", 1)
            self._zmq_kv_events_ports = allocate_ports(
                num_mockers * dp_size, BASE_PORT_ZMQ
            )
            bases = [self._zmq_kv_events_ports[i * dp_size] for i in range(num_mockers)]
230
231
            if not standalone_indexer:
                mocker_args["zmq_kv_events_ports"] = ",".join(str(p) for p in bases)
232
233
234
235
236
            logger.info(
                f"Allocated ZMQ KV event ports {self._zmq_kv_events_ports} "
                f"(bases: {bases}) for {num_mockers} workers"
            )

237
        if standalone_indexer:
238
239
240
241
242
            # 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))
243
244
245
246
247
248
249
250
251
            # Don't build a single mocker command — we'll launch per-mocker in launch_mockers_with_indexer
            self._process = None
        else:
            command = _build_mocker_command(
                endpoint=self.endpoint,
                store_backend=store_backend,
                num_workers=num_mockers,
                mocker_args=mocker_args,
            )
252

253
254
255
256
257
258
259
260
261
262
263
264
265
            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,
            )
266
267
        logger.info(
            f"Created mocker process with {num_mockers} worker(s), endpoint: {self.endpoint}"
268
            f"{', standalone_indexer=True' if standalone_indexer else ''}"
269
270
        )

271
272
273
274
275
276
    @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

277
278
279
280
281
282
    @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

283
    def __enter__(self):
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        if self._standalone_indexer:
            # Launch the standalone indexer binary
            block_size = self._mocker_args_orig.get("block_size", BLOCK_SIZE)
            indexer_cmd = [
                "cargo",
                "run",
                "-p",
                "dynamo-kv-router",
                "--features",
                "indexer-bin",
                "--bin",
                "dynamo-kv-indexer",
                "--",
                "--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__()
            # Don't start mocker processes yet — launch_mockers_with_indexer will do it
        else:
            logger.info(f"Starting mocker process with {self.num_workers} worker(s)")
            self._process.__enter__()
320
321
        return self

322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
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
383
384
385
386
387
388
389
390
391
392
393
394
395
396
    async def launch_mockers_with_indexer(self, endpoint):
        """Launch mockers one-by-one and register each with the standalone indexer.

        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 /workers 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)

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

            # Register each dp_rank endpoint with the standalone indexer
            zmq_addresses = {}
            register_url = f"{self.standalone_indexer_url}/register"
            async with aiohttp.ClientSession() as session:
                for dp_rank in range(dp_size):
                    port = self._zmq_kv_events_ports[i * dp_size + dp_rank]
                    endpoint = f"tcp://127.0.0.1:{port}"
                    zmq_addresses[str(dp_rank)] = endpoint

                    payload = {
                        "instance_id": new_worker_id,
                        "endpoint": endpoint,
                        "dp_rank": dp_rank,
397
398
399
400
                        "model_name": self.model_name,
                        "block_size": self._mocker_args_orig.get(
                            "block_size", BLOCK_SIZE
                        ),
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
                    }
                    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}"
            )

        logger.info(
            f"All {self.num_workers} mockers launched and registered with indexer"
        )

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
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
    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_str, zmq_endpoint in zmq_addresses.items():
                worker_entries.append(f"{worker_id}:{dp_rank_str}={zmq_endpoint}")
        workers_arg = ",".join(worker_entries)

        indexer_b_cmd = [
            "cargo",
            "run",
            "-p",
            "dynamo-kv-router",
            "--features",
            "indexer-bin",
            "--bin",
            "dynamo-kv-indexer",
            "--",
            "--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,
            display_output=True,
            health_check_ports=[self._standalone_indexer_b_port],
            health_check_urls=[],
            log_dir=self._request.node.name,
            terminate_all_matching_process_names=False,
            display_name="dynamo-kv-indexer-b",
        )
        logger.info(
            f"Starting standalone indexer B on port {self._standalone_indexer_b_port} "
            f"with peer http://localhost:{self._standalone_indexer_port}"
        )
        self._indexer_b_process.__enter__()

479
    def __exit__(self, exc_type, exc_val, exc_tb):
480
481
482
483
484
485
486
487
        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()
488
489
490
491
492
493
494
495
        # 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
496
497
498
499
500
501
502
503
504
        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)
        if self._process is not None:
            self._process.__exit__(exc_type, exc_val, exc_tb)
505
506
507
508
        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 = []
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528


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",
529
        request_plane: str = "nats",
530
        enable_bootstrap: bool = False,
531
532
533
534
    ):
        if worker_type not in ("prefill", "decode"):
            raise ValueError(
                f"worker_type must be 'prefill' or 'decode', got {worker_type}"
535
            )
536
537
538
539

        self.namespace = namespace
        self.worker_type = worker_type
        self.num_workers = num_mockers
540
        self._bootstrap_ports: list[int] = []
541
542
543
544
545
546
547
548
549

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

550
551
552
553
        mocker_args = (mocker_args or {}).copy()

        # Allocate bootstrap ports for prefill workers if enabled (one per worker)
        if enable_bootstrap and worker_type == "prefill":
554
            self._bootstrap_ports = allocate_ports(num_mockers, BASE_PORT_BOOTSTRAP)
555
556
557
558
559
560
            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"
            )
561
562
563
564
565
566
567
568
569

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

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

573
574
        self._process = ManagedProcess(
            command=command,
575
            env=env,
576
577
578
579
580
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
581
            terminate_all_matching_process_names=False,
582
583
584
585
586
        )
        logger.info(
            f"Created {worker_type} mocker process with {num_mockers} worker(s), "
            f"endpoint: {self.endpoint}"
        )
587

588
589
590
591
592
    @property
    def bootstrap_ports(self) -> list[int]:
        """Return the allocated bootstrap ports, if any."""
        return self._bootstrap_ports

593
    def __enter__(self):
594
595
596
597
        logger.info(
            f"Starting {self.worker_type} mocker process with {self.num_workers} worker(s)"
        )
        self._process.__enter__()
598
        return self
599

600
    def __exit__(self, exc_type, exc_val, exc_tb):
601
602
        logger.info(f"Stopping {self.worker_type} mocker process")
        self._process.__exit__(exc_type, exc_val, exc_tb)
603
604
605
606
607
        # 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 = []
608
609


610
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
611
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
612
@pytest.mark.parametrize(
613
    "durable_kv_events", [False], indirect=True
614
)  # Use NATS Core (local indexer)
615
def test_mocker_kv_router(
616
617
618
619
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
620
    durable_kv_events,
621
):
622
623
624
    """
    Test KV router with multiple mocker engine instances.
    This test doesn't require GPUs and runs quickly for pre-merge validation.
625
    Tests both NATS and TCP request planes.
626
627
    """

628
629
    # runtime_services starts etcd and optionally nats based on request_plane
    logger.info(f"Starting mocker KV router test with request_plane={request_plane}")
630

631
632
633
634
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
635
        "durable_kv_events": durable_kv_events,
636
    }
637

638
639
640
641
642
643
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
644
        # Start mocker instances with the new CLI interface
645
646
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
647

648
        # Get unique port for this test
649
650
651
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
652
653

        # Run basic router test (starts router internally and waits for workers to be ready)
654
655
656
657
        _test_router_basic(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
658
            frontend_port=frontend_port,
659
660
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
661
            request_plane=request_plane,
662
663
664
        )


665
@pytest.mark.parametrize("store_backend", ["etcd", "file"])
666
@pytest.mark.parametrize(
667
    "durable_kv_events", [False], indirect=True
668
)  # Use NATS Core (local indexer)
669
@pytest.mark.timeout(180)  # bumped for xdist contention (was 60s; ~19.86s serial avg)
670
671
def test_mocker_two_kv_router(
    request,
672
    runtime_services_dynamic_ports,
673
674
675
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
676
    durable_kv_events,
677
):
678
679
680
    """
    Test with two KV routers and multiple mocker engine instances.
    Alternates requests between the two routers to test load distribution.
681
    Tests with both etcd and file storage backends.
682
683
684
    """

    # runtime_services starts etcd and nats
685
686
687
    logger.info(
        f"Starting mocker two KV router test with {store_backend} storage backend"
    )
688

689
690
691
692
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
693
        "durable_kv_events": durable_kv_events,
694
    }
695

696
697
698
699
700
701
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
    ) as mockers:
702
        # Start mocker instances with the new CLI interface
703
704
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
705

706
707
708
709
710
        # Get unique ports for this test (2 ports for two routers)
        router_ports = get_unique_ports(
            request, num_ports=2, store_backend=store_backend
        )

711
712
713
714
715
        # 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,
716
            router_ports=router_ports,
717
718
719
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            store_backend=store_backend,
720
            skip_consumer_verification=not durable_kv_events,  # Skip JetStream checks in NATS Core mode
721
722
723
        )


724
@pytest.mark.skip(reason="Flaky, temporarily disabled")
725
@pytest.mark.parametrize(
726
    "durable_kv_events", [False], indirect=True
727
)  # Use NATS Core (local indexer)
728
@pytest.mark.timeout(60)  # ~3x average (~19.86s), rounded up (when enabled)
Alec's avatar
Alec committed
729
def test_mocker_kv_router_overload_503(
730
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
731
):
732
    """Test that KV router returns 503 when mocker workers are overloaded."""
733
    logger.info("Starting mocker KV router overload test for 503 status")
734
    # Create mocker args dictionary with limited resources - use local indexer (NATS Core mode)
735
736
737
738
    mocker_args = {
        "speedup_ratio": 10,
        "block_size": 4,  # Smaller block size
        "num_gpu_blocks": 64,  # Limited GPU blocks to exhaust quickly
739
        "durable_kv_events": durable_kv_events,
740
    }
741

742
    with MockerProcess(request, mocker_args=mocker_args, num_mockers=1) as mockers:
743
        # Start single mocker instance with limited resources
744
745
        logger.info("Starting single mocker instance with limited resources")
        logger.info(f"Mocker using endpoint: {mockers.endpoint}")
746

747
748
749
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

750
751
752
753
754
755
756
        # 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,
757
            blocks_threshold=0.2,
758
        )
759

760

761
@pytest.mark.timeout(90)  # bumped for xdist contention (was 22s; ~7.10s serial avg)
762
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
763
@pytest.mark.parametrize(
764
    "durable_kv_events", [False], indirect=True
765
)  # Use NATS Core (local indexer)
766
def test_kv_router_bindings(
767
768
769
770
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
771
    durable_kv_events,
772
):
773
774
    """Test KvRouter Python bindings with mocker engines."""
    logger.info("Starting KvRouter bindings test")
775
776
777
778
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
779
        "durable_kv_events": durable_kv_events,
780
    }
781

782
783
784
785
786
787
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
788
        # Start mocker instances
789
790
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
791

792
        # Get runtime and create endpoint
793
        runtime = get_runtime(request_plane=request_plane)
794
795
796
        endpoint = runtime.endpoint(
            f"{mockers.namespace}.{mockers.component_name}.generate"
        )
797

798
799
800
        # Run Python router bindings test
        _test_python_router_bindings(
            engine_workers=mockers,
801
802
            endpoint=endpoint,
            block_size=BLOCK_SIZE,
803
804
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
805
        )
806

807

808
@pytest.mark.parametrize(
809
    "store_backend,durable_kv_events,request_plane",
810
    [
811
812
813
        ("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
814
    ],
815
816
    ids=[
        "jetstream",
817
        "nats_core",
818
        "file",
819
    ],
820
    indirect=["request_plane", "durable_kv_events"],
821
)
822
@pytest.mark.timeout(180)
823
824
def test_indexers_sync(
    request,
825
    runtime_services_dynamic_ports,
826
    predownload_tokenizers,
827
    build_kv_indexer,
828
829
    file_storage_backend,
    store_backend,
830
    durable_kv_events,
831
    request_plane,
832
):
833
834
835
836
    """
    Test that two KV routers have synchronized indexer states after processing requests.
    This test verifies that both routers converge to the same internal state.

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

847
848
849
850
    # 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
851
    # Use 2 DP ranks to test per-dp_rank event ID tracking and recovery
852
853
854
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
855
        "durable_kv_events": durable_kv_events,
856
        "dp_size": 2,
857
858
    }

859
860
861
862
863
864
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
        request_plane=request_plane,
865
866
867
        zmq_kv_events=True,
        standalone_indexer=True,
        model_name=MODEL_NAME,
868
    ) as mockers:
869
870
        # 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")
871
872
873
874
        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
875
        # When using durable_kv_events=True, use JetStream mode for the router
876
877
878
879
880
881
882
        _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,
883
884
885
            test_nats_interruption=not durable_kv_events,
            nats_server=nats_process if not durable_kv_events else None,
            durable_kv_events=durable_kv_events,
886
            standalone_indexer_url=mockers.standalone_indexer_url,
887
            standalone_indexer_b_url=mockers.standalone_indexer_b_url,
888
889
890
891
892
        )

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


893
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
894
@pytest.mark.parametrize(
895
    "durable_kv_events", [False], indirect=True
896
)  # Use NATS Core (local indexer)
Alec's avatar
Alec committed
897
def test_query_instance_id_returns_worker_and_tokens(
898
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
899
):
900
    """Test query_instance_id annotation with mocker engines."""
901
    logger.info("Starting KV router query_instance_id annotation test")
902
903
904
905
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
906
        "durable_kv_events": durable_kv_events,
907
    }
908

909
910
911
    with MockerProcess(
        request, mocker_args=mocker_args, num_mockers=NUM_MOCKERS
    ) as mockers:
912
        # Start mocker instances
913
914
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
915

916
917
918
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

919
920
921
922
923
924
925
926
        # 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,
        )
927

928

929
@pytest.mark.timeout(90)  # bumped for xdist contention (was 29s; ~9.55s serial avg)
Yan Ru Pei's avatar
Yan Ru Pei committed
930
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
931
@pytest.mark.parametrize(
932
    "durable_kv_events,use_kv_events,zmq_kv_events",
933
    [
934
935
936
937
        (True, True, False),  # JetStream mode with KV events
        (False, True, False),  # NATS Core mode with local indexer (default)
        (False, False, False),  # Approximate mode (--no-kv-events) - no KV events
        (False, True, True),  # ZMQ mode: mocker → ZMQ PUB → relay → NATS
938
    ],
939
    ids=["jetstream", "nats_core", "no_kv_events", "zmq"],
940
    indirect=["durable_kv_events"],
941
)
942
def test_router_decisions(
943
944
945
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
946
    build_kv_indexer,
947
    durable_kv_events,
948
    use_kv_events,
949
    request_plane,
950
    zmq_kv_events,
951
952
953
):
    """Validate KV cache prefix reuse and dp_rank routing by sending progressive requests with overlapping prefixes.

954
    Parameterized to test:
955
956
    - JetStream mode: KV events via NATS JetStream (durable)
    - NATS Core mode (default): KV events via NATS Core with local indexer on workers
957
958
    - Approximate mode (--no-kv-events): No KV events, router predicts cache state
      based on routing decisions with TTL-based expiration and pruning
959
    """
960
    # runtime_services_dynamic_ports handles NATS and etcd startup
961
    logger.info(
962
        f"Starting test router decisions: durable_kv_events={durable_kv_events}, use_kv_events={use_kv_events}"
963
    )
964

Yan Ru Pei's avatar
Yan Ru Pei committed
965
    # Create mocker args dictionary with dp_size=4
966
    # durable_kv_events=True enables JetStream mode; False (default) uses NATS Core with local indexer
Yan Ru Pei's avatar
Yan Ru Pei committed
967
968
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
969
        "block_size": 8,
Yan Ru Pei's avatar
Yan Ru Pei committed
970
        "dp_size": 4,
971
        "durable_kv_events": durable_kv_events and use_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
972
    }
973

974
975
976
977
978
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
979
        zmq_kv_events=zmq_kv_events,
980
        standalone_indexer=zmq_kv_events,
981
        model_name=MODEL_NAME,
982
    ) as mockers:
983
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
984

985
986
        # Initialize mockers
        # Get runtime and create endpoint
987
        runtime = get_runtime(request_plane=request_plane)
988
        # Use the namespace from the mockers
989
        endpoint = runtime.endpoint(f"{mockers.namespace}.mocker.generate")
990

991
        _test_router_decisions(
992
993
994
995
996
997
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=use_kv_events,
998
            durable_kv_events=durable_kv_events,
999
            standalone_indexer_url=mockers.standalone_indexer_url,
1000
1001
        )

1002

1003
@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"])
1004
1005
1006
@pytest.mark.parametrize(
    "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"]
)
1007
@pytest.mark.timeout(180)  # bumped for xdist contention (was 59s; ~19.51s serial avg)
1008
def test_router_decisions_disagg(
1009
1010
1011
1012
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    registration_order,
1013
    enable_disagg_bootstrap,
1014
1015
1016
1017
1018
):
    """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.
1019

1020
1021
1022
    Parameterized to test:
    - registration_order: prefill_first vs decode_first
    - enable_disagg_bootstrap: without vs with bootstrap rendezvous
1023
    """
1024
    # runtime_services_dynamic_ports handles NATS and etcd startup
1025
1026
    logger.info(
        f"Starting disaggregated router prefix reuse test "
1027
        f"(registration_order={registration_order}, bootstrap={enable_disagg_bootstrap})"
1028
    )
1029
1030
1031
1032
1033

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

1034
    # Create mocker args - use NATS Core with local indexer (default mode)
1035
1036
1037
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1038
        # durable_kv_events defaults to False (NATS Core mode)
1039
    }
1040

1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
    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:
1053
1054
1055
1056
            logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

            # Then start decode workers
            logger.info("Starting 4 decode mocker instances (second)")
1057
            with DisaggMockerProcess(
1058
1059
1060
1061
1062
                request,
                namespace=shared_namespace,
                worker_type="decode",
                mocker_args=mocker_args,
                num_mockers=4,
1063
                request_plane="nats",
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
            ) 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,
            request_plane="nats",
        ) as decode_workers:
1093
1094
1095
1096
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")

            # Then start prefill workers
            logger.info("Starting 4 prefill mocker instances (second)")
1097
            with DisaggMockerProcess(
1098
1099
1100
1101
1102
                request,
                namespace=shared_namespace,
                worker_type="prefill",
                mocker_args=mocker_args,
                num_mockers=4,
1103
1104
                request_plane="nats",
                enable_bootstrap=enable_disagg_bootstrap,
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
            ) 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",
                )
1125
1126


1127
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
1128
@pytest.mark.parametrize(
1129
    "durable_kv_events", [False], indirect=True
1130
)  # Use NATS Core (local indexer)
1131
@pytest.mark.timeout(120)  # bumped for xdist contention (was 39s; ~12.84s serial avg)
1132
def test_busy_threshold_endpoint(
1133
1134
1135
1136
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
1137
    durable_kv_events,
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
):
    """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.
    """
1148
    # runtime_services_dynamic_ports handles NATS and etcd startup
1149
1150
1151
    logger.info(
        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
    )
1152

1153
1154
1155
1156
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1157
        "durable_kv_events": durable_kv_events,
1158
    }
1159

1160
1161
1162
1163
1164
1165
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
1166
1167
1168
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")

1169
1170
1171
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
1172
1173
1174
1175
1176
1177
1178

        _test_busy_threshold_endpoint(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
1179
            request_plane=request_plane,
1180
        )