test_router_e2e_with_vllm.py 22.8 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
import logging
import os
import time
from typing import Any, Dict, Optional

import pytest

16
from tests.router.common import (
17
18
    _test_router_basic,
    _test_router_decisions,
19
    _test_router_indexers_sync,
20
)
21
from tests.router.helper import 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

        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
150
        # - --kv-transfer-config enables KV cache transfer between ranks
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

        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
            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
209
                        # "--kv-transfer-config", '{"kv_connector":"NixlConnector","kv_role":"kv_both"}',  # Required for KV transfer between DP ranks
210
211
212
                    ]
                )

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
223
224
225
226
227
228
229
230
231
232
            # 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])

233
            env = os.environ.copy()  # Copy parent environment
234
235
236
237
            env_vars = {
                "CUDA_VISIBLE_DEVICES": gpu_device,
                "DYN_NAMESPACE": self.namespace,
                "DYN_REQUEST_PLANE": request_plane,
238
239
                "DYN_SYSTEM_PORT": str(system_port),
                "VLLM_NIXL_SIDE_CHANNEL_PORT": str(nixl_port),
240
241
242
243
244
245
246
247
                "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)
248
249
250
251
252
253
254
255
256
257

            # 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,
258
                terminate_all_matching_process_names=False,
259
260
261
262
263
            )
            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} "
264
                    f"(gpu_mem={gpu_memory_utilization}, system_port={system_port}) "
265
266
267
268
269
                    f"with endpoint: {self.endpoint}"
                )
            else:
                logger.info(
                    f"Created vLLM worker {worker_idx} on GPU {gpu_device} "
270
                    f"(gpu_mem={gpu_memory_utilization}, system_port={system_port}) "
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
                    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]
293
                process.log_dir = resolve_test_output_path(process.log_dir)
294
295
296
297
298
299
300
                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)

301
                process._terminate_all_matching_process_names()
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
355
356
357
358
359
360
361
362
363
364
                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)


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

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

387
388
389
390
391
392
393
    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:
394
395
396
397
        # Start vLLM workers
        logger.info(f"Starting {N_VLLM_WORKERS} vLLM workers")
        logger.info(f"All vLLM workers using namespace: {vllm_workers.namespace}")

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


413
@pytest.mark.pre_merge
414
@pytest.mark.gpu_1
415
@pytest.mark.timeout(150)  # ~3x average (~43s/test), rounded up
416
@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True)
417
def test_router_decisions_vllm_multiple_workers(
418
419
420
421
422
    request,
    runtime_services_dynamic_ports,
    predownload_models,
    set_ucx_tls_no_mm,
    request_plane,
423
424
425
426
427
):
    # runtime_services starts etcd and nats
    logger.info("Starting vLLM router prefix reuse test with two workers")
    N_WORKERS = 2

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

        # Get runtime and create endpoint
440
        runtime = get_runtime(request_plane=request_plane)
441
        endpoint = runtime.endpoint(f"{vllm_workers.namespace}.backend.generate")
442
443

        _test_router_decisions(
Yan Ru Pei's avatar
Yan Ru Pei committed
444
445
446
447
448
            vllm_workers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=False,
449
            block_size=BLOCK_SIZE,
450
451
452
453
        )


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

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

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

        _test_router_decisions(
493
494
495
496
497
498
            vllm_workers,
            endpoint,
            MODEL_NAME,
            request,
            test_dp_rank=True,
            block_size=BLOCK_SIZE,
499
500
        )

501
502
503

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

    Tests with configuration:
528
529
    - nats_core: etcd backend, local indexer with NATS Core, TCP request plane
                 (includes NATS interruption/recovery testing)
530
    """
531
    # runtime_services_dynamic_ports handles NATS and etcd startup
532
533
    nats_process, _etcd_process = runtime_services_dynamic_ports

534
535
    logger.info(
        f"Starting vLLM indexers sync test: store_backend={store_backend}, "
536
        f"durable_kv_events={durable_kv_events}, request_plane={request_plane}"
537
538
    )

539
540
    N_VLLM_WORKERS = 2

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

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