test_router_e2e_with_mockers.py 32.1 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
BASE_PORT = 9100  # Base port for all tests (high port to avoid conflicts)
51
NUM_REQUESTS = 100
52
BLOCK_SIZE = 16
53
54


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

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

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


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

92

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
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"])])
161
162
163
    # 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")
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
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
211
            terminate_all_matching_process_names=False,
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
        )
        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
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
297
            terminate_all_matching_process_names=False,
298
299
300
301
302
        )
        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(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
327
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
328
@pytest.mark.parametrize(
329
    "durable_kv_events", [False], indirect=True
330
)  # Use NATS Core (local indexer)
331
def test_mocker_kv_router(
332
333
334
335
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
336
    durable_kv_events,
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
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
351
        "durable_kv_events": durable_kv_events,
352
    }
353

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

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

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


381
@pytest.mark.parametrize("store_backend", ["etcd", "file"])
382
@pytest.mark.parametrize(
383
    "durable_kv_events", [False], indirect=True
384
)  # Use NATS Core (local indexer)
385
@pytest.mark.timeout(180)  # bumped for xdist contention (was 60s; ~19.86s serial avg)
386
387
def test_mocker_two_kv_router(
    request,
388
    runtime_services_dynamic_ports,
389
390
391
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
392
    durable_kv_events,
393
):
394
395
396
    """
    Test with two KV routers and multiple mocker engine instances.
    Alternates requests between the two routers to test load distribution.
397
    Tests with both etcd and file storage backends.
398
399
400
    """

    # runtime_services starts etcd and nats
401
402
403
    logger.info(
        f"Starting mocker two KV router test with {store_backend} storage backend"
    )
404

405
406
407
408
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
409
        "durable_kv_events": durable_kv_events,
410
    }
411

412
413
414
415
416
417
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
    ) as mockers:
418
        # Start mocker instances with the new CLI interface
419
420
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
421

422
423
424
425
426
        # Get unique ports for this test (2 ports for two routers)
        router_ports = get_unique_ports(
            request, num_ports=2, store_backend=store_backend
        )

427
428
429
430
431
        # 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,
432
            router_ports=router_ports,
433
434
435
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            store_backend=store_backend,
436
            skip_consumer_verification=not durable_kv_events,  # Skip JetStream checks in NATS Core mode
437
438
439
        )


440
@pytest.mark.skip(reason="Flaky, temporarily disabled")
441
@pytest.mark.parametrize(
442
    "durable_kv_events", [False], indirect=True
443
)  # Use NATS Core (local indexer)
444
@pytest.mark.timeout(60)  # ~3x average (~19.86s), rounded up (when enabled)
Alec's avatar
Alec committed
445
def test_mocker_kv_router_overload_503(
446
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
447
):
448
    """Test that KV router returns 503 when mocker workers are overloaded."""
449
    logger.info("Starting mocker KV router overload test for 503 status")
450
    # Create mocker args dictionary with limited resources - use local indexer (NATS Core mode)
451
452
453
454
    mocker_args = {
        "speedup_ratio": 10,
        "block_size": 4,  # Smaller block size
        "num_gpu_blocks": 64,  # Limited GPU blocks to exhaust quickly
455
        "durable_kv_events": durable_kv_events,
456
    }
457

458
    with MockerProcess(request, mocker_args=mocker_args, num_mockers=1) as mockers:
459
        # Start single mocker instance with limited resources
460
461
        logger.info("Starting single mocker instance with limited resources")
        logger.info(f"Mocker using endpoint: {mockers.endpoint}")
462

463
464
465
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

466
467
468
469
470
471
472
        # 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,
473
            blocks_threshold=0.2,
474
        )
475

476

477
@pytest.mark.timeout(90)  # bumped for xdist contention (was 22s; ~7.10s serial avg)
478
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
479
@pytest.mark.parametrize(
480
    "durable_kv_events", [False], indirect=True
481
)  # Use NATS Core (local indexer)
482
def test_kv_router_bindings(
483
484
485
486
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
487
    durable_kv_events,
488
):
489
490
    """Test KvRouter Python bindings with mocker engines."""
    logger.info("Starting KvRouter bindings test")
491
492
493
494
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
495
        "durable_kv_events": durable_kv_events,
496
    }
497

498
499
500
501
502
503
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
504
        # Start mocker instances
505
506
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
507

508
        # Get runtime and create endpoint
509
        runtime = get_runtime(request_plane=request_plane)
510
        namespace = runtime.namespace(mockers.namespace)
511
        component = namespace.component(mockers.component_name)
512
513
        endpoint = component.endpoint("generate")

514
515
516
        # Run Python router bindings test
        _test_python_router_bindings(
            engine_workers=mockers,
517
518
            endpoint=endpoint,
            block_size=BLOCK_SIZE,
519
520
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
521
        )
522

523

524
@pytest.mark.parametrize(
Yan Ru Pei's avatar
Yan Ru Pei committed
525
    "store_backend,durable_kv_events,request_plane,router_event_threads",
526
    [
Yan Ru Pei's avatar
Yan Ru Pei committed
527
528
529
530
        ("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
531
    ],
532
533
    ids=[
        "jetstream",
534
        "nats_core",
535
        "file",
Yan Ru Pei's avatar
Yan Ru Pei committed
536
        "nats_core_multi_thread",
537
    ],
538
    indirect=["request_plane", "durable_kv_events"],
539
)
540
@pytest.mark.timeout(180)  # bumped for xdist contention (was 90s; up to 33s under load)
541
542
def test_indexers_sync(
    request,
543
    runtime_services_dynamic_ports,
544
545
546
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
547
    durable_kv_events,
548
    request_plane,
Yan Ru Pei's avatar
Yan Ru Pei committed
549
    router_event_threads,
550
):
551
552
553
554
    """
    Test that two KV routers have synchronized indexer states after processing requests.
    This test verifies that both routers converge to the same internal state.

555
556
557
558
559
560
561
562
    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}, "
563
        f"durable_kv_events={durable_kv_events}, request_plane={request_plane}"
564
    )
565

566
567
568
569
    # 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
570
    # Use 2 DP ranks to test per-dp_rank event ID tracking and recovery
571
572
573
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
574
        "durable_kv_events": durable_kv_events,
575
        "dp_size": 2,
576
577
    }

578
579
580
581
582
583
584
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
        request_plane=request_plane,
    ) as mockers:
585
586
        # 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")
587
588
589
590
        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
591
        # When using durable_kv_events=True, use JetStream mode for the router
592
593
594
595
596
597
598
        _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,
599
600
601
            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
602
            router_event_threads=router_event_threads,
603
604
605
606
607
        )

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


608
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
609
@pytest.mark.parametrize(
610
    "durable_kv_events", [False], indirect=True
611
)  # Use NATS Core (local indexer)
Alec's avatar
Alec committed
612
def test_query_instance_id_returns_worker_and_tokens(
613
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
614
):
615
    """Test query_instance_id annotation with mocker engines."""
616
    logger.info("Starting KV router query_instance_id annotation test")
617
618
619
620
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
621
        "durable_kv_events": durable_kv_events,
622
    }
623
624
    os.makedirs(request.node.name, exist_ok=True)

625
626
627
    with MockerProcess(
        request, mocker_args=mocker_args, num_mockers=NUM_MOCKERS
    ) as mockers:
628
        # Start mocker instances
629
630
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
631

632
633
634
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

635
636
637
638
639
640
641
642
        # 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,
        )
643

644

645
@pytest.mark.timeout(90)  # bumped for xdist contention (was 29s; ~9.55s serial avg)
Yan Ru Pei's avatar
Yan Ru Pei committed
646
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
647
@pytest.mark.parametrize(
Yan Ru Pei's avatar
Yan Ru Pei committed
648
    "durable_kv_events,use_kv_events,router_event_threads",
649
    [
Yan Ru Pei's avatar
Yan Ru Pei committed
650
651
652
653
        (True, True, 1),  # JetStream mode with KV events
        (False, True, 1),  # NATS Core mode with local indexer (default)
        (False, False, 1),  # Approximate mode (--no-kv-events) - no KV events
        (False, True, 2),  # NATS Core mode - multi-threaded indexer
654
    ],
Yan Ru Pei's avatar
Yan Ru Pei committed
655
    ids=["jetstream", "nats_core", "no_kv_events", "nats_core_multi_thread"],
656
    indirect=["durable_kv_events"],
657
)
658
def test_router_decisions(
659
660
661
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
662
    durable_kv_events,
663
    use_kv_events,
664
    request_plane,
Yan Ru Pei's avatar
Yan Ru Pei committed
665
    router_event_threads,
666
667
668
):
    """Validate KV cache prefix reuse and dp_rank routing by sending progressive requests with overlapping prefixes.

669
    Parameterized to test:
670
671
    - JetStream mode: KV events via NATS JetStream (durable)
    - NATS Core mode (default): KV events via NATS Core with local indexer on workers
672
673
    - Approximate mode (--no-kv-events): No KV events, router predicts cache state
      based on routing decisions with TTL-based expiration and pruning
674
    """
675
    # runtime_services_dynamic_ports handles NATS and etcd startup
676
    logger.info(
677
        f"Starting test router decisions: durable_kv_events={durable_kv_events}, use_kv_events={use_kv_events}"
678
    )
679

Yan Ru Pei's avatar
Yan Ru Pei committed
680
    # Create mocker args dictionary with dp_size=4
681
    # durable_kv_events=True enables JetStream mode; False (default) uses NATS Core with local indexer
Yan Ru Pei's avatar
Yan Ru Pei committed
682
683
684
685
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
        "dp_size": 4,
686
        "durable_kv_events": durable_kv_events and use_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
687
    }
688

689
690
691
692
693
694
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
    ) as mockers:
695
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
696

697
698
        # Initialize mockers
        # Get runtime and create endpoint
699
        runtime = get_runtime(request_plane=request_plane)
700
701
702
703
704
        # Use the namespace from the mockers
        namespace = runtime.namespace(mockers.namespace)
        component = namespace.component("mocker")
        endpoint = component.endpoint("generate")

705
        _test_router_decisions(
706
707
708
709
710
711
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=use_kv_events,
712
            durable_kv_events=durable_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
713
            router_event_threads=router_event_threads,
714
715
        )

716

717
@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"])
718
719
720
@pytest.mark.parametrize(
    "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"]
)
721
@pytest.mark.timeout(180)  # bumped for xdist contention (was 59s; ~19.51s serial avg)
722
def test_router_decisions_disagg(
723
724
725
726
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    registration_order,
727
    enable_disagg_bootstrap,
728
729
730
731
732
):
    """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.
733

734
735
736
    Parameterized to test:
    - registration_order: prefill_first vs decode_first
    - enable_disagg_bootstrap: without vs with bootstrap rendezvous
737
    """
738
    # runtime_services_dynamic_ports handles NATS and etcd startup
739
740
    logger.info(
        f"Starting disaggregated router prefix reuse test "
741
        f"(registration_order={registration_order}, bootstrap={enable_disagg_bootstrap})"
742
    )
743
744
745
746
747

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

748
    # Create mocker args - use NATS Core with local indexer (default mode)
749
750
751
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
752
        # durable_kv_events defaults to False (NATS Core mode)
753
    }
754

755
756
757
758
759
760
761
762
763
764
765
766
    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:
767
768
769
770
            logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

            # Then start decode workers
            logger.info("Starting 4 decode mocker instances (second)")
771
            with DisaggMockerProcess(
772
773
774
775
776
                request,
                namespace=shared_namespace,
                worker_type="decode",
                mocker_args=mocker_args,
                num_mockers=4,
777
                request_plane="nats",
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
            ) 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:
807
808
809
810
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")

            # Then start prefill workers
            logger.info("Starting 4 prefill mocker instances (second)")
811
            with DisaggMockerProcess(
812
813
814
815
816
                request,
                namespace=shared_namespace,
                worker_type="prefill",
                mocker_args=mocker_args,
                num_mockers=4,
817
818
                request_plane="nats",
                enable_bootstrap=enable_disagg_bootstrap,
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
            ) 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",
                )
839
840


841
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
842
@pytest.mark.parametrize(
843
    "durable_kv_events", [False], indirect=True
844
)  # Use NATS Core (local indexer)
845
@pytest.mark.timeout(120)  # bumped for xdist contention (was 39s; ~12.84s serial avg)
846
def test_busy_threshold_endpoint(
847
848
849
850
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
851
    durable_kv_events,
852
853
854
855
856
857
858
859
860
861
):
    """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.
    """
862
    # runtime_services_dynamic_ports handles NATS and etcd startup
863
864
865
    logger.info(
        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
    )
866

867
868
869
870
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
871
        "durable_kv_events": durable_kv_events,
872
    }
873

874
875
876
877
878
879
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
880
881
882
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")

883
884
885
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
886
887
888
889
890
891
892

        _test_busy_threshold_endpoint(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
893
            request_plane=request_plane,
894
        )