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

4
5
6
7
8
9
10
11
12
"""
Test Execution Times (Last Run: 2025-12-09):
- test_request_migration_sglang_worker_failure: ~58s (gpu_1)
- test_request_migration_sglang_graceful_shutdown: ~58s (gpu_1, skipped)
- test_no_request_migration_sglang_worker_failure: ~38s (gpu_1)
- test_no_request_migration_sglang_graceful_shutdown: ~38s (gpu_1, skipped)
- Total: 115.71s (0:01:55) for enabled tests
"""

13
14
15
16
17
18
19
20
21
import logging
import os
import shutil

import pytest

from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
from tests.utils.managed_process import ManagedProcess, terminate_process_tree
from tests.utils.payloads import check_models_api
22
from tests.utils.port_utils import allocate_port, deallocate_port
23
24
25
26
27
28
29

# Import utilities from the refactored utils module
from .utils import (
    DynamoFrontendProcess,
    determine_request_receiving_worker,
    start_completion_request,
    validate_completion_response,
30
    verify_migration_metrics,
31
32
33
34
35
36
37
38
39
40
    verify_migration_occurred,
)

logger = logging.getLogger(__name__)

pytestmark = [
    pytest.mark.sglang,
    pytest.mark.gpu_1,
    pytest.mark.e2e,
    pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME),
41
    pytest.mark.post_merge,  # post_merge to pinpoint failure commit
42
    pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True),
43
44
45
46
47
48
]


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

49
50
51
52
53
54
55
56
    def __init__(
        self,
        request,
        worker_id: str,
        system_port: int,
        frontend_port: int,
        migration_limit: int = 3,
    ):
57
        self.worker_id = worker_id
58
        self.system_port = system_port
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

        command = [
            "python3",
            "-m",
            "dynamo.sglang",
            "--model-path",
            FAULT_TOLERANCE_MODEL_NAME,
            "--served-model-name",
            FAULT_TOLERANCE_MODEL_NAME,
            "--trust-remote-code",
            "--skip-tokenizer-init",
            "--mem-fraction-static",
            "0.45",
            "--context-length",
            "8192",
            "--migration-limit",
            str(migration_limit),
        ]

78
        # Set environment variables
79
        env = os.environ.copy()
80
        env["DYN_REQUEST_PLANE"] = request.getfixturevalue("request_plane")
81
        env["DYN_LOG"] = "debug"
82
83
84
85
86
        # 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"
87
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
88
89
        env["DYN_SYSTEM_PORT"] = str(system_port)
        env["DYN_HTTP_PORT"] = str(frontend_port)
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

        # TODO: Have the managed process take a command name explicitly to distinguish
        #       between processes started with the same command.
        log_dir = f"{request.node.name}_{worker_id}"

        # 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,
            health_check_urls=[
107
108
                (f"http://localhost:{frontend_port}/v1/models", check_models_api),
                (f"http://localhost:{system_port}/health", self.is_ready),
109
110
111
112
113
114
115
116
117
            ],
            timeout=300,
            display_output=True,
            terminate_existing=False,
            stragglers=["SGLANG:EngineCore"],
            straggler_commands=["-m dynamo.sglang"],
            log_dir=log_dir,
        )

118
119
120
121
122
123
124
125
126
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Release allocated port when worker exits."""
        try:
            # system_port is a required parameter, always set in __init__
            deallocate_port(self.system_port)
        except Exception as e:
            logging.warning(f"Failed to release SGLang worker port: {e}")

        return super().__exit__(exc_type, exc_val, exc_tb)
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

    def is_ready(self, response) -> bool:
        """Check the health of the worker process"""
        try:
            data = response.json()
            if data.get("status") == "ready":
                logger.info(f"{self.worker_id} status is ready")
                return True
            logger.warning(
                f"{self.worker_id} status is not ready: {data.get('status')}"
            )
        except ValueError:
            logger.warning(f"{self.worker_id} health response is not valid JSON")
        return False


143
@pytest.mark.timeout(235)  # 3x average
144
def test_request_migration_sglang_worker_failure(
145
    request, runtime_services_dynamic_ports, set_ucx_tls_no_mm, predownload_models
146
147
148
149
150
151
152
):
    """
    End-to-end test for worker fault tolerance with migration support using SGLang.

    This test verifies that when a worker is killed during request processing,
    the system can handle the failure gracefully and migrate the request to
    another worker.
153
154
155
156
157

    Timing (Last Run: 2025-12-09): ~58s total
    - Engine initialization: ~22s (Worker1: 12s, Worker2: 10s)
    - Test execution (request + migration): ~21s
    - Teardown: ~15s
158
159
    """

160
161
162
163
164
    # Allocate ports to avoid conflicts with parallel tests
    worker1_system_port = allocate_port(9100)
    worker2_system_port = allocate_port(9200)

    # Step 1: Start the frontend (allocates its own port)
165
166
167
168
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start 2 workers sequentially
169
170
171
172
173
174
        with DynamoWorkerProcess(
            request,
            "worker1",
            system_port=worker1_system_port,
            frontend_port=frontend.frontend_port,
        ) as worker1:
175
176
            logger.info(f"Worker 1 PID: {worker1.get_pid()}")

177
178
179
180
181
182
            with DynamoWorkerProcess(
                request,
                "worker2",
                system_port=worker2_system_port,
                frontend_port=frontend.frontend_port,
            ) as worker2:
183
184
185
                logger.info(f"Worker 2 PID: {worker2.get_pid()}")

                # Step 3: Send the request
186
187
188
                request_thread, response_list = start_completion_request(
                    frontend.frontend_port
                )
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206

                # Step 4: Use polling to determine which worker received the request
                worker, worker_name = determine_request_receiving_worker(
                    worker1, worker2, receiving_pattern="New Request ID: "
                )

                # Step 5: Kill the worker that has the request
                logger.info(
                    f"Killing {worker_name} with PID {worker.get_pid()} processing the request"
                )
                terminate_process_tree(worker.get_pid(), immediate_kill=True, timeout=0)

                # Step 6: Validate the completion response
                validate_completion_response(request_thread, response_list)

                # Step 7: Verify migration occurred
                verify_migration_occurred(frontend)

207
208
209
210
211
                # Step 8: Verify migration metrics
                verify_migration_metrics(
                    frontend.frontend_port, expected_ongoing_request_count=1
                )

212

213
@pytest.mark.timeout(235)  # 3x average
214
215
@pytest.mark.skip(reason="SGLang graceful shutdown not yet implemented")
def test_request_migration_sglang_graceful_shutdown(
216
    request, runtime_services_dynamic_ports, set_ucx_tls_no_mm, predownload_models
217
218
219
220
221
222
223
224
225
):
    """
    End-to-end test for worker fault tolerance with graceful shutdown and migration support using SGLang.

    This test verifies that when a worker receives a graceful shutdown signal (SIGTERM)
    during request processing, the system can handle the shutdown gracefully and migrate
    the request to another worker. Unlike the abrupt kill test, this simulates a more
    controlled shutdown scenario where the worker has time to clean up and notify the
    system about its shutdown.
226
227
228
229
230

    Timing (Last Run: 2025-12-09): ~58s total (estimated, similar to worker_failure)
    - Engine initialization: ~22s (Worker1: 12s, Worker2: 10s)
    - Test execution (request + graceful shutdown + migration): ~21s
    - Teardown: ~15s
231
232
    """

233
234
235
236
237
    # Allocate ports to avoid conflicts with parallel tests
    worker1_system_port = allocate_port(9100)
    worker2_system_port = allocate_port(9200)

    # Step 1: Start the frontend (allocates its own port)
238
239
240
241
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start 2 workers sequentially
242
243
244
245
246
247
        with DynamoWorkerProcess(
            request,
            "worker1",
            system_port=worker1_system_port,
            frontend_port=frontend.frontend_port,
        ) as worker1:
248
249
            logger.info(f"Worker 1 PID: {worker1.get_pid()}")

250
251
252
253
254
255
            with DynamoWorkerProcess(
                request,
                "worker2",
                system_port=worker2_system_port,
                frontend_port=frontend.frontend_port,
            ) as worker2:
256
257
258
                logger.info(f"Worker 2 PID: {worker2.get_pid()}")

                # Step 3: Send the request
259
260
261
                request_thread, response_list = start_completion_request(
                    frontend.frontend_port
                )
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281

                # Step 4: Use polling to determine which worker received the request
                worker, worker_name = determine_request_receiving_worker(
                    worker1, worker2, receiving_pattern="New Request ID: "
                )

                # Step 5: Gracefully shutdown the worker that has the request
                logger.info(
                    f"Gracefully shutting down {worker_name} with PID {worker.get_pid()} processing the request"
                )
                terminate_process_tree(
                    worker.get_pid(), immediate_kill=False, timeout=10
                )

                # Step 6: Validate the completion response
                validate_completion_response(request_thread, response_list)

                # Step 7: Verify migration occurred during graceful shutdown
                verify_migration_occurred(frontend)

282
283
284
285
286
                # Step 8: Verify migration metrics
                verify_migration_metrics(
                    frontend.frontend_port, expected_ongoing_request_count=1
                )

287

288
@pytest.mark.timeout(135)  # 3x average
289
def test_no_request_migration_sglang_worker_failure(
290
    request, runtime_services_dynamic_ports, set_ucx_tls_no_mm, predownload_models
291
292
293
294
295
296
297
):
    """
    End-to-end test for worker fault tolerance with migration disabled using SGLang.

    This test verifies that when migration is disabled (migration_limit=0) and a worker
    is killed during request processing, the request fails as expected without migration.
    This is the opposite behavior of test_request_migration_sglang_worker_failure.
298
299
300
301
302

    Timing (Last Run: 2025-12-09): ~38s total
    - Engine initialization: ~23s (Worker1: 13s, Worker2: 10s)
    - Test execution (failure validation): <1s
    - Teardown: ~15s
303
304
    """

305
306
307
308
309
    # Allocate ports to avoid conflicts with parallel tests
    worker1_system_port = allocate_port(9100)
    worker2_system_port = allocate_port(9200)

    # Step 1: Start the frontend (allocates its own port)
310
311
312
313
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start 2 workers sequentially with migration disabled
314
315
316
317
318
319
320
        with DynamoWorkerProcess(
            request,
            "worker1",
            system_port=worker1_system_port,
            frontend_port=frontend.frontend_port,
            migration_limit=0,
        ) as worker1:
321
322
            logger.info(f"Worker 1 PID: {worker1.get_pid()}")

323
324
325
326
327
328
329
            with DynamoWorkerProcess(
                request,
                "worker2",
                system_port=worker2_system_port,
                frontend_port=frontend.frontend_port,
                migration_limit=0,
            ) as worker2:
330
331
332
                logger.info(f"Worker 2 PID: {worker2.get_pid()}")

                # Step 3: Send the request
333
334
335
                request_thread, response_list = start_completion_request(
                    frontend.frontend_port
                )
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

                # Step 4: Use polling to determine which worker received the request
                worker, worker_name = determine_request_receiving_worker(
                    worker1, worker2, receiving_pattern="New Request ID: "
                )

                # Step 5: Kill the worker that has the request
                logger.info(
                    f"Killing {worker_name} with PID {worker.get_pid()} processing the request"
                )
                terminate_process_tree(worker.get_pid(), immediate_kill=True, timeout=0)

                # Step 6: Validate the completion response - should fail without migration
                try:
                    validate_completion_response(request_thread, response_list)
                    pytest.fail(
                        "Request succeeded unexpectedly when migration was disabled"
                    )
                except AssertionError as e:
                    assert "Request failed with status 500: " in str(
                        e
                    ), f"Unexpected request error message: {e}"

                # Step 7: Verify migration did NOT occur - should fail
                try:
                    verify_migration_occurred(frontend)
                    pytest.fail(
                        "Migration verification unexpectedly passed when migration was disabled"
                    )
                except AssertionError as e:
                    assert "'Cannot recreate stream: ...' error found in logs" in str(
                        e
                    ), f"Unexpected migration message: {e}"


371
@pytest.mark.timeout(135)  # 3x average
372
373
@pytest.mark.skip(reason="SGLang graceful shutdown not yet implemented")
def test_no_request_migration_sglang_graceful_shutdown(
374
    request, runtime_services_dynamic_ports, set_ucx_tls_no_mm, predownload_models
375
376
377
378
379
380
381
382
):
    """
    End-to-end test for worker fault tolerance with graceful shutdown and migration disabled using SGLang.

    This test verifies that when migration is disabled (migration_limit=0) and a worker
    receives a graceful shutdown signal (SIGTERM) during request processing, the request
    fails as expected without migration. This is the opposite behavior of
    test_request_migration_sglang_graceful_shutdown.
383
384
385
386
387

    Timing (Last Run: 2025-12-09): ~38s total (estimated, similar to no_migration_worker_failure)
    - Engine initialization: ~23s (Worker1: 13s, Worker2: 10s)
    - Test execution (graceful shutdown + failure validation): <1s
    - Teardown: ~15s
388
389
    """

390
391
392
393
394
    # Allocate ports to avoid conflicts with parallel tests
    worker1_system_port = allocate_port(9100)
    worker2_system_port = allocate_port(9200)

    # Step 1: Start the frontend (allocates its own port)
395
396
397
398
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start 2 workers sequentially with migration disabled
399
400
401
402
403
404
405
        with DynamoWorkerProcess(
            request,
            "worker1",
            system_port=worker1_system_port,
            frontend_port=frontend.frontend_port,
            migration_limit=0,
        ) as worker1:
406
407
            logger.info(f"Worker 1 PID: {worker1.get_pid()}")

408
409
410
411
412
413
414
            with DynamoWorkerProcess(
                request,
                "worker2",
                system_port=worker2_system_port,
                frontend_port=frontend.frontend_port,
                migration_limit=0,
            ) as worker2:
415
416
417
                logger.info(f"Worker 2 PID: {worker2.get_pid()}")

                # Step 3: Send the request
418
419
420
                request_thread, response_list = start_completion_request(
                    frontend.frontend_port
                )
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455

                # Step 4: Use polling to determine which worker received the request
                worker, worker_name = determine_request_receiving_worker(
                    worker1, worker2, receiving_pattern="New Request ID: "
                )

                # Step 5: Gracefully shutdown the worker that has the request
                logger.info(
                    f"Gracefully shutting down {worker_name} with PID {worker.get_pid()} processing the request"
                )
                terminate_process_tree(
                    worker.get_pid(), immediate_kill=False, timeout=10
                )

                # Step 6: Validate the completion response - should fail without migration
                try:
                    validate_completion_response(request_thread, response_list)
                    pytest.fail(
                        "Request succeeded unexpectedly when migration was disabled"
                    )
                except AssertionError as e:
                    assert "Request failed with status 500: " in str(
                        e
                    ), f"Unexpected request error message: {e}"

                # Step 7: Verify migration did NOT occur - should fail
                try:
                    verify_migration_occurred(frontend)
                    pytest.fail(
                        "Migration verification unexpectedly passed when migration was disabled"
                    )
                except AssertionError as e:
                    assert "'Cannot recreate stream: ...' error found in logs" in str(
                        e
                    ), f"Unexpected migration message: {e}"