test_vllm.py 14 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
69
70
        # Set debug logging environment
        env = os.environ.copy()
        env["DYN_LOG"] = "debug"
71
72
73
74
75
        # 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"
76
77
78
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
        env["DYN_SYSTEM_PORT"] = port

79
80
81
82
83
84
        # 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"

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


136
@pytest.mark.timeout(110)  # 3x average
137
138
139
def test_request_cancellation_vllm_aggregated(
    request, runtime_services, predownload_models
):
140
    """
141
    End-to-end test for request cancellation functionality in aggregated mode.
142
143
144

    This test verifies that when a request is cancelled by the client,
    the system properly handles the cancellation and cleans up resources
145
    on the worker side in aggregated (single worker) mode. Tests three scenarios:
146
147
148
149
150
151
152
153
154
155
    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
156
        with DynamoWorkerProcess(request) as worker:
157
158
            logger.info(f"Worker PID: {worker.get_pid()}")

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

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

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

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

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

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


210
@pytest.mark.timeout(150)  # 3x average
211
def test_request_cancellation_vllm_decode_cancel(
Alec's avatar
Alec committed
212
    request, runtime_services, predownload_models, set_ucx_tls_no_mm
Alec's avatar
Alec committed
213
):
214
    """
215
    End-to-end test for request cancellation during decode phase.
216

217
    This test verifies that when a request is cancelled by the client during the decode phase,
218
219
220
221
222
223
224
225
226
    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
227
        with DynamoWorkerProcess(request, is_prefill=True) as prefill_worker:
228
229
230
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

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

234
                # Step 4: Test request cancellation for streaming scenario
235
                logger.info(
236
237
238
239
240
241
                    "Testing chat completion stream request cancellation in decode worker (decode phase)..."
                )

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

242
                # Poll for "Decode Request ID" pattern in decode worker (vLLM v2 pattern)
243
244
                request_id, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
245
                    pattern="Decode Request ID: ",
246
247
248
                    match_type="contains",
                )

249
                # Verify same request ID reached prefill worker (as "Prefill Request ID")
250
251
                _, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
252
                    pattern=f"Prefill Request ID: {request_id}",
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
                )

                # 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",
273
274
275
                )

                logger.info(
276
                    "Chat completion stream cancellation in decode phase detected successfully"
277
278
279
                )


280
@pytest.mark.timeout(150)  # 3x average
281
def test_request_cancellation_vllm_prefill_cancel(
Alec's avatar
Alec committed
282
    request, runtime_services, predownload_models, set_ucx_tls_no_mm
283
):
284
    """
285
    End-to-end test for request cancellation during prefill phase.
286

287
    This test verifies that when a request is cancelled by the client during the prefill phase,
288
289
    the system properly handles the cancellation and cleans up resources
    on both the decode and prefill workers in a disaggregated setup.
290
    """
291
292
293
294
295
296

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

        # Step 2: Start the prefill worker
297
        with DynamoWorkerProcess(request, is_prefill=True) as prefill_worker:
298
299
300
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

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

304
305
306
                # 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
307
                logger.info(
308
                    "Testing completion request cancellation during prefill phase..."
309
310
                )

311
312
313
                # Send request with long prompt (non-blocking)
                cancellable_req = send_cancellable_request(
                    "completion", use_long_prompt=True
314
                )
315

316
317
318
                # 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(
319
                    process=prefill_worker,
320
321
                    pattern="Prefill Request ID: ",
                    match_type="contains",
322
323
324
325
                )

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

328
                # Poll for "Aborted Prefill Request ID" in prefill worker (where cancellation happens)
329
330
331
332
333
334
335
336
337
338
339
340
                _, 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",
                )

341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
                # 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}"

358
                logger.info(
359
                    "Completion request cancellation during prefill phase detected successfully"
360
                )