test_router_e2e_with_vllm.py 23 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

# Timing notes (measured locally):
# - GPU-1 subset (`-m "gpu_1 and not gpu_2"`): 130.43s total for 3 tests.
# These tests load a real model and can be slow/flaky when GPU resources are contended,
# so we set explicit pytest timeouts to fail fast on hangs (see per-test markers below).
8
import json
9
10
11
12
13
14
15
16
17
18
import logging
import os
import time
from typing import Any, Dict, Optional

import pytest

from tests.router.common import (  # utilities
    _test_router_basic,
    _test_router_decisions,
19
    _test_router_indexers_sync,
20
21
22
    generate_random_suffix,
    get_runtime,
)
23
from tests.utils.constants import DefaultPort
24
from tests.utils.managed_process import ManagedProcess
25
from tests.utils.port_utils import allocate_ports, deallocate_ports
26
from tests.utils.test_output import resolve_test_output_path
27
28
29
30

logger = logging.getLogger(__name__)

MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
31
32
33

pytestmark = [
    pytest.mark.e2e,
34
    pytest.mark.router,
35
36
37
    pytest.mark.vllm,
    pytest.mark.model(MODEL_NAME),
]
38
39
40
41
SPEEDUP_RATIO = 10.0
NUM_REQUESTS = 10
BLOCK_SIZE = 16

42
43
44
45
46
47
48
49

def allocate_frontend_ports(request, count: int) -> list[int]:
    """Allocate random free frontend ports for xdist-safe execution."""
    ports = allocate_ports(count, DefaultPort.FRONTEND.value)
    request.addfinalizer(lambda: deallocate_ports(ports))
    return ports


50
51
52
53
54
55
56
57
58
59
60
61
62
# 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,
}

63
64
65
66
67
68
69
70
71
72
# Shared vLLM configuration for all tests
# gpu_memory_utilization limits actual VRAM allocation (required for multi-worker on same GPU)
VLLM_ARGS: Dict[str, Any] = {
    "block_size": BLOCK_SIZE,
    "model": MODEL_NAME,
    "gpu_memory_utilization": 0.4,  # Limit VRAM allocation per worker
    "max_model_len": 1024,  # Limit context length to reduce KV cache size
    "enforce_eager": True,  # Disable CUDA graphs for faster startup & lower memory
}

73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

class VLLMProcess:
    """Manages vLLM workers using dynamo.vllm (HTTP API + KV events).

    This is a drop-in replacement for MockerProcess that uses real vLLM workers.
    The key difference: dynamo.vllm automatically handles:
    - HTTP API serving
    - KV cache event publishing (ZMQ → NATS bridge)
    - Integration with dynamo.frontend router
    """

    def __init__(
        self,
        request,
        vllm_args: Optional[Dict[str, Any]] = None,
        num_workers: int = 2,
        single_gpu: bool = False,
        data_parallel_size: Optional[int] = None,
91
92
        request_plane: str = "tcp",
        store_backend: str = "etcd",
93
        durable_kv_events: bool = False,
94
95
96
97
98
99
100
101
    ):
        """Initialize vLLM workers with dynamo integration.

        Args:
            request: pytest request fixture for log directory
            vllm_args: Configuration dict with keys:
                - block_size: KV cache block size (default: 16)
                - model: Model name/path (default: TinyLlama-1.1B)
102
103
                - gpu_memory_utilization: Fraction of GPU memory to allocate (optional)
                - num_gpu_blocks_override: Cap on number of KV cache blocks (optional)
104
                - max_model_len: Maximum sequence length (optional)
105
                - enforce_eager: Disable CUDA graphs (default: False)
106
            num_workers: Number of vLLM worker processes
107
            single_gpu: If True, all workers share GPU 0
108
            data_parallel_size: If set, enables data parallelism with this many ranks (num_workers must equal data_parallel_size)
109
110
            request_plane: Request plane to use ("nats", "tcp", or "http"). Defaults to "tcp".
            store_backend: Storage backend to use ("etcd" or "file"). Defaults to "etcd".
111
            durable_kv_events: If True, use JetStream for durable KV events. Defaults to False (NATS Core mode).
112
113
114
115
116
117
118
        """
        # Generate unique namespace for isolation
        namespace_suffix = generate_random_suffix()
        self.namespace = f"test-namespace-{namespace_suffix}"
        self.component_name = "backend"
        self.endpoint = f"dyn://{self.namespace}.{self.component_name}.generate"
        self.num_workers = num_workers
119
        self.data_parallel_size = data_parallel_size
120
        self.worker_processes = []
121
        self.store_backend = store_backend
122

123
124
125
126
127
128
129
130
131
132
133
        # Dynamically allocate unique system, KV event, and NIXL side-channel
        # ports (one of each per worker) to avoid conflicts in parallel test runs.
        self._system_ports = allocate_ports(num_workers, DefaultPort.SYSTEM1.value)
        self._kv_event_ports = allocate_ports(num_workers, DefaultPort.SYSTEM1.value)
        self._nixl_ports = allocate_ports(num_workers, DefaultPort.SYSTEM1.value)
        request.addfinalizer(
            lambda: deallocate_ports(
                self._system_ports + self._kv_event_ports + self._nixl_ports
            )
        )

134
135
136
137
138
        if vllm_args is None:
            vllm_args = {}

        block_size = vllm_args.get("block_size", BLOCK_SIZE)
        model = vllm_args.get("model", MODEL_NAME)
139
140
        gpu_memory_utilization = vllm_args.get("gpu_memory_utilization")
        num_gpu_blocks_override = vllm_args.get("num_gpu_blocks_override")
141
        max_model_len = vllm_args.get("max_model_len")
142
        enforce_eager = vllm_args.get("enforce_eager", False)
143
144
145
146
147
148
149
150

        self.model_name = model

        # Create vLLM worker processes
        # Matches test.sh behavior:
        # - When data_parallel_size is set, launch one process per DP rank
        # - Each process gets --data-parallel-rank and --data-parallel-size
        # - Each process runs on its own GPU via CUDA_VISIBLE_DEVICES
151
        # - --kv-transfer-config enables KV cache transfer between ranks
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180

        for worker_idx in range(num_workers):
            # Calculate GPU device for this process
            if single_gpu:
                # Force all processes to GPU 0 (for single-GPU testing)
                gpu_device = "0"
            elif data_parallel_size is not None:
                # Worker sees dp_rank GPUs (each DP rank gets its own GPU)
                worker_start_gpu = worker_idx * data_parallel_size
                gpu_device = ",".join(
                    str(i)
                    for i in range(
                        worker_start_gpu, worker_start_gpu + data_parallel_size
                    )
                )
            else:
                # No DP; worker sees one GPU
                gpu_device = str(worker_idx)

            command = [
                "python3",
                "-m",
                "dynamo.vllm",
                "--model",
                model,
                "--block-size",
                str(block_size),
            ]

181
182
183
184
185
186
187
188
189
190
            # Disable CUDA graphs for faster startup & lower memory
            if enforce_eager:
                command.append("--enforce-eager")

            # Limit VRAM allocation (required for multi-worker on same GPU)
            if gpu_memory_utilization is not None:
                command.extend(
                    ["--gpu-memory-utilization", str(gpu_memory_utilization)]
                )

191
192
193
194
            # Add optional max_model_len if specified
            if max_model_len is not None:
                command.extend(["--max-model-len", str(max_model_len)])

195
196
197
198
199
200
            # Cap block count for predictable KV cache behavior
            if num_gpu_blocks_override is not None:
                command.extend(
                    ["--num-gpu-blocks-override", str(num_gpu_blocks_override)]
                )

201
202
203
204
205
206
207
208
209
            if data_parallel_size is not None:
                # Add DP configuration for external load balancing
                # See: https://docs.vllm.ai/en/v0.10.0/serving/data_parallel_deployment.html#external-load-balancing
                command.extend(
                    [
                        "--data-parallel-size",
                        str(data_parallel_size),
                        # "--data-parallel-address", "127.0.0.1",  # Required for DP coordination
                        # "--data-parallel-rpc-port", "13345",  # RPC port for DP coordination
210
                        # "--kv-transfer-config", '{"kv_connector":"NixlConnector","kv_role":"kv_both"}',  # Required for KV transfer between DP ranks
211
212
213
                    ]
                )

214
215
216
217
            # Use --durable-kv-events to enable JetStream mode (local indexer disabled)
            if durable_kv_events:
                command.append("--durable-kv-events")

218
219
220
221
222
            # Ports are dynamically allocated for xdist-safe parallel execution.
            system_port = self._system_ports[worker_idx]
            kv_event_port = self._kv_event_ports[worker_idx]
            nixl_port = self._nixl_ports[worker_idx]

223
224
225
226
227
228
229
230
231
232
233
            # Pass KV events config explicitly via CLI
            kv_events_cfg = json.dumps(
                {
                    "publisher": "zmq",
                    "topic": "kv-events",
                    "endpoint": f"tcp://*:{kv_event_port}",
                    "enable_kv_cache_events": True,
                }
            )
            command.extend(["--kv-events-config", kv_events_cfg])

234
            env = os.environ.copy()  # Copy parent environment
235
236
237
238
            env_vars = {
                "CUDA_VISIBLE_DEVICES": gpu_device,
                "DYN_NAMESPACE": self.namespace,
                "DYN_REQUEST_PLANE": request_plane,
239
240
                "DYN_SYSTEM_PORT": str(system_port),
                "VLLM_NIXL_SIDE_CHANNEL_PORT": str(nixl_port),
241
242
243
244
245
246
247
248
                "PYTHONHASHSEED": "0",  # for deterministic event id's
            }

            # Add DYN_FILE_KV if using file storage backend
            if self.store_backend == "file" and "DYN_FILE_KV" in os.environ:
                env_vars["DYN_FILE_KV"] = os.environ["DYN_FILE_KV"]

            env.update(env_vars)
249
250
251
252
253
254
255
256
257
258

            # Create managed process for the worker
            process = ManagedProcess(
                command=command,
                env=env,
                timeout=120,  # Allow time for model loading
                display_output=True,
                health_check_ports=[],
                health_check_urls=[],
                log_dir=request.node.name,
259
                terminate_all_matching_process_names=False,
260
261
262
263
264
            )
            self.worker_processes.append(process)
            if data_parallel_size is not None:
                logger.info(
                    f"Created {data_parallel_size} DP ranks per worker on GPU(s) {gpu_device} "
265
                    f"(gpu_mem={gpu_memory_utilization}, system_port={system_port}) "
266
267
268
269
270
                    f"with endpoint: {self.endpoint}"
                )
            else:
                logger.info(
                    f"Created vLLM worker {worker_idx} on GPU {gpu_device} "
271
                    f"(gpu_mem={gpu_memory_utilization}, system_port={system_port}) "
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
                    f"with endpoint: {self.endpoint}"
                )

    def __enter__(self):
        """Start all vLLM worker processes with sequential initialization.

        Workers are started sequentially with a delay between each to avoid
        NIXL/UCX resource contention during initialization. This prevents
        UCX shared memory handle allocation failures when multiple workers
        try to initialize simultaneously on the same GPU.
        """
        logger.info(
            f"[VLLMProcess] Starting {len(self.worker_processes)} worker processes sequentially..."
        )

        # Start each process sequentially, waiting for NIXL initialization before next
        for i, process in enumerate(self.worker_processes):
            logger.info(f"[VLLMProcess] Starting vLLM worker {i}...")
            try:
                # Manually initialize the process without blocking on health checks
                process._logger = logging.getLogger(process.__class__.__name__)
                process._command_name = process.command[0]
294
                process.log_dir = resolve_test_output_path(process.log_dir)
295
296
297
298
299
300
301
                os.makedirs(process.log_dir, exist_ok=True)
                log_name = f"{process._command_name}.log.txt"
                process._log_path = os.path.join(process.log_dir, log_name)

                if process.data_dir:
                    process._remove_directory(process.data_dir)

302
                process._terminate_all_matching_process_names()
303
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
                logger.info(
                    f"[VLLMProcess] Launching process {i} (pid will be assigned)..."
                )
                process._start_process()  # Start the process but don't wait
                logger.info(
                    f"[VLLMProcess] Worker {i} launched with PID: {process.proc.pid if process.proc else 'unknown'}"
                )
                time.sleep(process.delayed_start)

                # Wait for NIXL initialization before starting next worker
                # This prevents UCX shared memory contention
                if i < len(self.worker_processes) - 1:
                    nixl_init_delay = 5  # seconds
                    logger.info(
                        f"[VLLMProcess] Waiting {nixl_init_delay}s for worker {i} to initialize NIXL before starting next worker..."
                    )
                    time.sleep(nixl_init_delay)

            except Exception:
                logger.exception(f"[VLLMProcess] Failed to start worker {i}")
                # Clean up on failure
                try:
                    process.__exit__(None, None, None)
                except Exception as cleanup_err:
                    logger.warning(f"[VLLMProcess] Error during cleanup: {cleanup_err}")
                raise

        logger.info(
            f"[VLLMProcess] All {len(self.worker_processes)} workers launched with sequential initialization."
        )
        logger.info("[VLLMProcess] Waiting for health checks to complete...")

        # Now wait for health checks for all processes
        for i, process in enumerate(self.worker_processes):
            logger.info(f"[VLLMProcess] Checking health for worker {i}...")
            try:
                elapsed = process._check_ports(process.timeout)
                process._check_urls(process.timeout - elapsed)
                process._check_funcs(process.timeout - elapsed)
                logger.info(f"[VLLMProcess] Worker {i} health checks passed")
            except Exception:
                logger.error(f"[VLLMProcess] Worker {i} health check failed")
                # Clean up all processes on failure
                self.__exit__(None, None, None)
                raise

        logger.info(
            "[VLLMProcess] All workers started successfully and passed health checks!"
        )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Stop all vLLM worker processes gracefully."""
        for i, process in enumerate(self.worker_processes):
            logger.info(f"Stopping vLLM worker {i}")
            process.__exit__(exc_type, exc_val, exc_tb)

        # Add delay to ensure full cleanup of NATS/ETCD/ZMQ resources
        # This prevents test isolation issues when running multiple tests
        logger.info("Waiting for vLLM worker resources to fully clean up...")
        time.sleep(2)


366
@pytest.mark.pre_merge
367
@pytest.mark.gpu_1
368
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
369
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
370
def test_vllm_kv_router_basic(
371
372
373
374
375
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
376
):
377
378
    """
    Quick e2e sanity test for KV router with vLLM engine instances.
379
    Tests both NATS and TCP request planes.
380
381
382
383
    """

    # runtime_services starts etcd and nats
    N_VLLM_WORKERS = 2
384
385
386
    logger.info(
        f"Starting vLLM KV router test with {N_VLLM_WORKERS} workers using request_plane={request_plane}"
    )
387

388
389
390
391
392
393
394
    with VLLMProcess(
        request,
        vllm_args=VLLM_ARGS,
        num_workers=N_VLLM_WORKERS,
        single_gpu=True,  # fit workers into one GPU
        request_plane=request_plane,
    ) as vllm_workers:
395
396
397
398
        # Start vLLM workers
        logger.info(f"Starting {N_VLLM_WORKERS} vLLM workers")
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

399
        # Run basic router test (starts router internally and waits for workers to be ready)
400
        frontend_port = allocate_frontend_ports(request, 1)[0]
401
402
403
404
        _test_router_basic(
            engine_workers=vllm_workers,
            block_size=BLOCK_SIZE,
            request=request,
405
            frontend_port=frontend_port,
406
407
408
409
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            frontend_timeout=180,  # 3 minutes should be plenty for TinyLlama
            store_backend="etcd",  # Explicit for clarity
410
            request_plane=request_plane,
411
412
413
        )


414
@pytest.mark.pre_merge
415
@pytest.mark.gpu_1
416
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
417
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
Yan Ru Pei's avatar
Yan Ru Pei committed
418
419
420
421
422
@pytest.mark.parametrize(
    "router_event_threads",
    [1, 2],
    ids=["single_thread", "multi_thread"],
)
423
def test_router_decisions_vllm_multiple_workers(
424
425
426
427
428
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
Yan Ru Pei's avatar
Yan Ru Pei committed
429
    router_event_threads,
430
431
432
433
434
):
    # runtime_services starts etcd and nats
    logger.info("Starting vLLM router prefix reuse test with two workers")
    N_WORKERS = 2

435
436
437
438
439
440
441
    with VLLMProcess(
        request,
        vllm_args=VLLM_ARGS,
        num_workers=N_WORKERS,
        single_gpu=True,  # Worker uses GPU 0
        request_plane=request_plane,
    ) as vllm_workers:
442
443
        # Start 2 worker processes on the same GPU
        logger.info("Starting 2 vLLM worker processes on single GPU (gpu_mem=0.4)")
444
445
446
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

        # Get runtime and create endpoint
447
        runtime = get_runtime(request_plane=request_plane)
448
        endpoint = runtime.endpoint(f"{vllm_workers.namespace}.backend.generate")
449
450

        _test_router_decisions(
Yan Ru Pei's avatar
Yan Ru Pei committed
451
452
453
454
455
            vllm_workers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=False,
456
            block_size=BLOCK_SIZE,
Yan Ru Pei's avatar
Yan Ru Pei committed
457
            router_event_threads=router_event_threads,
458
459
460
461
        )


@pytest.mark.gpu_2
462
@pytest.mark.nightly
463
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
464
@pytest.mark.timeout(600)  # 10 min max (multi-GPU + DP startup variance)
465
def test_router_decisions_vllm_dp(
466
467
468
469
470
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
471
):
472
473
474
475
476
477
478
479
480
481
    """Validate KV cache prefix reuse with vLLM by sending progressive requests with overlapping prefixes.
    Same flow as test_router_decisions_vllm_multiple_workers; force first request to (worker_id, dp_rank=1).
    Dump events from router and verify:
        * All but one (worker_id, dp_rank) should have no events (due to prefix reuse)
        * The (worker_id, dp_rank) with events should have exactly 4 events (one per request)
        * All events should be on the forced (worker_id, dp_rank=1) (verifying forced routing and prefix reuse)
    """
    N_WORKERS = 1
    DP_SIZE = 2

482
483
484
485
486
487
488
489
    with VLLMProcess(
        request,
        vllm_args=VLLM_ARGS,
        num_workers=N_WORKERS,  # Ignored when data_parallel_size is set
        single_gpu=False,
        data_parallel_size=DP_SIZE,  # Creates DP_SIZE processes (one per rank)
        request_plane=request_plane,
    ) as vllm_workers:
490
        logger.info("Starting 2 vLLM DP ranks (dp_size=2) (gpu_mem=0.4)")
491
492
493
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

        # Get runtime and create endpoint
494
        runtime = get_runtime(request_plane=request_plane)
495
        # Use the namespace from the vLLM workers
496
497
498
        endpoint = runtime.endpoint(
            f"{vllm_workers.namespace}.backend.generate"
        )  # endpoint is backend.generate
499
500

        _test_router_decisions(
501
502
503
504
505
506
            vllm_workers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            block_size=BLOCK_SIZE,
507
508
        )

509
510
511

@pytest.mark.pre_merge
@pytest.mark.gpu_1
512
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
513
@pytest.mark.parametrize(
514
    "store_backend,durable_kv_events,request_plane",
515
    [
516
        ("etcd", False, "tcp"),
517
    ],
518
519
    ids=["nats_core"],
    indirect=["durable_kv_events", "request_plane"],
520
)
521
def test_vllm_indexers_sync(
522
523
524
525
526
527
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    file_storage_backend,
    set_ucx_tls_no_mm,
    store_backend,
528
    durable_kv_events,
529
    request_plane,
530
531
532
533
):
    """
    Test that two KV routers have synchronized indexer states after processing requests
    with vLLM workers. This test verifies that both routers converge to the same internal state.
534
535

    Tests with configuration:
536
537
    - nats_core: etcd backend, local indexer with NATS Core, TCP request plane
                 (includes NATS interruption/recovery testing)
538
    """
539
    # runtime_services_dynamic_ports handles NATS and etcd startup
540
541
    nats_process, _etcd_process = runtime_services_dynamic_ports

542
543
    logger.info(
        f"Starting vLLM indexers sync test: store_backend={store_backend}, "
544
        f"durable_kv_events={durable_kv_events}, request_plane={request_plane}"
545
546
    )

547
548
    N_VLLM_WORKERS = 2

549
550
551
552
553
554
555
556
557
    with VLLMProcess(
        request,
        vllm_args=VLLM_ARGS,
        num_workers=N_VLLM_WORKERS,
        single_gpu=True,  # fit workers into one GPU
        request_plane=request_plane,
        store_backend=store_backend,
        durable_kv_events=durable_kv_events,
    ) as vllm_workers:
558
559
560
561
562
563
        # Start vLLM workers
        logger.info(f"Starting {N_VLLM_WORKERS} vLLM workers")
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

        # 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
564
        # When using durable_kv_events=True, use JetStream mode for the router
565
566
567
568
569
        _test_router_indexers_sync(
            engine_workers=vllm_workers,
            block_size=BLOCK_SIZE,
            model_name=MODEL_NAME,
            num_workers=N_VLLM_WORKERS,
570
571
            store_backend=store_backend,
            request_plane=request_plane,
572
573
574
            test_nats_interruption=not durable_kv_events,
            nats_server=nats_process if not durable_kv_events else None,
            durable_kv_events=durable_kv_events,
575
576
577
        )

        logger.info("vLLM indexers sync test completed successfully")