test_vllm.py 14.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import logging
import os
import shutil

import pytest

10
11
from tests.fault_tolerance.cancellation.utils import (
    DynamoFrontendProcess,
12
13
14
    poll_for_pattern,
    read_streaming_responses,
    send_cancellable_request,
15
)
Alec's avatar
Alec committed
16
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
17
from tests.utils.engine_process import FRONTEND_PORT
18
from tests.utils.managed_process import ManagedProcess
19
from tests.utils.payloads import check_health_generate, check_models_api
20
21
22

logger = logging.getLogger(__name__)

23
24
25
26
27
pytestmark = [
    pytest.mark.vllm,
    pytest.mark.gpu_1,
    pytest.mark.e2e,
    pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME),
28
    pytest.mark.post_merge,  # post_merge to pinpoint failure commit
29
30
]

31
32
33
34
35
36
37
38
39
40

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

    def __init__(self, request, is_prefill: bool = False):
        command = [
            "python3",
            "-m",
            "dynamo.vllm",
            "--model",
Alec's avatar
Alec committed
41
            FAULT_TOLERANCE_MODEL_NAME,
42
43
44
45
            "--enforce-eager",
            "--gpu-memory-utilization",
            "0.45",
            "--max-model-len",
46
            "16384",
47
48
49
50
51
52
53
            "--migration-limit",
            "3",
        ]

        # Set port based on worker type
        port = "8082" if is_prefill else "8081"

54
        # Configure health check based on worker type
55
        if is_prefill:
56
            # Prefill workers check their own status endpoint
57
58
            command.append("--is-prefill-worker")
            health_check_urls = [(f"http://localhost:{port}/health", self.is_ready)]
59
60
61
62
63
64
65
66
        else:
            # Decode workers should also check their own status endpoint first,
            # then verify the frontend sees the model
            health_check_urls = [
                (f"http://localhost:{port}/health", self.is_ready),
                (f"http://localhost:{FRONTEND_PORT}/v1/models", check_models_api),
                (f"http://localhost:{FRONTEND_PORT}/health", check_health_generate),
            ]
67

68
        # Set environment variables
69
        env = os.environ.copy()
70
71
        env["DYN_REQUEST_PLANE"] = request.getfixturevalue("request_plane")

72
        env["DYN_LOG"] = "debug"
73
74
75
76
77
        # 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"
78
79
80
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
        env["DYN_SYSTEM_PORT"] = port

81
82
83
84
85
86
        # Set KV event port and NIXL side channel port only for prefill worker
        # to avoid conflicts with decode worker
        if is_prefill:
            env["DYN_VLLM_KV_EVENT_PORT"] = "20082"
            env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = "5601"

87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
        # 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,
102
            health_check_urls=health_check_urls,
103
104
105
            timeout=300,
            display_output=True,
            terminate_existing=False,
106
107
108
109
110
111
112
            # Ensure any orphaned vLLM engine cores or child helpers are cleaned up
            stragglers=[
                "VLLM::EngineCore",
            ],
            straggler_commands=[
                "-m dynamo.vllm",
            ],
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
            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

    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


138
@pytest.mark.timeout(110)  # 3x average
139
@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
140
def test_request_cancellation_vllm_aggregated(request, runtime_services):
141
    """
142
    End-to-end test for request cancellation functionality in aggregated mode.
143
144
145

    This test verifies that when a request is cancelled by the client,
    the system properly handles the cancellation and cleans up resources
146
    on the worker side in aggregated (single worker) mode. Tests three scenarios:
147
148
149
150
151
152
153
154
155
156
    1. Completion request
    2. Chat completion request (non-streaming)
    3. Chat completion request (streaming)
    """

    # Step 1: Start the frontend
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start a single worker
157
        with DynamoWorkerProcess(request) as worker:
158
159
            logger.info(f"Worker PID: {worker.get_pid()}")

160
            # Step 3: Test request cancellation with polling approach
161
162
163
164
165
166
167
168
169
170
171
            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",
                ),
            ]

172
            for request_type, description in test_scenarios:
173
174
                logger.info(f"Testing {description.lower()}...")

175
176
177
                # Send the request (non-blocking)
                cancellable_req = send_cancellable_request(request_type)

178
                # Poll for "Decode Request ID" pattern (vLLM v2 pattern)
179
180
                request_id, worker_log_offset = poll_for_pattern(
                    process=worker,
181
                    pattern="Decode Request ID: ",
182
183
                    log_offset=worker_log_offset,
                    match_type="contains",
184
                )
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205

                # 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,
206
207
208
209
210
                )

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


211
@pytest.mark.timeout(150)  # 3x average
212
213
214
215
216
217
218
219
220
221
222
@pytest.mark.parametrize(
    "request_plane",
    [
        "nats",
        pytest.param(
            "tcp",
            marks=pytest.mark.xfail(reason="Multi-worker TCP unstable", strict=False),
        ),
    ],
    indirect=True,
)
223
def test_request_cancellation_vllm_decode_cancel(
224
    request, runtime_services, set_ucx_tls_no_mm
Alec's avatar
Alec committed
225
):
226
    """
227
    End-to-end test for request cancellation during decode phase.
228

229
    This test verifies that when a request is cancelled by the client during the decode phase,
230
231
232
233
234
235
236
237
238
    the system properly handles the cancellation and cleans up resources
    on the decode worker side in a disaggregated setup.
    """

    # Step 1: Start the frontend
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start the prefill worker
239
        with DynamoWorkerProcess(request, is_prefill=True) as prefill_worker:
240
241
242
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

            # Step 3: Start the decode worker
243
            with DynamoWorkerProcess(request, is_prefill=False) as decode_worker:
244
245
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

246
                # Step 4: Test request cancellation for streaming scenario
247
                logger.info(
248
249
250
251
252
253
                    "Testing chat completion stream request cancellation in decode worker (decode phase)..."
                )

                # Send streaming request (non-blocking)
                cancellable_req = send_cancellable_request("chat_completion_stream")

254
                # Poll for "Decode Request ID" pattern in decode worker (vLLM v2 pattern)
255
256
                request_id, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
257
                    pattern="Decode Request ID: ",
258
259
260
                    match_type="contains",
                )

261
                # Verify same request ID reached prefill worker (as "Prefill Request ID")
262
263
                _, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
264
                    pattern=f"Prefill Request ID: {request_id}",
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
                )

                # 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",
285
286
287
                )

                logger.info(
288
                    "Chat completion stream cancellation in decode phase detected successfully"
289
290
291
                )


292
@pytest.mark.timeout(150)  # 3x average
293
294
295
296
297
298
299
300
301
302
303
@pytest.mark.parametrize(
    "request_plane",
    [
        "nats",
        pytest.param(
            "tcp",
            marks=pytest.mark.xfail(reason="Multi-worker TCP unstable", strict=False),
        ),
    ],
    indirect=True,
)
304
def test_request_cancellation_vllm_prefill_cancel(
305
    request, runtime_services, set_ucx_tls_no_mm
306
):
307
    """
308
    End-to-end test for request cancellation during prefill phase.
309

310
    This test verifies that when a request is cancelled by the client during the prefill phase,
311
312
    the system properly handles the cancellation and cleans up resources
    on both the decode and prefill workers in a disaggregated setup.
313
    """
314
315
316
317
318
319

    # Step 1: Start the frontend
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start the prefill worker
320
        with DynamoWorkerProcess(request, is_prefill=True) as prefill_worker:
321
322
323
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

            # Step 3: Start the decode worker
324
            with DynamoWorkerProcess(request, is_prefill=False) as decode_worker:
325
326
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

327
328
329
                # 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
330
                logger.info(
331
                    "Testing completion request cancellation during prefill phase..."
332
333
                )

334
335
336
                # Send request with long prompt (non-blocking)
                cancellable_req = send_cancellable_request(
                    "completion", use_long_prompt=True
337
                )
338

339
340
341
                # 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(
342
                    process=prefill_worker,
343
344
                    pattern="Prefill Request ID: ",
                    match_type="contains",
345
346
347
348
                )

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

351
                # Poll for "Aborted Prefill Request ID" in prefill worker (where cancellation happens)
352
353
354
355
356
357
358
359
360
361
362
363
                _, 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",
                )

364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
                # 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}"

381
                logger.info(
382
                    "Completion request cancellation during prefill phase detected successfully"
383
                )