test_vllm.py 15.7 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
# 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

from tests.conftest import NatsServer
from tests.fault_tolerance.etcd_ha.utils import (
    DynamoFrontendProcess,
    EtcdCluster,
    send_inference_request,
    wait_for_processes_to_terminate,
)
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
from tests.utils.engine_process import FRONTEND_PORT
from tests.utils.managed_process import ManagedProcess
from tests.utils.payloads import check_health_generate, check_models_api

logger = logging.getLogger(__name__)


class DynamoWorkerProcess(ManagedProcess):
    """Process manager for Dynamo worker with vLLM backend and ETCD HA support"""

28
    def __init__(self, request, etcd_endpoints: list, is_prefill: bool = False):
29
30
31
32
33
34
35
36
37
38
39
40
41
        command = [
            "python3",
            "-m",
            "dynamo.vllm",
            "--model",
            FAULT_TOLERANCE_MODEL_NAME,
            "--enforce-eager",
            "--gpu-memory-utilization",
            "0.45",
            "--max-model-len",
            "8192",
        ]

42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
        # Set port based on worker type
        port = "8082" if is_prefill else "8081"

        # Configure health check based on worker type
        if is_prefill:
            # Prefill workers check their own status endpoint
            command.append("--is-prefill-worker")
            health_check_urls = [(f"http://localhost:{port}/health", self.is_ready)]
        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),
            ]
58
59
60
61
62

        # Set debug logging and ETCD endpoints
        env = os.environ.copy()
        env["DYN_LOG"] = "debug"
        env["ETCD_ENDPOINTS"] = ",".join(etcd_endpoints)
63
64
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
        env["DYN_SYSTEM_PORT"] = port
65

Alec's avatar
Alec committed
66
67
68
69
        if is_prefill:
            env["DYN_VLLM_KV_EVENT_PORT"] = "20082"
            env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = "5601"

70
71
72
        # Set log directory based on worker type
        worker_type = "prefill_worker" if is_prefill else "worker"
        log_dir = f"{request.node.name}_{worker_type}"
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

        # 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:
            pass

        super().__init__(
            command=command,
            env=env,
            health_check_urls=health_check_urls,
            timeout=120,
            display_output=True,
            terminate_existing=False,
            stragglers=[
                "VLLM::EngineCore",
            ],
            straggler_commands=[
                "-m dynamo.vllm",
            ],
            log_dir=log_dir,
        )

97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
        self.is_prefill = is_prefill

    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

114
115
116
117
118

@pytest.mark.vllm
@pytest.mark.gpu_1
@pytest.mark.e2e
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
119
@pytest.mark.nightly
120
@pytest.mark.timeout(600)
121
122
def test_etcd_ha_failover_vllm_aggregated(request, predownload_models):
    """
123
    Test ETCD High Availability with repeated node failures and recoveries.
124
125
126
127

    This test:
    1. Starts a 3-node ETCD cluster
    2. Starts NATS, frontend, and a vLLM worker
128
129
130
131
132
133
134
135
    3. Cycles through each of the 3 replicas:
       - Terminate the replica by index
       - Send inference request to verify system still works
       - Restart the terminated node

    This ensures testing of:
    - ETCD leader termination
    - Frontend/worker disconnection from their connected ETCD replica
136
137
138
139
140
141
    """
    # Step 1: Start NATS server
    with NatsServer(request):
        logger.info("NATS server started successfully")

        # Step 2: Start 3-node ETCD cluster
142
143
        num_replicas = 3
        with EtcdCluster(request, num_replicas=num_replicas) as etcd_cluster:
144
145
146
147
148
149
150
151
152
153
154
155
156
157
            logger.info("3-node ETCD cluster started successfully")

            # Get the endpoints for all ETCD nodes
            etcd_endpoints = etcd_cluster.get_client_endpoints()
            logger.info(f"ETCD endpoints: {etcd_endpoints}")

            # Step 3: Start the frontend with ETCD endpoints
            with DynamoFrontendProcess(request, etcd_endpoints):
                logger.info("Frontend started successfully")

                # Step 4: Start a vLLM worker
                with DynamoWorkerProcess(request, etcd_endpoints):
                    logger.info("Worker started successfully")

158
159
160
                    # Step 5: Send initial inference request to verify system is working
                    logger.info("Sending initial inference request")
                    result = send_inference_request("What is 2+2? The answer is")
161
                    assert (
162
163
                        "4" in result.lower() or "four" in result.lower()
                    ), f"Expected '4' or 'four' in response, got: '{result}'"
164

165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
                    # Step 6: Cycle through each replica to terminate/verify/restart
                    for i in range(num_replicas):
                        # Terminate a replica
                        logger.info(f"Iteration {i}: Terminating replica etcd-{i}")
                        etcd_cluster.terminate_replica(i)

                        # Send inference request to verify system still works
                        logger.info(
                            f"Iteration {i}: Sending inference request after termination"
                        )
                        result = send_inference_request(
                            "The capital of France is", max_tokens=20
                        )
                        assert (
                            "paris" in result.lower()
                        ), f"Iteration {i}: Expected 'Paris' in response, got: '{result}'"
181

182
183
184
                        # Restart the terminated replica
                        logger.info(f"Iteration {i}: Restarting replica etcd-{i}")
                        etcd_cluster.restart_replica(i)
185
186


187
188
189
@pytest.mark.vllm
@pytest.mark.gpu_1
@pytest.mark.e2e
190
@pytest.mark.nightly
191
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
192
@pytest.mark.timeout(600)
193
194
195
196
def test_etcd_ha_failover_vllm_disaggregated(
    request, predownload_models, set_ucx_tls_no_mm
):
    """
197
    Test ETCD High Availability with repeated node failures and recoveries in disaggregated mode.
198
199
200
201

    This test:
    1. Starts a 3-node ETCD cluster
    2. Starts NATS, frontend, and both prefill and decode vLLM workers
202
203
204
205
206
207
208
209
    3. Cycles through each of the 3 replicas:
       - Terminate the replica by index
       - Send inference request to verify system still works
       - Restart the terminated node

    This ensures testing of:
    - ETCD leader termination
    - Frontend/worker disconnection from their connected ETCD replica
210
211
212
213
214
215
    """
    # Step 1: Start NATS server
    with NatsServer(request):
        logger.info("NATS server started successfully")

        # Step 2: Start 3-node ETCD cluster
216
217
        num_replicas = 3
        with EtcdCluster(request, num_replicas=num_replicas) as etcd_cluster:
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
            logger.info("3-node ETCD cluster started successfully")

            # Get the endpoints for all ETCD nodes
            etcd_endpoints = etcd_cluster.get_client_endpoints()
            logger.info(f"ETCD endpoints: {etcd_endpoints}")

            # Step 3: Start the frontend with ETCD endpoints
            with DynamoFrontendProcess(request, etcd_endpoints):
                logger.info("Frontend started successfully")

                # Step 4: Start the prefill worker
                with DynamoWorkerProcess(request, etcd_endpoints, is_prefill=True):
                    logger.info("Prefill worker started successfully")

                    # Step 5: Start the decode worker
                    with DynamoWorkerProcess(request, etcd_endpoints, is_prefill=False):
                        logger.info("Decode worker started successfully")

236
237
238
                        # Step 6: Send initial inference request to verify system is working
                        logger.info("Sending initial inference request")
                        result = send_inference_request("What is 2+2? The answer is")
239
                        assert (
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
                            "4" in result.lower() or "four" in result.lower()
                        ), f"Expected '4' or 'four' in response, got: '{result}'"

                        # Step 7: Cycle through each replica to terminate/verify/restart
                        for i in range(num_replicas):
                            # Terminate a replica
                            logger.info(f"Iteration {i}: Terminating replica etcd-{i}")
                            etcd_cluster.terminate_replica(i)

                            # Send inference request to verify system still works
                            logger.info(
                                f"Iteration {i}: Sending inference request after termination"
                            )
                            result = send_inference_request(
                                "The capital of France is", max_tokens=20
                            )
                            assert (
                                "paris" in result.lower()
                            ), f"Iteration {i}: Expected 'Paris' in response, got: '{result}'"

                            # Restart the terminated replica
                            logger.info(f"Iteration {i}: Restarting replica etcd-{i}")
                            etcd_cluster.restart_replica(i)
263
264


265
266
267
@pytest.mark.vllm
@pytest.mark.gpu_1
@pytest.mark.e2e
268
@pytest.mark.nightly
269
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
270
@pytest.mark.timeout(600)
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
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
def test_etcd_non_ha_shutdown_vllm_aggregated(request, predownload_models):
    """
    Test that frontend and worker shut down when single ETCD node is terminated.

    This test:
    1. Starts a single ETCD node (no cluster)
    2. Starts NATS, frontend, and a vLLM worker
    3. Sends an inference request to verify the system works
    4. Terminates the single ETCD node
    5. Verifies that frontend and worker shut down gracefully
    """
    # Step 1: Start NATS server
    with NatsServer(request):
        logger.info("NATS server started successfully")

        # Step 2: Start single ETCD node using EtcdCluster with num_replicas=1
        with EtcdCluster(request, num_replicas=1) as etcd_cluster:
            logger.info("Single ETCD node started successfully")

            # Get the endpoint for the single ETCD node
            etcd_endpoints = etcd_cluster.get_client_endpoints()
            logger.info(f"ETCD endpoint: {etcd_endpoints}")

            # Step 3: Start the frontend with ETCD endpoint
            with DynamoFrontendProcess(request, etcd_endpoints) as frontend:
                logger.info("Frontend started successfully")

                # Step 4: Start a vLLM worker
                with DynamoWorkerProcess(request, etcd_endpoints) as worker:
                    logger.info("Worker started successfully")

                    # Step 5: Send inference request to verify system is working
                    logger.info("Sending inference request")
                    result = send_inference_request("What is 2+2? The answer is")
                    assert (
                        "4" in result.lower() or "four" in result.lower()
                    ), f"Expected '4' or 'four' in response, got: '{result}'"

                    logger.info("System is working correctly with single ETCD node")

                    # Step 6: Terminate the ETCD node
                    logger.info("Terminating single ETCD node")
                    etcd_cluster.stop()

                    # Step 7: Wait and verify frontend and worker detect the loss
                    wait_for_processes_to_terminate(
                        {"Worker": worker, "Frontend": frontend}
                    )
319
320
321
322
323


@pytest.mark.vllm
@pytest.mark.gpu_1
@pytest.mark.e2e
324
@pytest.mark.nightly
325
@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME)
326
@pytest.mark.timeout(600)
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
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
def test_etcd_non_ha_shutdown_vllm_disaggregated(
    request, predownload_models, set_ucx_tls_no_mm
):
    """
    Test that frontend and workers shut down when single ETCD node is terminated in disaggregated mode.

    This test:
    1. Starts a single ETCD node (no cluster)
    2. Starts NATS, frontend, and both prefill and decode vLLM workers
    3. Sends an inference request to verify the system works
    4. Terminates the single ETCD node
    5. Verifies that frontend and both workers shut down gracefully
    """
    # Step 1: Start NATS server
    with NatsServer(request):
        logger.info("NATS server started successfully")

        # Step 2: Start single ETCD node using EtcdCluster with num_replicas=1
        with EtcdCluster(request, num_replicas=1) as etcd_cluster:
            logger.info("Single ETCD node started successfully")

            # Get the endpoint for the single ETCD node
            etcd_endpoints = etcd_cluster.get_client_endpoints()
            logger.info(f"ETCD endpoint: {etcd_endpoints}")

            # Step 3: Start the frontend with ETCD endpoint
            with DynamoFrontendProcess(request, etcd_endpoints) as frontend:
                logger.info("Frontend started successfully")

                # Step 4: Start the prefill worker
                with DynamoWorkerProcess(
                    request, etcd_endpoints, is_prefill=True
                ) as prefill_worker:
                    logger.info("Prefill worker started successfully")

                    # Step 5: Start the decode worker
                    with DynamoWorkerProcess(
                        request, etcd_endpoints, is_prefill=False
                    ) as decode_worker:
                        logger.info("Decode worker started successfully")

                        # Step 6: Send inference request to verify system is working
                        logger.info("Sending inference request")
                        result = send_inference_request("What is 2+2? The answer is")
                        assert (
                            "4" in result.lower() or "four" in result.lower()
                        ), f"Expected '4' or 'four' in response, got: '{result}'"

                        logger.info(
                            "System is working correctly with single ETCD node in disaggregated mode"
                        )

                        # Step 7: Terminate the ETCD node
                        logger.info("Terminating single ETCD node")
                        etcd_cluster.stop()

                        # Step 8: Wait and verify frontend and both workers detect the loss
                        wait_for_processes_to_terminate(
                            {
                                "Decode Worker": decode_worker,
                                "Prefill Worker": prefill_worker,
                                "Frontend": frontend,
                            }
                        )