feat: register shutdown option as part of @dynamo_worker() signature
#5273
Conversation
Signed-off-by: michaelfeil <me@michaelfeil.eu>
Signed-off-by: michaelfeil <me@michaelfeil.eu>
|
👋 Hi michaelfeil! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
bot
added
external-contribution
michaelfeil
changed the title
@dynamo_worker() signature
WalkthroughThis change refactors shutdown handling across worker components by transitioning from manual signal management to framework-provided automatic shutdown. The Dynamo runtime decorator gains a Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Poem
🚥 Pre-merge checks | ❌ 3❌ Failed checks (1 warning, 2 inconclusive)
✏️ 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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
examples/multimodal/components/video_encode_worker.py (1)
237-246: Remove unusedgraceful_shutdownfunction.This function is no longer called since manual signal handling was removed. The framework now handles shutdown automatically via
_register_shutdown_signals(runtime)whenregister_shutdown=True.🧹 Proposed cleanup
-async def graceful_shutdown(runtime): - """ - By calling `runtime.shutdown()`, the endpoints will immediately be unavailable. - However, in-flight requests will still be processed until they are finished. - After all in-flight requests are finished, the `serve_endpoint` functions will return - and the engine will be shutdown by Python's garbage collector. - """ - logging.info("Received shutdown signal, shutting down DistributedRuntime") - runtime.shutdown() - logging.info("DistributedRuntime shutdown complete") - -examples/multimodal/components/worker.py (1)
397-406: Remove unusedgraceful_shutdownfunction.This function is dead code. The framework now handles shutdown automatically when
register_shutdown=Trueis set on the decorator.🧹 Proposed cleanup
-async def graceful_shutdown(runtime): - """ - By calling `runtime.shutdown()`, the endpoints will immediately be unavailable. - However, in-flight requests will still be processed until they are finished. - After all in-flight requests are finished, the `serve_endpoint` functions will return - and the engine will be shutdown by Python's garbage collector. - """ - logging.info("Received shutdown signal, shutting down DistributedRuntime") - runtime.shutdown() - logging.info("DistributedRuntime shutdown complete") - -examples/multimodal/components/processor.py (1)
271-280: Remove unusedgraceful_shutdownfunction.This function is no longer invoked. The framework automatically handles shutdown via
_register_shutdown_signals(runtime)whenregister_shutdown=True.🧹 Proposed cleanup
-async def graceful_shutdown(runtime): - """ - By calling `runtime.shutdown()`, the endpoints will immediately be unavailable. - However, in-flight requests will still be processed until they are finished. - After all in-flight requests are finished, the `serve_endpoint` functions will return - and the engine will be shutdown by Python's garbage collector. - """ - logging.info("Received shutdown signal, shutting down DistributedRuntime") - runtime.shutdown() - logging.info("DistributedRuntime shutdown complete") - -examples/multimodal/components/audio_encode_worker.py (1)
233-242: Remove unusedgraceful_shutdownfunction.This function is dead code since manual signal handling was removed. The framework handles shutdown automatically when
register_shutdown=True.🧹 Proposed cleanup
-async def graceful_shutdown(runtime): - """ - By calling `runtime.shutdown()`, the endpoints will immediately be unavailable. - However, in-flight requests will still be processed until they are finished. - After all in-flight requests are finished, the `serve_endpoint` functions will return - and the engine will be shutdown by Python's garbage collector. - """ - logging.info("Received shutdown signal, shutting down DistributedRuntime") - runtime.shutdown() - logging.info("DistributedRuntime shutdown complete") - -
🤖 Fix all issues with AI agents
In @lib/bindings/python/src/dynamo/runtime/__init__.py:
- Around line 31-32: The task created in _signal_handler is not retained and may
be garbage-collected; modify _signal_handler to store the
asyncio.create_task(...) result (e.g., assign it to a durable attribute like
runtime._shutdown_task or a module-level variable) so the Task returned by
asyncio.create_task(_shutdown_handler(runtime)) is kept until completion and can
be cleaned up automatically.
- Around line 25-28: The shutdown handler is defined synchronously but is passed
to asyncio.create_task(), causing a TypeError; change the function signature of
_shutdown_handler to be async (use async def _shutdown_handler(runtime)) so it
returns a coroutine, and keep the call to runtime.shutdown() inside it (it's
safe to call a synchronous method from an async function), ensuring any callers
like asyncio.create_task(_shutdown_handler(runtime)) receive a coroutine.
🧹 Nitpick comments (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
25-28: Consider using logging instead of print statements.For better observability and consistency with the Rust layer (which uses
tracing), consider replacing print statements with Python's logging module.♻️ Optional: Use logging
Add logging import at the top:
import logging logger = logging.getLogger(__name__)Then update the function:
async def _shutdown_handler(runtime): - print("Shutdown signal received, cancelling runtime...") + logger.info("Shutdown signal received, cancelling runtime...") runtime.shutdown() - print("Runtime shutdown complete.") + logger.info("Runtime shutdown complete.")
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
examples/multimodal/components/audio_encode_worker.pyexamples/multimodal/components/encode_worker.pyexamples/multimodal/components/processor.pyexamples/multimodal/components/video_encode_worker.pyexamples/multimodal/components/worker.pylib/bindings/python/examples/hello_world/server.pylib/bindings/python/src/dynamo/runtime/__init__.py
🧰 Additional context used
🧠 Learnings (6)
📓 Common learnings
Learnt from: keivenchang
Repo: ai-dynamo/dynamo PR: 4323
File: components/src/dynamo/vllm/main.py:197-205
Timestamp: 2025-11-14T01:09:35.244Z
Learning: keivenchang prefers using graceful_shutdown() over atexit for cleanup logic in components/src/dynamo/vllm/main.py, as it provides better control and consistency with existing shutdown patterns.
Learnt from: nnshah1
Repo: ai-dynamo/dynamo PR: 1444
File: tests/fault_tolerance/utils/metrics.py:30-32
Timestamp: 2025-07-01T13:55:03.940Z
Learning: The `dynamo_worker()` decorator in the dynamo codebase returns a wrapper that automatically injects the `runtime` parameter before calling the wrapped function. This means callers only need to provide the non-runtime parameters, while the decorator handles injecting the runtime argument automatically. For example, a function with signature `async def get_metrics(runtime, log_dir)` decorated with `dynamo_worker()` can be called as `get_metrics(log_dir)` because the decorator wrapper injects the runtime parameter.
📚 Learning: 2025-07-01T13:55:03.940Z
Learnt from: nnshah1
Repo: ai-dynamo/dynamo PR: 1444
File: tests/fault_tolerance/utils/metrics.py:30-32
Timestamp: 2025-07-01T13:55:03.940Z
Learning: The `dynamo_worker()` decorator in the dynamo codebase returns a wrapper that automatically injects the `runtime` parameter before calling the wrapped function. This means callers only need to provide the non-runtime parameters, while the decorator handles injecting the runtime argument automatically. For example, a function with signature `async def get_metrics(runtime, log_dir)` decorated with `dynamo_worker()` can be called as `get_metrics(log_dir)` because the decorator wrapper injects the runtime parameter.
Applied to files:
lib/bindings/python/src/dynamo/runtime/__init__.pyexamples/multimodal/components/encode_worker.pyexamples/multimodal/components/audio_encode_worker.pyexamples/multimodal/components/worker.pyexamples/multimodal/components/processor.pylib/bindings/python/examples/hello_world/server.pyexamples/multimodal/components/video_encode_worker.py
📚 Learning: 2025-11-14T01:09:35.244Z
Learnt from: keivenchang
Repo: ai-dynamo/dynamo PR: 4323
File: components/src/dynamo/vllm/main.py:197-205
Timestamp: 2025-11-14T01:09:35.244Z
Learning: keivenchang prefers using graceful_shutdown() over atexit for cleanup logic in components/src/dynamo/vllm/main.py, as it provides better control and consistency with existing shutdown patterns.
Applied to files:
lib/bindings/python/src/dynamo/runtime/__init__.pyexamples/multimodal/components/encode_worker.pyexamples/multimodal/components/audio_encode_worker.pyexamples/multimodal/components/worker.pylib/bindings/python/examples/hello_world/server.pyexamples/multimodal/components/video_encode_worker.py
📚 Learning: 2025-09-21T01:40:52.456Z
Learnt from: PeaBrane
Repo: ai-dynamo/dynamo PR: 3155
File: components/backends/vllm/src/dynamo/vllm/main.py:228-233
Timestamp: 2025-09-21T01:40:52.456Z
Learning: In the dynamo codebase, error handling for distributed runtime client initialization (like runtime.namespace().component().endpoint().client()) is handled at the Rust level in the distributed runtime bindings, so Python-level try/catch blocks are not needed and would be redundant.
Applied to files:
lib/bindings/python/src/dynamo/runtime/__init__.pylib/bindings/python/examples/hello_world/server.py
📚 Learning: 2026-01-03T01:27:08.763Z
Learnt from: tanmayv25
Repo: ai-dynamo/dynamo PR: 5143
File: examples/multimodal/components/processor.py:207-215
Timestamp: 2026-01-03T01:27:08.763Z
Learning: In examples/multimodal/components/processor.py, ensure the Dynamo multimodal processor supports flexible content ordering in multimodal requests (audio, image, or video content before text), reflecting the OpenAI multimodal API specification where content items can appear in any order. Do not enforce a text-first ordering as mandatory. Implement and verify logic that checks for required content types regardless of position, and update tests to cover all permutations of content ordering.
Applied to files:
examples/multimodal/components/processor.py
📚 Learning: 2026-01-03T01:29:46.547Z
Learnt from: tanmayv25
Repo: ai-dynamo/dynamo PR: 5143
File: examples/multimodal/components/processor.py:207-214
Timestamp: 2026-01-03T01:29:46.547Z
Learning: In examples/multimodal/components/processor.py, the processor currently supports only a single text field per request. Do not add support for multiple text fields; this is out of scope for the current implementation. When reviewing or modifying this file, ensure changes do not introduce multi-text-field handling and, if relevant, update tests and documentation to reflect the single-text-field limitation.
Applied to files:
examples/multimodal/components/processor.py
🧬 Code graph analysis (7)
lib/bindings/python/src/dynamo/runtime/__init__.py (5)
lib/runtime/src/distributed.rs (3)
runtime(288-290)shutdown(323-326)request_plane(382-384)lib/bindings/python/rust/lib.rs (1)
shutdown(616-618)lib/bindings/python/src/dynamo/_core.pyi (5)
shutdown(49-53)shutdown(809-813)shutdown(898-902)DistributedRuntime(36-85)get(1658-1659)lib/runtime/src/runtime.rs (1)
shutdown(302-333)tests/conftest.py (1)
request_plane(515-524)
examples/multimodal/components/encode_worker.py (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
dynamo_worker(41-73)
examples/multimodal/components/audio_encode_worker.py (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
dynamo_worker(41-73)
examples/multimodal/components/worker.py (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
dynamo_worker(41-73)
examples/multimodal/components/processor.py (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
dynamo_worker(41-73)
lib/bindings/python/examples/hello_world/server.py (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
dynamo_worker(41-73)
examples/multimodal/components/video_encode_worker.py (1)
lib/bindings/python/src/dynamo/runtime/__init__.py (1)
dynamo_worker(41-73)
🪛 Ruff (0.14.10)
lib/bindings/python/src/dynamo/runtime/__init__.py
32-32: Store a reference to the return value of asyncio.create_task
(RUF006)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Build and Test - dynamo
🔇 Additional comments (9)
examples/multimodal/components/video_encode_worker.py (1)
249-249: LGTM! Framework-managed shutdown enabled.The decorator change correctly enables automatic SIGINT/SIGTERM handling via the framework's built-in shutdown mechanism, eliminating the need for manual signal handling.
examples/multimodal/components/encode_worker.py (1)
190-190: LGTM! Framework-managed shutdown enabled.The decorator change correctly enables automatic SIGINT/SIGTERM handling, aligning with the project-wide transition to framework-provided shutdown behavior.
examples/multimodal/components/worker.py (1)
409-409: LGTM! Framework-managed shutdown enabled.The decorator change correctly transitions shutdown handling to the framework's built-in mechanism, replacing manual signal management.
examples/multimodal/components/processor.py (1)
283-283: LGTM! Framework-managed shutdown enabled.The decorator change correctly enables automatic shutdown handling, consistent with the project-wide refactoring to centralize shutdown logic in the framework.
examples/multimodal/components/audio_encode_worker.py (1)
245-245: LGTM! Framework-managed shutdown enabled.The decorator change correctly enables framework-provided shutdown handling, eliminating manual signal management and aligning with the project-wide transition.
lib/bindings/python/src/dynamo/runtime/__init__.py (3)
35-38: Signal registration implementation looks correct.The signal handling setup properly integrates with the asyncio event loop for SIGINT and SIGTERM. The lambda correctly captures the runtime reference.
41-44: Good API design with backward-compatible defaults.The new parameters enable opt-in shutdown handling and configurable KV store selection while maintaining backward compatibility.
50-52: Implementation correctly uses configurable parameters and shutdown handling.The wrapper properly applies the
kv_storeparameter and conditionally registers signal handlers via the standard asyncioloop.add_signal_handler()pattern whenregister_shutdown=True. Verification shows selective adoption in examples—onlyhello_world/server.pyenables signal handling, while other examples rely on defaults, which is appropriate for their use cases.lib/bindings/python/examples/hello_world/server.py (1)
23-23: Excellent simplification by delegating shutdown to the framework.The transition to
register_shutdown=Trueremoves manual signal handling boilerplate while maintaining graceful shutdown behavior. The decorator properly registers SIGINT and SIGTERM handlers through_register_shutdown_signals()when enabled.
| def _shutdown_handler(runtime): | ||
| print("Shutdown signal received, cancelling runtime...") | ||
| runtime.shutdown() | ||
| print("Runtime shutdown complete.") |
There was a problem hiding this comment.
Critical: _shutdown_handler must be an async function.
Line 32 passes _shutdown_handler(runtime) to asyncio.create_task(), which requires a coroutine. Since _shutdown_handler is not defined with async def, this will raise a TypeError at runtime.
🐛 Fix: Define as async function
-def _shutdown_handler(runtime):
+async def _shutdown_handler(runtime):
print("Shutdown signal received, cancelling runtime...")
runtime.shutdown()
print("Runtime shutdown complete.")Note: runtime.shutdown() is synchronous but can be safely called from an async function.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _shutdown_handler(runtime): | |
| print("Shutdown signal received, cancelling runtime...") | |
| runtime.shutdown() | |
| print("Runtime shutdown complete.") | |
| async def _shutdown_handler(runtime): | |
| print("Shutdown signal received, cancelling runtime...") | |
| runtime.shutdown() | |
| print("Runtime shutdown complete.") |
🤖 Prompt for AI Agents
In @lib/bindings/python/src/dynamo/runtime/__init__.py around lines 25 - 28, The
shutdown handler is defined synchronously but is passed to
asyncio.create_task(), causing a TypeError; change the function signature of
_shutdown_handler to be async (use async def _shutdown_handler(runtime)) so it
returns a coroutine, and keep the call to runtime.shutdown() inside it (it's
safe to call a synchronous method from an async function), ensuring any callers
like asyncio.create_task(_shutdown_handler(runtime)) receive a coroutine.
| def _signal_handler(runtime): | ||
| asyncio.create_task(_shutdown_handler(runtime)) |
There was a problem hiding this comment.
Store the task reference to prevent premature garbage collection.
The task created by asyncio.create_task() is not stored, which can lead to the task being garbage collected before completion. While unlikely during shutdown, this violates best practices and triggers RUF006.
♻️ Proposed fix: Store task reference
+_shutdown_tasks = set()
+
def _signal_handler(runtime):
- asyncio.create_task(_shutdown_handler(runtime))
+ task = asyncio.create_task(_shutdown_handler(runtime))
+ _shutdown_tasks.add(task)
+ task.add_done_callback(_shutdown_tasks.discard)This ensures the task reference is held until completion and automatically cleaned up.
🧰 Tools
🪛 Ruff (0.14.10)
32-32: Store a reference to the return value of asyncio.create_task
(RUF006)
🤖 Prompt for AI Agents
In @lib/bindings/python/src/dynamo/runtime/__init__.py around lines 31 - 32, The
task created in _signal_handler is not retained and may be garbage-collected;
modify _signal_handler to store the asyncio.create_task(...) result (e.g.,
assign it to a durable attribute like runtime._shutdown_task or a module-level
variable) so the Task returned by
asyncio.create_task(_shutdown_handler(runtime)) is kept until completion and can
be cleaned up automatically.
1 participant
Overview:
I noticed a lot of manual setup's with
runtime.shutdown(). Ideally, it would be best to just automatically shutdown the runtime. Some other examples, such as vllm claim that there is a graceful shutdown, but sometimes there is not.Details:
Where should the reviewer start?
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
New Features
Refactor
✏️ Tip: You can customize this high-level summary in your review settings.