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

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

21
import aiohttp
22
23
import pytest

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

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

MODEL_NAME = ROUTER_MODEL_NAME

54
55
56
57
pytestmark = [
    pytest.mark.pre_merge,
    pytest.mark.gpu_0,
    pytest.mark.integration,
58
    pytest.mark.model(MODEL_NAME),
59
]
60
61
NUM_MOCKERS = 2
SPEEDUP_RATIO = 10.0
62
63
64
BASE_PORT = 9100  # Base port for general test allocations (frontend, system, etc.)
BASE_PORT_BOOTSTRAP = 10100  # Base port for disagg bootstrap rendezvous
BASE_PORT_ZMQ = 11100  # Base port for ZMQ KV event publishing
65
NUM_REQUESTS = 100
66
BLOCK_SIZE = 16
67
68


69
def get_unique_ports(
70
71
72
73
    request,
    num_ports: int = 1,
    store_backend: str = "etcd",
    request_plane: str = "nats",
74
    registration_order: str = "prefill_first",
75
) -> list[int]:
76
    """Allocate random free ports for xdist-safe router tests.
77

78
79
80
    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.
81

82
83
84
85
    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.
86
    """
87
88
89
    _ = (store_backend, request_plane, registration_order)
    ports = allocate_ports(num_ports, BASE_PORT)
    request.addfinalizer(lambda: deallocate_ports(ports))
90
91
92
    return ports


93
94
95
96
97
98
99
100
101
102
103
104
105
# 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,
}

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
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,
134
        "--discovery-backend",
135
136
137
138
139
140
141
        store_backend,
        "--num-workers",
        str(num_workers),
    ]

    # Add worker type flag for disaggregated mode
    if worker_type == "prefill":
142
        command.extend(["--disaggregation-mode", "prefill"])
143
    elif worker_type == "decode":
144
        command.extend(["--disaggregation-mode", "decode"])
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

    # 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")
171
172
    if "preemption_mode" in mocker_args:
        command.extend(["--preemption-mode", str(mocker_args["preemption_mode"])])
173
174
    if "dp_size" in mocker_args:
        command.extend(["--data-parallel-size", str(mocker_args["dp_size"])])
175
176
177
    # 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")
178
179
    if "bootstrap_ports" in mocker_args:
        command.extend(["--bootstrap-ports", mocker_args["bootstrap_ports"]])
180
181
    if "zmq_kv_events_ports" in mocker_args:
        command.extend(["--zmq-kv-events-ports", mocker_args["zmq_kv_events_ports"]])
182
183
    if "zmq_replay_ports" in mocker_args:
        command.extend(["--zmq-replay-ports", mocker_args["zmq_replay_ports"]])
184
185
186
187

    return command


188
class MockerProcess:
189
190
191
192
193
194
    """Manages mocker engine instances with shared tokio runtime via --num-workers.

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

196
197
198
199
200
    def __init__(
        self,
        request,
        mocker_args: Optional[Dict[str, Any]] = None,
        num_mockers: int = 1,
201
        store_backend: str = "etcd",
202
        request_plane: str = "nats",
203
        zmq_kv_events: bool = False,
204
        standalone_indexer: bool = False,
205
        model_name: str = "mocker",
206
        zmq_replay: bool = False,
207
    ):
208
209
        namespace_suffix = generate_random_suffix()
        self.namespace = f"test-namespace-{namespace_suffix}"
210
        self.component_name = "mocker"
211
        self.model_name = model_name
212
        self.endpoint = f"dyn://{self.namespace}.{self.component_name}.generate"
213
        self.num_workers = num_mockers
214
        self._zmq_kv_events_ports: list[int] = []
215
        self._zmq_replay_ports: list[int] = []
216
217
218
219
220
221
        self._standalone_indexer = standalone_indexer
        self._standalone_indexer_port: Optional[int] = None
        self._standalone_indexer_b_port: Optional[int] = None
        self._indexer_process: Optional[ManagedProcess] = None
        self._indexer_b_process: Optional[ManagedProcess] = None
        self._mocker_processes: list[ManagedProcess] = []
222
223
224
225
226
227
228
        self._request = request
        self._store_backend = store_backend
        self._request_plane = request_plane
        self._mocker_args_orig: Dict[str, Any] = (mocker_args or {}).copy()
        self.worker_id_to_zmq_ports: dict[int, dict[int, str]] = {}

        mocker_args = self._mocker_args_orig.copy()
229
230
        # Store dp_size for DP-aware test functions
        self.dp_size = mocker_args.get("dp_size")
231
232
        # Alias for consistency with vLLM/SGLang workers
        self.data_parallel_size = self.dp_size
233

234
235
        # Allocate contiguous ZMQ port blocks for KV event publishing because
        # the mocker binds base_port + dp_rank for each DP rank.
236
237
        if zmq_kv_events:
            dp_size = mocker_args.get("dp_size", 1)
238
239
            self._zmq_kv_events_ports = allocate_contiguous_ports(
                num_mockers, dp_size, BASE_PORT_ZMQ
240
241
            )
            bases = [self._zmq_kv_events_ports[i * dp_size] for i in range(num_mockers)]
242
243
            if not standalone_indexer:
                mocker_args["zmq_kv_events_ports"] = ",".join(str(p) for p in bases)
244
245
246
247
248
            logger.info(
                f"Allocated ZMQ KV event ports {self._zmq_kv_events_ports} "
                f"(bases: {bases}) for {num_mockers} workers"
            )

249
        # Allocate contiguous ZMQ replay port blocks with the same layout.
250
251
        if zmq_replay and zmq_kv_events:
            dp_size = mocker_args.get("dp_size", 1)
252
253
            self._zmq_replay_ports = allocate_contiguous_ports(
                num_mockers, dp_size, BASE_PORT_ZMQ + 1000
254
255
256
257
            )
            replay_bases = [
                self._zmq_replay_ports[i * dp_size] for i in range(num_mockers)
            ]
258
259
            if not standalone_indexer:
                mocker_args["zmq_replay_ports"] = ",".join(str(p) for p in replay_bases)
260
261
262
263
264
            logger.info(
                f"Allocated ZMQ replay ports {self._zmq_replay_ports} "
                f"(bases: {replay_bases}) for {num_mockers} workers"
            )

265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
        if standalone_indexer:
            # Allocate ports for standalone indexer A and B (P2P recovery peer)
            indexer_ports = allocate_ports(2, BASE_PORT)
            self._standalone_indexer_port = indexer_ports[0]
            self._standalone_indexer_b_port = indexer_ports[1]
            request.addfinalizer(lambda: deallocate_ports(indexer_ports))
            # Don't build a single mocker command — we'll launch per-mocker in launch_mockers_with_indexer
            self._process = None
        else:
            command = _build_mocker_command(
                endpoint=self.endpoint,
                store_backend=store_backend,
                num_workers=num_mockers,
                mocker_args=mocker_args,
            )

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

            self._process = ManagedProcess(
                command=command,
                env=env,
                timeout=60,
                display_output=True,
                health_check_ports=[],
                health_check_urls=[],
                log_dir=request.node.name,
                terminate_all_matching_process_names=False,
            )
        logger.info(
            f"Created mocker process with {num_mockers} worker(s), endpoint: {self.endpoint}"
            f"{', standalone_indexer=True' if standalone_indexer else ''}"
297
298
        )

299
300
301
302
303
    @property
    def standalone_indexer_url(self) -> Optional[str]:
        if self._standalone_indexer_port is not None:
            return f"http://localhost:{self._standalone_indexer_port}"
        return None
304

305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
    @property
    def standalone_indexer_b_url(self) -> Optional[str]:
        if self._standalone_indexer_b_port is not None:
            return f"http://localhost:{self._standalone_indexer_b_port}"
        return None

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

    async def launch_mockers_with_indexer(self, endpoint):
        """Launch mockers one-by-one and register each with the standalone indexer.

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

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

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

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

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

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

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

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

407
408
409
410
            # Register each dp_rank endpoint with the standalone indexer.
            # The mocker binds on base_port + dp_rank (contiguous), so we must
            # use the same formula here rather than indexing into the allocated
            # port list, which may contain gaps when intervening ports are busy.
411
412
            zmq_addresses = {}
            register_url = f"{self.standalone_indexer_url}/register"
413
414
415
            replay_base = (
                self._zmq_replay_ports[i * dp_size] if self._zmq_replay_ports else None
            )
416
417
            async with aiohttp.ClientSession() as session:
                for dp_rank in range(dp_size):
418
                    port = base_port + dp_rank
419
420
421
422
423
424
425
426
427
428
429
430
                    endpoint = f"tcp://127.0.0.1:{port}"
                    zmq_addresses[dp_rank] = endpoint

                    payload = {
                        "instance_id": new_worker_id,
                        "endpoint": endpoint,
                        "dp_rank": dp_rank,
                        "model_name": self.model_name,
                        "block_size": self._mocker_args_orig.get(
                            "block_size", BLOCK_SIZE
                        ),
                    }
431
432
433
434
                    if replay_base is not None:
                        payload[
                            "replay_endpoint"
                        ] = f"tcp://127.0.0.1:{replay_base + dp_rank}"
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
                    async with session.post(register_url, json=payload) as resp:
                        if resp.status != 201:
                            body = await resp.text()
                            raise RuntimeError(
                                f"Failed to register instance {new_worker_id} "
                                f"dp_rank {dp_rank}: {resp.status} {body}"
                            )

            self.worker_id_to_zmq_ports[new_worker_id] = zmq_addresses

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

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

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

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

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

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

        indexer_b_cmd = [
            *get_kv_indexer_command(),
            "--block-size",
            str(block_size),
            "--port",
            str(self._standalone_indexer_b_port),
            "--peers",
            f"http://localhost:{self._standalone_indexer_port}",
            "--workers",
            workers_arg,
            "--model-name",
            self.model_name,
        ]
        self._indexer_b_process = ManagedProcess(
            command=indexer_b_cmd,
            timeout=120,
494
            display_output=True,
495
            health_check_ports=[self._standalone_indexer_b_port],
496
            health_check_urls=[],
497
            log_dir=self._request.node.name,
498
            terminate_all_matching_process_names=False,
499
            display_name="dynamo-kv-indexer-b",
500
501
        )
        logger.info(
502
503
            f"Starting standalone indexer B on port {self._standalone_indexer_b_port} "
            f"with peer http://localhost:{self._standalone_indexer_port}"
504
        )
505
        self._indexer_b_process.__enter__()
506

507
    def __exit__(self, exc_type, exc_val, exc_tb):
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        logger.info("Stopping mocker process(es)")
        # Stop individual mocker processes (standalone_indexer mode)
        for proc in self._mocker_processes:
            try:
                proc.__exit__(exc_type, exc_val, exc_tb)
            except Exception as e:
                logger.warning(f"Error stopping mocker process: {e}")
        self._mocker_processes.clear()
        # Stop standalone indexer B (P2P recovery peer)
        if self._indexer_b_process is not None:
            try:
                self._indexer_b_process.__exit__(exc_type, exc_val, exc_tb)
            except Exception as e:
                logger.warning(f"Error stopping indexer B process: {e}")
            self._indexer_b_process = None
        # Stop standalone indexer A
        if self._indexer_process is not None:
            try:
                self._indexer_process.__exit__(exc_type, exc_val, exc_tb)
            except Exception as e:
                logger.warning(f"Error stopping indexer process: {e}")
            self._indexer_process = None
        # Stop single mocker process (non-standalone mode)
531
532
        if self._process is not None:
            self._process.__exit__(exc_type, exc_val, exc_tb)
533
534
535
536
        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 = []
537
538
539
540
        if self._zmq_replay_ports:
            deallocate_ports(self._zmq_replay_ports)
            logger.info(f"Deallocated ZMQ replay ports {self._zmq_replay_ports}")
            self._zmq_replay_ports = []
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560


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",
561
        request_plane: str = "nats",
562
        enable_bootstrap: bool = False,
563
564
565
566
    ):
        if worker_type not in ("prefill", "decode"):
            raise ValueError(
                f"worker_type must be 'prefill' or 'decode', got {worker_type}"
567
            )
568
569
570
571

        self.namespace = namespace
        self.worker_type = worker_type
        self.num_workers = num_mockers
572
        self._bootstrap_ports: list[int] = []
573
574
575
576
577
578
579
580
581

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

582
583
584
585
        mocker_args = (mocker_args or {}).copy()

        # Allocate bootstrap ports for prefill workers if enabled (one per worker)
        if enable_bootstrap and worker_type == "prefill":
586
            self._bootstrap_ports = allocate_ports(num_mockers, BASE_PORT_BOOTSTRAP)
587
588
589
590
591
592
            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"
            )
593
594
595
596
597
598
599
600
601

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

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

605
606
        self._process = ManagedProcess(
            command=command,
607
            env=env,
608
609
610
611
612
            timeout=60,
            display_output=True,
            health_check_ports=[],
            health_check_urls=[],
            log_dir=request.node.name,
613
            terminate_all_matching_process_names=False,
614
615
616
617
618
        )
        logger.info(
            f"Created {worker_type} mocker process with {num_mockers} worker(s), "
            f"endpoint: {self.endpoint}"
        )
619

620
621
622
623
624
    @property
    def bootstrap_ports(self) -> list[int]:
        """Return the allocated bootstrap ports, if any."""
        return self._bootstrap_ports

625
    def __enter__(self):
626
627
628
629
        logger.info(
            f"Starting {self.worker_type} mocker process with {self.num_workers} worker(s)"
        )
        self._process.__enter__()
630
        return self
631

632
    def __exit__(self, exc_type, exc_val, exc_tb):
633
634
        logger.info(f"Stopping {self.worker_type} mocker process")
        self._process.__exit__(exc_type, exc_val, exc_tb)
635
636
637
638
639
        # 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 = []
640
641


642
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
643
@pytest.mark.parametrize(
644
645
646
647
648
649
650
651
652
653
654
    "router_mode,durable_kv_events",
    [
        pytest.param("kv", False, id="kv-nondurable"),
        pytest.param("kv", True, id="kv-durable"),
        pytest.param("round-robin", False, id="roundrobin"),
        pytest.param("random", False, id="random"),
    ],
    indirect=["durable_kv_events"],
)
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
def test_mocker_router(
655
656
657
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
658
    router_mode,
659
    request_plane,
660
    durable_kv_events,
661
):
662
    """Test router with multiple mocker engine instances across all router modes.
663

664
665
    Covers kv, round-robin, and random routing. Tests both NATS and TCP request planes.
    """
666
    # runtime_services starts etcd and optionally nats based on request_plane
667
668
669
    logger.info(
        f"Starting mocker router test: router_mode={router_mode}, request_plane={request_plane}"
    )
670

671
672
673
674
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
675
        "durable_kv_events": durable_kv_events,
676
    }
677

678
679
680
681
682
683
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
684
        # Start mocker instances with the new CLI interface
685
686
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
687

688
        # Get unique port for this test
689
690
691
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
692
693

        # Run basic router test (starts router internally and waits for workers to be ready)
694
695
696
697
        _test_router_basic(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
698
            frontend_port=frontend_port,
699
700
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
701
            request_plane=request_plane,
702
            router_mode=router_mode,
703
704
705
        )


706
@pytest.mark.parametrize("store_backend", ["etcd", "file"])
707
@pytest.mark.parametrize(
708
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
709
)  # Use NATS Core (local indexer)
710
@pytest.mark.timeout(180)  # bumped for xdist contention (was 60s; ~19.86s serial avg)
711
712
def test_mocker_two_kv_router(
    request,
713
    runtime_services_dynamic_ports,
714
715
716
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
717
    durable_kv_events,
718
):
719
720
721
    """
    Test with two KV routers and multiple mocker engine instances.
    Alternates requests between the two routers to test load distribution.
722
    Tests with both etcd and file storage backends.
723
724
725
    """

    # runtime_services starts etcd and nats
726
727
728
    logger.info(
        f"Starting mocker two KV router test with {store_backend} storage backend"
    )
729

730
731
732
733
    # Create mocker args dictionary - use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
734
        "durable_kv_events": durable_kv_events,
735
    }
736

737
738
739
740
741
742
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        store_backend=store_backend,
    ) as mockers:
743
        # Start mocker instances with the new CLI interface
744
745
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
746

747
748
749
750
751
        # Get unique ports for this test (2 ports for two routers)
        router_ports = get_unique_ports(
            request, num_ports=2, store_backend=store_backend
        )

752
753
754
755
756
        # 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,
757
            router_ports=router_ports,
758
759
760
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            store_backend=store_backend,
761
            skip_consumer_verification=not durable_kv_events,  # Skip JetStream checks in NATS Core mode
762
763
764
        )


765
@pytest.mark.skip(reason="Flaky, temporarily disabled")
766
@pytest.mark.parametrize(
767
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
768
)  # Use NATS Core (local indexer)
769
@pytest.mark.timeout(60)  # ~3x average (~19.86s), rounded up (when enabled)
Alec's avatar
Alec committed
770
def test_mocker_kv_router_overload_503(
771
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
772
):
773
    """Test that KV router returns 503 when mocker workers are overloaded."""
774
    logger.info("Starting mocker KV router overload test for 503 status")
775
    # Create mocker args dictionary with limited resources - use local indexer (NATS Core mode)
776
777
778
779
    mocker_args = {
        "speedup_ratio": 10,
        "block_size": 4,  # Smaller block size
        "num_gpu_blocks": 64,  # Limited GPU blocks to exhaust quickly
780
        "durable_kv_events": durable_kv_events,
781
    }
782

783
    with MockerProcess(request, mocker_args=mocker_args, num_mockers=1) as mockers:
784
        # Start single mocker instance with limited resources
785
786
        logger.info("Starting single mocker instance with limited resources")
        logger.info(f"Mocker using endpoint: {mockers.endpoint}")
787

788
789
790
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

791
792
793
794
795
796
797
        # 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,
798
            blocks_threshold=0.2,
799
        )
800

801

802
@pytest.mark.timeout(90)  # bumped for xdist contention (was 22s; ~7.10s serial avg)
803
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
804
@pytest.mark.parametrize(
805
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
806
)  # Use NATS Core (local indexer)
807
def test_kv_router_bindings(
808
809
810
811
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
812
    durable_kv_events,
813
):
814
815
    """Test KvRouter Python bindings with mocker engines."""
    logger.info("Starting KvRouter bindings test")
816
817
818
819
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
820
        "durable_kv_events": durable_kv_events,
821
    }
822

823
824
825
826
827
828
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
829
        # Start mocker instances
830
831
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
832

833
        # Get runtime and create endpoint
834
        runtime = get_runtime(request_plane=request_plane)
835
836
837
        endpoint = runtime.endpoint(
            f"{mockers.namespace}.{mockers.component_name}.generate"
        )
838

839
840
841
        # Run Python router bindings test
        _test_python_router_bindings(
            engine_workers=mockers,
842
843
            endpoint=endpoint,
            block_size=BLOCK_SIZE,
844
845
            model_name=MODEL_NAME,
            num_workers=NUM_MOCKERS,
846
        )
847

848

849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
@pytest.mark.parametrize(
    "store_backend,durable_kv_events,request_plane",
    [
        ("etcd", True, "nats"),  # JetStream mode - uses JetStream
        ("etcd", False, "tcp"),  # NATS core mode (with gap detection) - no JetStream
        ("file", True, "nats"),  # File backend - uses JetStream
    ],
    ids=[
        "jetstream",
        "nats_core",
        "file",
    ],
    indirect=["request_plane", "durable_kv_events"],
)
@pytest.mark.timeout(300)
def test_indexers_sync(
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    file_storage_backend,
    store_backend,
    durable_kv_events,
    request_plane,
):
    """
    Test that two KV routers have synchronized indexer states after processing requests.
    This test verifies that both routers converge to the same internal state.

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

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

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

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

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

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


935
@pytest.mark.timeout(120)  # bumped for xdist contention (was 42s; ~13.80s serial avg)
936
@pytest.mark.parametrize(
937
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
938
)  # Use NATS Core (local indexer)
Alec's avatar
Alec committed
939
def test_query_instance_id_returns_worker_and_tokens(
940
    request, runtime_services_dynamic_ports, predownload_tokenizers, durable_kv_events
Alec's avatar
Alec committed
941
):
942
    """Test query_instance_id annotation with mocker engines."""
943
    logger.info("Starting KV router query_instance_id annotation test")
944
945
946
947
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
948
        "durable_kv_events": durable_kv_events,
949
    }
950

951
952
953
    with MockerProcess(
        request, mocker_args=mocker_args, num_mockers=NUM_MOCKERS
    ) as mockers:
954
        # Start mocker instances
955
956
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
957

958
959
960
        # Get unique port for this test
        frontend_port = get_unique_ports(request, num_ports=1)[0]

961
962
963
964
965
966
967
968
        # 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,
        )
969

970

971
@pytest.mark.timeout(300)  # bumped for xdist contention (was 29s; ~9.55s serial avg)
Yan Ru Pei's avatar
Yan Ru Pei committed
972
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
973
@pytest.mark.parametrize(
974
    "durable_kv_events,use_kv_events,zmq_kv_events",
975
    [
976
977
978
        (True, True, False),  # JetStream mode with KV events
        (False, True, False),  # NATS Core mode with local indexer (default)
        (False, False, False),  # Approximate mode (--no-kv-events) - no KV events
979
        (False, True, True),  # ZMQ mode: mocker → ZMQ PUB → relay → NATS
980
    ],
981
    ids=["jetstream", "nats_core", "no_kv_events", "zmq"],
982
    indirect=["durable_kv_events"],
983
)
984
def test_router_decisions(
985
986
987
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
988
    durable_kv_events,
989
    use_kv_events,
990
    request_plane,
991
    zmq_kv_events,
992
993
994
):
    """Validate KV cache prefix reuse and dp_rank routing by sending progressive requests with overlapping prefixes.

995
    Parameterized to test:
996
997
    - JetStream mode: KV events via NATS JetStream (durable)
    - NATS Core mode (default): KV events via NATS Core with local indexer on workers
998
999
    - Approximate mode (--no-kv-events): No KV events, router predicts cache state
      based on routing decisions with TTL-based expiration and pruning
1000
    """
1001
    # runtime_services_dynamic_ports handles NATS and etcd startup
1002
    logger.info(
1003
        f"Starting test router decisions: durable_kv_events={durable_kv_events}, use_kv_events={use_kv_events}"
1004
    )
1005

Yan Ru Pei's avatar
Yan Ru Pei committed
1006
    # Create mocker args dictionary with dp_size=4
1007
    # durable_kv_events=True enables JetStream mode; False (default) uses NATS Core with local indexer
Yan Ru Pei's avatar
Yan Ru Pei committed
1008
1009
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
1010
        "block_size": 8,
Yan Ru Pei's avatar
Yan Ru Pei committed
1011
        "dp_size": 4,
1012
        "durable_kv_events": durable_kv_events and use_kv_events,
Yan Ru Pei's avatar
Yan Ru Pei committed
1013
    }
1014

1015
1016
1017
1018
1019
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=2,
        request_plane=request_plane,
1020
        zmq_kv_events=zmq_kv_events,
1021
        standalone_indexer=zmq_kv_events,
1022
        model_name=MODEL_NAME,
1023
    ) as mockers:
1024
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")
1025

1026
1027
        # Initialize mockers
        # Get runtime and create endpoint
1028
        runtime = get_runtime(request_plane=request_plane)
1029
        # Use the namespace from the mockers
1030
        endpoint = runtime.endpoint(f"{mockers.namespace}.mocker.generate")
1031

1032
        _test_router_decisions(
1033
1034
1035
1036
1037
1038
            mockers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            use_kv_events=use_kv_events,
1039
            durable_kv_events=durable_kv_events,
1040
            standalone_indexer_url=mockers.standalone_indexer_url,
1041
1042
        )

1043

1044
@pytest.mark.parametrize("registration_order", ["prefill_first", "decode_first"])
1045
1046
1047
@pytest.mark.parametrize(
    "enable_disagg_bootstrap", [False, True], ids=["no_bootstrap", "with_bootstrap"]
)
1048
@pytest.mark.timeout(180)  # bumped for xdist contention (was 59s; ~19.51s serial avg)
1049
def test_router_decisions_disagg(
1050
1051
1052
1053
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    registration_order,
1054
    enable_disagg_bootstrap,
1055
1056
1057
1058
1059
):
    """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.
1060

1061
1062
1063
    Parameterized to test:
    - registration_order: prefill_first vs decode_first
    - enable_disagg_bootstrap: without vs with bootstrap rendezvous
1064
    """
1065
    # runtime_services_dynamic_ports handles NATS and etcd startup
1066
1067
    logger.info(
        f"Starting disaggregated router prefix reuse test "
1068
        f"(registration_order={registration_order}, bootstrap={enable_disagg_bootstrap})"
1069
    )
1070
1071
1072
1073
1074

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

1075
    # Create mocker args - use NATS Core with local indexer (default mode)
1076
1077
1078
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1079
        # durable_kv_events defaults to False (NATS Core mode)
1080
    }
1081

1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
    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:
1094
1095
1096
1097
            logger.info(f"Prefill workers using endpoint: {prefill_workers.endpoint}")

            # Then start decode workers
            logger.info("Starting 4 decode mocker instances (second)")
1098
            with DisaggMockerProcess(
1099
1100
1101
1102
1103
                request,
                namespace=shared_namespace,
                worker_type="decode",
                mocker_args=mocker_args,
                num_mockers=4,
1104
                request_plane="nats",
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
            ) 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:
1134
1135
1136
1137
            logger.info(f"Decode workers using endpoint: {decode_workers.endpoint}")

            # Then start prefill workers
            logger.info("Starting 4 prefill mocker instances (second)")
1138
            with DisaggMockerProcess(
1139
1140
1141
1142
1143
                request,
                namespace=shared_namespace,
                worker_type="prefill",
                mocker_args=mocker_args,
                num_mockers=4,
1144
1145
                request_plane="nats",
                enable_bootstrap=enable_disagg_bootstrap,
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
            ) 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",
                )
1166
1167


1168
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
1169
@pytest.mark.parametrize(
1170
    "durable_kv_events", [False], ids=["nondurable"], indirect=True
1171
)  # Use NATS Core (local indexer)
1172
@pytest.mark.timeout(120)  # bumped for xdist contention (was 39s; ~12.84s serial avg)
1173
def test_busy_threshold_endpoint(
1174
1175
1176
1177
    request,
    runtime_services_dynamic_ports,
    predownload_tokenizers,
    request_plane,
1178
    durable_kv_events,
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
):
    """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.
    """
1189
    # runtime_services_dynamic_ports handles NATS and etcd startup
1190
1191
1192
    logger.info(
        f"Starting busy_threshold endpoint test with request_plane={request_plane}"
    )
1193

1194
1195
1196
1197
    # Use local indexer (NATS Core mode)
    mocker_args = {
        "speedup_ratio": SPEEDUP_RATIO,
        "block_size": BLOCK_SIZE,
1198
        "durable_kv_events": durable_kv_events,
1199
    }
1200

1201
1202
1203
1204
1205
1206
    with MockerProcess(
        request,
        mocker_args=mocker_args,
        num_mockers=NUM_MOCKERS,
        request_plane=request_plane,
    ) as mockers:
1207
1208
1209
        logger.info(f"Starting {NUM_MOCKERS} mocker instances")
        logger.info(f"All mockers using endpoint: {mockers.endpoint}")

1210
1211
1212
        frontend_port = get_unique_ports(
            request, num_ports=1, request_plane=request_plane
        )[0]
1213
1214
1215
1216
1217
1218
1219

        _test_busy_threshold_endpoint(
            engine_workers=mockers,
            block_size=BLOCK_SIZE,
            request=request,
            frontend_port=frontend_port,
            test_payload=TEST_PAYLOAD,
1220
            request_plane=request_plane,
1221
        )
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279


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

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

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

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

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

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

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

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

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