feat: request migration for trtllm

[jh-nv]

Overview:

Implement request migration when worker shuts down for trtllm

Details:

1. Created shutdown_event (asyncio.Event) to track shutdown signals, set the shutdown event before shutting down runtime during graceful shutdown.
2. Refactored _handle_cancellation() → _handle_cancellation_and_shutdown() to monitor both cancellation and shutdown events simultaneously
3. Enhanced _cancellation_monitor context manager to detect which event triggered (shutdown vs cancellation) and raise GeneratorExit specifically for shutdown cases, which triggers the request migration.

Where should the reviewer start?

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx

Summary by CodeRabbit

• Bug Fixes
  • Improved graceful shutdown coordination across system components for more reliable termination.
  • Enhanced request cancellation handling during shutdown to prevent incomplete operations.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Walkthrough

The changes implement graceful shutdown event coordination across the TensorRT LLM runtime. An asyncio.Event is threaded through initialization and handler configuration to enable coordinated shutdown signaling during request generation and cancellation handling.

Changes

Cohort / File(s)	Summary
Core Shutdown Event Initialization components/src/dynamo/trtllm/main.py	Updated graceful_shutdown() signature to accept shutdown_event parameter. Modified worker() to instantiate shutdown_event and pass it through init() call. Expanded init() signature to accept shutdown_event and wired it into RequestHandlerConfig dataclass initialization.
Handler Shutdown Integration components/src/dynamo/trtllm/request_handlers/handler_base.py	Added shutdown_event optional field to RequestHandlerConfig dataclass. Wired shutdown_event into HandlerBase via self.shutdown_event. Renamed _handle_cancellation() to _handle_cancellation_and_shutdown() with extended logic to monitor both cancellation and shutdown signals. Enhanced _cancellation_monitor() context manager to return monitoring task and propagate GeneratorExit when shutdown fires.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Poem

🐰✨ A shutdown event hops through the code,
Graceful coordination on the runtime road,
From main to handler, the signal takes flight,
Cancellation and shutdown, now unified right! 🌙

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 68.75% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat: request migration for trtllm' accurately reflects the main objective of the changeset, which implements request migration when a worker shuts down for the trtllm backend.
Description check	✅ Passed	The description includes all required sections (Overview, Details, Where to start, Related Issues) and adequately explains the key changes made to implement request migration, though the 'Where should the reviewer start?' section is empty.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[kthui]

LGTM! Minor comment on the redundant try...except block.

Will let @tanmayv25 to give the final check mark on TRT-LLM, as updated offline that he will be able to review by the EOD.

[kthui]

@coderabbitai full review

[coderabbitai]

✅ Actions performed

Full review triggered.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@components/src/dynamo/trtllm/request_handlers/handler_base.py`:
- Around line 182-199: The comment incorrectly claims
context.async_killed_or_stopped() is "already a Future"; it is an async def and
returns a coroutine, so change the comment and wrap that coroutine in a Task
before passing to asyncio.wait: replace the current cancellation_future =
context.async_killed_or_stopped() with creating a task via
asyncio.create_task(context.async_killed_or_stopped()) (keep the variable name
cancellation_future or rename to cancellation_task if you prefer), ensure
shutdown_event handling still creates shutdown_task via
asyncio.create_task(self.shutdown_event.wait()), and pass these tasks to
asyncio.wait; update the misleading comment to state that the coroutine is
wrapped into a Task for compatibility with asyncio.wait.

🧹 Nitpick comments (2)

components/src/dynamo/trtllm/main.py (1)

141-146: Store the task reference to prevent potential garbage collection.

The task created by asyncio.create_task() should be stored to prevent it from being garbage collected before completion. While this is in a shutdown context, storing the reference is best practice.

♻️ Suggested fix

+    _shutdown_task = None
+
     # Set up signal handler for graceful shutdown
     def signal_handler():
         # Schedule the shutdown coroutine instead of calling it directly
-        asyncio.create_task(graceful_shutdown(runtime, shutdown_event))
+        nonlocal _shutdown_task
+        _shutdown_task = asyncio.create_task(graceful_shutdown(runtime, shutdown_event))

components/src/dynamo/trtllm/request_handlers/handler_base.py (1)

252-259: Consider logging unexpected exceptions instead of silently ignoring them.

The except Exception: pass block will silently swallow any unexpected errors from the monitoring task, which could hide bugs. Adding debug-level logging would aid troubleshooting without changing the control flow.

♻️ Suggested improvement

             else:
                 # Task completed, check if it was due to shutdown
                 try:
                     monitor_task.result()
                 except GeneratorExit:
                     raise
                 except Exception:
-                    pass
+                    logging.debug(
+                        f"Monitor task completed with unexpected exception for context {context.id()}",
+                        exc_info=True,
+                    )

[tanmayv25]

Are these shutdown changes being tested somewhere?

[jh-nv]

Are these shutdown changes being tested somewhere?

Right now, we rely on the fault_tolerance/migration to test. The change yield same result as before for trtllm, we will need to add better simulation and test when the migration story is fully implemented.