feat(runtime): add Bulwark failover replay support

Signed-off-by: mkhadkevich <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,9 +17,38 @@
 _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()
 
 
+def is_shutdown_in_progress() -> bool:
+    """Return whether this process is already handling graceful shutdown."""
+    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 == "":
@@ -106,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()

components/src/dynamo/frontend/frontend_args.py

@@ -25,6 +25,7 @@
     add_negatable_bool_argument,
     env_or_default,
 )
+from dynamo.common.utils.namespace import get_worker_namespace
 
 from . import __version__
 
@@ -62,6 +63,7 @@ class FrontendConfig(RouterConfigBase, KvRouterConfigBase, AicPerfConfigBase):
 
     namespace: Optional[str] = None
     namespace_prefix: Optional[str] = None
+    bulwark_gateway_endpoint: Optional[str] = None
 
     migration_limit: int
     migration_max_seq_len: Optional[int]
@@ -94,6 +96,26 @@ def validate(self) -> None:
             self.router_mode = "kv"
         self.apply_load_aware_preset()
 
+        if self.bulwark_gateway_endpoint:
+            if self.namespace and not self.namespace_prefix:
+                self.namespace = get_worker_namespace(self.namespace)
+            if not self.bulwark_gateway_endpoint.startswith("dyn://"):
+                raise ValueError(
+                    "--bulwark-gateway-endpoint must be a dyn://namespace.component.endpoint path"
+                )
+            if not self.namespace and not self.namespace_prefix:
+                raise ValueError(
+                    "--bulwark-gateway-endpoint requires --namespace or --namespace-prefix "
+                    "to select the private primary/shadow worker namespace"
+                )
+            if self.interactive:
+                raise ValueError(
+                    "--bulwark-gateway-endpoint cannot be combined with --interactive"
+                )
+            if self.kserve_grpc_server:
+                raise ValueError(
+                    "--bulwark-gateway-endpoint cannot be combined with --kserve-grpc-server"
+                )
         if bool(self.tls_cert_path) ^ bool(self.tls_key_path):  # ^ is XOR
             raise ValueError(
                 "--tls-cert-path and --tls-key-path must be provided together"
@@ -265,6 +287,19 @@ def add_arguments(self, parser) -> None:
             ),
         )
 
+        add_argument(
+            g,
+            flag_name="--bulwark-gateway-endpoint",
+            env_var="DYN_BULWARK_GATEWAY_ENDPOINT",
+            default=None,
+            help=(
+                "Expose this frontend as a stable request-plane worker endpoint for Bulwark "
+                "compound workers. The value is a public dyn://namespace.component.endpoint "
+                "path; private primary/shadow workers are discovered using --namespace or "
+                "--namespace-prefix."
+            ),
+        )
+
         add_argument(
             g,
             flag_name="--migration-limit",

components/src/dynamo/frontend/main.py

@@ -171,14 +171,18 @@ async def async_main():
     Initializes the distributed runtime, configures routing, and starts
     the HTTP server or interactive mode based on command-line arguments.
     """
-    # The system status server port is a worker concern.
+    # The system status server port is usually a worker concern.
     #
     # Serve tests set DYN_SYSTEM_PORT for the worker, but aggregated launch scripts
-    # start `dynamo.frontend` first. If the frontend inherits DYN_SYSTEM_PORT, it can
-    # bind that port before the worker, causing port conflicts and/or scraping the
-    # wrong metrics endpoint.
+    # start `dynamo.frontend` first. If the normal frontend inherits DYN_SYSTEM_PORT,
+    # it can bind that port before the worker, causing port conflicts and/or scraping
+    # the wrong metrics endpoint. Bulwark gateway endpoint mode intentionally keeps
+    # the system server so Kubernetes readiness can gate the stable worker entity.
+    frontend_system_port = os.environ.get("DYN_SYSTEM_PORT")
     os.environ.pop("DYN_SYSTEM_PORT", None)
     config, vllm_flags, sglang_flags = parse_args()
+    if config.bulwark_gateway_endpoint and frontend_system_port:
+        os.environ["DYN_SYSTEM_PORT"] = frontend_system_port
     dump_config(config.dump_config_to, config)
     max_seq_info = (
         f", max_seq_len: {config.migration_max_seq_len}"
@@ -189,8 +193,9 @@ async def async_main():
         f"Request migration {'enabled' if config.migration_limit > 0 else 'disabled'} "
         f"(limit: {config.migration_limit}{max_seq_info})"
     )
-    # Warn if DYN_SYSTEM_PORT is set (frontend doesn't use system metrics server)
-    if os.environ.get("DYN_SYSTEM_PORT"):
+    # Warn if DYN_SYSTEM_PORT is set for a normal frontend. Gateway endpoint mode
+    # uses it for Kubernetes readiness.
+    if os.environ.get("DYN_SYSTEM_PORT") and not config.bulwark_gateway_endpoint:
         logger.warning(
             "=" * 80 + "\n"
             "WARNING: DYN_SYSTEM_PORT is set but NOT used by the frontend!\n"
@@ -297,7 +302,15 @@ def signal_handler():
     engine = await make_engine(runtime, e)
 
     try:
-        if config.interactive:
+        if config.bulwark_gateway_endpoint:
+            logger.info(
+                "Starting Bulwark frontend gateway at %s; private workers selected by namespace=%s namespace_prefix=%s",
+                config.bulwark_gateway_endpoint,
+                config.namespace,
+                config.namespace_prefix,
+            )
+            await run_input(runtime, config.bulwark_gateway_endpoint, engine)
+        elif config.interactive:
             await run_input(runtime, "text", engine)
         elif config.kserve_grpc_server:
             await run_input(runtime, "grpc", engine)