test_router_e2e_with_vllm.py 22.4 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
26
27
28

logger = logging.getLogger(__name__)

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

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

40
41
42
43
44
45
46
47

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


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

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

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

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

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

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

        block_size = vllm_args.get("block_size", BLOCK_SIZE)
        model = vllm_args.get("model", MODEL_NAME)
137
138
        gpu_memory_utilization = vllm_args.get("gpu_memory_utilization")
        num_gpu_blocks_override = vllm_args.get("num_gpu_blocks_override")
139
        max_model_len = vllm_args.get("max_model_len")
140
        enforce_eager = vllm_args.get("enforce_eager", False)
141
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

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

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

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

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

199
200
201
202
203
204
205
206
207
208
209
210
211
            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
                    ]
                )

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

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

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

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

289
                process._terminate_all_matching_process_names()
290
291
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
                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)


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

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

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

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


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

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

        # Get runtime and create endpoint
434
        runtime = get_runtime(request_plane=request_plane)
435
436
437
438
439
        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
440
441
442
443
444
445
            vllm_workers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=False,
            router_event_threads=router_event_threads,
446
447
448
449
        )


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

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

        # Get runtime and create endpoint
482
        runtime = get_runtime(request_plane=request_plane)
483
484
485
486
487
488
489
490
491
        # 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
        )

492
493
494

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

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

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

530
531
    N_VLLM_WORKERS = 2

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

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