feat: Backend accept new requests during shutdown grace period

components/src/dynamo/common/utils/graceful_shutdown.py

@@ -0,0 +1,134 @@
+# 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
+
+logger = logging.getLogger(__name__)
+
+# TODO: make this using cli flag
+_DEFAULT_GRACE_PERIOD_SECS = 5.0
+_GRACE_PERIOD_ENV = "DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS"
+_shutdown_started = asyncio.Event()
+
+
+def get_grace_period_seconds() -> float:
+    value = os.getenv(_GRACE_PERIOD_ENV)
+    if value is None or value == "":
+        return _DEFAULT_GRACE_PERIOD_SECS
+    try:
+        parsed = float(value)
+    except ValueError:
+        logger.warning(
+            "Invalid %s=%r; using default %s",
+            _GRACE_PERIOD_ENV,
+            value,
+            _DEFAULT_GRACE_PERIOD_SECS,
+        )
+        return _DEFAULT_GRACE_PERIOD_SECS
+    if parsed < 0:
+        logger.warning(
+            "Negative %s=%r; using 0",
+            _GRACE_PERIOD_ENV,
+            value,
+        )
+        return 0.0
+    return parsed
+
+
+async def _unregister_endpoints(endpoints: Iterable) -> None:
+    seen = set()
+    tasks = []
+    for endpoint in 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 graceful_shutdown_with_discovery(
+    runtime,
+    endpoints: Iterable,
+    shutdown_event: Optional[asyncio.Event] = None,
+    grace_period_s: Optional[float] = None,
+) -> None:
+    if _shutdown_started.is_set():
+        return
+    _shutdown_started.set()
+
+    if grace_period_s is None:
+        grace_period_s = get_grace_period_seconds()
+
+    logger.info("Received shutdown signal; unregistering endpoints from discovery")
+    await _unregister_endpoints(list(endpoints))
+
+    if grace_period_s > 0:
+        logger.info("Grace period %.2fs before stopping endpoints", grace_period_s)
+        await asyncio.sleep(grace_period_s)
+
+    if shutdown_event is not None:
+        shutdown_event.set()
+
+    logger.info("Initiating runtime shutdown")
+    runtime.shutdown()
+
+
+def install_signal_handlers(
+    loop: asyncio.AbstractEventLoop,
+    runtime,
+    endpoints: Iterable,
+    shutdown_event: Optional[asyncio.Event] = None,
+    grace_period_s: Optional[float] = None,
+) -> None:
+    shutdown_task: Optional[asyncio.Task[None]] = None
+
+    def _on_shutdown_done(task: asyncio.Task[None]) -> None:
+        nonlocal shutdown_task
+        try:
+            task.result()
+        except asyncio.CancelledError:
+            logger.info("Graceful shutdown task cancelled")
+        except Exception:
+            logger.exception("Graceful shutdown task failed")
+        finally:
+            if shutdown_task is task:
+                shutdown_task = None
+
+    def signal_handler() -> None:
+        nonlocal shutdown_task
+        if shutdown_task is not None and not shutdown_task.done():
+            logger.debug("Shutdown already in progress; ignoring duplicate signal")
+            return
+
+        shutdown_task = asyncio.create_task(
+            graceful_shutdown_with_discovery(
+                runtime,
+                endpoints,
+                shutdown_event=shutdown_event,
+                grace_period_s=grace_period_s,
+            )
+        )
+        shutdown_task.add_done_callback(_on_shutdown_done)
+
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        loop.add_signal_handler(sig, signal_handler)
+
+    logger.info(
+        "Signal handlers set up for graceful shutdown "
+        "(discovery unregister + grace period)"
+    )

components/src/dynamo/common/utils/runtime.py

@@ -7,14 +7,12 @@
 Provides:
     - parse_endpoint: Parse 'dyn://namespace.component.endpoint' strings
     - graceful_shutdown: Shutdown DistributedRuntime with optional event signaling
-    - create_runtime: Create DistributedRuntime with signal handlers
+    - create_runtime: Create DistributedRuntime.
 """
 
 import asyncio
-import logging
 import os
-import signal
-from typing import Optional, Tuple
+from typing import Tuple
 
 from dynamo.runtime import DistributedRuntime
 
@@ -43,42 +41,22 @@ def parse_endpoint(endpoint: str) -> Tuple[str, str, str]:
     return namespace, component, endpoint_name
 
 
-async def graceful_shutdown(
-    runtime: DistributedRuntime,
-    shutdown_event: Optional[asyncio.Event] = None,
-) -> None:
-    """Shutdown DistributedRuntime with optional event signaling.
-
-    Args:
-        runtime: The DistributedRuntime instance to shut down.
-        shutdown_event: Optional event to set before shutting down,
-            signaling in-flight handlers to finish.
-    """
-    logging.info("Received shutdown signal, shutting down DistributedRuntime")
-    if shutdown_event is not None:
-        shutdown_event.set()
-    runtime.shutdown()
-    logging.info("DistributedRuntime shutdown complete")
-
-
 def create_runtime(
     discovery_backend: str,
     request_plane: str,
     event_plane: str,
     use_kv_events: bool,
-    shutdown_event: Optional[asyncio.Event] = None,
 ) -> Tuple[DistributedRuntime, asyncio.AbstractEventLoop]:
-    """Create a DistributedRuntime and register signal handlers for graceful shutdown.
+    """Create a DistributedRuntime.
 
     Sets DYN_EVENT_PLANE in the environment, computes whether NATS is needed,
-    creates the runtime, and installs SIGTERM/SIGINT handlers.
+    and creates the runtime.
 
     Args:
         discovery_backend: Discovery backend type (kubernetes, etcd, file, mem).
         request_plane: Request distribution method (nats, http, tcp).
         event_plane: Event publishing method (nats, zmq).
         use_kv_events: Whether KV events are enabled.
-        shutdown_event: Optional event to set on shutdown signal.
 
     Returns:
         Tuple of (runtime, event_loop).
@@ -91,12 +69,4 @@ def create_runtime(
 
     runtime = DistributedRuntime(loop, discovery_backend, request_plane, enable_nats)
 
-    def signal_handler():
-        asyncio.create_task(graceful_shutdown(runtime, shutdown_event))
-
-    for sig in (signal.SIGTERM, signal.SIGINT):
-        loop.add_signal_handler(sig, signal_handler)
-
-    logging.debug("Signal handlers set up for graceful shutdown")
-
     return runtime, loop