test_router_e2e_with_vllm.py 22.6 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
9
10
11
12
13
14
15
16
17
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,
18
    _test_router_indexers_sync,
19
20
21
    generate_random_suffix,
    get_runtime,
)
22
from tests.utils.constants import DefaultPort
23
from tests.utils.managed_process import ManagedProcess
24
from tests.utils.port_utils import allocate_ports, deallocate_ports
25
from tests.utils.test_output import resolve_test_output_path
26
27
28
29

logger = logging.getLogger(__name__)

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

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

41
42
43
44
45
46
47
48

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


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

62
63
64
65
66
67
68
69
70
71
# 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
}

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

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,
90
91
        request_plane: str = "tcp",
        store_backend: str = "etcd",
92
        durable_kv_events: bool = False,
93
94
95
96
97
98
99
100
    ):
        """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)
101
102
                - gpu_memory_utilization: Fraction of GPU memory to allocate (optional)
                - num_gpu_blocks_override: Cap on number of KV cache blocks (optional)
103
                - max_model_len: Maximum sequence length (optional)
104
                - enforce_eager: Disable CUDA graphs (default: False)
105
            num_workers: Number of vLLM worker processes
106
            single_gpu: If True, all workers share GPU 0
107
            data_parallel_size: If set, enables data parallelism with this many ranks (num_workers must equal data_parallel_size)
108
109
            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".
110
            durable_kv_events: If True, use JetStream for durable KV events. Defaults to False (NATS Core mode).
111
112
113
114
115
116
117
        """
        # 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
118
        self.data_parallel_size = data_parallel_size
119
        self.worker_processes = []
120
        self.store_backend = store_backend
121

122
123
124
125
126
127
128
129
130
131
132
        # 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
            )
        )

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

        block_size = vllm_args.get("block_size", BLOCK_SIZE)
        model = vllm_args.get("model", MODEL_NAME)
138
139
        gpu_memory_utilization = vllm_args.get("gpu_memory_utilization")
        num_gpu_blocks_override = vllm_args.get("num_gpu_blocks_override")
140
        max_model_len = vllm_args.get("max_model_len")
141
        enforce_eager = vllm_args.get("enforce_eager", False)
142
143
144
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
171
172
173
174
175
176
177
178
179

        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
        # - --connector nixl enables KV cache transfer between ranks

        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),
            ]

180
181
182
183
184
185
186
187
188
189
            # 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)]
                )

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

194
195
196
197
198
199
            # 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)]
                )

200
201
202
203
204
205
206
207
208
209
210
211
212
            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
                        # "--connector", "nixl",  # Required for KV transfer between DP ranks
                    ]
                )

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

217
218
219
220
221
            # 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]

222
            env = os.environ.copy()  # Copy parent environment
223
224
225
226
            env_vars = {
                "CUDA_VISIBLE_DEVICES": gpu_device,
                "DYN_NAMESPACE": self.namespace,
                "DYN_REQUEST_PLANE": request_plane,
227
228
229
                "DYN_SYSTEM_PORT": str(system_port),
                "DYN_VLLM_KV_EVENT_PORT": str(kv_event_port),
                "VLLM_NIXL_SIDE_CHANNEL_PORT": str(nixl_port),
230
231
232
233
234
235
236
237
                "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)
238
239
240
241
242
243
244
245
246
247

            # 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,
248
                terminate_all_matching_process_names=False,
249
250
251
252
253
            )
            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} "
254
                    f"(gpu_mem={gpu_memory_utilization}, system_port={system_port}) "
255
256
257
258
259
                    f"with endpoint: {self.endpoint}"
                )
            else:
                logger.info(
                    f"Created vLLM worker {worker_idx} on GPU {gpu_device} "
260
                    f"(gpu_mem={gpu_memory_utilization}, system_port={system_port}) "
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
                    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]
283
                process.log_dir = resolve_test_output_path(process.log_dir)
284
285
286
287
288
289
290
                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)

291
                process._terminate_all_matching_process_names()
292
293
294
295
296
297
298
299
300
301
302
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
                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)


355
@pytest.mark.pre_merge
356
@pytest.mark.gpu_1
357
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
358
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
359
def test_vllm_kv_router_basic(
360
361
362
363
364
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
365
):
366
367
    """
    Quick e2e sanity test for KV router with vLLM engine instances.
368
    Tests both NATS and TCP request planes.
369
370
371
372
    """

    # runtime_services starts etcd and nats
    N_VLLM_WORKERS = 2
373
374
375
    logger.info(
        f"Starting vLLM KV router test with {N_VLLM_WORKERS} workers using request_plane={request_plane}"
    )
376

377
378
379
380
381
382
383
    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:
384
385
386
387
        # Start vLLM workers
        logger.info(f"Starting {N_VLLM_WORKERS} vLLM workers")
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

388
        # Run basic router test (starts router internally and waits for workers to be ready)
389
        frontend_port = allocate_frontend_ports(request, 1)[0]
390
391
392
393
        _test_router_basic(
            engine_workers=vllm_workers,
            block_size=BLOCK_SIZE,
            request=request,
394
            frontend_port=frontend_port,
395
396
397
398
            test_payload=TEST_PAYLOAD,
            num_requests=NUM_REQUESTS,
            frontend_timeout=180,  # 3 minutes should be plenty for TinyLlama
            store_backend="etcd",  # Explicit for clarity
399
            request_plane=request_plane,
400
401
402
        )


403
@pytest.mark.pre_merge
404
@pytest.mark.gpu_1
405
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
406
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
Yan Ru Pei's avatar
Yan Ru Pei committed
407
408
409
410
411
@pytest.mark.parametrize(
    "router_event_threads",
    [1, 2],
    ids=["single_thread", "multi_thread"],
)
412
def test_router_decisions_vllm_multiple_workers(
413
414
415
416
417
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
Yan Ru Pei's avatar
Yan Ru Pei committed
418
    router_event_threads,
419
420
421
422
423
):
    # runtime_services starts etcd and nats
    logger.info("Starting vLLM router prefix reuse test with two workers")
    N_WORKERS = 2

424
425
426
427
428
429
430
    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:
431
432
        # Start 2 worker processes on the same GPU
        logger.info("Starting 2 vLLM worker processes on single GPU (gpu_mem=0.4)")
433
434
435
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

        # Get runtime and create endpoint
436
        runtime = get_runtime(request_plane=request_plane)
437
438
439
440
441
        namespace = runtime.namespace(vllm_workers.namespace)
        component = namespace.component("backend")
        endpoint = component.endpoint("generate")

        _test_router_decisions(
Yan Ru Pei's avatar
Yan Ru Pei committed
442
443
444
445
446
447
            vllm_workers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=False,
            router_event_threads=router_event_threads,
448
449
450
451
        )


@pytest.mark.gpu_2
452
@pytest.mark.nightly
453
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
454
@pytest.mark.timeout(600)  # 10 min max (multi-GPU + DP startup variance)
455
def test_router_decisions_vllm_dp(
456
457
458
459
460
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
461
):
462
463
464
465
466
467
468
469
470
471
    """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

472
473
474
475
476
477
478
479
    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:
480
        logger.info("Starting 2 vLLM DP ranks (dp_size=2) (gpu_mem=0.4)")
481
482
483
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

        # Get runtime and create endpoint
484
        runtime = get_runtime(request_plane=request_plane)
485
486
487
488
489
490
491
492
493
        # Use the namespace from the vLLM workers
        namespace = runtime.namespace(vllm_workers.namespace)
        component = namespace.component("backend")  # endpoint is backend.generate
        endpoint = component.endpoint("generate")

        _test_router_decisions(
            vllm_workers, endpoint, MODEL_NAME, request, test_dp_rank=True
        )

494
495
496

@pytest.mark.pre_merge
@pytest.mark.gpu_1
497
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
498
@pytest.mark.parametrize(
499
    "store_backend,durable_kv_events,request_plane",
500
    [
501
        ("etcd", False, "tcp"),
502
    ],
503
504
    ids=["nats_core"],
    indirect=["durable_kv_events", "request_plane"],
505
)
506
def test_vllm_indexers_sync(
507
508
509
510
511
512
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    file_storage_backend,
    set_ucx_tls_no_mm,
    store_backend,
513
    durable_kv_events,
514
    request_plane,
515
516
517
518
):
    """
    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.
519
520

    Tests with configuration:
521
522
    - nats_core: etcd backend, local indexer with NATS Core, TCP request plane
                 (includes NATS interruption/recovery testing)
523
    """
524
    # runtime_services_dynamic_ports handles NATS and etcd startup
525
526
    nats_process, _etcd_process = runtime_services_dynamic_ports

527
528
    logger.info(
        f"Starting vLLM indexers sync test: store_backend={store_backend}, "
529
        f"durable_kv_events={durable_kv_events}, request_plane={request_plane}"
530
531
    )

532
533
    N_VLLM_WORKERS = 2

534
535
536
537
538
539
540
541
542
    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:
543
544
545
546
547
548
        # 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
549
        # When using durable_kv_events=True, use JetStream mode for the router
550
551
552
553
554
        _test_router_indexers_sync(
            engine_workers=vllm_workers,
            block_size=BLOCK_SIZE,
            model_name=MODEL_NAME,
            num_workers=N_VLLM_WORKERS,
555
556
            store_backend=store_backend,
            request_plane=request_plane,
557
558
559
            test_nats_interruption=not durable_kv_events,
            nats_server=nats_process if not durable_kv_events else None,
            durable_kv_events=durable_kv_events,
560
561
562
        )

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