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

# 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.parallel,
47
    pytest.mark.model(MODEL_NAME),
48
]
49
50
NUM_MOCKERS = 2
SPEEDUP_RATIO = 10.0
51
BASE_PORT = 9100  # Base port for all tests (high port to avoid conflicts)
52
NUM_REQUESTS = 100
53
BLOCK_SIZE = 16
54
55


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

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

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


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

93

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
122
123
124
125
126
127
128
129
130
131
132
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
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,
        "--store-kv",
        store_backend,
        "--num-workers",
        str(num_workers),
    ]

    # Add worker type flag for disaggregated mode
    if worker_type == "prefill":
        command.append("--is-prefill-worker")
    elif worker_type == "decode":
        command.append("--is-decode-worker")

    # 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"])])
162
163
    if mocker_args.get("enable_local_indexer"):
        command.append("--enable-local-indexer")
164
165
    if "bootstrap_ports" in mocker_args:
        command.extend(["--bootstrap-ports", mocker_args["bootstrap_ports"]])
166
167
168
169

    return command


170
class MockerProcess:
171
    """Manages mocker engine instances with shared tokio runtime via --num-workers."""
172

173
174
175
176
177
    def __init__(
        self,
        request,
        mocker_args: Optional[Dict[str, Any]] = None,
        num_mockers: int = 1,
178
        store_backend: str = "etcd",
179
        request_plane: str = "nats",
180
    ):
181
182
        namespace_suffix = generate_random_suffix()
        self.namespace = f"test-namespace-{namespace_suffix}"
183
184
        self.component_name = "mocker"
        self.endpoint = f"dyn://{self.namespace}.{self.component_name}.generate"
185
186
187
        self.num_workers = num_mockers

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

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

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

203
204
        self._process = ManagedProcess(
            command=command,
205
            env=env,
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
            terminate_existing=False,
        )
        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)


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",
245
        request_plane: str = "nats",
246
        enable_bootstrap: bool = False,
247
248
249
250
    ):
        if worker_type not in ("prefill", "decode"):
            raise ValueError(
                f"worker_type must be 'prefill' or 'decode', got {worker_type}"
251
            )
252
253
254
255

        self.namespace = namespace
        self.worker_type = worker_type
        self.num_workers = num_mockers
256
        self._bootstrap_ports: list[int] = []
257
258
259
260
261
262
263
264
265

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

266
267
268
269
270
271
272
273
274
275
276
        mocker_args = (mocker_args or {}).copy()

        # Allocate bootstrap ports for prefill workers if enabled (one per worker)
        if enable_bootstrap and worker_type == "prefill":
            self._bootstrap_ports = allocate_ports(num_mockers, BASE_PORT)
            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"
            )
277
278
279
280
281
282
283
284
285

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

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

289
290
        self._process = ManagedProcess(
            command=command,
291
            env=env,
292
293
294
295
296
297
298
299
300
301
302
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
            terminate_existing=False,
        )
        logger.info(
            f"Created {worker_type} mocker process with {num_mockers} worker(s), "
            f"endpoint: {self.endpoint}"
        )
303

304
305
306
307
308
    @property
    def bootstrap_ports(self) -> list[int]:
        """Return the allocated bootstrap ports, if any."""
        return self._bootstrap_ports

309
    def __enter__(self):
310
311
312
313
        logger.info(
            f"Starting {self.worker_type} mocker process with {self.num_workers} worker(s)"
        )
        self._process.__enter__()
314
        return self
315

316
    def __exit__(self, exc_type, exc_val, exc_tb):
317
318
        logger.info(f"Stopping {self.worker_type} mocker process")
        self._process.__exit__(exc_type, exc_val, exc_tb)
319
320
321
322
323
        # 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 = []
324
325


326
@pytest.mark.timeout(42)  # ~3x average (~13.80s), rounded up
327
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
328
329
330
@pytest.mark.parametrize(
    "use_nats_core", [True], indirect=True
)  # Use NATS Core (local indexer)
331
def test_mocker_kv_router(
332
333
334
335
336
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
    use_nats_core,
337
):
338
339
340
    """
    Test KV router with multiple mocker engine instances.
    This test doesn't require GPUs and runs quickly for pre-merge validation.
341
    Tests both NATS and TCP request planes.
342
343
    """

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

347
348
349
350
351
352
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": use_nats_core,
    }
353
354

    try:
355
        # Start mocker instances with the new CLI interface
356
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
357
        mockers = MockerProcess(
358
359
360
361
            request,
            mocker_args=mocker_args,
            num_mockers=NUM_MOCKERS,
            request_plane=request_plane,
362
        )
363
364
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
        mockers.__enter__()
365

366
        # Get unique port for this test
367
368
369
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
370
371

        # Run basic router test (starts router internally and waits for workers to be ready)
372
373
374
375
        _test_router_basic(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
376
            frontend_port=frontend_port,
377
378
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
379
            request_plane=request_plane,
380
381
382
        )

    finally:
383
384
        if "mockers" in locals():
            mockers.__exit__(None, None, None)
385
386


387
@pytest.mark.parametrize("store_backend", ["etcd", "file"])
388
389
390
@pytest.mark.parametrize(
    "use_nats_core", [True], indirect=True
)  # Use NATS Core (local indexer)
391
@pytest.mark.timeout(60)  # ~3x average (~19.86s), rounded up
392
393
def test_mocker_two_kv_router(
    request,
394
    runtime_services_dynamic_ports,
395
396
397
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
398
    use_nats_core,
399
):
400
401
402
    """
    Test with two KV routers and multiple mocker engine instances.
    Alternates requests between the two routers to test load distribution.
403
    Tests with both etcd and file storage backends.
404
405
406
    """

    # runtime_services starts etcd and nats
407
408
409
    logger.info(
        f"Starting mocker two KV router test with {store_backend} storage backend"
    )
410

411
412
413
414
415
416
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": use_nats_core,
    }
417
418

    try:
419
        # Start mocker instances with the new CLI interface
420
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
421
        mockers = MockerProcess(
422
423
424
425
            request,
            mocker_args=mocker_args,
            num_mockers=NUM_MOCKERS,
            store_backend=store_backend,
426
        )
427
428
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
        mockers.__enter__()
429

430
431
432
433
434
        # Get unique ports for this test (2 ports for two routers)
        router_ports = get_unique_ports(
            request, num_ports=2, store_backend=store_backend
        )

435
436
437
438
439
        # 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,
440
            router_ports=router_ports,
441
442
443
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            store_backend=store_backend,
444
            skip_consumer_verification=use_nats_core,  # Skip JetStream checks in NATS Core mode
445
446
447
        )

    finally:
448
449
        if "mockers" in locals():
            mockers.__exit__(None, None, None)
450
451


452
@pytest.mark.skip(reason="Flaky, temporarily disabled")
453
454
455
@pytest.mark.parametrize(
    "use_nats_core", [True], indirect=True
)  # Use NATS Core (local indexer)
456
@pytest.mark.timeout(60)  # ~3x average (~19.86s), rounded up (when enabled)
Alec's avatar
Alec committed
457
def test_mocker_kv_router_overload_503(
458
    request, runtime_services_dynamic_ports, predownload_tokenizers, use_nats_core
Alec's avatar
Alec committed
459
):
460
    """Test that KV router returns 503 when mocker workers are overloaded."""
461
    logger.info("Starting mocker KV router overload test for 503 status")
462
    # Create mocker args dictionary with limited resources - use local indexer (NATS Core mode)
463
464
465
466
    mocker_args = {
        "speedup_ratio": 10,
        "block_size": 4,  # Smaller block size
        "num_gpu_blocks": 64,  # Limited GPU blocks to exhaust quickly
467
        "enable_local_indexer": use_nats_core,
468
    }
469

470
    try:
471
        # Start single mocker instance with limited resources
472
        logger.info("Starting single mocker instance with limited resources")
473
        mockers = MockerProcess(request, mocker_args=mocker_args, num_mockers=1)
474
475
        logger.info(f"Mocker using endpoint: {mockers.endpoint}")
        mockers.__enter__()
476

477
478
479
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

480
481
482
483
484
485
486
        # 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,
487
            blocks_threshold=0.2,
488
        )
489
490

    finally:
491
492
        if "mockers" in locals():
            mockers.__exit__(None, None, None)
493

494

495
@pytest.mark.timeout(22)  # ~3x average (~7.10s), rounded up
496
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
497
498
499
@pytest.mark.parametrize(
    "use_nats_core", [True], indirect=True
)  # Use NATS Core (local indexer)
500
def test_kv_push_router_bindings(
501
502
503
504
505
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
    use_nats_core,
506
):
507
    """Test KvPushRouter Python bindings with mocker engines."""
508
    logger.info("Starting KvPushRouter bindings test")
509
510
511
512
513
514
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": use_nats_core,
    }
515
516

    try:
517
        # Start mocker instances
518
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
519
        mockers = MockerProcess(
520
521
522
523
            request,
            mocker_args=mocker_args,
            num_mockers=NUM_MOCKERS,
            request_plane=request_plane,
524
        )
525
526
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
        mockers.__enter__()
527

528
        # Get runtime and create endpoint
529
        runtime = get_runtime(request_plane=request_plane)
530
        namespace = runtime.namespace(mockers.namespace)
531
        component = namespace.component(mockers.component_name)
532
533
        endpoint = component.endpoint("generate")

534
535
536
        # Run Python router bindings test
        _test_python_router_bindings(
            engine_workers=mockers,
537
538
            endpoint=endpoint,
            block_size=BLOCK_SIZE,
539
540
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
541
        )
542
543

    finally:
544
545
546
547
        if "mockers" in locals():
            mockers.__exit__(None, None, None)


548
549
550
@pytest.mark.parametrize(
    "store_backend,use_nats_core,request_plane",
    [
551
552
553
        ("etcd", False, "nats"),  # JetStream mode - uses JetStream (default)
        ("etcd", True, "tcp"),  # NATS core mode (with gap detection) - no JetStream
        ("file", False, "nats"),  # File backend - uses JetStream (default)
554
    ],
555
556
    ids=[
        "jetstream",
557
        "nats_core",
558
        "file",
559
    ],
560
    indirect=["request_plane", "use_nats_core"],
561
)
562
@pytest.mark.timeout(90)  # TODO: figure out a timeout
563
564
def test_indexers_sync(
    request,
565
    runtime_services_dynamic_ports,
566
567
568
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
569
570
    use_nats_core,
    request_plane,
571
):
572
573
574
575
    """
    Test that two KV routers have synchronized indexer states after processing requests.
    This test verifies that both routers converge to the same internal state.

576
577
578
579
580
581
582
583
584
585
    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}, "
        f"use_nats_core={use_nats_core}, request_plane={request_plane}"
    )
586

587
588
589
590
    # 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
591
    # Use 2 DP ranks to test per-dp_rank event ID tracking and recovery
592
593
594
595
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": use_nats_core,
596
        "dp_size": 2,
597
598
599
    }

    try:
600
601
        # 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")
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
        mockers = MockerProcess(
            request,
            mocker_args=mocker_args,
            num_mockers=NUM_MOCKERS,
            store_backend=store_backend,
            request_plane=request_plane,
        )
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
        mockers.__enter__()

        # 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
        _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=use_nats_core,
            nats_server=nats_process if use_nats_core else None,
        )

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

    finally:
        if "mockers" in locals():
            mockers.__exit__(None, None, None)


@pytest.mark.timeout(42)  # ~3x average (~13.80s), rounded up
633
634
635
@pytest.mark.parametrize(
    "use_nats_core", [True], indirect=True
)  # Use NATS Core (local indexer)
Alec's avatar
Alec committed
636
def test_query_instance_id_returns_worker_and_tokens(
637
    request, runtime_services_dynamic_ports, predownload_tokenizers, use_nats_core
Alec's avatar
Alec committed
638
):
639
    """Test query_instance_id annotation with mocker engines."""
640
    logger.info("Starting KV router query_instance_id annotation test")
641
642
643
644
645
646
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": use_nats_core,
    }
647
648
649
    os.makedirs(request.node.name, exist_ok=True)

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

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

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

    finally:
671
672
        if "mockers" in locals():
            mockers.__exit__(None, None, None)
673
674


675
@pytest.mark.timeout(29)  # ~3x average (~9.55s), rounded up
676
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
677
678
679
@pytest.mark.parametrize(
    "use_nats_core,use_kv_events",
    [
680
681
682
        (False, True),  # JetStream mode (default) - uses JetStream
        (True, True),  # NATS Core + local indexer mode - no JetStream
        (False, False),  # Approximate mode (--no-kv-events) - uses JetStream
683
684
    ],
    ids=["jetstream", "nats_core", "no_kv_events"],
685
    indirect=["use_nats_core"],
686
)
687
def test_router_decisions(
688
689
690
691
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    use_nats_core,
692
    use_kv_events,
693
    request_plane,
694
695
696
):
    """Validate KV cache prefix reuse and dp_rank routing by sending progressive requests with overlapping prefixes.

697
698
699
700
701
    Parameterized to test:
    - JetStream mode (default): KV events via JetStream
    - NATS Core mode: KV events via NATS Core with local indexer on workers
    - Approximate mode (--no-kv-events): No KV events, router predicts cache state
      based on routing decisions with TTL-based expiration and pruning
702
    """
703
    # runtime_services_dynamic_ports handles NATS and etcd startup
704
705
706
707
708
709
    if not use_kv_events:
        mode = "Approximate (no-kv-events)"
    elif use_nats_core:
        mode = "NATS Core (local indexer)"
    else:
        mode = "JetStream"
710
711
712
    logger.info(
        f"Starting test router prefix reuse and KV events synchronization ({mode})"
    )
713

Yan Ru Pei's avatar
Yan Ru Pei committed
714
    # Create mocker args dictionary with dp_size=4
715
    # Note: enable_local_indexer only applies when use_kv_events=True and use_nats_core=True
Yan Ru Pei's avatar
Yan Ru Pei committed
716
717
718
719
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "dp_size": 4,
720
        "enable_local_indexer": use_nats_core and use_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
721
    }
722
723

    try:
Yan Ru Pei's avatar
Yan Ru Pei committed
724
        logger.info(
725
            f"Starting 2 mocker instances with dp_size=4 each (8 total dp ranks), {mode}"
726
        )
727
728
729
730
731
732
        mockers = MockerProcess(
            request,
            mocker_args=mocker_args,
            num_mockers=2,
            request_plane=request_plane,
        )
733
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
734

735
736
737
738
        # Initialize mockers
        mockers.__enter__()

        # Get runtime and create endpoint
739
        runtime = get_runtime(request_plane=request_plane)
740
741
742
743
744
        # Use the namespace from the mockers
        namespace = runtime.namespace(mockers.namespace)
        component = namespace.component("mocker")
        endpoint = component.endpoint("generate")

745
        _test_router_decisions(
746
747
748
749
750
751
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=use_kv_events,
752
753
754
755
756
        )

    finally:
        if "mockers" in locals():
            mockers.__exit__(None, None, None)
757
758


759
@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"])
760
761
762
@pytest.mark.parametrize(
    "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"]
)
763
@pytest.mark.timeout(59)  # ~3x average (~19.51s), rounded up
764
def test_router_decisions_disagg(
765
766
767
768
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    registration_order,
769
    enable_disagg_bootstrap,
770
771
772
773
774
):
    """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.
775

776
777
778
    Parameterized to test:
    - registration_order: prefill_first vs decode_first
    - enable_disagg_bootstrap: without vs with bootstrap rendezvous
779
    """
780
    # runtime_services_dynamic_ports handles NATS and etcd startup
781
782
    logger.info(
        f"Starting disaggregated router prefix reuse test "
783
        f"(registration_order={registration_order}, bootstrap={enable_disagg_bootstrap})"
784
    )
785
786
787
788
789

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

790
791
792
793
794
795
    # Create mocker args - use JetStream for KV events (more reliable than NATS Core)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": False,
    }
796
797
798
799
800

    prefill_workers = None
    decode_workers = None

    try:
801
802
803
804
805
806
807
808
809
        if registration_order == "prefill_first":
            # Start prefill workers first
            logger.info("Starting 4 prefill mocker instances (first)")
            prefill_workers = DisaggMockerProcess(
                request,
                namespace=shared_namespace,
                worker_type="prefill",
                mocker_args=mocker_args,
                num_mockers=4,
810
811
                request_plane="nats",
                enable_bootstrap=enable_disagg_bootstrap,
812
813
814
815
816
817
818
819
820
821
822
823
            )
            prefill_workers.__enter__()
            logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

            # Then start decode workers
            logger.info("Starting 4 decode mocker instances (second)")
            decode_workers = DisaggMockerProcess(
                request,
                namespace=shared_namespace,
                worker_type="decode",
                mocker_args=mocker_args,
                num_mockers=4,
824
                request_plane="nats",
825
826
827
828
829
830
831
832
833
834
835
836
            )
            decode_workers.__enter__()
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")
        else:
            # Start decode workers first
            logger.info("Starting 4 decode mocker instances (first)")
            decode_workers = DisaggMockerProcess(
                request,
                namespace=shared_namespace,
                worker_type="decode",
                mocker_args=mocker_args,
                num_mockers=4,
837
                request_plane="nats",
838
839
840
841
842
843
844
845
846
847
848
849
            )
            decode_workers.__enter__()
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")

            # Then start prefill workers
            logger.info("Starting 4 prefill mocker instances (second)")
            prefill_workers = DisaggMockerProcess(
                request,
                namespace=shared_namespace,
                worker_type="prefill",
                mocker_args=mocker_args,
                num_mockers=4,
850
851
                request_plane="nats",
                enable_bootstrap=enable_disagg_bootstrap,
852
853
854
            )
            prefill_workers.__enter__()
            logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")
855
856

        # Get unique port for this test
857
858
859
        frontend_port = get_unique_ports(
            request, num_ports=1, registration_order=registration_order
        )[0]
860
861

        # Run disagg routing test
862
        _test_router_decisions_disagg(
863
864
865
866
867
868
            prefill_workers=prefill_workers,
            decode_workers=decode_workers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
869
            request_plane="nats",
870
871
872
873
874
875
876
        )

    finally:
        if decode_workers is not None:
            decode_workers.__exit__(None, None, None)
        if prefill_workers is not None:
            prefill_workers.__exit__(None, None, None)
877
878


879
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
880
881
882
@pytest.mark.parametrize(
    "use_nats_core", [True], indirect=True
)  # Use NATS Core (local indexer)
883
@pytest.mark.timeout(39)  # ~3x average (~12.84s), rounded up
884
def test_busy_threshold_endpoint(
885
886
887
888
889
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
    use_nats_core,
890
891
892
893
894
895
896
897
898
899
):
    """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.
    """
900
    # runtime_services_dynamic_ports handles NATS and etcd startup
901
902
903
    logger.info(
        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
    )
904

905
906
907
908
909
910
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "enable_local_indexer": use_nats_core,
    }
911
912
913
914

    try:
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        mockers = MockerProcess(
915
916
917
918
            request,
            mocker_args=mocker_args,
            num_mockers=NUM_MOCKERS,
            request_plane=request_plane,
919
920
921
922
        )
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
        mockers.__enter__()

923
924
925
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
926
927
928
929
930
931
932

        _test_busy_threshold_endpoint(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
933
            request_plane=request_plane,
934
935
936
937
938
        )

    finally:
        if "mockers" in locals():
            mockers.__exit__(None, None, None)