test_router_e2e_with_mockers.py 33.5 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
16
import logging
import os
17
from typing import Any, Dict, Optional
18
19
20

import pytest

21
from tests.router.common import (  # utilities
22
    _test_busy_threshold_endpoint,
23
24
25
    _test_python_router_bindings,
    _test_router_basic,
    _test_router_decisions,
26
    _test_router_decisions_disagg,
27
28
29
30
31
32
33
    _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
34
from tests.utils.constants import ROUTER_MODEL_NAME
35
from tests.utils.managed_process import ManagedProcess
36
from tests.utils.port_utils import allocate_ports, deallocate_ports
37

38
39
40
41
logger = logging.getLogger(__name__)

MODEL_NAME = ROUTER_MODEL_NAME

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


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

66
67
68
    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.
69

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


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

94

95
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
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,
122
        "--discovery-backend",
123
124
125
126
127
128
129
        store_backend,
        "--num-workers",
        str(num_workers),
    ]

    # Add worker type flag for disaggregated mode
    if worker_type == "prefill":
130
        command.extend(["--disaggregation-mode", "prefill"])
131
    elif worker_type == "decode":
132
        command.extend(["--disaggregation-mode", "decode"])
133
134
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
161
162

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

    return command


174
class MockerProcess:
175
    """Manages mocker engine instances with shared tokio runtime via --num-workers."""
176

177
178
179
180
181
    def __init__(
        self,
        request,
        mocker_args: Optional[Dict[str, Any]] = None,
        num_mockers: int = 1,
182
        store_backend: str = "etcd",
183
        request_plane: str = "nats",
184
        zmq_kv_events: bool = False,
185
    ):
186
187
        namespace_suffix = generate_random_suffix()
        self.namespace = f"test-namespace-{namespace_suffix}"
188
189
        self.component_name = "mocker"
        self.endpoint = f"dyn://{self.namespace}.{self.component_name}.generate"
190
        self.num_workers = num_mockers
191
        self._zmq_kv_events_ports: list[int] = []
192

193
        mocker_args = (mocker_args or {}).copy()
194
195
        # Store dp_size for DP-aware test functions
        self.dp_size = mocker_args.get("dp_size")
196
197
        # Alias for consistency with vLLM/SGLang workers
        self.data_parallel_size = self.dp_size
198

199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
        # 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)]
            mocker_args["zmq_kv_events_ports"] = ",".join(str(p) for p in bases)
            logger.info(
                f"Allocated ZMQ KV event ports {self._zmq_kv_events_ports} "
                f"(bases: {bases}) for {num_mockers} workers"
            )

215
216
217
218
219
220
221
        command = _build_mocker_command(
            endpoint=self.endpoint,
            store_backend=store_backend,
            num_workers=num_mockers,
            mocker_args=mocker_args,
        )

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

225
226
        self._process = ManagedProcess(
            command=command,
227
            env=env,
228
229
230
231
232
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
233
            terminate_all_matching_process_names=False,
234
235
236
237
238
239
240
241
242
243
244
245
246
        )
        logger.info(
            f"Created mocker process with {num_mockers} worker(s), endpoint: {self.endpoint}"
        )

    def __enter__(self):
        logger.info(f"Starting mocker process with {self.num_workers} worker(s)")
        self._process.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        logger.info("Stopping mocker process")
        self._process.__exit__(exc_type, exc_val, exc_tb)
247
248
249
250
        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 = []
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270


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",
271
        request_plane: str = "nats",
272
        enable_bootstrap: bool = False,
273
274
275
276
    ):
        if worker_type not in ("prefill", "decode"):
            raise ValueError(
                f"worker_type must be 'prefill' or 'decode', got {worker_type}"
277
            )
278
279
280
281

        self.namespace = namespace
        self.worker_type = worker_type
        self.num_workers = num_mockers
282
        self._bootstrap_ports: list[int] = []
283
284
285
286
287
288
289
290
291

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

292
293
294
295
        mocker_args = (mocker_args or {}).copy()

        # Allocate bootstrap ports for prefill workers if enabled (one per worker)
        if enable_bootstrap and worker_type == "prefill":
296
            self._bootstrap_ports = allocate_ports(num_mockers, BASE_PORT_BOOTSTRAP)
297
298
299
300
301
302
            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"
            )
303
304
305
306
307
308
309
310
311

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

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

315
316
        self._process = ManagedProcess(
            command=command,
317
            env=env,
318
319
320
321
322
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
323
            terminate_all_matching_process_names=False,
324
325
326
327
328
        )
        logger.info(
            f"Created {worker_type} mocker process with {num_mockers} worker(s), "
            f"endpoint: {self.endpoint}"
        )
329

330
331
332
333
334
    @property
    def bootstrap_ports(self) -> list[int]:
        """Return the allocated bootstrap ports, if any."""
        return self._bootstrap_ports

335
    def __enter__(self):
336
337
338
339
        logger.info(
            f"Starting {self.worker_type} mocker process with {self.num_workers} worker(s)"
        )
        self._process.__enter__()
340
        return self
341

342
    def __exit__(self, exc_type, exc_val, exc_tb):
343
344
        logger.info(f"Stopping {self.worker_type} mocker process")
        self._process.__exit__(exc_type, exc_val, exc_tb)
345
346
347
348
349
        # 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 = []
350
351


352
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
353
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
354
@pytest.mark.parametrize(
355
    "durable_kv_events", [False], indirect=True
356
)  # Use NATS Core (local indexer)
357
def test_mocker_kv_router(
358
359
360
361
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
362
    durable_kv_events,
363
):
364
365
366
    """
    Test KV router with multiple mocker engine instances.
    This test doesn't require GPUs and runs quickly for pre-merge validation.
367
    Tests both NATS and TCP request planes.
368
369
    """

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

373
374
375
376
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
377
        "durable_kv_events": durable_kv_events,
378
    }
379

380
381
382
383
384
385
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
386
        # Start mocker instances with the new CLI interface
387
388
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
389

390
        # Get unique port for this test
391
392
393
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
394
395

        # Run basic router test (starts router internally and waits for workers to be ready)
396
397
398
399
        _test_router_basic(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
400
            frontend_port=frontend_port,
401
402
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
403
            request_plane=request_plane,
404
405
406
        )


407
@pytest.mark.parametrize("store_backend", ["etcd", "file"])
408
@pytest.mark.parametrize(
409
    "durable_kv_events", [False], indirect=True
410
)  # Use NATS Core (local indexer)
411
@pytest.mark.timeout(180)  # bumped for xdist contention (was 60s; ~19.86s serial avg)
412
413
def test_mocker_two_kv_router(
    request,
414
    runtime_services_dynamic_ports,
415
416
417
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
418
    durable_kv_events,
419
):
420
421
422
    """
    Test with two KV routers and multiple mocker engine instances.
    Alternates requests between the two routers to test load distribution.
423
    Tests with both etcd and file storage backends.
424
425
426
    """

    # runtime_services starts etcd and nats
427
428
429
    logger.info(
        f"Starting mocker two KV router test with {store_backend} storage backend"
    )
430

431
432
433
434
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
435
        "durable_kv_events": durable_kv_events,
436
    }
437

438
439
440
441
442
443
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
    ) as mockers:
444
        # Start mocker instances with the new CLI interface
445
446
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
447

448
449
450
451
452
        # Get unique ports for this test (2 ports for two routers)
        router_ports = get_unique_ports(
            request, num_ports=2, store_backend=store_backend
        )

453
454
455
456
457
        # 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,
458
            router_ports=router_ports,
459
460
461
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            store_backend=store_backend,
462
            skip_consumer_verification=not durable_kv_events,  # Skip JetStream checks in NATS Core mode
463
464
465
        )


466
@pytest.mark.skip(reason="Flaky, temporarily disabled")
467
@pytest.mark.parametrize(
468
    "durable_kv_events", [False], indirect=True
469
)  # Use NATS Core (local indexer)
470
@pytest.mark.timeout(60)  # ~3x average (~19.86s), rounded up (when enabled)
Alec's avatar
Alec committed
471
def test_mocker_kv_router_overload_503(
472
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
473
):
474
    """Test that KV router returns 503 when mocker workers are overloaded."""
475
    logger.info("Starting mocker KV router overload test for 503 status")
476
    # Create mocker args dictionary with limited resources - use local indexer (NATS Core mode)
477
478
479
480
    mocker_args = {
        "speedup_ratio": 10,
        "block_size": 4,  # Smaller block size
        "num_gpu_blocks": 64,  # Limited GPU blocks to exhaust quickly
481
        "durable_kv_events": durable_kv_events,
482
    }
483

484
    with MockerProcess(request, mocker_args=mocker_args, num_mockers=1) as mockers:
485
        # Start single mocker instance with limited resources
486
487
        logger.info("Starting single mocker instance with limited resources")
        logger.info(f"Mocker using endpoint: {mockers.endpoint}")
488

489
490
491
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

492
493
494
495
496
497
498
        # 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,
499
            blocks_threshold=0.2,
500
        )
501

502

503
@pytest.mark.timeout(90)  # bumped for xdist contention (was 22s; ~7.10s serial avg)
504
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
505
@pytest.mark.parametrize(
506
    "durable_kv_events", [False], indirect=True
507
)  # Use NATS Core (local indexer)
508
def test_kv_router_bindings(
509
510
511
512
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
513
    durable_kv_events,
514
):
515
516
    """Test KvRouter Python bindings with mocker engines."""
    logger.info("Starting KvRouter bindings test")
517
518
519
520
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
521
        "durable_kv_events": durable_kv_events,
522
    }
523

524
525
526
527
528
529
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
530
        # Start mocker instances
531
532
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
533

534
        # Get runtime and create endpoint
535
        runtime = get_runtime(request_plane=request_plane)
536
537
538
        endpoint = runtime.endpoint(
            f"{mockers.namespace}.{mockers.component_name}.generate"
        )
539

540
541
542
        # Run Python router bindings test
        _test_python_router_bindings(
            engine_workers=mockers,
543
544
            endpoint=endpoint,
            block_size=BLOCK_SIZE,
545
546
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
547
        )
548

549

550
@pytest.mark.parametrize(
Yan Ru Pei's avatar
Yan Ru Pei committed
551
    "store_backend,durable_kv_events,request_plane,router_event_threads",
552
    [
Yan Ru Pei's avatar
Yan Ru Pei committed
553
554
555
556
        ("etcd", True, "nats", 1),  # JetStream mode - uses JetStream
        ("etcd", False, "tcp", 1),  # NATS core mode (with gap detection) - no JetStream
        ("file", True, "nats", 1),  # File backend - uses JetStream
        ("etcd", False, "tcp", 2),  # NATS core mode - multi-threaded indexer
557
    ],
558
559
    ids=[
        "jetstream",
560
        "nats_core",
561
        "file",
Yan Ru Pei's avatar
Yan Ru Pei committed
562
        "nats_core_multi_thread",
563
    ],
564
    indirect=["request_plane", "durable_kv_events"],
565
)
566
@pytest.mark.timeout(180)
567
568
def test_indexers_sync(
    request,
569
    runtime_services_dynamic_ports,
570
571
572
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
573
    durable_kv_events,
574
    request_plane,
Yan Ru Pei's avatar
Yan Ru Pei committed
575
    router_event_threads,
576
):
577
578
579
580
    """
    Test that two KV routers have synchronized indexer states after processing requests.
    This test verifies that both routers converge to the same internal state.

581
582
583
584
585
586
587
588
    Tests with three configurations:
    - jetstream: etcd backend, JetStream for KV events, NATS request plane
    - nats_core: etcd backend, local indexer with NATS Core, TCP request plane
                 (includes NATS interruption/recovery testing)
    - file: file backend, JetStream for KV events, NATS request plane
    """
    logger.info(
        f"Starting indexers sync test: store_backend={store_backend}, "
589
        f"durable_kv_events={durable_kv_events}, request_plane={request_plane}"
590
    )
591

592
593
594
595
    # 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
596
    # Use 2 DP ranks to test per-dp_rank event ID tracking and recovery
597
598
599
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
600
        "durable_kv_events": durable_kv_events,
601
        "dp_size": 2,
602
603
    }

604
605
606
607
608
609
610
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
        request_plane=request_plane,
    ) as mockers:
611
612
        # 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")
613
614
615
616
        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
617
        # When using durable_kv_events=True, use JetStream mode for the router
618
619
620
621
622
623
624
        _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,
625
626
627
            test_nats_interruption=not durable_kv_events,
            nats_server=nats_process if not durable_kv_events else None,
            durable_kv_events=durable_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
628
            router_event_threads=router_event_threads,
629
630
631
632
633
        )

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


634
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
635
@pytest.mark.parametrize(
636
    "durable_kv_events", [False], indirect=True
637
)  # Use NATS Core (local indexer)
Alec's avatar
Alec committed
638
def test_query_instance_id_returns_worker_and_tokens(
639
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
640
):
641
    """Test query_instance_id annotation with mocker engines."""
642
    logger.info("Starting KV router query_instance_id annotation test")
643
644
645
646
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
647
        "durable_kv_events": durable_kv_events,
648
    }
649

650
651
652
    with MockerProcess(
        request, mocker_args=mocker_args, num_mockers=NUM_MOCKERS
    ) as mockers:
653
        # Start mocker instances
654
655
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
656

657
658
659
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

660
661
662
663
664
665
666
667
        # 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,
        )
668

669

670
@pytest.mark.timeout(90)  # bumped for xdist contention (was 29s; ~9.55s serial avg)
Yan Ru Pei's avatar
Yan Ru Pei committed
671
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
672
@pytest.mark.parametrize(
673
    "durable_kv_events,use_kv_events,router_event_threads,zmq_kv_events",
674
    [
675
676
677
678
679
        (True, True, 1, False),  # JetStream mode with KV events
        (False, True, 1, False),  # NATS Core mode with local indexer (default)
        (False, False, 1, False),  # Approximate mode (--no-kv-events) - no KV events
        (False, True, 2, False),  # NATS Core mode - multi-threaded indexer
        (False, True, 1, True),  # ZMQ mode: mocker → ZMQ PUB → relay → NATS
680
    ],
681
    ids=["jetstream", "nats_core", "no_kv_events", "nats_core_multi_thread", "zmq"],
682
    indirect=["durable_kv_events"],
683
)
684
def test_router_decisions(
685
686
687
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
688
    durable_kv_events,
689
    use_kv_events,
690
    request_plane,
Yan Ru Pei's avatar
Yan Ru Pei committed
691
    router_event_threads,
692
    zmq_kv_events,
693
694
695
):
    """Validate KV cache prefix reuse and dp_rank routing by sending progressive requests with overlapping prefixes.

696
    Parameterized to test:
697
698
    - JetStream mode: KV events via NATS JetStream (durable)
    - NATS Core mode (default): KV events via NATS Core with local indexer on workers
699
700
    - Approximate mode (--no-kv-events): No KV events, router predicts cache state
      based on routing decisions with TTL-based expiration and pruning
701
    """
702
    # runtime_services_dynamic_ports handles NATS and etcd startup
703
    logger.info(
704
        f"Starting test router decisions: durable_kv_events={durable_kv_events}, use_kv_events={use_kv_events}"
705
    )
706

Yan Ru Pei's avatar
Yan Ru Pei committed
707
    # Create mocker args dictionary with dp_size=4
708
    # durable_kv_events=True enables JetStream mode; False (default) uses NATS Core with local indexer
Yan Ru Pei's avatar
Yan Ru Pei committed
709
710
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
711
        "block_size": 8,
Yan Ru Pei's avatar
Yan Ru Pei committed
712
        "dp_size": 4,
713
        "durable_kv_events": durable_kv_events and use_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
714
    }
715

716
717
718
719
720
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
721
        zmq_kv_events=zmq_kv_events,
722
    ) as mockers:
723
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
724

725
726
        # Initialize mockers
        # Get runtime and create endpoint
727
        runtime = get_runtime(request_plane=request_plane)
728
        # Use the namespace from the mockers
729
        endpoint = runtime.endpoint(f"{mockers.namespace}.mocker.generate")
730

731
        _test_router_decisions(
732
733
734
735
736
737
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=use_kv_events,
738
            durable_kv_events=durable_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
739
            router_event_threads=router_event_threads,
740
741
        )

742

743
@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"])
744
745
746
@pytest.mark.parametrize(
    "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"]
)
747
@pytest.mark.timeout(180)  # bumped for xdist contention (was 59s; ~19.51s serial avg)
748
def test_router_decisions_disagg(
749
750
751
752
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    registration_order,
753
    enable_disagg_bootstrap,
754
755
756
757
758
):
    """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.
759

760
761
762
    Parameterized to test:
    - registration_order: prefill_first vs decode_first
    - enable_disagg_bootstrap: without vs with bootstrap rendezvous
763
    """
764
    # runtime_services_dynamic_ports handles NATS and etcd startup
765
766
    logger.info(
        f"Starting disaggregated router prefix reuse test "
767
        f"(registration_order={registration_order}, bootstrap={enable_disagg_bootstrap})"
768
    )
769
770
771
772
773

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

774
    # Create mocker args - use NATS Core with local indexer (default mode)
775
776
777
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
778
        # durable_kv_events defaults to False (NATS Core mode)
779
    }
780

781
782
783
784
785
786
787
788
789
790
791
792
    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:
793
794
795
796
            logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

            # Then start decode workers
            logger.info("Starting 4 decode mocker instances (second)")
797
            with DisaggMockerProcess(
798
799
800
801
802
                request,
                namespace=shared_namespace,
                worker_type="decode",
                mocker_args=mocker_args,
                num_mockers=4,
803
                request_plane="nats",
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
            ) 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:
833
834
835
836
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")

            # Then start prefill workers
            logger.info("Starting 4 prefill mocker instances (second)")
837
            with DisaggMockerProcess(
838
839
840
841
842
                request,
                namespace=shared_namespace,
                worker_type="prefill",
                mocker_args=mocker_args,
                num_mockers=4,
843
844
                request_plane="nats",
                enable_bootstrap=enable_disagg_bootstrap,
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
            ) 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",
                )
865
866


867
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
868
@pytest.mark.parametrize(
869
    "durable_kv_events", [False], indirect=True
870
)  # Use NATS Core (local indexer)
871
@pytest.mark.timeout(120)  # bumped for xdist contention (was 39s; ~12.84s serial avg)
872
def test_busy_threshold_endpoint(
873
874
875
876
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
877
    durable_kv_events,
878
879
880
881
882
883
884
885
886
887
):
    """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.
    """
888
    # runtime_services_dynamic_ports handles NATS and etcd startup
889
890
891
    logger.info(
        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
    )
892

893
894
895
896
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
897
        "durable_kv_events": durable_kv_events,
898
    }
899

900
901
902
903
904
905
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
906
907
908
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")

909
910
911
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
912
913
914
915
916
917
918

        _test_busy_threshold_endpoint(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
919
            request_plane=request_plane,
920
        )