feat(runtime): add Bulwark failover replay support

Signed-off-by: Maksim Khadkevich <mkhadkevich@nvidia.com>

Author: hutm

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

@@ -5,6 +5,7 @@
 import logging
 import os
 import signal
+import sys
 from typing import Any, Callable, Coroutine, Iterable, Optional
 
 from dynamo._core import DistributedRuntime
@@ -16,6 +17,7 @@
 _DEFAULT_DRAIN_TIMEOUT_SECS = 30.0
 _DEFAULT_CLEANUP_TIMEOUT_SECS = 30.0
 _GRACE_PERIOD_ENV = "DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS"
+_FAST_FAILOVER_EXIT_ENV = "DYN_GMS_FAILOVER_FAST_EXIT_ON_SIGTERM"
 _shutdown_started = asyncio.Event()
 
 
@@ -24,6 +26,29 @@ def is_shutdown_in_progress() -> bool:
     return _shutdown_started.is_set()
 
 
+def fast_failover_exit_enabled() -> bool:
+    value = os.getenv(_FAST_FAILOVER_EXIT_ENV)
+    if value is None:
+        return False
+    return value.strip().lower() not in {"", "0", "false", "no", "off"}
+
+
+def _fast_exit_after_failover_unregister(
+    shutdown_event: Optional[asyncio.Event],
+) -> None:
+    if shutdown_event is not None:
+        shutdown_event.set()
+    logger.warning(
+        "GMS failover fast-exit enabled; exiting after discovery unregister "
+        "so the kernel releases failover ownership immediately"
+    )
+    try:
+        sys.stdout.flush()
+        sys.stderr.flush()
+    finally:
+        os._exit(0)
+
+
 def get_grace_period_seconds() -> float:
     value = os.getenv(_GRACE_PERIOD_ENV)
     if value is None or value == "":
@@ -111,6 +136,9 @@ async def graceful_shutdown_with_discovery(
     logger.info("Received shutdown signal; unregistering endpoints from discovery")
     await _unregister_endpoints(list(endpoints))
 
+    if fast_failover_exit_enabled():
+        _fast_exit_after_failover_unregister(shutdown_event)
+
     if grace_period_s > 0:
         logger.info("Grace period %.2fs before stopping endpoints", grace_period_s)
         await asyncio.sleep(grace_period_s)

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

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Compatibility helpers for forwarding OTel trace context to engines."""
+
+from __future__ import annotations
+
+from dynamo._core import Context
+
+
+def build_trace_headers(
+    context: Context | None,
+    *,
+    enabled: bool = True,
+) -> dict[str, str] | None:
+    """Return W3C trace headers for an engine request when available."""
+    if not enabled or context is None:
+        return None
+    return context.trace_headers()

components/src/dynamo/common/utils/tests/test_graceful_shutdown.py

@@ -36,12 +36,20 @@
 
 _GRACEFUL_SHUTDOWN_PATH = Path(__file__).parent.parent / "graceful_shutdown.py"
 
-# Provide a minimal dynamo._core stub so the module can be loaded
-_dynamo_stub = types.ModuleType("dynamo")
+# Provide a minimal dynamo._core stub while loading graceful_shutdown.py without
+# importing the native extension. The stub is restored immediately after module
+# load so pytest package collection can still import the real runtime bindings.
 _dynamo_core_stub = types.ModuleType("dynamo._core")
 _dynamo_core_stub.DistributedRuntime = object
-sys.modules.setdefault("dynamo", _dynamo_stub)
-sys.modules.setdefault("dynamo._core", _dynamo_core_stub)
+_previous_dynamo = sys.modules.get("dynamo")
+_previous_dynamo_core = sys.modules.get("dynamo._core")
+_added_dynamo_stub = False
+if importlib.util.find_spec("dynamo") is None:
+    _dynamo_stub = types.ModuleType("dynamo")
+    _dynamo_stub.__path__ = []  # mark as package for dynamo._core imports
+    sys.modules["dynamo"] = _dynamo_stub
+    _added_dynamo_stub = True
+sys.modules["dynamo._core"] = _dynamo_core_stub
 
 
 def _load_graceful_shutdown():
@@ -55,8 +63,20 @@ def _load_graceful_shutdown():
 
 
 _gs = _load_graceful_shutdown()
+if _previous_dynamo_core is None:
+    sys.modules.pop("dynamo._core", None)
+else:
+    sys.modules["dynamo._core"] = _previous_dynamo_core
+if _added_dynamo_stub:
+    if _previous_dynamo is None:
+        sys.modules.pop("dynamo", None)
+    else:
+        sys.modules["dynamo"] = _previous_dynamo
+
 graceful_shutdown_with_discovery = _gs.graceful_shutdown_with_discovery
 install_signal_handlers = _gs.install_signal_handlers
+is_shutdown_in_progress = _gs.is_shutdown_in_progress
+fast_failover_exit_enabled = _gs.fast_failover_exit_enabled
 
 
 # ---------------------------------------------------------------------------
@@ -309,3 +329,75 @@ async def _run():
 
     asyncio.run(_run())
     mock_runtime.shutdown.assert_called_once()
+
+
+def test_shutdown_in_progress_reflects_active_shutdown():
+    """The shutdown-in-progress flag is visible before shutdown_event is set."""
+    assert is_shutdown_in_progress() is False
+
+    async def _run():
+        mock_runtime = MagicMock()
+        mock_endpoint = AsyncMock()
+        mock_endpoint.unregister_endpoint_instance = AsyncMock(return_value=None)
+        await graceful_shutdown_with_discovery(
+            runtime=mock_runtime,
+            endpoints=[mock_endpoint],
+            shutdown_event=None,
+            grace_period_s=0,
+        )
+
+    asyncio.run(_run())
+
+    assert is_shutdown_in_progress() is True
+
+
+def test_fast_failover_exit_env_parsing(monkeypatch):
+    monkeypatch.delenv("DYN_GMS_FAILOVER_FAST_EXIT_ON_SIGTERM", raising=False)
+    assert fast_failover_exit_enabled() is False
+
+    for value in ("1", "true", "yes", "on", "enabled"):
+        monkeypatch.setenv("DYN_GMS_FAILOVER_FAST_EXIT_ON_SIGTERM", value)
+        assert fast_failover_exit_enabled() is True
+
+    for value in ("", "0", "false", "no", "off"):
+        monkeypatch.setenv("DYN_GMS_FAILOVER_FAST_EXIT_ON_SIGTERM", value)
+        assert fast_failover_exit_enabled() is False
+
+
+def test_fast_failover_exit_runs_after_unregister_before_runtime_shutdown(monkeypatch):
+    class FastExit(Exception):
+        pass
+
+    call_order = []
+    shutdown_event = asyncio.Event()
+    mock_runtime = MagicMock()
+    mock_runtime.shutdown = MagicMock(side_effect=lambda: call_order.append("shutdown"))
+
+    async def _run():
+        mock_endpoint = AsyncMock()
+        mock_endpoint.unregister_endpoint_instance = AsyncMock(
+            side_effect=lambda: call_order.append("unregister")
+        )
+
+        def fake_fast_exit(event):
+            assert event is shutdown_event
+            event.set()
+            call_order.append("fast_exit")
+            raise FastExit
+
+        monkeypatch.setenv("DYN_GMS_FAILOVER_FAST_EXIT_ON_SIGTERM", "1")
+        monkeypatch.setattr(_gs, "_fast_exit_after_failover_unregister", fake_fast_exit)
+
+        await graceful_shutdown_with_discovery(
+            runtime=mock_runtime,
+            endpoints=[mock_endpoint],
+            shutdown_event=shutdown_event,
+            grace_period_s=0,
+        )
+
+    with pytest.raises(FastExit):
+        asyncio.run(_run())
+
+    assert call_order == ["unregister", "fast_exit"]
+    assert shutdown_event.is_set()
+    mock_runtime.shutdown.assert_not_called()

docs/fault-tolerance/request-migration.md

@@ -34,6 +34,15 @@ The migration limit is configured at the **frontend** level and applies globally
 - Set via `--migration-limit` flag on the frontend
 - Applies to all models served by the frontend
 
+### Retry Concurrency Configuration
+
+During worker failover, many in-flight streams can disconnect at the same time and retry against the newly active worker. `DYN_MIGRATION_RETRY_CONCURRENCY` optionally caps the number of migrated retry streams that can run concurrently in a frontend process:
+
+- Default behavior: unlimited retry concurrency (`DYN_MIGRATION_RETRY_CONCURRENCY=0` or unset)
+- Set via `DYN_MIGRATION_RETRY_CONCURRENCY` on the frontend or Bulwark gateway
+- The cap applies only after a request has entered migration/retry; first-attempt traffic is not throttled
+- A migrated stream holds a permit until that stream completes, which prevents a failover replay burst from overwhelming a warm shadow worker
+
 ### Max Sequence Length Configuration
 
 The max sequence length setting controls how long the migration system will cache token state for a request. Once the total sequence length (prompt + generated tokens) exceeds this limit, migration is disabled for that request and token tracking stops: