feat: Backend accept new requests during shutdown grace period - #6093
Conversation
github-actions
Bot
added
feat
documentation
WalkthroughThis pull request introduces a centralized graceful shutdown utility module with signal handling and endpoint cleanup, integrates it into TRT-LLM and vLLM worker processes, extends migration tests with grace period parameterization, updates shutdown documentation, and adds utility functions for migration verification and logging. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. 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
🤖 Fix all issues with AI agents
In `@components/src/dynamo/common/utils/graceful_shutdown.py`:
- Around line 98-106: The signal_handler currently fires
asyncio.create_task(...) without storing the returned Task, risking GC
cancellation; modify signal_handler (and the module) to capture and retain the
Task returned by asyncio.create_task(graceful_shutdown_with_discovery(...)) —
e.g., append it to a module-level list or assign to a persistent variable so the
Task has a strong reference until it completes; ensure you still pass runtime,
endpoints, shutdown_event and grace_period_s into
graceful_shutdown_with_discovery and do not change the coroutine call semantics.
In `@docs/fault_tolerance/graceful_shutdown.md`:
- Line 56: The docs claim a default grace period of 5s while the actual default
constant _DEFAULT_GRACE_PERIOD_SECS in graceful_shutdown.py is 40.0, so update
one to match the other: either change the documentation line to state "default
40s" or change _DEFAULT_GRACE_PERIOD_SECS to 5.0 and ensure the environment
variable DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS default handling (and any
references to that constant in functions/classes inside graceful_shutdown.py)
reflects the chosen value so doc and runtime behavior are consistent.
🧹 Nitpick comments (4)
tests/fault_tolerance/migration/utils.py (3)
452-477: Remove commented-out code.This entire block is dead code that duplicates the now-refactored
verify_migration_occurred. It adds noise and hampers readability.
681-685: Typo:AssertionError→AssertionError.Wait — actually, let me re-examine this. The code reads:
except AssertionError as e:This is the correct Python built-in
AssertionError, so this is fine.However, the logic here is fragile: it catches the assertion error from
verify_migration_requestedand then asserts a specific substring of the error message to confirm it was the expected assertion failure. If the assertion message inverify_migration_requested(Line 438) ever changes, this check silently passes when it shouldn't.Consider a more direct approach:
Proposed simplification
else: - try: - verify_migration_requested(frontend) - pytest.fail("Migration requested unexpectedly.") - except AssertionError as e: - assert ( - "migration request expected, but 'Stream disconnected... recreating stream...' message not found in logs" - in str(e) - ) + log_content = frontend.read_logs() + assert ( + "Stream disconnected... recreating stream..." not in log_content + ), "Migration was requested unexpectedly"
590-593:expect_request_success: bool | None = Nonedefaults to treatingNoneasFalse.Line 655 uses
if expect_request_success:, soNonefalls into the else branch which expects the request to fail. All current callers set this explicitly, but a future caller omitting it would silently expect failure. Consider defaulting toTrueor usingif expect_request_success is not False:if the intent is to succeed by default.tests/fault_tolerance/migration/test_vllm.py (1)
80-92: Docstring does not documentgrace_period_sparameter.The class docstring lists
is_prefillbut not the newgrace_period_sconstructor parameter.
|
/coderabbitai review |
kthui
left a comment
kthui
left a comment
There was a problem hiding this comment.
Thanks for implementing the complete graceful shutdown logic and unifying it across all three backends! I have two suggestions for this PR.
1. Maintainability of the Graceful Shutdown Logic
Many individual variables (e.g. endpoints, shutdown_event) currently contribute to the shutdown behavior. Passing these explicitly as function arguments makes the code harder to read. For example, init_llm_worker in TRT-LLM takes 4 arguments: the first two are essential for the worker, but the last two are solely for graceful shutdown. This means 50% of the arguments are unrelated to the core worker functionality.
Would it be more maintainable if the entire graceful_shutdown.py were refactored into a class? For example:
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import asyncio
import logging
import os
import signal
from typing import Iterable, Optional, Set
logger = logging.getLogger(__name__)
_DEFAULT_GRACE_PERIOD_SECS = 40.0
_GRACE_PERIOD_ENV = "DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS"
class GracefulShutdownManager:
"""
Manages the graceful shutdown process for Dynamo components.
It handles:
1. Listening for SIGINT/SIGTERM.
2. Unregistering endpoints from discovery to stop new traffic.
3. Waiting for a grace period to allow in-flight requests to complete.
4. Shutting down the distributed runtime.
"""
def __init__(
self,
runtime,
endpoints: Iterable,
shutdown_event: Optional[asyncio.Event] = None,
grace_period_s: Optional[float] = None,
):
self.runtime = runtime
self.endpoints = list(endpoints)
self.shutdown_event = shutdown_event
# Store the explicit argument, but defer env var lookup until shutdown
self._explicit_grace_period_s = grace_period_s
self._shutdown_task: Optional[asyncio.Task] = None
self._shutdown_started = asyncio.Event()
def _get_effective_grace_period(self) -> float:
"""
Determines the grace period to use.
Priority:
1. Explicit argument passed to __init__ (if not None)
2. Current value of environment variable
3. Default constant
"""
if self._explicit_grace_period_s is not None:
return self._explicit_grace_period_s
value = os.getenv(_GRACE_PERIOD_ENV)
if value is None or value == "":
return _DEFAULT_GRACE_PERIOD_SECS
try:
parsed = float(value)
if parsed < 0:
logger.warning(
"Negative %s=%r; using 0", _GRACE_PERIOD_ENV, value
)
return 0.0
return parsed
except ValueError:
logger.warning(
"Invalid %s=%r; using default %s",
_GRACE_PERIOD_ENV,
value,
_DEFAULT_GRACE_PERIOD_SECS,
)
return _DEFAULT_GRACE_PERIOD_SECS
async def _unregister_endpoints(self) -> None:
"""Unregisters all managed endpoints from discovery."""
seen: Set[int] = set()
tasks = []
for endpoint in self.endpoints:
endpoint_id = id(endpoint)
if endpoint_id in seen:
continue
seen.add(endpoint_id)
tasks.append(endpoint.unregister_endpoint_instance())
if not tasks:
return
results = await asyncio.gather(*tasks, return_exceptions=True)
for result in results:
if isinstance(result, Exception):
logger.warning(
"Failed to unregister endpoint instance from discovery: %s",
result,
)
async def shutdown(self) -> None:
"""Executes the graceful shutdown workflow."""
if self._shutdown_started.is_set():
return
self._shutdown_started.set()
# Calculate grace period at the time of shutdown
grace_period = self._get_effective_grace_period()
logger.info("Received shutdown signal; unregistering endpoints from discovery")
# 1. Unregister endpoints to stop routing new requests
await self._unregister_endpoints()
# 2. Wait for grace period
if grace_period > 0:
logger.info("Grace period %.2fs before stopping endpoints", grace_period)
await asyncio.sleep(grace_period)
# 3. Signal external event if provided
if self.shutdown_event is not None:
self.shutdown_event.set()
# 4. Stop the runtime
logger.info("Initiating runtime shutdown")
self.runtime.shutdown()
def _on_shutdown_done(self, task: asyncio.Task) -> None:
"""Callback for when the shutdown task completes."""
try:
task.result()
except asyncio.CancelledError:
logger.info("Graceful shutdown task cancelled")
except Exception:
logger.exception("Graceful shutdown task failed")
finally:
if self._shutdown_task is task:
self._shutdown_task = None
def _signal_handler(self, loop: asyncio.AbstractEventLoop) -> None:
"""Internal handler triggered by the OS signal."""
if self._shutdown_task and not self._shutdown_task.done():
logger.debug("Shutdown already in progress; ignoring duplicate signal")
return
self._shutdown_task = loop.create_task(self.shutdown())
self._shutdown_task.add_done_callback(self._on_shutdown_done)
def install_signal_handlers(self, loop: asyncio.AbstractEventLoop) -> None:
"""Registers signal handlers for SIGINT and SIGTERM."""
for sig in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(sig, lambda: self._signal_handler(loop))
logger.info(
"Signal handlers set up for graceful shutdown (discovery unregister + grace period)"
)
def install_signal_handlers(
loop: asyncio.AbstractEventLoop,
runtime,
endpoints: Iterable,
shutdown_event: Optional[asyncio.Event] = None,
grace_period_s: Optional[float] = None,
) -> GracefulShutdownManager:
"""
Convenience wrapper to maintain backward compatibility or simple usage.
Creates and installs the manager in one go.
"""
manager = GracefulShutdownManager(runtime, endpoints, shutdown_event, grace_period_s)
manager.install_signal_handlers(loop)
return managerInstead of explicitly passing shutdown_event and shutdown_endpoints as arguments, we can pass an instance of the GracefulShutdownManager, where those variables can be access via the manager interface.
2. Test Coverage
It is great that you are adding tests! Let us take a step back and clarify the goal: we want to handle the corner case where new requests arrive at the worker after it begins to shutdown. The timeout essentially "delays" the shutdown to allow those requests to finish.
The new tests should check whether new requests are still accepted after shutdown begins, but before the grace period ends. For example:
- Start a frontend.
- Start a worker, with grace period set to 10 seconds.
- Send a request to ensure the setup works.
- Isolate the frontend from receiving discovery plane updates (e.g. via firewall rules).
- Initialize graceful shutdown of the worker.
- Within 10 seconds, send another request, ensuring the worker is still accepting requests.
A second variant of the test should set the grace period to 0 seconds, then assert that the worker stops receiving requests immediately after initializing graceful shutdown.
The new tests will NOT work on the current migration test suite, since we are not testing for migration here. I would recommend:
- Set the grace period to 0 on existing migration graceful shutdown tests. This resembles the scenario where requests are migrated immediately once the grace period ends.
- There is no need for E2E migration tests with grace period > 0, because the code path after the grace period ends is the same regardless of the duration.
- Add the two new E2E graceful shutdown scenarios (without migration) to the K8s test suite at
/tests/fault_tolerance/deploy. More control over the environment to isolate the frontend from receiving discovery plane updates can be found there.- Alternatively, a unit test for
GracefulShutdownManagercould be written. By mockingruntime,endpointsandshutdown_event, the correct methods are called at the right time with different grace periods can be asserted.
- Alternatively, a unit test for
Signed-off-by: Jacky <18255193+kthui@users.noreply.github.com>
Signed-off-by: Jacky <18255193+kthui@users.noreply.github.com>
…wn_DIS-1233 Signed-off-by: Jacky <18255193+kthui@users.noreply.github.com>
Signed-off-by: Jacky <18255193+kthui@users.noreply.github.com>
kthui
changed the title
Signed-off-by: Jacky <18255193+kthui@users.noreply.github.com> Co-authored-by: Jacky <18255193+kthui@users.noreply.github.com> Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
…namo#6093) Signed-off-by: Jacky <18255193+kthui@users.noreply.github.com> Co-authored-by: Jacky <18255193+kthui@users.noreply.github.com>
3 participants
Summary
Changes
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Closes DIS-1233
Summary by CodeRabbit
New Features
Documentation
Tests