test_trtllm.py 20.7 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
Test Execution Times (Last Run: 2025-12-13):
6
- test_request_cancellation_trtllm_aggregated: ~45s (gpu_1)
7
8
9
10
- test_request_cancellation_trtllm_decode_cancel: ~65s (gpu_1)
- test_request_cancellation_trtllm_prefill_cancel: ~65s (gpu_1)
- test_request_cancellation_trtllm_kv_transfer_cancel: ~65s (gpu_1)
- Total: ~240s x2 request planes = ~480s (0:08:00)
11
12
"""

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

import pytest

from tests.fault_tolerance.cancellation.utils import (
    DynamoFrontendProcess,
22
23
24
    poll_for_pattern,
    read_streaming_responses,
    send_cancellable_request,
25
26
27
28
)
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
from tests.utils.managed_process import ManagedProcess
from tests.utils.payloads import check_health_generate, check_models_api
29
from tests.utils.port_utils import allocate_port, deallocate_port
30
31
32

logger = logging.getLogger(__name__)

33
34
35
36
pytestmark = [
    pytest.mark.trtllm,
    pytest.mark.gpu_1,
    pytest.mark.e2e,
37
    pytest.mark.post_merge,
38
    pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME),
39
    pytest.mark.post_merge,  # post_merge to pinpoint failure commit
40
    pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True),
41
42
]

43
44
45
46

class DynamoWorkerProcess(ManagedProcess):
    """Process manager for Dynamo worker with TensorRT-LLM backend"""

47
48
49
50
51
52
    def __init__(
        self,
        request,
        frontend_port: int,
        mode: str = "prefill_and_decode",
    ):
53
54
55
56
57
        """
        Initialize TensorRT-LLM worker process.

        Args:
            request: pytest request object
58
            frontend_port: Port for the frontend server
59
60
            mode: One of "prefill_and_decode", "prefill", "decode"
        """
61
62
63
64
        # Allocate system port for this worker
        system_port = allocate_port(9100)
        self.system_port = system_port
        self.frontend_port = frontend_port
65
66
67
        # Prefill workers require migration_limit=0 (no KV cache migration support)
        migration_limit = "0" if mode == "prefill" else "3"

68
69
70
71
72
73
74
75
76
        command = [
            "python3",
            "-m",
            "dynamo.trtllm",
            "--model",
            FAULT_TOLERANCE_MODEL_NAME,
            "--disaggregation-mode",
            mode,
            "--max-seq-len",
77
78
79
            "16384",
            "--max-num-tokens",
            "16384",
80
            "--migration-limit",
81
            migration_limit,
82
83
84
        ]
        if mode != "prefill_and_decode":
            with open("test_request_cancellation_trtllm_config.yaml", "w") as f:
85
86
87
                f.write(
                    "cache_transceiver_config:\n  backend: DEFAULT\n  max_tokens_in_buffer: 16384\n"
                )
88
                f.write("disable_overlap_scheduler: true\n")
89
                f.write("kv_cache_config:\n  max_tokens: 16384\n")
90
91
92
93
94
95
            command += [
                "--extra-engine-args",
                "test_request_cancellation_trtllm_config.yaml",
            ]

        health_check_urls = [
96
97
            (f"http://localhost:{frontend_port}/v1/models", check_models_api),
            (f"http://localhost:{frontend_port}/health", check_health_generate),
98
99
        ]

100
101
102
103
104
        # Set health check based on worker type
        if mode in ["prefill", "decode"]:
            health_check_urls = [
                (f"http://localhost:{system_port}/health", self.is_ready)
            ]
105

106
        # Set environment variables
107
        env = os.environ.copy()
108
109
        env["DYN_REQUEST_PLANE"] = request.getfixturevalue("request_plane")

110
        env["DYN_LOG"] = "debug"
111
112
113
114
115
        # 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"
116
        env["DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"] = '["generate"]'
117
        env["DYN_SYSTEM_PORT"] = str(system_port)
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

        # Set log directory based on worker type
        log_dir = f"{request.node.name}_{mode}_worker"

        # 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=health_check_urls,
            timeout=300,
            display_output=True,
            terminate_existing=False,
            log_dir=log_dir,
        )

        self.mode = mode

    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.mode.capitalize()} worker status is ready")
                return True
            logger.warning(
                f"{self.mode.capitalize()} worker status is not ready: {data.get('status')}"
            )
        except ValueError:
            logger.warning(
                f"{self.mode.capitalize()} worker health response is not valid JSON"
            )
        return False

158
159
160
161
162
163
164
165
166
167
    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)

168

169
@pytest.mark.timeout(135)  # 3x average
170
def test_request_cancellation_trtllm_aggregated(
171
    request, runtime_services_dynamic_ports, predownload_models
172
):
173
174
175
176
177
    """
    End-to-end test for request cancellation functionality in aggregated mode.

    This test verifies that when a request is cancelled by the client,
    the system properly handles the cancellation and cleans up resources
178
179
180
181
182
183
184
185
186
    on the worker side in aggregated (prefill_and_decode) mode. Tests three scenarios:
    1. Completion request
    2. Chat completion request (non-streaming)
    3. Chat completion request (streaming)

    Timing (Last Run: 2025-12-09): ~45s total
    - Engine initialization: ~27s (frontend + worker)
    - Testing 3 scenarios: ~15s (~5s each)
    - Teardown: ~3s
187
188
    """

189
    # Step 1: Start the frontend (allocates its own frontend_port)
190
191
192
193
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

        # Step 2: Start an aggregated worker
194
195
196
197
        # Step 2: Start a single worker (allocates its own system_port)
        with DynamoWorkerProcess(
            request, frontend.frontend_port, mode="prefill_and_decode"
        ) as worker:
198
199
200
201
202
            logger.info(f"Aggregated Worker PID: {worker.get_pid()}")

            # TODO: Why wait after worker ready fixes frontend 404 / 500 flakiness?
            time.sleep(2)

203
            # Step 3: Test request cancellation with polling approach
204
205
206
207
208
209
210
211
212
213
214
            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",
                ),
            ]

215
            for request_type, description in test_scenarios:
216
217
                logger.info(f"Testing {description.lower()}...")

218
                # Send the request (non-blocking)
219
220
221
                cancellable_req = send_cancellable_request(
                    frontend.frontend_port, request_type
                )
222
223
224
225
226
227
228

                # Poll for "New Request ID" pattern
                request_id, worker_log_offset = poll_for_pattern(
                    process=worker,
                    pattern="New Request ID: ",
                    log_offset=worker_log_offset,
                    match_type="contains",
229
                )
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250

                # 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,
251
252
253
254
255
                )

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


256
@pytest.mark.timeout(195)  # 3x average
257
def test_request_cancellation_trtllm_decode_cancel(
258
    request, runtime_services_dynamic_ports, predownload_models
259
):
260
    """
261
    End-to-end test for request cancellation during decode phase with unified frontend.
262
263
264

    This test verifies that when a request is cancelled by the client during the decode phase,
    the system properly handles the cancellation and cleans up resources
265
    on the decode worker side in a disaggregated setup.
266
267
268
269
270

    Timing (Last Run: 2025-12-09): ~115s total (2 workers at 45% GPU each)
    - Engine initialization: ~92s (frontend: 2s, prefill worker: 45s, decode worker: 45s sequential)
    - Testing stream cancellation during decode: ~20s
    - Teardown: ~3s
271
272
    """

273
    # Step 1: Start the frontend (allocates its own frontend_port)
274
275
276
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

277
278
279
280
        # Step 2: Start the prefill worker (allocates its own system_port)
        with DynamoWorkerProcess(
            request, frontend.frontend_port, mode="prefill"
        ) as prefill_worker:
281
282
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

283
284
285
286
            # Step 3: Start the decode worker (allocates its own system_port)
            with DynamoWorkerProcess(
                request, frontend.frontend_port, mode="decode"
            ) as decode_worker:
287
288
289
290
291
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

                # TODO: Why wait after worker ready fixes frontend 404 / 500 flakiness?
                time.sleep(2)

292
                # Step 4: Test request cancellation for streaming scenario
293
                logger.info(
294
                    "Testing chat completion stream request cancellation in decode worker (decode phase)..."
295
296
                )

297
                # Send streaming request (non-blocking)
298
299
300
                cancellable_req = send_cancellable_request(
                    frontend.frontend_port, "chat_completion_stream"
                )
301

302
303
304
305
                # Poll for "Prefill Request ID" pattern in prefill worker (frontend routes here first)
                request_id, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
                    pattern="Prefill Request ID: ",
306
307
308
                    match_type="contains",
                )

309
310
311
312
                # Verify same request ID reached decode worker (after prefill completes)
                _, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
                    pattern=f"Decode Request ID: {request_id}",
313
                )
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336

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

                logger.info(
                    "Chat completion stream cancellation in decode phase detected successfully"
337
338
339
                )


340
@pytest.mark.timeout(195)  # 3x average
341
def test_request_cancellation_trtllm_prefill_cancel(
342
    request, runtime_services_dynamic_ports, predownload_models
343
):
344
    """
345
    End-to-end test for request cancellation during prefill phase with unified frontend.
346

347
348
349
    This test verifies that when a request is cancelled by the client during the prefill phase,
    the system properly handles the cancellation and cleans up resources on the prefill worker.
    Since the request is cancelled before prefill completes, the decode worker never receives it.
350
351
352
353
354

    Timing (Last Run: 2025-12-09): ~115s total (2 workers at 45% GPU each)
    - Engine initialization: ~92s (frontend: 2s, prefill worker: 45s, decode worker: 45s sequential)
    - Testing cancellation during prefill: ~20s
    - Teardown: ~3s
355
356
    """

357
    # Step 1: Start the frontend (allocates its own frontend_port)
358
359
360
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

361
362
363
364
        # Step 2: Start the prefill worker (allocates its own system_port)
        with DynamoWorkerProcess(
            request, frontend.frontend_port, mode="prefill"
        ) as prefill_worker:
365
366
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

367
368
369
370
            # Step 3: Start the decode worker (allocates its own system_port)
            with DynamoWorkerProcess(
                request, frontend.frontend_port, mode="decode"
            ) as decode_worker:
371
372
373
374
375
376
377
378
379
380
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

                # TODO: Why wait after worker ready fixes frontend 404 / 500 flakiness?
                time.sleep(2)

                # Step 4: Test request cancellation during prefill phase
                logger.info(
                    "Testing completion request cancellation during prefill phase..."
                )

381
382
                # Send request with long prompt (non-blocking)
                cancellable_req = send_cancellable_request(
383
                    frontend.frontend_port, "completion", use_long_prompt=True
384
385
                )

386
                # Poll for "Prefill Request ID" pattern in prefill worker (frontend routes here first)
387
388
                request_id, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
389
                    pattern="Prefill Request ID: ",
390
391
392
                    match_type="contains",
                )

393
                # Cancel during prefill phase
394
                cancellable_req.cancel()
395
                logger.info(f"Cancelled request ID: {request_id} during prefill")
396

397
                # Poll for "Aborted Request ID" in prefill worker (where cancellation happens)
398
399
400
401
402
403
404
405
406
407
                _, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
                    pattern=f"Aborted 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",
408
                )
409

410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
                # 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}"

427
428
                logger.info(
                    "Completion request cancellation during prefill phase detected successfully"
429
                )
430
431


432
433
@pytest.mark.xfail(reason="Test fails only on CI", strict=False)
@pytest.mark.timeout(195)  # 3x average
434
def test_request_cancellation_trtllm_kv_transfer_cancel(
435
    request, runtime_services_dynamic_ports, predownload_models
436
):
437
438
439
440
441
    """
    End-to-end test for request cancellation during prefill to decode KV transfer phase.

    This test verifies that when a request is cancelled by the client during the KV transfer phase,
    the system properly handles the cancellation and cleans up resources on the workers.
442
443
444
445
446

    Timing (Last Run: 2025-12-09): ~115s total (2 workers at 45% GPU each)
    - Engine initialization: ~92s (frontend: 2s, prefill worker: 45s, decode worker: 45s sequential)
    - Testing KV transfer cancellation: ~20s
    - Teardown: ~3s
447
448
    """

449
    # Step 1: Start the frontend (allocates its own frontend_port)
450
451
452
    with DynamoFrontendProcess(request) as frontend:
        logger.info("Frontend started successfully")

453
454
455
456
        # Step 2: Start the prefill worker (allocates its own system_port)
        with DynamoWorkerProcess(
            request, frontend.frontend_port, mode="prefill"
        ) as prefill_worker:
457
458
            logger.info(f"Prefill Worker PID: {prefill_worker.get_pid()}")

459
460
461
462
            # Step 3: Start the decode worker (allocates its own system_port)
            with DynamoWorkerProcess(
                request, frontend.frontend_port, mode="decode"
            ) as decode_worker:
463
464
465
466
467
468
469
470
471
472
473
474
                logger.info(f"Decode Worker PID: {decode_worker.get_pid()}")

                # TODO: Why wait after worker ready fixes frontend 404 / 500 flakiness?
                time.sleep(2)

                # Step 4: Test request cancellation during KV transfer phase
                logger.info(
                    "Testing completion request cancellation during KV transfer phase..."
                )

                # Send request with long prompt
                cancellable_req = send_cancellable_request(
475
                    frontend.frontend_port, "completion", use_long_prompt=True
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
                )

                # Poll for "Prefill Request ID" pattern in prefill worker
                request_id, prefill_log_offset = poll_for_pattern(
                    process=prefill_worker,
                    pattern="Prefill Request ID: ",
                    match_type="contains",
                )

                # Poll for decode worker entry signaling start of KV transfer phase
                _, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
                    pattern=f"Decode Request ID: {request_id}",
                    poll_interval_ms=2,
                )

                # Cancel during KV transfer phase in decode worker
                cancellable_req.cancel()
                logger.info(
                    f"Cancelled request ID: {request_id} at beginning of decode"
                )

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

                logger.info(
                    "Completion request cancellation at beginning of decode detected successfully"
                )

                # Verify the workers are still functional
516
517
518
                cancellable_req = send_cancellable_request(
                    frontend.frontend_port, "chat_completion_stream"
                )
519
520
521
522
523
524
525
526
527
528
529
                _, decode_log_offset = poll_for_pattern(
                    process=decode_worker,
                    pattern="Decode Request ID: ",
                    log_offset=decode_log_offset,
                    match_type="contains",
                )
                read_streaming_responses(cancellable_req, expected_count=5)

                logger.info(
                    "Workers are functional after cancellation during KV transfer"
                )