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

4
5
6
7
8
9
10
11
"""
Test Execution Times (Last Run: 2025-12-09):
- test_request_cancellation_vllm_aggregated: ~55s (gpu_1)
- test_request_cancellation_vllm_decode_cancel: ~53s (gpu_2)
- test_request_cancellation_vllm_prefill_cancel: ~53s (gpu_2)
- Total: 161.65s (0:02:41)
"""

12
import json
13
14
15
16
17
18
import logging
import os
import shutil

import pytest

19
20
from tests.fault_tolerance.cancellation.utils import (
    DynamoFrontendProcess,
21
22
23
    poll_for_pattern,
    read_streaming_responses,
    send_cancellable_request,
24
)
Alec's avatar
Alec committed
25
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
26
from tests.utils.managed_process import ManagedProcess
27
from tests.utils.payloads import check_health_generate, check_models_api
28
from tests.utils.port_utils import allocate_port, deallocate_port
29
30
31

logger = logging.getLogger(__name__)

32
pytestmark = [
33
    pytest.mark.fault_tolerance,
34
35
36
37
    pytest.mark.vllm,
    pytest.mark.gpu_1,
    pytest.mark.e2e,
    pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME),
38
    pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True),
39
40
]

41
42
43
44

class DynamoWorkerProcess(ManagedProcess):
    """Process manager for Dynamo worker with vLLM backend"""

45
46
47
48
49
50
51
52
53
54
55
    def __init__(
        self,
        request,
        frontend_port: int,
        is_prefill: bool = False,
    ):
        # Allocate system port for this worker
        system_port = allocate_port(9100)
        self.system_port = system_port
        self.frontend_port = frontend_port

56
57
58
59
60
        command = [
            "python3",
            "-m",
            "dynamo.vllm",
            "--model",
Alec's avatar
Alec committed
61
            FAULT_TOLERANCE_MODEL_NAME,
62
63
64
65
            "--enforce-eager",
            "--gpu-memory-utilization",
            "0.45",
            "--max-model-len",
66
            "16384",
67
68
        ]

69
        # Configure health check based on worker type
70
        if is_prefill:
71
            # Prefill workers check their own status endpoint
72
            command.extend(["--disaggregation-mode", "prefill"])
73
74
75
            health_check_urls = [
                (f"http://localhost:{system_port}/health", self.is_ready)
            ]
76
77
78
79
        else:
            # Decode workers should also check their own status endpoint first,
            # then verify the frontend sees the model
            health_check_urls = [
80
81
82
                (f"http://localhost:{system_port}/health", self.is_ready),
                (f"http://localhost:{frontend_port}/v1/models", check_models_api),
                (f"http://localhost:{frontend_port}/health", check_health_generate),
83
            ]
84

85
        # Set environment variables
86
        env = os.environ.copy()
87
88
        env["DYN_REQUEST_PLANE"] = request.getfixturevalue("request_plane")

89
        env["DYN_LOG"] = "debug"
90
91
92
93
94
        # Disable canary health check - these tests expect full control over requests
        # sent to the workers where canary health check intermittently sends dummy
        # requests to workers interfering with the test process which may cause
        # intermittent failures
        env["DYN_HEALTH_CHECK_ENABLED"] = "false"
95
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
96
97
        env["DYN_SYSTEM_PORT"] = str(system_port)
        env["DYN_HTTP_PORT"] = str(frontend_port)
98

99
        # Set KV events config and NIXL side channel port only for prefill worker
100
101
        # to avoid conflicts with decode worker
        if is_prefill:
102
103
104
105
106
107
108
109
110
111
112
113
114
            command.extend(
                [
                    "--kv-events-config",
                    json.dumps(
                        {
                            "publisher": "zmq",
                            "topic": "kv-events",
                            "endpoint": "tcp://*:20082",
                            "enable_kv_cache_events": True,
                        }
                    ),
                ]
            )
115
116
117
            env[
                "VLLM_NIXL_SIDE_CHANNEL_PORT"
            ] = "5601"  # TODO: use dynamic port allocation
118

119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
        # Set log directory based on worker type
        worker_type = "prefill_worker" if is_prefill else "worker"
        log_dir = f"{request.node.name}_{worker_type}"

        # Clean up any existing log directory from previous runs
        try:
            shutil.rmtree(log_dir)
            logger.info(f"Cleaned up existing log directory: {log_dir}")
        except FileNotFoundError:
            # Directory doesn't exist, which is fine
            pass

        super().__init__(
            command=command,
            env=env,
134
            health_check_urls=health_check_urls,
135
136
            timeout=300,
            display_output=True,
137
            terminate_all_matching_process_names=False,
138
139
140
141
142
143
144
            # Ensure any orphaned vLLM engine cores or child helpers are cleaned up
            stragglers=[
                "VLLM::EngineCore",
            ],
            straggler_commands=[
                "-m dynamo.vllm",
            ],
145
146
147
148
149
150
151
152
153
            log_dir=log_dir,
        )

        self.is_prefill = is_prefill

    def get_pid(self):
        """Get the PID of the worker process"""
        return self.proc.pid if self.proc else None

154
155
156
157
158
159
160
161
162
163
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Release allocated port when worker exits."""
        try:
            # system_port is always allocated in __init__
            deallocate_port(self.system_port)
        except Exception as e:
            logging.warning(f"Failed to release vLLM worker port: {e}")

        return super().__exit__(exc_type, exc_val, exc_tb)

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    def is_ready(self, response) -> bool:
        """Check the health of the worker process"""
        try:
            data = response.json()
            if data.get("status") == "ready":
                worker_type = "Prefill worker" if self.is_prefill else "Worker"
                logger.info(f"{worker_type} status is ready")
                return True
            worker_type = "Prefill worker" if self.is_prefill else "Worker"
            logger.warning(f"{worker_type} status is not ready: {data.get('status')}")
        except ValueError:
            worker_type = "Prefill worker" if self.is_prefill else "Worker"
            logger.warning(f"{worker_type} health response is not valid JSON")
        return False


180
@pytest.mark.timeout(110)  # 3x average
181
@pytest.mark.post_merge
182
183
184
def test_request_cancellation_vllm_aggregated(
    request, runtime_services_dynamic_ports, predownload_models
):
185
    """
186
    End-to-end test for request cancellation functionality in aggregated mode.
187
188
189

    This test verifies that when a request is cancelled by the client,
    the system properly handles the cancellation and cleans up resources
190
    on the worker side in aggregated (single worker) mode. Tests three scenarios:
191
192
193
    1. Completion request
    2. Chat completion request (non-streaming)
    3. Chat completion request (streaming)
194
195
196
197
198

    Timing (Last Run: 2025-12-09): ~55s total
    - Engine initialization: ~15s
    - Testing 3 scenarios: ~38s (~12s each)
    - Teardown: ~2s
199
200
    """

201
    # Step 1: Start the frontend (allocates its own frontend_port)
202
203
204
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

205
206
        # Step 2: Start a single worker (allocates its own system_port)
        with DynamoWorkerProcess(request, frontend.frontend_port) as worker:
207
208
            logger.info(f"Worker PID: {worker.get_pid()}")

209
            # Step 3: Test request cancellation with polling approach
210
211
212
213
214
215
216
217
218
219
220
            frontend_log_offset, worker_log_offset = 0, 0

            test_scenarios = [
                ("completion", "Completion request cancellation"),
                ("chat_completion", "Chat completion request cancellation"),
                (
                    "chat_completion_stream",
                    "Chat completion stream request cancellation",
                ),
            ]

221
            for request_type, description in test_scenarios:
222
223
                logger.info(f"Testing {description.lower()}...")

224
                # Send the request (non-blocking)
225
226
227
                cancellable_req = send_cancellable_request(
                    frontend.frontend_port, request_type
                )
228

229
                # Poll for "Decode Request ID" pattern (vLLM v2 pattern)
230
231
                request_id, worker_log_offset = poll_for_pattern(
                    process=worker,
232
                    pattern="Decode Request ID: ",
233
234
                    log_offset=worker_log_offset,
                    match_type="contains",
235
                )
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256

                # For streaming, read 5 responses before cancelling
                if request_type == "chat_completion_stream":
                    read_streaming_responses(cancellable_req, expected_count=5)

                # Now cancel the request
                cancellable_req.cancel()
                logger.info(f"Cancelled request ID: {request_id}")

                # Poll for "Aborted Request ID" with matching ID
                _, worker_log_offset = poll_for_pattern(
                    process=worker,
                    pattern=f"Aborted Request ID: {request_id}",
                    log_offset=worker_log_offset,
                )

                # Verify frontend log has kill message
                _, frontend_log_offset = poll_for_pattern(
                    process=frontend,
                    pattern="issued control message Kill to sender",
                    log_offset=frontend_log_offset,
257
258
259
260
261
                )

                logger.info(f"{description} detected successfully")


262
@pytest.mark.timeout(150)  # 3x average
263
@pytest.mark.nightly
264
def test_request_cancellation_vllm_decode_cancel(
265
    request, runtime_services_dynamic_ports, set_ucx_tls_no_mm, predownload_models
Alec's avatar
Alec committed
266
):
267
    """
268
    End-to-end test for request cancellation during decode phase.
269

270
    This test verifies that when a request is cancelled by the client during the decode phase,
271
272
    the system properly handles the cancellation and cleans up resources
    on the decode worker side in a disaggregated setup.
273
274
275
276
277

    Timing (Last Run: 2025-12-09): ~53s total (requires 2 GPUs)
    - Engine initialization: ~23s (decode + prefill workers)
    - Testing stream cancellation during decode: ~28s
    - Teardown: ~2s
278
279
    """

280
    # Step 1: Start the frontend (allocates its own frontend_port)
281
282
283
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

284
285
286
287
        # Step 2: Start the prefill worker (allocates its own system_port)
        with DynamoWorkerProcess(
            request, frontend.frontend_port, is_prefill=True
        ) as prefill_worker:
288
289
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

290
291
292
293
            # Step 3: Start the decode worker (allocates its own system_port)
            with DynamoWorkerProcess(
                request, frontend.frontend_port, is_prefill=False
            ) as decode_worker:
294
295
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

296
                # Step 4: Test request cancellation for streaming scenario
297
                logger.info(
298
299
300
301
                    "Testing chat completion stream request cancellation in decode worker (decode phase)..."
                )

                # Send streaming request (non-blocking)
302
303
304
                cancellable_req = send_cancellable_request(
                    frontend.frontend_port, "chat_completion_stream"
                )
305

306
                # Poll for "Decode Request ID" pattern in decode worker (vLLM v2 pattern)
307
308
                request_id, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
309
                    pattern="Decode Request ID: ",
310
311
312
                    match_type="contains",
                )

313
                # Verify same request ID reached prefill worker (as "Prefill Request ID")
314
315
                _, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
316
                    pattern=f"Prefill Request ID: {request_id}",
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
                )

                # Read 5 streaming responses (decode phase)
                read_streaming_responses(cancellable_req, expected_count=5)

                # Now cancel the request
                cancellable_req.cancel()
                logger.info(f"Cancelled request ID: {request_id}")

                # Poll for "Aborted Request ID" in decode worker
                _, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
                    pattern=f"Aborted Request ID: {request_id}",
                    log_offset=decode_log_offset,
                )

                # Verify frontend log has kill message
                _, frontend_log_offset = poll_for_pattern(
                    process=frontend,
                    pattern="issued control message Kill to sender",
337
338
339
                )

                logger.info(
340
                    "Chat completion stream cancellation in decode phase detected successfully"
341
342
343
                )


344
@pytest.mark.timeout(150)  # 3x average
345
@pytest.mark.nightly
346
def test_request_cancellation_vllm_prefill_cancel(
347
    request, runtime_services_dynamic_ports, set_ucx_tls_no_mm, predownload_models
348
):
349
    """
350
    End-to-end test for request cancellation during prefill phase.
351

352
    This test verifies that when a request is cancelled by the client during the prefill phase,
353
354
    the system properly handles the cancellation and cleans up resources
    on both the decode and prefill workers in a disaggregated setup.
355
356
357
358
359

    Timing (Last Run: 2025-12-09): ~53s total (requires 2 GPUs)
    - Engine initialization: ~23s (decode + prefill workers)
    - Testing cancellation during prefill: ~28s
    - Teardown: ~2s
360
    """
361

362
    # Step 1: Start the frontend (allocates its own frontend_port)
363
364
365
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

366
367
368
369
        # Step 2: Start the prefill worker (allocates its own system_port)
        with DynamoWorkerProcess(
            request, frontend.frontend_port, is_prefill=True
        ) as prefill_worker:
370
371
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

372
373
374
375
            # Step 3: Start the decode worker (allocates its own system_port)
            with DynamoWorkerProcess(
                request, frontend.frontend_port, is_prefill=False
            ) as decode_worker:
376
377
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

378
379
380
                # Step 4: Test request cancellation during prefill phase
                # Note: With the new architecture, prefill routing happens in the frontend,
                # so the request goes directly to the prefill worker first
381
                logger.info(
382
                    "Testing completion request cancellation during prefill phase..."
383
384
                )

385
386
                # Send request with long prompt (non-blocking)
                cancellable_req = send_cancellable_request(
387
                    frontend.frontend_port, "completion", use_long_prompt=True
388
                )
389

390
391
392
                # Poll for "Prefill Request ID" pattern in prefill worker (vLLM v2 pattern)
                # With new architecture, prefill is routed by frontend's internal router
                request_id, prefill_log_offset = poll_for_pattern(
393
                    process=prefill_worker,
394
395
                    pattern="Prefill Request ID: ",
                    match_type="contains",
396
397
398
399
                )

                # Cancel during prefill phase
                cancellable_req.cancel()
400
                logger.info(f"Cancelled request ID: {request_id} during prefill")
401

402
                # Poll for "Aborted Prefill Request ID" in prefill worker (where cancellation happens)
403
404
405
406
407
408
409
410
411
412
413
414
                _, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
                    pattern=f"Aborted Prefill Request ID: {request_id}",
                    log_offset=prefill_log_offset,
                )

                # Verify frontend log has kill message
                _, frontend_log_offset = poll_for_pattern(
                    process=frontend,
                    pattern="issued control message Kill to sender",
                )

415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
                # Verify decode worker never received the request
                pattern = "Request ID: "
                try:
                    _, decode_log_offset = poll_for_pattern(
                        process=decode_worker,
                        pattern=pattern,
                        max_wait_ms=10,
                        match_type="contains",
                    )
                    pytest.fail(
                        "Decode worker received request cancelled during prefill phase"
                    )
                except AssertionError as e:
                    assert str(e).startswith(
                        f"Failed to find '{pattern}' pattern after 2 iterations "
                    ), f"Unexpected error: {e}"

432
                logger.info(
433
                    "Completion request cancellation during prefill phase detected successfully"
434
                )