test_trtllm.py 16.2 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
Test Execution Times (Last Run: 2026-01-12):
- test_request_migration_trtllm_aggregated: ~95s
- test_request_migration_trtllm_prefill: N/A
- test_request_migration_trtllm_kv_transfer: N/A
- test_request_migration_trtllm_decode: N/A
10
11
"""

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

import pytest

from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
19
from tests.utils.managed_process import ManagedProcess
20
from tests.utils.payloads import check_models_api
21
from tests.utils.port_utils import allocate_port, deallocate_port
22

23
24
# Customized utils for migration tests
from .utils import DynamoFrontendProcess, run_migration_test
25
26
27
28
29
30
31
32

logger = logging.getLogger(__name__)

pytestmark = [
    pytest.mark.trtllm,
    pytest.mark.gpu_1,
    pytest.mark.e2e,
    pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME),
33
    pytest.mark.post_merge,  # post_merge to pinpoint failure commit
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
    pytest.mark.parametrize(
        "migration_limit", [3, 0], ids=["migration_enabled", "migration_disabled"]
    ),
    pytest.mark.parametrize(
        "immediate_kill",
        [
            pytest.param(True, id="worker_failure"),
            pytest.param(
                False,
                id="graceful_shutdown",
                marks=pytest.mark.xfail(
                    strict=False, reason="TRT-LLM graceful shutdown not yet implemented"
                ),
            ),
        ],
    ),
    pytest.mark.parametrize(
        "request_api",
        [
            pytest.param("chat"),
            pytest.param(
                "completion",
                marks=pytest.mark.skip(reason="Behavior unverified yet"),
            ),
        ],
    ),
    pytest.mark.parametrize(
        "stream",
        [
            pytest.param(True, id="stream"),
            pytest.param(
                False,
                id="unary",
                marks=pytest.mark.skip(reason="Behavior unverified yet"),
            ),
        ],
    ),
71
    pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True),
72
73
74
75
]


class DynamoWorkerProcess(ManagedProcess):
76
77
78
79
80
81
82
83
84
85
86
87
    """Process manager for Dynamo worker with TRT-LLM backend

    Supports both aggregated mode (single worker) and disaggregated mode
    (separate prefill and decode workers).

    Args:
        request: pytest request fixture
        worker_id: Unique identifier for the worker (e.g., "worker1", "prefill1")
        frontend_port: Port where the frontend is running
        migration_limit: Maximum number of migration attempts (default: 3)
        mode: "prefill_and_decode" for aggregated, "prefill" or "decode" for disaggregated
    """
88

89
90
91
92
93
94
    def __init__(
        self,
        request,
        worker_id: str,
        frontend_port: int,
        migration_limit: int = 3,
95
        mode: str = "prefill_and_decode",
96
    ):
97
        self.worker_id = worker_id
98
99
        self.system_port = allocate_port(9100)
        self.mode = mode
100

101
102
103
104
        # Prefill workers require migration_limit=0 (no KV cache migration support)
        if mode == "prefill":
            logging.info("Prefill worker - setting migration_limit to 0")
            migration_limit = 0
105
106
107
108
109
110
111
112

        command = [
            "python3",
            "-m",
            "dynamo.trtllm",
            "--model",
            FAULT_TOLERANCE_MODEL_NAME,
            "--disaggregation-mode",
113
            mode,
114
115
            "--max-seq-len",
            "8192",
116
117
118
119
            "--max-num-tokens",
            "8192",
            "--free-gpu-memory-fraction",
            "0.15",  # avoid validation error on TRT-LLM available memory checks
120
121
122
            "--migration-limit",
            str(migration_limit),
        ]
123
124
125
126
127
128
129
130
131
132
133
        if mode != "prefill_and_decode":
            config_file = (
                f"test_request_migration_trtllm_config_{self.system_port}.yaml"
            )
            with open(config_file, "w") as f:
                f.write(
                    "cache_transceiver_config:\n  backend: DEFAULT\n  max_tokens_in_buffer: 8192\n"
                )
                f.write("disable_overlap_scheduler: true\n")
                f.write("kv_cache_config:\n  max_tokens: 8192\n")
            command += ["--extra-engine-args", config_file]
134

135
        # Set environment variables
136
        env = os.environ.copy()
137
        env["DYN_REQUEST_PLANE"] = request.getfixturevalue("request_plane")
138

139
        env["DYN_LOG"] = "debug"
140
141
142
143
144
        # 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"
145
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
146
147
148
149
150
151
152
153
154
155
156
        env["DYN_SYSTEM_PORT"] = str(self.system_port)
        env["DYN_HTTP_PORT"] = str(frontend_port)

        # Configure health check based on worker type
        health_check_urls = [
            (f"http://localhost:{self.system_port}/health", self.is_ready)
        ]
        if mode in ["decode", "prefill_and_decode"]:
            health_check_urls.append(
                (f"http://localhost:{frontend_port}/v1/models", check_models_api)
            )
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

        # 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,
173
            health_check_urls=health_check_urls,
174
175
176
177
            timeout=300,
            display_output=True,
            terminate_existing=False,
            log_dir=log_dir,
178
            display_name=worker_id,
179
180
        )

181
182
183
184
185
186
187
188
189
    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 TRT-LLM worker port: {e}")

        return super().__exit__(exc_type, exc_val, exc_tb)
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205

    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


206
@pytest.mark.timeout(290)  # 3x average
207
208
209
210
211
212
213
214
215
def test_request_migration_trtllm_aggregated(
    request,
    runtime_services_dynamic_ports,
    set_ucx_tls_no_mm,
    predownload_models,
    migration_limit,
    immediate_kill,
    request_api,
    stream,
216
217
):
    """
218
    End-to-end test for aggregated worker request migration.
219

220
221
222
223
224
    Parameters:
        immediate_kill: True for abrupt kill (SIGKILL), False for graceful shutdown (SIGTERM)
        migration_limit: > 0 to verify migration succeeds, 0 to verify request fails
        request_api: "chat" for chat completion API, "completion" for completion API
        stream: True for streaming, False for non-streaming
225
226
    """

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

231
232
233
234
        # Step 2: Start 2 workers
        with DynamoWorkerProcess(
            request, "worker1", frontend.frontend_port, migration_limit=migration_limit
        ) as worker1:
235
236
            logger.info(f"Worker 1 PID: {worker1.get_pid()}")

237
            with DynamoWorkerProcess(
238
239
240
241
                request,
                "worker2",
                frontend.frontend_port,
                migration_limit=migration_limit,
242
            ) as worker2:
243
244
                logger.info(f"Worker 2 PID: {worker2.get_pid()}")

245
246
247
248
249
250
251
252
253
254
                # Step 3: Run migration test
                run_migration_test(
                    frontend,
                    worker1,
                    worker2,
                    receiving_pattern="New Request ID: ",
                    migration_limit=migration_limit,
                    immediate_kill=immediate_kill,
                    use_chat_completion=(request_api == "chat"),
                    stream=stream,
255
256
                )

257

258
259
260
261
262
263
264
265
266
267
268
@pytest.mark.xfail(strict=False, reason="Prefill migration not yet supported")
@pytest.mark.timeout(350)  # 3x average
def test_request_migration_trtllm_prefill(
    request,
    runtime_services_dynamic_ports,
    set_ucx_tls_no_mm,
    predownload_models,
    migration_limit,
    immediate_kill,
    request_api,
    stream,
269
270
):
    """
271
272
273
274
275
276
277
278
279
    End-to-end test for prefill worker request migration in disaggregated mode.

    Setup: 1 decode worker + 2 prefill workers

    Parameters:
        immediate_kill: True for abrupt kill (SIGKILL), False for graceful shutdown (SIGTERM)
        migration_limit: > 0 to verify migration succeeds, 0 to verify request fails
        request_api: "chat" for chat completion API, "completion" for completion API
        stream: True for streaming, False for non-streaming
280
281
    """

282
283
    # Step 1: Start the frontend
    with DynamoFrontendProcess(request, enforce_disagg=True) as frontend:
284
285
        logger.info("Frontend started successfully")

286
287
288
289
290
291
292
293
294
        # Step 2: Start decode worker first (required for prefill workers to connect)
        with DynamoWorkerProcess(
            request,
            "worker0",
            frontend.frontend_port,
            migration_limit=migration_limit,
            mode="decode",
        ) as decode_worker:
            logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")
295

296
            # Step 3: Start 2 prefill workers
297
            with DynamoWorkerProcess(
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
                request,
                "worker1",
                frontend.frontend_port,
                migration_limit=migration_limit,
                mode="prefill",
            ) as prefill1:
                logger.info(f"Prefill Worker 1 PID: {prefill1.get_pid()}")

                with DynamoWorkerProcess(
                    request,
                    "worker2",
                    frontend.frontend_port,
                    migration_limit=migration_limit,
                    mode="prefill",
                ) as prefill2:
                    logger.info(f"Prefill Worker 2 PID: {prefill2.get_pid()}")

                    # Step 4: Run migration test
                    run_migration_test(
                        frontend,
                        prefill1,
                        prefill2,
                        receiving_pattern="Prefill Request ID: ",
                        migration_limit=migration_limit,
                        immediate_kill=immediate_kill,
                        use_chat_completion=(request_api == "chat"),
                        stream=stream,
                        use_long_prompt=True,
                    )
327

328

329
330
331
332
333
334
335
336
337
338
339
@pytest.mark.skip(reason="Decode worker can get stuck downloading kv cache")
@pytest.mark.timeout(350)  # 3x average
def test_request_migration_trtllm_kv_transfer(
    request,
    runtime_services_dynamic_ports,
    set_ucx_tls_no_mm,
    predownload_models,
    migration_limit,
    immediate_kill,
    request_api,
    stream,
340
341
):
    """
342
    End-to-end test for request migration during KV transfer in disaggregated mode.
343

344
    Setup: 1 prefill worker + 2 decode workers
345

346
347
348
349
350
    Parameters:
        immediate_kill: True for abrupt kill (SIGKILL), False for graceful shutdown (SIGTERM)
        migration_limit: > 0 to verify migration succeeds, 0 to verify request fails
        request_api: "chat" for chat completion API, "completion" for completion API
        stream: True for streaming, False for non-streaming
351
352
    """

353
354
    # Step 1: Start the frontend
    with DynamoFrontendProcess(request, enforce_disagg=True) as frontend:
355
356
        logger.info("Frontend started successfully")

357
        # Step 2: Start prefill worker first
358
359
        with DynamoWorkerProcess(
            request,
360
            "worker0",
361
            frontend.frontend_port,
362
363
364
365
            migration_limit=migration_limit,
            mode="prefill",
        ) as prefill_worker:
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")
366

367
            # Step 3: Start 2 decode workers
368
369
            with DynamoWorkerProcess(
                request,
370
                "worker1",
371
                frontend.frontend_port,
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
                migration_limit=migration_limit,
                mode="decode",
            ) as decode1:
                logger.info(f"Decode Worker 1 PID: {decode1.get_pid()}")

                with DynamoWorkerProcess(
                    request,
                    "worker2",
                    frontend.frontend_port,
                    migration_limit=migration_limit,
                    mode="decode",
                ) as decode2:
                    logger.info(f"Decode Worker 2 PID: {decode2.get_pid()}")

                    # Step 4: Run migration test
                    run_migration_test(
                        frontend,
                        decode1,
                        decode2,
                        receiving_pattern="Decode Request ID: ",
                        migration_limit=migration_limit,
                        immediate_kill=immediate_kill,
                        use_chat_completion=(request_api == "chat"),
                        stream=stream,
                        use_long_prompt=True,
397
398
399
                    )


400
401
402
403
404
405
406
407
408
409
@pytest.mark.timeout(350)  # 3x average
def test_request_migration_trtllm_decode(
    request,
    runtime_services_dynamic_ports,
    set_ucx_tls_no_mm,
    predownload_models,
    migration_limit,
    immediate_kill,
    request_api,
    stream,
410
411
):
    """
412
    End-to-end test for decode worker request migration in disaggregated mode.
413

414
    Setup: 1 prefill worker + 2 decode workers
415

416
417
418
419
420
    Parameters:
        immediate_kill: True for abrupt kill (SIGKILL), False for graceful shutdown (SIGTERM)
        migration_limit: > 0 to verify migration succeeds, 0 to verify request fails
        request_api: "chat" for chat completion API, "completion" for completion API
        stream: True for streaming, False for non-streaming
421
    """
422
423
424
425
    if not stream:
        pytest.skip(
            "Decode test requires streaming to wait for response before stopping worker"
        )
426

427
428
    # Step 1: Start the frontend
    with DynamoFrontendProcess(request, enforce_disagg=True) as frontend:
429
430
        logger.info("Frontend started successfully")

431
        # Step 2: Start prefill worker first
432
433
        with DynamoWorkerProcess(
            request,
434
            "worker0",
435
            frontend.frontend_port,
436
437
438
439
            migration_limit=migration_limit,
            mode="prefill",
        ) as prefill_worker:
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")
440

441
            # Step 3: Start 2 decode workers
442
443
            with DynamoWorkerProcess(
                request,
444
                "worker1",
445
                frontend.frontend_port,
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
                migration_limit=migration_limit,
                mode="decode",
            ) as decode1:
                logger.info(f"Decode Worker 1 PID: {decode1.get_pid()}")

                with DynamoWorkerProcess(
                    request,
                    "worker2",
                    frontend.frontend_port,
                    migration_limit=migration_limit,
                    mode="decode",
                ) as decode2:
                    logger.info(f"Decode Worker 2 PID: {decode2.get_pid()}")

                    # Step 4: Run migration test
                    run_migration_test(
                        frontend,
                        decode1,
                        decode2,
                        receiving_pattern="Decode Request ID: ",
                        migration_limit=migration_limit,
                        immediate_kill=immediate_kill,
                        use_chat_completion=(request_api == "chat"),
                        stream=stream,
                        wait_for_new_response_before_stop=True,
471
                    )