feat: request migration for SGLang

[jh-nv]

Overview:
Implement request migration when worker shuts down for sglang

Details:

1. Created shutdown_event (asyncio.Event) to track shutdown signals, set the shutdown event before shutting down runtime during graceful shutdown.
2. use signals to chain the signint and sigterm processing so that SGLang can also process it.
3. Refactored _handle_cancellation() to monitor both cancellation and shutdown events simultaneously
4. 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

• Chores
  • Implemented a centralized graceful shutdown mechanism to enable coordinated termination across all system components and services.
  • Enhanced signal handling infrastructure to ensure proper shutdown propagation, clean termination sequences, and complete resource cleanup throughout the system.
  • Improved coordination of shutdown procedures across all processing variants including embedding, diffusion, text generation, and multimodal handlers.

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

[coderabbitai]

Walkthrough

This PR introduces a centralized graceful shutdown mechanism using asyncio.Event across sglang components. A shutdown_event parameter is threaded through initialization functions and worker handlers to enable coordinated termination. Signal handlers chain to previous handlers and trigger both the shutdown_event and runtime shutdown.

Changes

Cohort / File(s)	Summary
Core Shutdown Infrastructure components/src/dynamo/sglang/main.py	Centralized shutdown mechanism with signal chaining; updated 7 async init functions (init, init_prefill, init_diffusion, init_embedding, init_multimodal_processor, init_multimodal_encode_worker, init_multimodal_prefill_worker) to accept and propagate shutdown_event; wired signal handlers to trigger shutdown_event and invoke runtime.shutdown(); updated handler instantiations to receive shutdown_event
Base Handler & Cancellation components/src/dynamo/sglang/request_handlers/handler_base.py	Added shutdown_event parameter to BaseWorkerHandler.__init__; enhanced _handle_cancellation to monitor both cancellation signals and shutdown_event; updated _cancellation_monitor to detect shutdown-triggered completion and raise GeneratorExit; adjusted logging to reflect "cancellation or shutdown" semantics
LLM Request Handlers components/src/dynamo/sglang/request_handlers/llm/{decode,diffusion,prefill}_handler.py	Added optional shutdown_event: Optional[asyncio.Event] parameter to DecodeWorkerHandler, DiffusionWorkerHandler, and PrefillWorkerHandler; propagated parameter to parent class via super().__init__()
Embedding & Multimodal Handlers components/src/dynamo/sglang/request_handlers/embedding/embedding_handler.py, components/src/dynamo/sglang/request_handlers/multimodal/{processor,encode_worker,worker}_handler.py	Added optional shutdown_event parameter to EmbeddingWorkerHandler, MultimodalProcessorHandler, MultimodalEncodeWorkerHandler, MultimodalWorkerHandler, and MultimodalPrefillWorkerHandler; forwarded parameter to superclass constructor

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 Shutdown signals now flow with grace,
Through handlers coordinated in place,
Events and chains align,
Graceful closure—divine!
No stray threads left to trace. ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% 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 SGLang' clearly describes the main feature being implemented and aligns with the core objective of enabling request migration during SGLang worker shutdown.
Description check	✅ Passed	The PR description covers Overview, Details with numbered implementation steps, and references a related GitHub issue. However, the 'Where should the reviewer start' section is empty and the issue reference shows placeholder text '#xxx' instead of an actual issue number.

_{✏️ 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.}

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

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

309-333: Shutdown can be missed before the request ID is available.

_handle_cancellation awaits request_id_future before watching shutdown_event, so a shutdown that occurs before the first response won't be detected. The task remains blocked at the await request_id_future line and never reaches the code that checks for shutdown. Consider racing request_id_future against shutdown_event concurrently, exiting early with GeneratorExit if shutdown fires first.

🐛 Suggested fix

-            # Always wait for the request ID to ensure we can abort the request
-            sglang_request_id = await request_id_future
+            shutdown_task = None
+            if self.shutdown_event:
+                shutdown_task = asyncio.create_task(self.shutdown_event.wait())
+
+            # If shutdown happens before we have a request ID, bail out early
+            if shutdown_task:
+                done, _ = await asyncio.wait(
+                    [request_id_future, shutdown_task],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                if shutdown_task in done:
+                    raise GeneratorExit(
+                        "Engine was shut down before request id became available"
+                    )
+
+            # Always wait for the request ID to ensure we can abort the request
+            sglang_request_id = await request_id_future
@@
-            shutdown_task = None
-
-            if self.shutdown_event:
-                # Create task for shutdown monitoring and add to wait list
-                shutdown_task = asyncio.create_task(self.shutdown_event.wait())
-                wait_for.append(shutdown_task)
+            if shutdown_task:
+                wait_for.append(shutdown_task)

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

39-67: Non‑leader nodes can hang on shutdown.

_handle_non_leader_node() waits on a new asyncio.Event() instead of the shared shutdown_event, so SIGTERM/SIGINT won’t unblock it and the process can hang indefinitely. Please accept the shared event and await it.

🔧 Proposed fix

-async def _handle_non_leader_node(
-    engine: sgl.Engine,
-    generate_endpoint,
-) -> None:
+async def _handle_non_leader_node(
+    engine: sgl.Engine,
+    generate_endpoint,
+    shutdown_event: asyncio.Event,
+) -> None:
@@
-    await asyncio.Event().wait()
+    await shutdown_event.wait()

-        await _handle_non_leader_node(engine, generate_endpoint)
+        await _handle_non_leader_node(engine, generate_endpoint, shutdown_event)

🤖 Fix all issues with AI agents

In `@components/src/dynamo/sglang/main.py`:
- Around line 83-99: signal_handler currently calls any previous handler
(old_handlers[signum]) which for SIGINT may be signal.default_int_handler and
raises KeyboardInterrupt before runtime.shutdown() runs; change signal_handler
(the function referencing shutdown_event, loop.call_soon_threadsafe,
old_handlers, and runtime.shutdown) so that runtime.shutdown() is guaranteed to
run first (e.g., schedule/await shutdown_event and call runtime.shutdown in a
try/finally) and only invoke the old handler afterwards, or explicitly skip
calling old_handler when it is signal.default_int_handler; if you must call
old_handler preserve it inside a try/except so a raised KeyboardInterrupt does
not prevent runtime.shutdown().

[kthui]

I see the strategy here - suppressing the overwrite does ensure Dynamo keeps control of the signal, which is great.

However, I have a concern about discarding the handler SGLang tries to register. By returning None in watching_add_signal_handler, we are effectively "silencing" SGLang's internal shutdown logic. If the sgl.Engine relies on that handler to flush buffers, detach from a Ray cluster, or clean up shared memory, those operations will never happen.

Instead of just suppressing it, we should capture the callback and invoke it manually within our own signal_handler.

[jh-nv]

I see the strategy here - suppressing the overwrite does ensure Dynamo keeps control of the signal, which is great.

However, I have a concern about discarding the handler SGLang tries to register. By returning None in watching_add_signal_handler, we are effectively "silencing" SGLang's internal shutdown logic. If the sgl.Engine relies on that handler to flush buffers, detach from a Ray cluster, or clean up shared memory, those operations will never happen.

Instead of just suppressing it, we should capture the callback and invoke it manually within our own signal_handler.

Updated:

1. cache the SGLang cancel handling
2. Use a flag to track shutdown state, so we don't trigger multiple times
3. move shutdown to the loop rather than in the signal handler.

All migration test pass

[kthui]

Thanks for preserving the SGLang handler.

The logic looks generally correct to me, but I have one shutdown flow item and one cleanup suggestion:

1. Ordering Confirmation: Could we hit a race condition if SGLang shutdown signal handler is called before runtime.shutdown()?
2. Code Cleanup: The install_graceful_shutdown implementation can be better refined - e.g. runtime.shutdown() returns None, so the asyncio.iscoroutine check that followed is unnecessary. The cleanup can be done in a follow-up PR.

[kthui]

@ishandhanani could you help double-check the monkey patch logic on Python shutdown signal handler?

High-level flow summary for verification:

1. Register SGLang backend shutdown signal handler.
2. Monkey-patch loop.add_signal_handler to intercept registration.
3. Capture and store SGLang's handler when it tries to register.
4. On shutdown signal:
   i. Run captured SGLang handler
   ii. Run runtime.shutdown()

[kthui]

Thanks for the updates, and cleaning up corner cases!

For the record, here is the summary of the SGLang backend shutdown flow:

1. Register SGLang backend shutdown signal handler.
2. Monkey-patch loop.add_signal_handler to intercept registration.
3. Capture and store SGLang's handler when it tries to register.
4. On shutdown signal:
   a. Stop accepting new requests - runtime.shutdown().
   b. Cancel and return incomplete for all ongoing requests - shutdown_event.set().
   c. Clean up the engine - handler.cleanup().
   d. Call the captured SGLang shutdown signal handler.