diff --git a/components/src/dynamo/common/utils/graceful_shutdown.py b/components/src/dynamo/common/utils/graceful_shutdown.py index 95de6452e4fa..d86703654186 100644 --- a/components/src/dynamo/common/utils/graceful_shutdown.py +++ b/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) diff --git a/components/src/dynamo/common/utils/otel_tracing.py b/components/src/dynamo/common/utils/otel_tracing.py new file mode 100644 index 000000000000..e3f5ddce903c --- /dev/null +++ b/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() diff --git a/components/src/dynamo/common/utils/tests/test_graceful_shutdown.py b/components/src/dynamo/common/utils/tests/test_graceful_shutdown.py index 2fcc8ed8a6a1..23b3e203915d 100644 --- a/components/src/dynamo/common/utils/tests/test_graceful_shutdown.py +++ b/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() diff --git a/components/src/dynamo/frontend/frontend_args.py b/components/src/dynamo/frontend/frontend_args.py index 43e94efabb56..c6dec60f4089 100644 --- a/components/src/dynamo/frontend/frontend_args.py +++ b/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", diff --git a/components/src/dynamo/frontend/main.py b/components/src/dynamo/frontend/main.py index 13b60007486b..3cdd34f43e55 100644 --- a/components/src/dynamo/frontend/main.py +++ b/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) diff --git a/lib/llm/src/entrypoint/input/endpoint.rs b/lib/llm/src/entrypoint/input/endpoint.rs index 1802266da41a..8fd92e8a9c62 100644 --- a/lib/llm/src/entrypoint/input/endpoint.rs +++ b/lib/llm/src/entrypoint/input/endpoint.rs @@ -1,13 +1,23 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{future::Future, pin::Pin, sync::Arc}; +use std::{collections::HashSet, future::Future, pin::Pin, sync::Arc}; + +use anyhow::Context as _; +use futures::{StreamExt, future::pending}; +use tokio::time::{Duration, Instant, sleep_until}; use crate::{ backend::Backend, + discovery::{KvWorkerMonitor, ModelManager, WORKER_TYPE_DECODE}, engines::StreamingEngineAdapter, + entrypoint::{EngineConfig, build_preprocessed_routing}, + http::service::metrics::Metrics, + model_card::ModelDeploymentCard, model_type::{ModelInput, ModelType}, + namespace::NamespaceFilter, preprocessor::{BackendOutput, PreprocessedRequest}, + protocols::common::llm_backend::LLMEngineOutput, types::{ Annotated, openai::chat_completions::{ @@ -18,12 +28,18 @@ use crate::{ }; use dynamo_runtime::engine::AsyncEngineStream; -use dynamo_runtime::pipeline::{ - Context, ManyOut, Operator, SegmentSource, ServiceBackend, SingleIn, Source, network::Ingress, +use dynamo_runtime::{ + DistributedRuntime, + config::HealthStatus, + discovery::{ + DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoverySpec, + }, + pipeline::{ + Context, ManyOut, Operator, RouterMode, SegmentSource, ServiceBackend, SingleIn, Source, + network::Ingress, + }, + protocols::EndpointId, }; -use dynamo_runtime::{DistributedRuntime, protocols::EndpointId}; - -use crate::entrypoint::EngineConfig; pub async fn run( distributed_runtime: DistributedRuntime, @@ -123,8 +139,115 @@ pub async fn run( let fut = endpoint.endpoint_builder().handler(ingress).start(); Box::pin(fut) } - EngineConfig::Dynamic { .. } => { - unreachable!("An endpoint input will never have a Dynamic engine"); + EngineConfig::Dynamic { + model: local_model, + chat_engine_factory, + prefill_load_estimator, + } => { + if chat_engine_factory.is_some() { + tracing::warn!( + "Bulwark gateway endpoint mode receives preprocessed request-plane traffic; \ + dyn-chat-processor factories are ignored in this mode" + ); + } + + let namespace_filter = NamespaceFilter::from_namespace_and_prefix( + local_model.namespace(), + local_model.namespace_prefix(), + ); + if namespace_filter.is_global() { + anyhow::bail!( + "Bulwark gateway endpoint mode requires --namespace or --namespace-prefix \ + to select private primary/shadow workers" + ); + } + + let (private_endpoint_id, private_card) = + wait_for_gateway_backend(&distributed_runtime, &endpoint_id, &namespace_filter) + .await?; + + let private_component = distributed_runtime + .namespace(&private_endpoint_id.namespace)? + .component(&private_endpoint_id.component)?; + let private_endpoint = private_component.endpoint(&private_endpoint_id.name); + let private_client = private_endpoint.client().await?; + + let model_manager = Arc::new(ModelManager::new()); + let metrics = Arc::new(Metrics::new()); + let router_config = private_card + .router_config + .as_ref() + .unwrap_or(local_model.router_config()); + + let kv_chooser = if router_config.router_mode == RouterMode::KV { + Some( + model_manager + .kv_chooser_for( + &private_endpoint, + private_card.kv_cache_block_size, + Some(router_config.kv_router_config.clone()), + prefill_load_estimator.clone(), + WORKER_TYPE_DECODE, + Some(private_card.display_name.clone()), + private_card.runtime_config.enable_eagle, + ) + .await?, + ) + } else { + None + }; + + let monitor_client = kv_chooser + .as_ref() + .map(|chooser| chooser.client().clone()) + .unwrap_or_else(|| private_client.clone()); + let worker_monitor = Some(KvWorkerMonitor::new( + monitor_client, + router_config.load_threshold_config.clone(), + )); + + let routing = build_preprocessed_routing( + &private_client, + model_manager, + router_config.router_mode, + worker_monitor, + kv_chooser, + None, + false, + router_config.session_affinity_ttl_secs, + ) + .await + .context("build Bulwark gateway preprocessed routing")?; + let pipeline = routing + .build_preprocessed_pipeline( + &private_card, + local_model.migration_limit(), + local_model.migration_max_seq_len(), + metrics, + ) + .context("build Bulwark gateway preprocessed pipeline")?; + let ingress = Ingress::< + SingleIn, + ManyOut>, + >::for_engine(pipeline)?; + + let public_card = gateway_public_card(&private_card); + register_gateway_model_card(&distributed_runtime, &endpoint_id, &public_card).await?; + spawn_gateway_backend_readiness_monitor( + distributed_runtime.clone(), + endpoint_id.clone(), + namespace_filter, + ); + + tracing::info!( + public_endpoint = %endpoint_id, + private_endpoint = %private_endpoint_id, + model = %public_card.name(), + "Bulwark gateway endpoint registered" + ); + + let fut = endpoint.endpoint_builder().handler(ingress).start(); + Box::pin(fut) } }; @@ -153,6 +276,401 @@ pub async fn run( } } +fn gateway_public_card(private_card: &ModelDeploymentCard) -> ModelDeploymentCard { + let mut card = private_card.clone(); + let mut model_type = ModelType::empty(); + if private_card.model_type.supports_chat() { + model_type |= ModelType::Chat; + } + if private_card.model_type.supports_completions() { + model_type |= ModelType::Completions; + } + card.model_type = model_type; + card.model_input = ModelInput::Tokens; + card.worker_type = Some(WorkerType::Aggregated); + card.needs.clear(); + card.router_config = None; + card +} + +fn is_gateway_backend_card(card: &ModelDeploymentCard) -> bool { + card.model_input == ModelInput::Tokens + && (card.model_type.supports_chat() || card.model_type.supports_completions()) + && !card.model_type.supports_prefill() + && matches!(card.worker_type, None | Some(WorkerType::Aggregated)) + && card.needs.is_empty() +} + +async fn wait_for_gateway_backend( + distributed_runtime: &DistributedRuntime, + public_endpoint_id: &EndpointId, + namespace_filter: &NamespaceFilter, +) -> anyhow::Result<(EndpointId, ModelDeploymentCard)> { + let discovery_stream = distributed_runtime + .discovery() + .list_and_watch( + DiscoveryQuery::AllModels, + Some(distributed_runtime.primary_token().clone()), + ) + .await?; + tokio::pin!(discovery_stream); + + tracing::info!( + public_endpoint = %public_endpoint_id, + namespace_filter = ?namespace_filter, + "Waiting for private aggregated token worker for Bulwark gateway" + ); + + while let Some(event) = discovery_stream.next().await { + let event = event?; + let DiscoveryEvent::Added(instance) = event else { + continue; + }; + let DiscoveryInstance::Model { + namespace, + component, + endpoint, + .. + } = &instance + else { + continue; + }; + + if !namespace_filter.matches(namespace) { + continue; + } + + let endpoint_id = EndpointId { + namespace: namespace.clone(), + component: component.clone(), + name: endpoint.clone(), + }; + if &endpoint_id == public_endpoint_id { + tracing::debug!(%endpoint_id, "Skipping Bulwark gateway public endpoint"); + continue; + } + + let card = match instance.deserialize_model::() { + Ok(card) => card, + Err(err) => { + tracing::warn!(%endpoint_id, %err, "Skipping unreadable model card"); + continue; + } + }; + if !is_gateway_backend_card(&card) { + tracing::debug!( + %endpoint_id, + model = %card.name(), + model_input = %card.model_input.as_str(), + model_type = %card.model_type.as_str(), + worker_type = ?card.worker_type, + needs = ?card.needs, + "Skipping model card that is not an aggregated token backend for Bulwark gateway" + ); + continue; + } + + tracing::info!( + private_endpoint = %endpoint_id, + model = %card.name(), + "Found private Bulwark backend for gateway" + ); + return Ok((endpoint_id, card)); + } + + anyhow::bail!( + "model discovery stream ended before a private aggregated token worker appeared for Bulwark gateway" + ) +} + +const BULWARK_GATEWAY_NOTREADY_GRACE_MS_ENV: &str = "DYN_BULWARK_GATEWAY_NOTREADY_GRACE_MS"; +const DEFAULT_BULWARK_GATEWAY_NOTREADY_GRACE_MS: u64 = 10_000; + +fn gateway_notready_grace() -> Duration { + match std::env::var(BULWARK_GATEWAY_NOTREADY_GRACE_MS_ENV) { + Ok(raw) => match raw.parse::() { + Ok(ms) => Duration::from_millis(ms), + Err(err) => { + tracing::warn!( + value = %raw, + %err, + default_ms = DEFAULT_BULWARK_GATEWAY_NOTREADY_GRACE_MS, + "Invalid DYN_BULWARK_GATEWAY_NOTREADY_GRACE_MS; using default gateway readiness grace" + ); + Duration::from_millis(DEFAULT_BULWARK_GATEWAY_NOTREADY_GRACE_MS) + } + }, + Err(_) => Duration::from_millis(DEFAULT_BULWARK_GATEWAY_NOTREADY_GRACE_MS), + } +} + +fn spawn_gateway_backend_readiness_monitor( + distributed_runtime: DistributedRuntime, + public_endpoint_id: EndpointId, + namespace_filter: NamespaceFilter, +) { + let cancel_token = distributed_runtime.primary_token().clone(); + let health_runtime = distributed_runtime.clone(); + + tokio::spawn(async move { + let discovery_stream = match distributed_runtime + .discovery() + .list_and_watch(DiscoveryQuery::AllModels, Some(cancel_token.clone())) + .await + { + Ok(stream) => stream, + Err(err) => { + tracing::error!( + %err, + public_endpoint = %public_endpoint_id, + "Bulwark gateway failed to start private backend readiness monitor" + ); + set_gateway_health(&health_runtime, HealthStatus::NotReady, 0); + return; + } + }; + tokio::pin!(discovery_stream); + + let notready_grace = gateway_notready_grace(); + tracing::info!( + public_endpoint = %public_endpoint_id, + notready_grace_ms = notready_grace.as_millis(), + "Bulwark gateway readiness monitor configured" + ); + + let mut active_backends: HashSet = HashSet::new(); + let mut last_ready = false; + let mut pending_notready_deadline: Option = None; + set_gateway_health(&health_runtime, HealthStatus::NotReady, 0); + + loop { + let pending_notready = async { + match pending_notready_deadline { + Some(deadline) => sleep_until(deadline).await, + None => pending::<()>().await, + } + }; + + tokio::select! { + event = discovery_stream.next() => { + let Some(event) = event else { + tracing::warn!( + public_endpoint = %public_endpoint_id, + "Bulwark gateway private backend readiness stream ended" + ); + set_gateway_health(&health_runtime, HealthStatus::NotReady, 0); + break; + }; + + match event { + Ok(DiscoveryEvent::Added(instance)) => { + let id = instance.id(); + if is_gateway_backend_instance( + &instance, + &public_endpoint_id, + &namespace_filter, + ) { + active_backends.insert(id); + } else { + active_backends.remove(&id); + } + } + Ok(DiscoveryEvent::Removed(id)) => { + active_backends.remove(&id); + } + Err(err) => { + tracing::warn!( + %err, + public_endpoint = %public_endpoint_id, + "Bulwark gateway private backend readiness watch event failed" + ); + continue; + } + } + + let ready = !active_backends.is_empty(); + if ready { + pending_notready_deadline = None; + if !last_ready { + set_gateway_health(&health_runtime, HealthStatus::Ready, active_backends.len()); + tracing::info!( + public_endpoint = %public_endpoint_id, + active_private_backends = active_backends.len(), + ready = true, + "Bulwark gateway readiness changed" + ); + last_ready = true; + } + } else if last_ready { + if notready_grace.is_zero() { + set_gateway_health(&health_runtime, HealthStatus::NotReady, 0); + tracing::info!( + public_endpoint = %public_endpoint_id, + active_private_backends = 0, + ready = false, + "Bulwark gateway readiness changed" + ); + last_ready = false; + } else if pending_notready_deadline.is_none() { + pending_notready_deadline = Some(Instant::now() + notready_grace); + tracing::info!( + public_endpoint = %public_endpoint_id, + notready_grace_ms = notready_grace.as_millis(), + "Bulwark gateway private backends empty; delaying NotReady" + ); + } + } + } + _ = pending_notready => { + pending_notready_deadline = None; + if active_backends.is_empty() && last_ready { + set_gateway_health(&health_runtime, HealthStatus::NotReady, 0); + tracing::info!( + public_endpoint = %public_endpoint_id, + active_private_backends = 0, + ready = false, + "Bulwark gateway readiness changed after grace" + ); + last_ready = false; + } + } + } + } + }); +} + +fn set_gateway_health( + distributed_runtime: &DistributedRuntime, + status: HealthStatus, + active_private_backends: usize, +) { + distributed_runtime + .system_health() + .lock() + .set_health_status(status.clone()); + tracing::debug!( + ?status, + active_private_backends, + "Bulwark gateway system health updated" + ); +} + +fn is_gateway_backend_instance( + instance: &DiscoveryInstance, + public_endpoint_id: &EndpointId, + namespace_filter: &NamespaceFilter, +) -> bool { + let DiscoveryInstance::Model { + namespace, + component, + endpoint, + .. + } = instance + else { + return false; + }; + + if !namespace_filter.matches(namespace) { + return false; + } + + let endpoint_id = EndpointId { + namespace: namespace.clone(), + component: component.clone(), + name: endpoint.clone(), + }; + if &endpoint_id == public_endpoint_id { + return false; + } + + instance + .deserialize_model::() + .is_ok_and(|card| is_gateway_backend_card(&card)) +} + +async fn register_gateway_model_card( + distributed_runtime: &DistributedRuntime, + endpoint_id: &EndpointId, + card: &ModelDeploymentCard, +) -> anyhow::Result<()> { + let spec = DiscoverySpec::from_model( + endpoint_id.namespace.clone(), + endpoint_id.component.clone(), + endpoint_id.name.clone(), + card, + )?; + distributed_runtime.discovery().register(spec).await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::entrypoint::RouterConfig; + + fn token_chat_card() -> ModelDeploymentCard { + let mut card = ModelDeploymentCard::default(); + card.model_input = ModelInput::Tokens; + card.model_type = ModelType::Chat; + card.worker_type = Some(WorkerType::Aggregated); + card + } + + #[test] + fn gateway_backend_card_accepts_aggregated_token_chat() { + let card = token_chat_card(); + assert!(is_gateway_backend_card(&card)); + } + + #[test] + fn gateway_backend_card_accepts_legacy_missing_worker_type() { + let mut card = token_chat_card(); + card.worker_type = None; + assert!(is_gateway_backend_card(&card)); + } + + #[test] + fn gateway_backend_card_rejects_disaggregated_decode() { + let mut card = token_chat_card(); + card.worker_type = Some(WorkerType::Decode); + card.needs = vec![vec![WorkerType::Prefill]]; + assert!(!is_gateway_backend_card(&card)); + } + + #[test] + fn gateway_backend_card_rejects_text_or_prefill_cards() { + let mut text_card = token_chat_card(); + text_card.model_input = ModelInput::Text; + assert!(!is_gateway_backend_card(&text_card)); + + let mut prefill_card = token_chat_card(); + prefill_card.model_type = ModelType::Prefill; + prefill_card.worker_type = Some(WorkerType::Prefill); + prefill_card.needs = vec![vec![WorkerType::Decode]]; + assert!(!is_gateway_backend_card(&prefill_card)); + } + + #[test] + fn gateway_public_card_hides_private_topology() { + let mut private = token_chat_card(); + private.model_type = ModelType::Chat | ModelType::Completions | ModelType::Images; + private.worker_type = Some(WorkerType::Decode); + private.needs = vec![vec![WorkerType::Prefill]]; + private.router_config = Some(RouterConfig::default()); + + let public = gateway_public_card(&private); + + assert!(public.model_type.supports_chat()); + assert!(public.model_type.supports_completions()); + assert!(!public.model_type.supports_images()); + assert!(!public.model_type.supports_prefill()); + assert_eq!(public.model_input, ModelInput::Tokens); + assert_eq!(public.worker_type, Some(WorkerType::Aggregated)); + assert!(public.needs.is_empty()); + assert!(public.router_config.is_none()); + } +} + #[cfg(test)] #[cfg(feature = "integration")] mod integration_tests { diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index d61cc9b3d6d2..63084d75a321 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -45,12 +45,14 @@ use super::{ }, service_v2, }; +use crate::discovery::ModelManagerError; use crate::engines::ValidateRequest; use crate::preprocessor::PRESERVE_OMITTED_MAX_TOKENS_CONTEXT_KEY; use crate::protocols::common::extensions::{ AGENT_CONTEXT_CONTEXT_KEY, AgentContext, SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId, agent_context_from_headers, apply_header_routing_overrides, session_affinity_from_headers, }; +use crate::protocols::openai::ParsingOptions; use crate::protocols::openai::chat_completions::aggregator::ChatCompletionAggregator; use crate::protocols::openai::{ audios::{NvAudioSpeechResponse, NvCreateAudioSpeechRequest}, @@ -75,7 +77,6 @@ use dynamo_runtime::logging::get_distributed_tracing_context; use tracing::Instrument; pub const DYNAMO_REQUEST_ID_HEADER: &str = "x-dynamo-request-id"; - /// Dynamo Annotation for the request ID pub const ANNOTATION_REQUEST_ID: &str = "request_id"; @@ -91,6 +92,55 @@ pub(super) fn rl_router( Ok(dynamo_rl::rl_router(state)) } +fn is_transient_model_lookup_error(error: &ModelManagerError) -> bool { + matches!( + error, + ModelManagerError::ModelNotFound(_) | ModelManagerError::ModelUnavailable(_) + ) +} + +async fn get_chat_completions_engine_with_failover_wait( + state: &Arc, + model: &str, +) -> Result< + ( + crate::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine, + ParsingOptions, + ), + ModelManagerError, +> { + service_v2::wait_for_failover( + "model", + model, + || { + state + .manager() + .get_chat_completions_engine_with_parsing(model) + }, + is_transient_model_lookup_error, + ) + .await +} + +async fn get_completions_engine_with_failover_wait( + state: &Arc, + model: &str, +) -> Result< + ( + crate::types::openai::completions::OpenAICompletionsStreamingEngine, + ParsingOptions, + ), + ModelManagerError, +> { + service_v2::wait_for_failover( + "model", + model, + || state.manager().get_completions_engine_with_parsing(model), + is_transient_model_lookup_error, + ) + .await +} + // Default axum max body limit without configuring is 2MB: https://docs.rs/axum/latest/axum/extract/struct.DefaultBodyLimit.html /// Default body limit in bytes (45MB) to support 500k+ token payloads. /// Can be configured at runtime using the DYN_HTTP_BODY_LIMIT_MB environment variable. @@ -740,9 +790,8 @@ async fn completions_single( let http_queue_guard = state.metrics_clone().create_http_queue_guard(&metric_model); // todo - error handling should be more robust - let (engine, parsing_options) = state - .manager() - .get_completions_engine_with_parsing(&model) + let (engine, parsing_options) = get_completions_engine_with_failover_wait(&state, &model) + .await .map_err(|e| { let err_response = ErrorMessage::from_model_error(&e); inflight_guard.mark_error(extract_error_type_from_response(&err_response)); @@ -890,9 +939,8 @@ async fn completions_batch( // Create http_queue_guard early - tracks time waiting to be processed let http_queue_guard = state.metrics_clone().create_http_queue_guard(&metric_model); - let (engine, parsing_options) = state - .manager() - .get_completions_engine_with_parsing(&model) + let (engine, parsing_options) = get_completions_engine_with_failover_wait(&state, &model) + .await .map_err(|e| { let err_response = ErrorMessage::from_model_error(&e); inflight_guard.mark_error(extract_error_type_from_response(&err_response)); @@ -1830,9 +1878,8 @@ async fn chat_completions( tracing::trace!("Getting chat completions engine for model: {}", model); - let (engine, parsing_options) = state - .manager() - .get_chat_completions_engine_with_parsing(&model) + let (engine, parsing_options) = get_chat_completions_engine_with_failover_wait(&state, &model) + .await .map_err(|e| { let err_response = ErrorMessage::from_model_error(&e); inflight_guard.mark_error(extract_error_type_from_response(&err_response)); @@ -2309,9 +2356,8 @@ async fn responses( tracing::trace!("Getting chat completions engine for model: {}", model); - let (engine, parsing_options) = state - .manager() - .get_chat_completions_engine_with_parsing(&model) + let (engine, parsing_options) = get_chat_completions_engine_with_failover_wait(&state, &model) + .await .map_err(|e| { let err_response = ErrorMessage::from_model_error(&e); inflight_guard.mark_error(extract_error_type_from_response(&err_response)); diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 4e314108940f..17f0df41e6fb 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use std::env::var; use std::path::PathBuf; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU8; use std::sync::atomic::AtomicU64; @@ -21,6 +22,7 @@ use super::metrics; use super::metrics::{register_lora_allocation_metrics, register_worker_timing_metrics}; use crate::discovery::ModelManager; use crate::endpoint_type::EndpointType; +use crate::frontend_config::{FrontendApiConfig, MetricsConfig}; use crate::kv_router::metrics::{ RoutingOverheadMetrics, register_router_queue_metrics, register_worker_load_metrics, }; @@ -42,10 +44,102 @@ use dynamo_runtime::metrics::{ use std::net::SocketAddr; use tokio::sync::Notify; use tokio::task::JoinHandle; +use tokio::time::{Instant, sleep}; use tokio_util::sync::CancellationToken; use tower_http::trace::TraceLayer; -use crate::frontend_config::{FrontendApiConfig, MetricsConfig}; +const DYN_HTTP_MODEL_FAILOVER_WAIT_MS: &str = "DYN_HTTP_MODEL_FAILOVER_WAIT_MS"; +pub(super) const MODEL_FAILOVER_WAIT_POLL_MS: u64 = 50; + +pub(super) fn model_failover_wait() -> Duration { + static WAIT: OnceLock = OnceLock::new(); + *WAIT.get_or_init(|| { + std::env::var(DYN_HTTP_MODEL_FAILOVER_WAIT_MS) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) + .unwrap_or_default() + }) +} + +fn should_wait_for_endpoint_failover(endpoint_type: EndpointType) -> bool { + matches!( + endpoint_type, + EndpointType::Chat | EndpointType::Completion | EndpointType::Responses + ) +} + +pub(super) async fn wait_for_failover( + subject_kind: &str, + subject: &str, + mut operation: F, + is_transient: P, +) -> Result +where + F: FnMut() -> Result, + P: Fn(&E) -> bool, +{ + let wait = model_failover_wait(); + let deadline = Instant::now() + wait; + let mut attempts = 0_u32; + + loop { + match operation() { + Ok(value) => { + if attempts > 0 { + tracing::info!( + subject_kind, + subject, + attempts, + wait_ms = wait.as_millis(), + "Failover subject became available" + ); + } + return Ok(value); + } + Err(error) if is_transient(&error) => { + let now = Instant::now(); + if wait.is_zero() || now >= deadline { + if attempts > 0 { + tracing::warn!( + subject_kind, + subject, + attempts, + wait_ms = wait.as_millis(), + "Failover subject remained unavailable" + ); + } + return Err(error); + } + attempts += 1; + sleep(std::cmp::min( + Duration::from_millis(MODEL_FAILOVER_WAIT_POLL_MS), + deadline - now, + )) + .await; + } + Err(error) => return Err(error), + } + } +} + +async fn endpoint_enabled_or_wait(state: &State, endpoint_type: EndpointType) -> bool { + if state.flags.get(&endpoint_type) { + return true; + } + if !should_wait_for_endpoint_failover(endpoint_type) { + return false; + } + + wait_for_failover( + "endpoint", + endpoint_type.as_str(), + || state.flags.get(&endpoint_type).then_some(()).ok_or(()), + |_| true, + ) + .await + .is_ok() +} /// Middleware that echoes `x-request-id` from request to response headers. async fn echo_request_id_header( @@ -892,13 +986,16 @@ impl HttpServiceConfigBuilder { let cancel_token = config.cancel_token.unwrap_or_default(); // Use the provided discovery client, or fall back to a no-op memory-backed one // (for in-process modes that don't need discovery) - let discovery_client = config.discovery.unwrap_or_else(|| { - use dynamo_runtime::discovery::KVStoreDiscovery; - Arc::new(KVStoreDiscovery::new( - dynamo_runtime::storage::kv::Manager::memory(), - cancel_token.child_token(), - )) as Arc - }); + let discovery_client = match config.discovery { + Some(client) => client, + None => { + use dynamo_runtime::discovery::KVStoreDiscovery; + Arc::new(KVStoreDiscovery::new( + dynamo_runtime::storage::kv::Manager::memory(), + cancel_token.child_token(), + )?) as Arc + } + }; // Both surfaces are on by default; an env-truthy DISABLE var turns them // off. The builder flag can also force off (e.g. tests), and wins. let nvext_enabled = @@ -1219,7 +1316,7 @@ impl HttpServiceConfigBuilder { let state: Arc = state_route.clone(); async move { // Check if the endpoint is enabled - let enabled = state.flags.get(&endpoint_type); + let enabled = endpoint_enabled_or_wait(&state, endpoint_type).await; if enabled { Ok(next.run(req).await) } else { @@ -1256,6 +1353,44 @@ mod tests { } } + #[tokio::test] + async fn failover_wait_returns_immediate_success() { + let result: Result = wait_for_failover("model", "test", || Ok(7), |_| true).await; + assert_eq!(result, Ok(7)); + } + + #[tokio::test] + async fn failover_wait_does_not_retry_non_transient_errors() { + let attempts = std::cell::Cell::new(0_u32); + let result: Result<(), &str> = wait_for_failover( + "model", + "test", + || { + attempts.set(attempts.get() + 1); + Err("permanent") + }, + |_| false, + ) + .await; + assert_eq!(result, Err("permanent")); + assert_eq!(attempts.get(), 1); + } + + #[test] + fn test_endpoint_failover_wait_scope() { + assert!(should_wait_for_endpoint_failover(EndpointType::Chat)); + assert!(should_wait_for_endpoint_failover(EndpointType::Completion)); + assert!(should_wait_for_endpoint_failover(EndpointType::Responses)); + + assert!(!should_wait_for_endpoint_failover(EndpointType::Embedding)); + assert!(!should_wait_for_endpoint_failover(EndpointType::Images)); + assert!(!should_wait_for_endpoint_failover(EndpointType::Videos)); + assert!(!should_wait_for_endpoint_failover(EndpointType::Audios)); + assert!(!should_wait_for_endpoint_failover( + EndpointType::AnthropicMessages + )); + } + #[tokio::test] #[serial_test::serial] async fn test_liveness_endpoint_stays_live_while_draining() { diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index b71533509fc1..c9b38ffa0c0e 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -106,6 +106,20 @@ pub struct DisaggregatedEndpoint { pub bootstrap_port: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +pub struct GmsRuntimeConfig { + #[serde(default)] + pub control_enabled: bool, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub nixl_ip: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub daemon_socket: Option, +} + +pub const GMS_RUNTIME_CONFIG_KEY: &str = "gms"; + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Validate)] #[validate(schema(function = "validate_model_runtime_config"))] /// Runtime-resolved metadata published by a worker after its engine starts. @@ -433,6 +447,31 @@ impl ModelRuntimeConfig { } } + pub fn set_gms_placement_enabled(&mut self) -> anyhow::Result<()> { + self.set_engine_specific(GMS_RUNTIME_CONFIG_KEY, { + let mut config = self.gms_runtime_config().unwrap_or_default(); + config.control_enabled = true; + config + }) + } + + pub fn set_gms_control_enabled(&mut self) -> anyhow::Result<()> { + self.set_gms_placement_enabled() + } + + pub fn set_gms_daemon_socket(&mut self, daemon_socket: String) -> anyhow::Result<()> { + let mut config = self.gms_runtime_config().unwrap_or_default(); + config.control_enabled = true; + config.daemon_socket = Some(daemon_socket); + self.set_engine_specific(GMS_RUNTIME_CONFIG_KEY, config) + } + + pub fn gms_runtime_config(&self) -> Option { + self.get_engine_specific(GMS_RUNTIME_CONFIG_KEY) + .ok() + .flatten() + } + pub fn effective_tokenizer_backend(&self) -> TokenizerBackend { self.tokenizer_backend .unwrap_or_else(TokenizerBackend::from_env_or_default) @@ -786,3 +825,42 @@ mod tests { } } } + +#[cfg(test)] +mod gms_runtime_config_tests { + use super::*; + + #[test] + fn gms_runtime_config_is_absent_by_default() { + let config = ModelRuntimeConfig::default(); + + assert!(config.gms_runtime_config().is_none()); + assert!(!config.runtime_data.contains_key(GMS_RUNTIME_CONFIG_KEY)); + } + + #[test] + fn set_gms_placement_enabled_is_optional_capability_metadata() { + let mut config = ModelRuntimeConfig::default(); + + config.set_gms_placement_enabled().unwrap(); + + let gms = config.gms_runtime_config().unwrap(); + assert!(gms.control_enabled); + assert_eq!(gms.nixl_ip, None); + assert_eq!(gms.daemon_socket, None); + } + + #[test] + fn set_gms_daemon_socket_is_preserved_when_enabling_placement() { + let mut config = ModelRuntimeConfig::default(); + + config + .set_gms_daemon_socket("/var/run/gms/kv.sock".to_string()) + .unwrap(); + config.set_gms_placement_enabled().unwrap(); + + let gms = config.gms_runtime_config().unwrap(); + assert!(gms.control_enabled); + assert_eq!(gms.daemon_socket.as_deref(), Some("/var/run/gms/kv.sock")); + } +} diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs index a0b51a474028..05cf30aaf330 100644 --- a/lib/llm/src/migration.rs +++ b/lib/llm/src/migration.rs @@ -3,10 +3,14 @@ use std::collections::BTreeMap; use std::error::Error as StdError; -use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::{Arc, Mutex}; +use std::time::Duration; use anyhow::{Error, Result}; use futures::{stream, stream::StreamExt}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::time::{Instant, sleep, timeout}; use crate::{ http::service::metrics::Metrics, @@ -55,6 +59,16 @@ impl HasTokenIds for LLMEngineOutput { } } +const DYN_MIGRATION_FAILOVER_WAIT_MS: &str = "DYN_MIGRATION_FAILOVER_WAIT_MS"; +const DYN_HTTP_MODEL_FAILOVER_WAIT_MS: &str = "DYN_HTTP_MODEL_FAILOVER_WAIT_MS"; +const DYN_MIGRATION_RETRY_CONCURRENCY: &str = "DYN_MIGRATION_RETRY_CONCURRENCY"; +const DYN_MIGRATION_FAILOVER_RETRY_CONCURRENCY: &str = "DYN_MIGRATION_FAILOVER_RETRY_CONCURRENCY"; +const DYN_MIGRATION_FIRST_CHUNK_TIMEOUT_MS: &str = "DYN_MIGRATION_FIRST_CHUNK_TIMEOUT_MS"; +const DYN_MIGRATION_RECENT_FAILOVER_WINDOW_MS: &str = "DYN_MIGRATION_RECENT_FAILOVER_WINDOW_MS"; +const DYN_MIGRATION_FAILOVER_RETRY_COOLDOWN_MS: &str = "DYN_MIGRATION_FAILOVER_RETRY_COOLDOWN_MS"; +const DYN_MIGRATION_FAILOVER_POLL_MS: &str = "DYN_MIGRATION_FAILOVER_POLL_MS"; +const DEFAULT_MIGRATION_FAILOVER_WAIT_POLL_MS: u64 = 50; + /// Check if an error chain indicates the request should be migrated. fn is_migratable(err: &(dyn StdError + 'static)) -> bool { const MIGRATABLE: &[ErrorType] = &[ @@ -64,7 +78,212 @@ fn is_migratable(err: &(dyn StdError + 'static)) -> bool { ErrorType::Backend(BackendError::EngineShutdown), ]; const NON_MIGRATABLE: &[ErrorType] = &[ErrorType::Cancelled, ErrorType::ResourceExhausted]; - error::match_error_chain(err, MIGRATABLE, NON_MIGRATABLE) + if error::match_error_chain(err, MIGRATABLE, NON_MIGRATABLE) { + return true; + } + + // Some request-plane routing failures still arrive as string-wrapped + // anyhow/pipeline errors before the typed DynamoError survives the full + // stack. Treat discovery misses as transient during migration so Bulwark + // failover can wait for the shadow endpoint to register. + let mut current: Option<&(dyn StdError + 'static)> = Some(err); + while let Some(e) = current { + let message = e.to_string(); + let message = message.to_ascii_lowercase(); + if message.contains("no instances found for endpoint") + || message.contains("no instances in selected device group for endpoint") + || message.contains("connection refused") + || message.contains("connection reset by peer") + || message.contains("connection aborted") + || message.contains("worker disconnected before response stream was established") + { + return true; + } + current = e.source(); + } + + false +} + +fn is_failover_backend_cancel_message(err: &(dyn StdError + 'static)) -> bool { + let mut current: Option<&(dyn StdError + 'static)> = Some(err); + while let Some(e) = current { + let message = e.to_string().to_ascii_lowercase(); + if message.contains("cancellederror") + || message.contains("backendcancelled") + || message.contains("event loop is closed") + { + return true; + } + current = e.source(); + } + + false +} + +#[cfg(test)] +fn is_failover_transient_cancel(err: &(dyn StdError + 'static)) -> bool { + is_failover_backend_cancel_message(err) + || error::match_error_chain(err, &[ErrorType::Cancelled], &[]) +} + +fn is_failover_cancel_replayable(err: &(dyn StdError + 'static)) -> bool { + if is_failover_backend_cancel_message(err) { + return !migration_failover_wait().is_zero() || recently_observed_failover(); + } + + // A plain typed Cancelled error can also mean a user/client cancellation. + // Only replay it after an independent failover signal has been observed. + error::match_error_chain(err, &[ErrorType::Cancelled], &[]) && recently_observed_failover() +} + +fn migration_failover_wait() -> Duration { + static WAIT: OnceLock = OnceLock::new(); + *WAIT.get_or_init(|| { + std::env::var(DYN_MIGRATION_FAILOVER_WAIT_MS) + .or_else(|_| std::env::var(DYN_HTTP_MODEL_FAILOVER_WAIT_MS)) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) + .unwrap_or_default() + }) +} + +fn migration_failover_retry_cooldown() -> Duration { + static COOLDOWN: OnceLock = OnceLock::new(); + *COOLDOWN.get_or_init(|| { + std::env::var(DYN_MIGRATION_FAILOVER_RETRY_COOLDOWN_MS) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) + .unwrap_or_default() + }) +} + +fn migration_retry_concurrency() -> usize { + static LIMIT: OnceLock = OnceLock::new(); + *LIMIT.get_or_init(|| { + std::env::var(DYN_MIGRATION_RETRY_CONCURRENCY) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or_default() + }) +} + +fn migration_retry_semaphore() -> Option> { + static SEMAPHORE: OnceLock>> = OnceLock::new(); + SEMAPHORE + .get_or_init(|| { + let limit = migration_retry_concurrency(); + if limit == 0 { + None + } else { + Some(Arc::new(Semaphore::new(limit))) + } + }) + .as_ref() + .cloned() +} + +fn migration_failover_retry_concurrency() -> usize { + static LIMIT: OnceLock = OnceLock::new(); + *LIMIT.get_or_init(|| { + std::env::var(DYN_MIGRATION_FAILOVER_RETRY_CONCURRENCY) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or_default() + }) +} + +fn migration_failover_retry_semaphore() -> Option> { + static SEMAPHORE: OnceLock>> = OnceLock::new(); + SEMAPHORE + .get_or_init(|| { + let limit = migration_failover_retry_concurrency(); + if limit == 0 { + None + } else { + Some(Arc::new(Semaphore::new(limit))) + } + }) + .as_ref() + .cloned() +} + +fn migration_failover_poll_interval() -> Duration { + static POLL: OnceLock = OnceLock::new(); + *POLL.get_or_init(|| { + std::env::var(DYN_MIGRATION_FAILOVER_POLL_MS) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) + .unwrap_or_else(|| Duration::from_millis(DEFAULT_MIGRATION_FAILOVER_WAIT_POLL_MS)) + }) +} + +fn migration_first_chunk_timeout() -> Duration { + static TIMEOUT: OnceLock = OnceLock::new(); + *TIMEOUT.get_or_init(|| { + std::env::var(DYN_MIGRATION_FIRST_CHUNK_TIMEOUT_MS) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) + .unwrap_or_default() + }) +} + +fn migration_recent_failover_window() -> Duration { + static WINDOW: OnceLock = OnceLock::new(); + *WINDOW.get_or_init(|| { + if let Some(window) = std::env::var(DYN_MIGRATION_RECENT_FAILOVER_WINDOW_MS) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) + { + return window; + } + + let first_chunk_timeout = migration_first_chunk_timeout(); + if first_chunk_timeout.is_zero() { + Duration::ZERO + } else { + std::cmp::max( + migration_failover_wait(), + first_chunk_timeout + .checked_mul(3) + .unwrap_or(first_chunk_timeout), + ) + } + }) +} + +fn observed_failover_at() -> &'static Mutex> { + static OBSERVED_AT: OnceLock>> = OnceLock::new(); + OBSERVED_AT.get_or_init(|| Mutex::new(None)) +} + +fn mark_failover_observed() { + if migration_first_chunk_timeout().is_zero() && migration_recent_failover_window().is_zero() { + return; + } + if let Ok(mut observed_at) = observed_failover_at().lock() { + *observed_at = Some(Instant::now()); + } +} + +pub(crate) fn observed_failover_elapsed() -> Option { + let Ok(observed_at) = observed_failover_at().lock() else { + return None; + }; + observed_at.as_ref().map(Instant::elapsed) +} + +fn recently_observed_failover() -> bool { + let window = migration_recent_failover_window(); + if window.is_zero() { + return false; + } + observed_failover_elapsed().is_some_and(|elapsed| elapsed <= window) } pub struct Migration { @@ -193,6 +412,34 @@ where /// Latest worker span pointer seen on the active stream; stamped as /// `migration_link` on the next retry. Populated by `track_response`. last_worker_link: Option, + /// Whether the current stream has produced backend output. Retries for + /// streams that failed before any output are treated as fresh request + /// retries instead of replayed in-flight streams. + current_stream_has_output: bool, + /// Optional process-wide retry permit. Held for the lifetime of a migrated + /// stream so failover does not stampede a newly active shadow. + migration_retry_permit: Option, + /// Optional failover-only replay permit. This is held until the replayed + /// stream finishes so Bulwark caps full replay pressure, not only the + /// first-token retry burst. + failover_retry_permit: Option, + /// Set only after an established stream fails before any backend output. + /// This permits one bounded replay even if the failed transport path has + /// already marked the original request context stopped. Initial requests + /// with a stopped context still fail fast. + allow_pre_output_stopped_retry: bool, + /// Set after a cancellation-like stream error while Bulwark failover + /// handling is active. Python backends can surface primary death as an + /// asyncio CancelledError even though replay is the right client-visible + /// behavior. + allow_failover_stopped_retry: bool, + /// Whether the current stream was created by the failover replay path. + /// Once such a stream is established, avoid repeatedly cancelling it on + /// first-token timeout; under a hot shadow that compounds queueing. + current_stream_from_failover_replay: bool, + /// Set after suppressing the timeout for the current replay stream so the + /// next poll waits on the backend stream normally. + current_stream_first_chunk_timeout_suppressed: bool, } impl RetryManager @@ -251,6 +498,13 @@ where model_name, metrics, last_worker_link: None, + current_stream_has_output: false, + migration_retry_permit: None, + failover_retry_permit: None, + allow_pre_output_stopped_retry: false, + allow_failover_stopped_retry: false, + current_stream_from_failover_replay: false, + current_stream_first_chunk_timeout_suppressed: false, }; slf.new_stream().await?; slf.exceed_max_seq_len(0); // disable migration if prompt len > max_seq_len @@ -266,63 +520,303 @@ where return Some(Annotated::from_error("next_stream is None")); } }; - if let Some(response) = response_stream.next().await { - // Check if this is a migratable error that should trigger stream recreation. - if let Some(err) = response.error.as_ref() - && is_migratable(err) - { - tracing::warn!(error = %err, "Stream disconnected, recreating stream"); - self.metrics.inc_migration_ongoing_request(&self.model_name); - if let Err(err) = self.new_stream().await { - tracing::warn!(error = ?err, "Cannot recreate stream"); - } else { + let response = if !self.current_stream_has_output + && self.retries_left > 0 + && !self.current_stream_first_chunk_timeout_suppressed + && !migration_first_chunk_timeout().is_zero() + { + match timeout(migration_first_chunk_timeout(), response_stream.next()).await { + Ok(response) => response, + Err(_) => { + if recently_observed_failover() { + if self.current_stream_from_failover_replay { + tracing::warn!( + timeout_ms = migration_first_chunk_timeout().as_millis(), + request_id = %self.context.id(), + "No first response from established failover replay stream; waiting without recreating" + ); + self.current_stream_first_chunk_timeout_suppressed = true; + continue; + } + tracing::warn!( + timeout_ms = migration_first_chunk_timeout().as_millis(), + request_id = %self.context.id(), + "No first response from stream during recent failover; recreating stream" + ); + self.metrics.inc_migration_new_request(&self.model_name); + self.allow_pre_output_stopped_retry = true; + self.allow_failover_stopped_retry = true; + if let Err(err) = self.new_stream().await { + tracing::warn!(error = ?err, "Cannot recreate timed-out stream"); + self.current_stream_has_output = true; + return Some(Annotated::from_error(err.to_string())); + } + } continue; } } + } else { + response_stream.next().await + }; + if let Some(response) = response { + // Check if this is a migratable error that should trigger stream recreation. + if let Some(err) = response.error.as_ref() { + let failover_cancel_replayable = is_failover_cancel_replayable(err); + if is_migratable(err) || failover_cancel_replayable { + if failover_cancel_replayable { + self.allow_failover_stopped_retry = true; + } + mark_failover_observed(); + if self.output_budget_exhausted() { + tracing::warn!( + request_id = %self.context.id(), + "Stream disconnected after output token budget was exhausted; completing without replay" + ); + self.retries_left = 0; + return None; + } + tracing::warn!(error = %err, "Stream disconnected, recreating stream"); + self.metrics.inc_migration_ongoing_request(&self.model_name); + if !self.current_stream_has_output { + self.allow_pre_output_stopped_retry = true; + } + if let Err(err) = self.new_stream().await { + tracing::warn!(error = ?err, "Cannot recreate stream"); + self.current_stream_has_output = true; + } else { + continue; + } + } + } self.track_response(&response); return Some(response); } + if !self.current_stream_has_output { + if self.retries_left > 0 { + tracing::warn!( + request_id = %self.context.id(), + "Stream ended before first response; recreating stream" + ); + self.metrics.inc_migration_new_request(&self.model_name); + self.allow_pre_output_stopped_retry = true; + if let Err(err) = self.new_stream().await { + tracing::warn!(error = ?err, "Cannot recreate empty stream"); + self.current_stream_has_output = true; + return Some(Annotated::from_error(err.to_string())); + } + continue; + } + tracing::warn!( + request_id = %self.context.id(), + "Stream ended before first response and migration budget is exhausted" + ); + self.current_stream_has_output = true; + return Some(Annotated::from_error( + "stream ended before first response".to_string(), + )); + } return None; } } + async fn ensure_migration_retry_permit(&mut self) -> Result<()> { + if self.migration_retry_permit.is_some() { + return Ok(()); + } + let Some(semaphore) = migration_retry_semaphore() else { + return Ok(()); + }; + let limit = migration_retry_concurrency(); + tracing::debug!( + limit, + request_id = %self.context.id(), + "Waiting for migration retry concurrency permit" + ); + let permit = semaphore + .acquire_owned() + .await + .map_err(|_| Error::msg("migration retry semaphore closed"))?; + tracing::debug!( + limit, + request_id = %self.context.id(), + "Acquired migration retry concurrency permit" + ); + self.migration_retry_permit = Some(permit); + Ok(()) + } + + async fn ensure_failover_retry_permit(&mut self) -> Result<()> { + if self.failover_retry_permit.is_some() { + return Ok(()); + } + let Some(semaphore) = migration_failover_retry_semaphore() else { + return Ok(()); + }; + let limit = migration_failover_retry_concurrency(); + tracing::debug!( + limit, + request_id = %self.context.id(), + "Waiting for failover replay concurrency permit" + ); + let permit = semaphore + .acquire_owned() + .await + .map_err(|_| Error::msg("failover replay semaphore closed"))?; + tracing::debug!( + limit, + request_id = %self.context.id(), + "Acquired failover replay concurrency permit" + ); + self.failover_retry_permit = Some(permit); + Ok(()) + } + + fn failover_retry_controls_active(&self) -> bool { + self.allow_failover_stopped_retry + || (self.next_stream.is_some() + && !self.current_stream_has_output + && recently_observed_failover()) + } + async fn new_stream(&mut self) -> Result<()> { let mut response_stream: Option>>> = None; + let mut stream_from_failover_replay = false; while self.retries_left > 0 { self.retries_left -= 1; // Once any chunks have arrived from a previous attempt, stamp - // that worker's span as `migration_link` so the next worker's + // that worker's span as migration_link so the next worker's // span renders an OTel Link back to it. Guarded so the initial - // attempt doesn't clobber a `migration_link` set upstream. + // attempt doesn't clobber a migration_link set upstream. if let Some(link) = self.last_worker_link.as_ref() { self.request.migration_link = Some(link.clone()); } - let mut request = Context::with_id_and_metadata( - self.request.clone(), - self.context.id().to_string(), - self.metadata.clone(), - ); - if let Some(session_affinity) = self.session_affinity.as_ref() { - request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_affinity.clone()); - } - self.context.link_child(request.context()); - if self.context.is_stopped() || self.context.is_killed() { - tracing::debug!("Abort creating new stream after context is stopped or killed"); - return Err(DynamoError::builder() - .error_type(ErrorType::Cancelled) - .message(format!( - "Context id {} is stopped or killed", - self.context.id() - )) - .build() - .into()); + + let wait = migration_failover_wait(); + let deadline = if wait.is_zero() { + None + } else { + Some(Instant::now() + wait) + }; + let mut counted_attempt = false; + let mut failover_polls = 0_u32; + let mut failover_cooldown_applied = false; + + loop { + let retry_stream = self.next_stream.is_some(); + let failover_retry = self.failover_retry_controls_active(); + if retry_stream { + self.ensure_migration_retry_permit().await?; + } + if failover_retry { + stream_from_failover_replay = true; + self.ensure_failover_retry_permit().await?; + + let cooldown = migration_failover_retry_cooldown(); + if !failover_cooldown_applied && !cooldown.is_zero() { + failover_cooldown_applied = true; + tracing::warn!( + cooldown_ms = cooldown.as_millis(), + request_id = %self.context.id(), + "Delaying failover replay before recreating stream" + ); + sleep(cooldown).await; + } + } + + let mut request = Context::with_id_and_metadata( + self.request.clone(), + self.context.id().to_string(), + self.metadata.clone(), + ); + if let Some(session_affinity) = self.session_affinity.as_ref() { + request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_affinity.clone()); + } + self.context.link_child(request.context()); + let context_stopped = self.context.is_stopped(); + let context_killed = self.context.is_killed(); + if context_killed { + tracing::debug!("Abort creating new stream after context is killed"); + return Err(DynamoError::builder() + .error_type(ErrorType::Cancelled) + .message(format!("Context id {} is killed", self.context.id())) + .build() + .into()); + } + if context_stopped && !context_killed { + let allow_stopped_retry = self.allow_pre_output_stopped_retry + && self.next_stream.is_some() + && !self.current_stream_has_output; + let allow_failover_stopped_retry = self.allow_failover_stopped_retry + && self.next_stream.is_some() + && (!migration_failover_wait().is_zero() || recently_observed_failover()); + if allow_stopped_retry || allow_failover_stopped_retry { + tracing::warn!( + request_id = %self.context.id(), + pre_output = allow_stopped_retry, + failover_cancel = allow_failover_stopped_retry, + "Creating failover replay stream after context stop" + ); + } else { + tracing::debug!("Abort creating new stream after context is stopped"); + return Err(DynamoError::builder() + .error_type(ErrorType::Cancelled) + .message(format!("Context id {} is stopped", self.context.id())) + .build() + .into()); + } + } + + response_stream = Some(self.next_generate.generate(request).await); + let Some(Err(err)) = response_stream.as_ref() else { + break; + }; + let failover_cancel_replayable = is_failover_cancel_replayable(err.as_ref()); + if !(is_migratable(err.as_ref()) || failover_cancel_replayable) { + break; + } + if failover_cancel_replayable { + self.allow_failover_stopped_retry = true; + } + stream_from_failover_replay = true; + mark_failover_observed(); + + if !counted_attempt { + self.metrics.inc_migration_new_request(&self.model_name); + counted_attempt = true; + } + + if let Some(deadline) = deadline { + let now = Instant::now(); + if now < deadline { + failover_polls += 1; + tracing::warn!( + polls = failover_polls, + wait_ms = wait.as_millis(), + "Creating new stream failed; waiting for failover endpoint: {}", + err + ); + sleep(std::cmp::min( + migration_failover_poll_interval(), + deadline - now, + )) + .await; + continue; + } + tracing::warn!( + polls = failover_polls, + wait_ms = wait.as_millis(), + "Creating new stream failed after failover wait; retrying if budget remains: {}", + err + ); + } else { + tracing::warn!(error = %err, "Creating new stream, retrying"); + } + break; } - response_stream = Some(self.next_generate.generate(request).await); + if let Some(err) = response_stream.as_ref().unwrap().as_ref().err() - && is_migratable(err.as_ref()) + && (is_migratable(err.as_ref()) || is_failover_cancel_replayable(err.as_ref())) { - tracing::warn!(error = %err, "Creating new stream, retrying"); - self.metrics.inc_migration_new_request(&self.model_name); continue; } break; @@ -330,6 +824,11 @@ where match response_stream { Some(Ok(next_stream)) => { self.next_stream = Some(next_stream); + self.current_stream_has_output = false; + self.current_stream_from_failover_replay = stream_from_failover_replay; + self.current_stream_first_chunk_timeout_suppressed = false; + self.allow_pre_output_stopped_retry = false; + self.allow_failover_stopped_retry = false; Ok(()) } Some(Err(err)) => Err(err), // should propagate original error if any @@ -340,13 +839,14 @@ where } fn track_response(&mut self, response: &Annotated) { - if self.retries_left == 0 { - return; - } let llm_engine_output = match response.data.as_ref() { Some(output) => output, None => return, }; + self.current_stream_has_output = true; + if self.retries_left == 0 { + return; + } // Capture the worker's engine.generate span pointer so a future // migration retry can render an OTel Link back to it. The adapter // stamps this on the first non-empty chunk; subsequent chunks may @@ -367,6 +867,10 @@ where } } + fn output_budget_exhausted(&self) -> bool { + matches!(self.request.stop_conditions.max_tokens, Some(0)) + } + /// Returns `true` if the tracked request token length plus `new_output_len` /// exceeds the configured max_seq_len, in which case migration is disabled. fn exceed_max_seq_len(&mut self, new_output_len: u32) -> bool { @@ -462,6 +966,8 @@ mod tests { MidStreamFailAlwaysStreamError { fail_after: usize, }, + /// Opens an empty stream on the first call, then succeeds on retry. + EmptyThenSuccess, /// Always fails with NoResponders error (same as FailThenSuccess first call) AlwaysFail, } @@ -697,6 +1203,20 @@ mod tests { )) } } + MockBehavior::EmptyThenSuccess => { + if call_num == 0 { + let (_tx, rx) = mpsc::channel(1); + let stream = tokio_stream::wrappers::ReceiverStream::new(rx); + let ctx = Arc::new(Controller::new(self.context_id.clone())); + Ok(dynamo_runtime::pipeline::ResponseStream::new( + Box::pin(stream), + ctx, + )) + } else { + self.send_responses(responses_already_generated, self.num_responses) + .await + } + } MockBehavior::AlwaysFail => { // Always fail with NoResponders error (same as FailThenSuccess first call) Err(anyhow::anyhow!( @@ -737,6 +1257,54 @@ mod tests { } } + #[test] + fn test_string_wrapped_discovery_miss_is_migratable() { + let err = anyhow::anyhow!( + "no instances found for endpoint mkhadkevich-dev-gms-bulwark-sglang/backend/generate" + ); + + assert!(is_migratable(err.as_ref())); + } + + #[test] + fn test_string_wrapped_request_plane_connect_errors_are_migratable() { + for message in [ + "Connection refused (os error 111)", + "error trying to connect: tcp connect error: Connection refused (os error 111)", + "connection reset by peer", + "connection aborted", + "worker disconnected before response stream was established", + ] { + let err = anyhow::anyhow!(message); + assert!(is_migratable(err.as_ref()), "{message}"); + } + } + + #[test] + fn test_python_cancel_errors_are_failover_transient_cancels() { + for message in [ + "CancelledError:", + "BackendCancelled: primary stopped", + "RuntimeError: Event loop is closed", + ] { + let err = anyhow::anyhow!(message); + assert!( + is_failover_backend_cancel_message(err.as_ref()), + "{message}" + ); + assert!(is_failover_transient_cancel(err.as_ref()), "{message}"); + assert!(!is_migratable(err.as_ref()), "{message}"); + } + + let typed = DynamoError::builder() + .error_type(ErrorType::Cancelled) + .message("Context was stopped") + .build(); + assert!(!is_failover_backend_cancel_message(&typed)); + assert!(is_failover_transient_cancel(&typed)); + assert!(!is_migratable(&typed)); + } + /// Test case 1: No migration needed /// Tests the normal case where the RetryManager successfully processes all responses /// from a single stream without any failures or need for retries/migration. @@ -865,6 +1433,103 @@ mod tests { assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); } + #[tokio::test] + async fn test_retry_manager_empty_first_stream_migration() { + dynamo_runtime::logging::init(); + let context_id = uuid::Uuid::new_v4().to_string(); + let request = create_mock_request(3); + let mock_engine = Arc::new(MockEngine::new( + MockBehavior::EmptyThenSuccess, + 3, + 100, + context_id.clone(), + )); + let next_generate: ServerStreamingEngine> = + mock_engine; + + let ctx = Arc::new(Controller::new(context_id.clone())); + let metrics = Arc::new(Metrics::new()); + let mut retry_manager = RetryManager::build( + ctx, + BTreeMap::new(), + request, + next_generate, + 1, + None, + Arc::new(TEST_MODEL.to_string()), + metrics.clone(), + None, + ) + .await + .expect("Failed to build RetryManager"); + + let mut responses = Vec::new(); + while let Some(response) = retry_manager.next().await { + responses.push(response); + } + + assert_eq!(responses.len(), 3); + for (i, response) in responses.iter().enumerate() { + assert!(response.err().is_none()); + if let Some(output) = &response.data { + assert_eq!(output.token_ids, vec![101 + i as u32]); + } + } + assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 1); + assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); + } + + #[tokio::test] + async fn test_retry_manager_empty_stream_retries_after_context_stop() { + dynamo_runtime::logging::init(); + let context_id = uuid::Uuid::new_v4().to_string(); + let request = create_mock_request(2); + let mock_engine = Arc::new(MockEngine::new( + MockBehavior::EmptyThenSuccess, + 2, + 100, + context_id.clone(), + )); + let next_generate: ServerStreamingEngine> = + mock_engine; + + let ctx = Arc::new(Controller::new(context_id.clone())); + let metrics = Arc::new(Metrics::new()); + let mut retry_manager = RetryManager::build( + ctx.clone(), + BTreeMap::new(), + request, + next_generate, + 1, + None, + Arc::new(TEST_MODEL.to_string()), + metrics.clone(), + None, + ) + .await + .expect("Failed to build RetryManager"); + + // The failover transport path can mark the parent context stopped as + // the failed first attempt closes. Because no output was published, + // the retry manager should still be allowed to replay once. + ctx.stop_generating(); + + let mut responses = Vec::new(); + while let Some(response) = retry_manager.next().await { + responses.push(response); + } + + assert_eq!(responses.len(), 2); + for (i, response) in responses.iter().enumerate() { + assert!(response.err().is_none()); + if let Some(output) = &response.data { + assert_eq!(output.token_ids, vec![101 + i as u32]); + } + } + assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 1); + assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); + } + /// Test case 3: Ongoing request migration /// Tests the scenario where a worker fails mid-stream during an ongoing request. /// This simulates a connection being lost after partial response delivery, requiring @@ -922,6 +1587,55 @@ mod tests { assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 1); } + #[tokio::test] + async fn test_retry_manager_does_not_replay_when_output_budget_exhausted() { + dynamo_runtime::logging::init(); + + let context_id = uuid::Uuid::new_v4().to_string(); + let request = create_mock_request(10); + let mock_engine = Arc::new(MockEngine::new( + MockBehavior::MidStreamFail { fail_after: 10 }, + 10, + 100, + context_id.clone(), + )); + let call_count = mock_engine.call_count.clone(); + let next_generate: ServerStreamingEngine> = + mock_engine; + + let ctx = Arc::new(Controller::new(context_id.clone())); + let metrics = Arc::new(Metrics::new()); + let mut retry_manager = RetryManager::build( + ctx, + BTreeMap::new(), + request, + next_generate, + 3, + None, + Arc::new(TEST_MODEL.to_string()), + metrics.clone(), + None, + ) + .await + .expect("Failed to build RetryManager"); + + let mut responses = Vec::new(); + while let Some(response) = retry_manager.next().await { + responses.push(response); + } + + assert_eq!(responses.len(), 10); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 0); + assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); + for (i, response) in responses.iter().enumerate() { + assert!(response.err().is_none()); + if let Some(output) = &response.data { + assert_eq!(output.token_ids, vec![101 + i as u32]); + } + } + } + /// Test case 4: New request migration - indefinite failure /// Tests the scenario where a worker becomes unreachable for new requests indefinitely. /// The RetryManager should exhaust all retries and return the original error from the first attempt. @@ -1130,7 +1844,7 @@ mod tests { assert!( error .to_string() - .contains(&format!("Context id {} is stopped or killed", context_id)) + .contains(&format!("Context id {} is stopped", context_id)) ); // Verify the error is a typed DynamoError with Cancelled type let dynamo_err = error diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index 0afad99a50b6..0b7f0b2f6583 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/lib/llm/src/protocols/common/preprocessor.rs @@ -108,6 +108,84 @@ pub struct TraceLink { pub span_id: String, } +/// One contiguous memory region inside a GMS block descriptor. +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct GmsMemoryRegion { + /// NIXL-registered remote pointer (interpreted by the source + /// daemon's NIXL agent). + pub remote_ptr: u64, + /// Region size in bytes. + pub size: u64, + /// Source tier: "hbm", "host", or "storage". + pub tier: String, + /// Source model layer for host-tier multi-range descriptors. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub layer: Option, + /// Byte offset inside the source layer allocation. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub offset: Option, +} + +/// One block's location inside a source GMS daemon's NIXL-registered +/// memory. Returned by the daemon's `get_bootstrap_info` RPC; the +/// receiver engine's connector uses it to issue `nixl_agent.read()` +/// directly against the source daemon. +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct GmsBlockDescriptor { + /// Backward-compatible first/only region pointer. New consumers + /// should prefer `ranges` when it is non-empty. + pub remote_ptr: u64, + /// Total descriptor size in bytes. For multi-range host blocks, + /// this is the sum of all ranges. + pub size: u64, + /// Source tier: "hbm", "host", or "storage". Engine connectors + /// may choose to fall back to `fetch_remote` if the source tier + /// can't be NIXL-read directly (e.g., disk without GDS plugin). + pub tier: String, + /// Vectored memory regions aligned in source copy order. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub ranges: Vec, + /// Source block generation when the engine/daemon can provide it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub generation: Option, + /// True when the source daemon considers the bytes stable. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sealed: Option, +} + +/// In-band metadata for GMS-managed KV transfer. Populated by the +/// router (or routing-time orchestration) when a request's prefix +/// overlap is served from a remote GMS daemon. Carries enough info +/// for the receiving engine's connector to issue an engine-direct +/// `nixl_agent.read()` without the router opening any GMS daemon +/// sockets. Worker-local daemon sockets remain private to workers. +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct GmsPlacementInfo { + /// Source daemon's NIXL agent name (peer identifier). + pub source_nixl_agent_name: String, + + /// Hex-encoded NIXL agent metadata blob for `add_remote_agent`. + /// Sized ~few-KB; first-time use only — engine caches per peer. + pub source_nixl_agent_metadata_hex: String, + + /// Source daemon NIXL listen address when the source worker exposes it. + /// The hot engine-direct path primarily uses the metadata blob above. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source_nixl_ip: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source_nixl_listen_port: Option, + + /// Per-hash descriptors aligned with `hashes`. `None` entries + /// mean the source daemon doesn't have that hash; the engine + /// should skip it and recompute that block locally. + pub descriptors: Vec>, + + /// Hex-encoded content hashes the engine is asking about (length + /// matches `descriptors`). + pub hashes: Vec, +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct PrefillResult { /// Disaggregated execution parameters. Engine-owned; the framework @@ -254,6 +332,16 @@ pub struct PreprocessedRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub bootstrap_info: Option, + /// GMS placement info for disaggregated serving when the source + /// daemon holds KV blocks that can be NIXL-read directly. Carried + /// in-band like [`BootstrapInfo`]; the receiving engine's connector + /// consults it and either issues an engine-direct read (default) + /// or asks its local daemon to `fetch_remote` (backup). The router + /// itself is GMS-agnostic — it just forwards this field. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub gms_placement: Option, + /// Additional arguments for extensibility #[builder(default)] #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/lib/runtime/src/config/environment_names.rs b/lib/runtime/src/config/environment_names.rs index 48425f08ce1b..3df34cc901cf 100644 --- a/lib/runtime/src/config/environment_names.rs +++ b/lib/runtime/src/config/environment_names.rs @@ -146,6 +146,8 @@ pub mod nats { /// NATS request/reply timeout in seconds. Unset = async-nats default (10 s). pub const DYN_NATS_REQUEST_TIMEOUT_SECS: &str = "DYN_NATS_REQUEST_TIMEOUT_SECS"; + /// Request-plane ack timeout in milliseconds for NATS requests. + pub const DYN_NATS_REQUEST_TIMEOUT_MS: &str = "DYN_NATS_REQUEST_TIMEOUT_MS"; /// NATS authentication environment variables (checked in priority order) pub mod auth { @@ -600,6 +602,15 @@ pub mod discovery { /// Kube discovery mode: "pod" (default) or "container" (each container registers independently) pub const DYN_KUBE_DISCOVERY_MODE: &str = "DYN_KUBE_DISCOVERY_MODE"; + + /// Kube discovery daemon debounce window in milliseconds (default: 500). + pub const DYN_KUBE_DISCOVERY_DEBOUNCE_MS: &str = "DYN_KUBE_DISCOVERY_DEBOUNCE_MS"; + + /// Explicit logical discovery instance id shared by a failover cohort. + pub const DYN_DISCOVERY_LOGICAL_INSTANCE_ID: &str = "DYN_DISCOVERY_LOGICAL_INSTANCE_ID"; + + /// Stable logical discovery key hashed into an instance id shared by a failover cohort. + pub const DYN_DISCOVERY_LOGICAL_INSTANCE_KEY: &str = "DYN_DISCOVERY_LOGICAL_INSTANCE_KEY"; } /// CUDA and GPU environment variables @@ -688,6 +699,7 @@ mod tests { // NATS nats::NATS_SERVER, nats::DYN_NATS_REQUEST_TIMEOUT_SECS, + nats::DYN_NATS_REQUEST_TIMEOUT_MS, nats::auth::NATS_AUTH_USERNAME, nats::auth::NATS_AUTH_PASSWORD, nats::auth::NATS_AUTH_TOKEN, diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index 611f2c2862c1..862a05a44ce0 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -17,7 +17,7 @@ use crate::CancellationToken; use crate::discovery::{ ClaimCloseOutcome, ClaimOutcome, ClaimPayloadFuture, Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryMetadata, DiscoveryQuery, DiscoverySpec, - DiscoveryStream, MetadataSnapshot, + DiscoveryStream, MetadataSnapshot, resolve_logical_instance_id, }; use anyhow::Result; use async_trait::async_trait; @@ -49,15 +49,17 @@ impl KubeDiscoveryClient { cancel_token: CancellationToken, ) -> Result { let pod_info = PodInfo::from_env()?; - let instance_id = pod_info.target.instance_id(); + let physical_instance_id = pod_info.target.instance_id(); + let instance_id = resolve_logical_instance_id(physical_instance_id)?; let cr_name = pod_info.target.cr_name(); tracing::info!( - "Initializing KubeDiscoveryClient: mode={:?}, target={:?}, cr_name={}, instance_id={:x}, namespace={}, pod_uid={}", + "Initializing KubeDiscoveryClient: mode={:?}, target={:?}, cr_name={}, instance_id={:x}, physical_instance_id={:x}, namespace={}, pod_uid={}", pod_info.mode, pod_info.target, cr_name, instance_id, + physical_instance_id, pod_info.pod_namespace, pod_info.pod_uid ); @@ -378,8 +380,10 @@ impl Discovery for KubeDiscoveryClient { } } - // Track known instances by their unique ID - let mut known: HashSet = initial.into_keys().collect(); + // Track known instances by ID and value so transport updates for a stable + // logical worker id are emitted as an upsert instead of being coalesced away. + let mut known: std::collections::HashMap = + initial; loop { tracing::trace!( @@ -428,9 +432,11 @@ impl Discovery for KubeDiscoveryClient { "Watch received snapshot update" ); - // Compute diff using keys + // Compute diff by ID, then treat value changes as upserts. This is + // important for failover cohorts where a logical worker id stays stable + // while the active pod's transport address changes. let current_keys: HashSet<&DiscoveryInstanceId> = current.keys().collect(); - let known_keys: HashSet<&DiscoveryInstanceId> = known.iter().collect(); + let known_keys: HashSet<&DiscoveryInstanceId> = known.keys().collect(); let added: Vec<&DiscoveryInstanceId> = current_keys.difference(&known_keys).copied().collect(); @@ -440,8 +446,14 @@ impl Discovery for KubeDiscoveryClient { .map(|&id| id.clone()) .collect(); + let updated: Vec<&DiscoveryInstanceId> = current_keys + .intersection(&known_keys) + .copied() + .filter(|id| current.get(*id) != known.get(*id)) + .collect(); + // Log diff results (even if empty, for debugging) - if added.is_empty() && removed.is_empty() { + if added.is_empty() && removed.is_empty() && updated.is_empty() { tracing::debug!( stream_id = %stream_id, seq = snapshot.sequence, @@ -453,13 +465,16 @@ impl Discovery for KubeDiscoveryClient { seq = snapshot.sequence, added = added.len(), removed = removed.len(), + updated = updated.len(), total = current.len(), "Watch detected changes" ); } - // Emit Added events - for id in added { + // Emit Added events for both new and changed instances. Consumers treat + // Added as an upsert, which preserves KV-router state for the same worker + // id while refreshing the request-plane transport. + for id in added.into_iter().chain(updated.into_iter()) { if let Some(instance) = current.get(id) { tracing::info!( stream_id = %stream_id, @@ -492,8 +507,8 @@ impl Discovery for KubeDiscoveryClient { } } - // Update known set - known = current.into_keys().collect(); + // Update known map + known = current; } Err(_) => { tracing::info!( diff --git a/lib/runtime/src/discovery/kube/daemon.rs b/lib/runtime/src/discovery/kube/daemon.rs index f60ee206ea91..eb97d50002b3 100644 --- a/lib/runtime/src/discovery/kube/daemon.rs +++ b/lib/runtime/src/discovery/kube/daemon.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::CancellationToken; +use crate::config::environment_names::discovery; use crate::discovery::{DiscoveryMetadata, MetadataSnapshot}; use anyhow::Result; use futures::StreamExt; @@ -19,7 +20,25 @@ use tokio::time::{Duration, timeout}; use super::crd::DynamoWorkerMetadata; use super::utils::{KubeDiscoveryMode, PodInfo, extract_endpoint_info, extract_ready_containers}; -const DEBOUNCE_DURATION: Duration = Duration::from_millis(500); +const DEFAULT_DEBOUNCE_DURATION: Duration = Duration::from_millis(500); + +fn discovery_debounce_duration() -> Duration { + match std::env::var(discovery::DYN_KUBE_DISCOVERY_DEBOUNCE_MS) { + Ok(raw) => match raw.parse::() { + Ok(ms) => Duration::from_millis(ms), + Err(err) => { + tracing::warn!( + value = %raw, + %err, + default_ms = DEFAULT_DEBOUNCE_DURATION.as_millis(), + "Invalid DYN_KUBE_DISCOVERY_DEBOUNCE_MS; using default" + ); + DEFAULT_DEBOUNCE_DURATION + } + }, + Err(_) => DEFAULT_DEBOUNCE_DURATION, + } +} #[derive(Clone)] struct CachedCrMetadata { @@ -197,6 +216,11 @@ impl DiscoveryDaemon { tokio::spawn(cr_reflector_stream); // Event-driven loop with debouncing + let debounce_duration = discovery_debounce_duration(); + tracing::info!( + debounce_ms = debounce_duration.as_millis(), + "Discovery daemon debounce configured" + ); let mut sequence = 0u64; let mut prev_snapshot = MetadataSnapshot::empty(); // Keeps transient invalid CR updates from looking like removals. @@ -205,7 +229,7 @@ impl DiscoveryDaemon { loop { tokio::select! { _ = notify.notified() => { - tokio::time::sleep(DEBOUNCE_DURATION).await; + tokio::time::sleep(debounce_duration).await; let _ = timeout(Duration::ZERO, notify.notified()).await; tracing::trace!("Debounce window elapsed, processing snapshot"); diff --git a/lib/runtime/src/discovery/kube/utils.rs b/lib/runtime/src/discovery/kube/utils.rs index 1d8218abc7fb..a54ac7a248a6 100644 --- a/lib/runtime/src/discovery/kube/utils.rs +++ b/lib/runtime/src/discovery/kube/utils.rs @@ -4,14 +4,12 @@ use anyhow::Result; use k8s_openapi::api::core::v1::Pod; use k8s_openapi::api::discovery::v1::EndpointSlice; -use std::collections::hash_map::DefaultHasher; use std::fs; -use std::hash::{Hash, Hasher}; use std::path::Path; use crate::config::environment_names::discovery; +use crate::discovery::{DISCOVERY_INSTANCE_ID_MASK, hash_logical_instance_key}; -const INSTANCE_ID_MASK: u64 = 0x001F_FFFF_FFFF_FFFFu64; const MAIN_CONTAINER_NAME: &str = "main"; /// Kube discovery mode. @@ -60,9 +58,7 @@ impl KubeDiscoveryTarget { /// Deterministic instance ID derived from cr_name. pub fn instance_id(&self) -> u64 { - let mut hasher = DefaultHasher::new(); - self.cr_name().hash(&mut hasher); - hasher.finish() & INSTANCE_ID_MASK + hash_logical_instance_key(&self.cr_name()) & DISCOVERY_INSTANCE_ID_MASK } pub fn pod_name(&self) -> &str { @@ -76,9 +72,7 @@ impl KubeDiscoveryTarget { /// /// Used by C bindings (EPP) for pod-level worker ID mapping. pub fn hash_pod_name(pod_name: &str) -> u64 { - let mut hasher = DefaultHasher::new(); - pod_name.hash(&mut hasher); - hasher.finish() & INSTANCE_ID_MASK + hash_logical_instance_key(pod_name) & DISCOVERY_INSTANCE_ID_MASK } /// Extract (instance_id, pod_name) tuples from an EndpointSlice for ready endpoints. diff --git a/lib/runtime/src/discovery/kv_store.rs b/lib/runtime/src/discovery/kv_store.rs index 528b20ff4c9f..6ce401f3f3c0 100644 --- a/lib/runtime/src/discovery/kv_store.rs +++ b/lib/runtime/src/discovery/kv_store.rs @@ -16,6 +16,7 @@ use super::{ ClaimCloseOutcome, ClaimEvent, ClaimOutcome, ClaimPayload, ClaimPayloadFuture, Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoverySpec, DiscoveryStream, EndpointInstanceId, EventChannelInstanceId, ModelCardInstanceId, + resolve_logical_instance_id, }; use crate::storage::kv; @@ -31,6 +32,7 @@ pub struct KVStoreDiscovery { store: Arc, cancel_token: CancellationToken, claims: ClaimState, + instance_id: u64, } /// Process-local invalidation relay for the shared claims bucket. @@ -126,12 +128,15 @@ impl Drop for ClaimWatcherActiveGuard { } impl KVStoreDiscovery { - pub fn new(store: kv::Manager, cancel_token: CancellationToken) -> Self { - Self { + pub fn new(store: kv::Manager, cancel_token: CancellationToken) -> Result { + let physical_instance_id = store.connection_id(); + let instance_id = resolve_logical_instance_id(physical_instance_id)?; + Ok(Self { store: Arc::new(store), cancel_token, claims: ClaimState::new(), - } + instance_id, + }) } async fn ensure_claim_watcher(&self) -> Result<()> { @@ -412,7 +417,7 @@ impl KVStoreDiscovery { #[async_trait] impl Discovery for KVStoreDiscovery { fn instance_id(&self) -> u64 { - self.store.connection_id() + self.instance_id } async fn register_internal(&self, spec: DiscoverySpec) -> Result { @@ -1160,7 +1165,7 @@ mod tests { async fn test_kv_store_discovery_register_endpoint() { let store = kv::Manager::memory(); let cancel_token = CancellationToken::new(); - let client = KVStoreDiscovery::new(store, cancel_token); + let client = KVStoreDiscovery::new(store, cancel_token).unwrap(); let spec = DiscoverySpec::Endpoint { namespace: "test".to_string(), @@ -1186,7 +1191,7 @@ mod tests { async fn test_kv_store_discovery_list() { let store = kv::Manager::memory(); let cancel_token = CancellationToken::new(); - let client = KVStoreDiscovery::new(store, cancel_token); + let client = KVStoreDiscovery::new(store, cancel_token).unwrap(); // Register multiple endpoints let spec1 = DiscoverySpec::Endpoint { @@ -1244,7 +1249,7 @@ mod tests { async fn test_kv_store_discovery_watch() { let store = kv::Manager::memory(); let cancel_token = CancellationToken::new(); - let client = Arc::new(KVStoreDiscovery::new(store, cancel_token.clone())); + let client = Arc::new(KVStoreDiscovery::new(store, cancel_token.clone()).unwrap()); // Start watching before registering let mut stream = client diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs index 91ec2b76c911..03b83a9fb6eb 100644 --- a/lib/runtime/src/discovery/mod.rs +++ b/lib/runtime/src/discovery/mod.rs @@ -5,7 +5,9 @@ use anyhow::{Context, Result}; use async_trait::async_trait; use futures::Stream; use serde::{Deserialize, Serialize}; +use std::collections::hash_map::DefaultHasher; use std::future::Future; +use std::hash::{Hash, Hasher}; use std::pin::Pin; use tokio::sync::broadcast; use tokio_util::sync::CancellationToken; @@ -47,6 +49,91 @@ pub enum ClaimEvent { Reset, } +pub(crate) const DISCOVERY_INSTANCE_ID_MASK: u64 = 0x001F_FFFF_FFFF_FFFFu64; + +pub(crate) fn hash_logical_instance_key(key: &str) -> u64 { + let mut hasher = DefaultHasher::new(); + key.hash(&mut hasher); + hasher.finish() & DISCOVERY_INSTANCE_ID_MASK +} + +fn parse_logical_instance_id(raw: &str) -> Result { + let value = raw.trim(); + if value.is_empty() { + anyhow::bail!("DYN_DISCOVERY_LOGICAL_INSTANCE_ID must not be empty"); + } + + if let Some(hex) = value + .strip_prefix("0x") + .or_else(|| value.strip_prefix("0X")) + { + return u64::from_str_radix(hex, 16) + .with_context(|| format!("invalid hex DYN_DISCOVERY_LOGICAL_INSTANCE_ID `{value}`")); + } + + value + .parse::() + .with_context(|| format!("invalid DYN_DISCOVERY_LOGICAL_INSTANCE_ID `{value}`")) +} + +fn expand_env_placeholders(value: &str) -> Result { + let mut out = String::with_capacity(value.len()); + let bytes = value.as_bytes(); + let mut i = 0; + + while i < bytes.len() { + if bytes[i] == b'$' && i + 1 < bytes.len() && bytes[i + 1] == b'(' { + let start = i + 2; + let Some(rel_end) = value[start..].find(')') else { + anyhow::bail!( + "unterminated env placeholder in DYN_DISCOVERY_LOGICAL_INSTANCE_KEY: `{value}`" + ); + }; + let end = start + rel_end; + let name = &value[start..end]; + if name.is_empty() || !name.chars().all(|c| c == '_' || c.is_ascii_alphanumeric()) { + anyhow::bail!( + "invalid env placeholder `$({name})` in DYN_DISCOVERY_LOGICAL_INSTANCE_KEY" + ); + } + let replacement = std::env::var(name).with_context(|| { + format!("DYN_DISCOVERY_LOGICAL_INSTANCE_KEY references unset env var `{name}`") + })?; + out.push_str(&replacement); + i = end + 1; + } else { + let ch = value[i..] + .chars() + .next() + .expect("valid utf-8 while expanding env placeholders"); + out.push(ch); + i += ch.len_utf8(); + } + } + + Ok(out) +} + +pub(crate) fn resolve_logical_instance_id(default_id: u64) -> Result { + use crate::config::environment_names::discovery::{ + DYN_DISCOVERY_LOGICAL_INSTANCE_ID, DYN_DISCOVERY_LOGICAL_INSTANCE_KEY, + }; + + if let Ok(raw_id) = std::env::var(DYN_DISCOVERY_LOGICAL_INSTANCE_ID) { + return parse_logical_instance_id(&raw_id); + } + + if let Ok(raw_key) = std::env::var(DYN_DISCOVERY_LOGICAL_INSTANCE_KEY) { + let expanded = expand_env_placeholders(&raw_key)?; + if expanded.trim().is_empty() { + anyhow::bail!("DYN_DISCOVERY_LOGICAL_INSTANCE_KEY must not expand to an empty string"); + } + return Ok(hash_logical_instance_key(&expanded)); + } + + Ok(default_id) +} + /// Transport kind for event plane - used for configuration and env var selection. /// /// This enum represents the *type* of transport without connection details. @@ -901,3 +988,64 @@ pub trait Discovery: Send + Sync { /// waiting for TTL expiry. Default is a no-op for backends that don't need cleanup. fn shutdown(&self) {} } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + static ENV_LOCK: Mutex<()> = Mutex::new(()); + + #[test] + fn parse_logical_instance_id_accepts_decimal_and_hex() { + assert_eq!(parse_logical_instance_id("123").unwrap(), 123); + assert_eq!(parse_logical_instance_id("0x7b").unwrap(), 123); + assert!(parse_logical_instance_id("").is_err()); + assert!(parse_logical_instance_id("not-a-number").is_err()); + } + + fn restore_env(name: &str, value: Option) { + // SAFETY: callers hold ENV_LOCK before mutating process environment. + unsafe { + if let Some(value) = value { + std::env::set_var(name, value); + } else { + std::env::remove_var(name); + } + } + } + + #[test] + fn expand_env_placeholders_rejects_missing_vars() { + let _guard = ENV_LOCK.lock().unwrap(); + let name = "DYN_TEST_MISSING_LOGICAL_INSTANCE_KEY"; + let old = std::env::var(name).ok(); + restore_env(name, None); + assert!(expand_env_placeholders("$(DYN_TEST_MISSING_LOGICAL_INSTANCE_KEY)").is_err()); + restore_env(name, old); + } + + #[test] + fn resolve_logical_instance_id_hashes_expanded_key() { + let _guard = ENV_LOCK.lock().unwrap(); + let old_id = std::env::var("DYN_DISCOVERY_LOGICAL_INSTANCE_ID").ok(); + let old_key = std::env::var("DYN_DISCOVERY_LOGICAL_INSTANCE_KEY").ok(); + let old_pcsg = std::env::var("DYN_TEST_PCSG").ok(); + + restore_env("DYN_DISCOVERY_LOGICAL_INSTANCE_ID", None); + restore_env("DYN_TEST_PCSG", Some("pcsg-a".to_string())); + restore_env( + "DYN_DISCOVERY_LOGICAL_INSTANCE_KEY", + Some("ns/$(DYN_TEST_PCSG)/role".to_string()), + ); + + let expected = hash_logical_instance_key("ns/pcsg-a/role"); + let actual = resolve_logical_instance_id(99).unwrap(); + + restore_env("DYN_DISCOVERY_LOGICAL_INSTANCE_ID", old_id); + restore_env("DYN_DISCOVERY_LOGICAL_INSTANCE_KEY", old_key); + restore_env("DYN_TEST_PCSG", old_pcsg); + + assert_eq!(actual, expected); + } +} diff --git a/lib/runtime/src/distributed.rs b/lib/runtime/src/distributed.rs index 45eeb921bf1f..dd344a54bc5a 100644 --- a/lib/runtime/src/distributed.rs +++ b/lib/runtime/src/distributed.rs @@ -176,7 +176,7 @@ impl DistributedRuntime { }; use crate::discovery::KVStoreDiscovery; ( - Arc::new(KVStoreDiscovery::new(store, runtime.primary_token())) + Arc::new(KVStoreDiscovery::new(store, runtime.primary_token())?) as Arc, None, ) diff --git a/lib/runtime/src/pipeline/network/egress/addressed_router.rs b/lib/runtime/src/pipeline/network/egress/addressed_router.rs index fcdcf3c0bf7e..89a88027d287 100644 --- a/lib/runtime/src/pipeline/network/egress/addressed_router.rs +++ b/lib/runtime/src/pipeline/network/egress/addressed_router.rs @@ -104,6 +104,17 @@ where } } else if is_complete_final { None + } else if first_response { + // The worker-side stream closed before emitting any frame. Even if the + // request context has already been stopped by failover, treating this as + // a clean EOF leaks an empty HTTP 200 to clients and prevents request + // migration from replaying the request. + let err = DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("Stream ended before first response") + .build(); + tracing::debug!("{err}"); + Some(U::from_err(err)) } else if engine_ctx_for_stream.is_stopped() { tracing::debug!("Request cancelled and then trying to read a response"); None diff --git a/lib/runtime/src/pipeline/network/egress/nats_client.rs b/lib/runtime/src/pipeline/network/egress/nats_client.rs index 34924cf881dc..5d2be15b33f5 100644 --- a/lib/runtime/src/pipeline/network/egress/nats_client.rs +++ b/lib/runtime/src/pipeline/network/egress/nats_client.rs @@ -7,11 +7,13 @@ //! providing a consistent interface across all transport types. use super::unified_client::{ClientStats, Headers, RequestPlaneClient}; +use crate::config::environment_names::nats::DYN_NATS_REQUEST_TIMEOUT_MS; use crate::error::{DynamoError, ErrorType}; use crate::metrics::transport_metrics::NATS_ERRORS_TOTAL; use anyhow::Result; use async_trait::async_trait; use bytes::Bytes; +use std::time::Duration; /// Client-facing message for oversized request payloads. /// @@ -20,6 +22,19 @@ use bytes::Bytes; /// diagnostics (subject, payload_bytes, max_payload_bytes) live in logs and /// the `nats_errors_total` metric instead. const MAX_PAYLOAD_USER_MESSAGE: &str = "Request payload is too large for this deployment. Reduce the input size or metadata size and retry."; +const DEFAULT_NATS_REQUEST_TIMEOUT_MS: u64 = 2_000; + +fn request_timeout_from_value(value: Option<&str>) -> Duration { + value + .and_then(|raw| raw.parse::().ok()) + .filter(|&millis| millis > 0) + .map(Duration::from_millis) + .unwrap_or_else(|| Duration::from_millis(DEFAULT_NATS_REQUEST_TIMEOUT_MS)) +} + +fn request_timeout_from_env() -> Duration { + request_timeout_from_value(std::env::var(DYN_NATS_REQUEST_TIMEOUT_MS).ok().as_deref()) +} /// NATS implementation of RequestPlaneClient /// @@ -28,6 +43,7 @@ const MAX_PAYLOAD_USER_MESSAGE: &str = "Request payload is too large for this de pub struct NatsRequestClient { client: async_nats::Client, max_payload: usize, + request_timeout: Duration, } impl NatsRequestClient { @@ -45,6 +61,7 @@ impl NatsRequestClient { Self { client, max_payload, + request_timeout: request_timeout_from_env(), } } @@ -90,25 +107,50 @@ impl RequestPlaneClient for NatsRequestClient { nats_headers.insert(key.as_str(), value.as_str()); } - // Send request with headers - let response = self - .client - .request_with_headers(address.clone(), nats_headers, payload) - .await - .map_err(|e| { + // Send request with headers. The response is only the request-plane + // acknowledgment; generation output arrives over the response stream. + let request_timeout = self.request_timeout; + let response = tokio::time::timeout( + request_timeout, + self.client + .request_with_headers(address.clone(), nats_headers, payload), + ) + .await; + + match response { + Ok(Ok(response)) => Ok(response.payload), + Ok(Err(e)) => { NATS_ERRORS_TOTAL .with_label_values(&["request_failed"]) .inc(); - anyhow::anyhow!( + Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::CannotConnect) .message(format!("NATS request to {address} failed")) .cause(e) .build() - ) - })?; - - Ok(response.payload) + )) + } + Err(_) => { + NATS_ERRORS_TOTAL + .with_label_values(&["request_timeout"]) + .inc(); + tracing::warn!( + address = %address, + timeout_ms = request_timeout.as_millis(), + "NATS request-plane ack timed out" + ); + Err(anyhow::anyhow!( + DynamoError::builder() + .error_type(ErrorType::CannotConnect) + .message(format!( + "NATS request to {address} timed out after {}ms", + request_timeout.as_millis() + )) + .build() + )) + } + } } fn transport_name(&self) -> &'static str { @@ -158,4 +200,28 @@ mod tests { assert!(!err.message().contains("payload_bytes")); assert!(!err.message().contains("nats-server")); } + + #[test] + fn request_timeout_value_uses_default_for_missing_or_invalid_values() { + assert_eq!( + request_timeout_from_value(None), + Duration::from_millis(DEFAULT_NATS_REQUEST_TIMEOUT_MS) + ); + assert_eq!( + request_timeout_from_value(Some("0")), + Duration::from_millis(DEFAULT_NATS_REQUEST_TIMEOUT_MS) + ); + assert_eq!( + request_timeout_from_value(Some("not-a-number")), + Duration::from_millis(DEFAULT_NATS_REQUEST_TIMEOUT_MS) + ); + } + + #[test] + fn request_timeout_value_accepts_positive_millis() { + assert_eq!( + request_timeout_from_value(Some("750")), + Duration::from_millis(750) + ); + } } diff --git a/lib/runtime/src/pipeline/network/egress/push_router.rs b/lib/runtime/src/pipeline/network/egress/push_router.rs index 42dddd5aa770..c366ffed39c8 100644 --- a/lib/runtime/src/pipeline/network/egress/push_router.rs +++ b/lib/runtime/src/pipeline/network/egress/push_router.rs @@ -399,12 +399,18 @@ fn spawn_instance_removal_watcher( } Some(Ok(_)) => {} Some(Err(e)) => { + if cancel_token.is_cancelled() { + break 'reconnect; + } tracing::warn!( endpoint = %endpoint_name, "Instance removal watcher stream error: {e}" ); } None => { + if cancel_token.is_cancelled() { + break 'reconnect; + } tracing::warn!( endpoint = %endpoint_name, "Instance removal watcher stream ended; reconnecting" diff --git a/lib/runtime/src/pipeline/network/ingress/push_handler.rs b/lib/runtime/src/pipeline/network/ingress/push_handler.rs index d796d93eb894..eb926b236a27 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_handler.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_handler.rs @@ -4,6 +4,7 @@ use super::*; use crate::engine::AsyncEngineContext; +use crate::error::{DynamoError, ErrorType}; use crate::metrics::prometheus_names::work_handler; use crate::metrics::work_handler_perf::{ WORK_HANDLER_NETWORK_TRANSIT_SECONDS, WORK_HANDLER_TIME_TO_FIRST_RESPONSE_SECONDS, @@ -161,6 +162,7 @@ impl Ingress { // TODO: Detect end-of-stream using Server-Sent Events (SSE) let mut send_complete_final = true; let mut saw_error_response = false; + let mut sent_response_frame = false; while let Some(resp) = stream.next().await { tracing::trace!("Sending response: {:?}", resp); let is_error = resp.err().is_some(); @@ -202,7 +204,10 @@ impl Ingress { .inc(); } break; - } else if !is_error { + } + + sent_response_frame = true; + if !is_error { // Only notify on non-error chunks — error responses don't prove // the engine is healthy and should not reset the canary timer. if let Some(notifier) = self.endpoint_health_check_notifier.get() { @@ -210,6 +215,43 @@ impl Ingress { } } } + + if send_complete_final && !sent_response_frame { + let err = DynamoError::builder() + .error_type(ErrorType::Disconnected) + .message("Stream ended before first response") + .build(); + tracing::warn!( + error = %err, + "Response stream ended before first frame for stream {}", + context.id() + ); + + let resp_wrapper = NetworkStreamWrapper { + data: Some(U::from_err(err)), + complete_final: false, + }; + let resp_bytes = serde_json::to_vec(&resp_wrapper) + .expect("fatal error: invalid response object - this should never happen"); + if let Some(m) = self.metrics() { + m.response_bytes.inc_by(resp_bytes.len() as u64); + } + if (publisher.send(resp_bytes.into()).await).is_err() { + send_complete_final = false; + tracing::error!( + "Failed to publish empty-stream error for stream {}", + context.id() + ); + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::PUBLISH_RESPONSE]) + .inc(); + } + } else { + saw_error_response = true; + } + } + if send_complete_final { let resp_wrapper = NetworkStreamWrapper:: { data: None, diff --git a/lib/runtime/src/storage/kv.rs b/lib/runtime/src/storage/kv.rs index 207142842b50..2ebd34d7521d 100644 --- a/lib/runtime/src/storage/kv.rs +++ b/lib/runtime/src/storage/kv.rs @@ -342,14 +342,27 @@ impl Manager { // Send all the existing keys for (key, bytes) in bucket.entries().await? { - if let Err(err) = tx + match tx .send_timeout( WatchEvent::Put(KeyValue::new(key, bytes)), WATCH_SEND_TIMEOUT, ) .await { - tracing::error!(bucket_name, %err, "KeyValueStoreManager.watch failed adding existing key to channel"); + Ok(()) => {} + Err(tokio::sync::mpsc::error::SendTimeoutError::Closed(_)) => { + tracing::debug!( + bucket_name, + "KeyValueStoreManager.watch receiver closed while adding existing key; stopping watch task" + ); + return Ok::<(), StoreError>(()); + } + Err(tokio::sync::mpsc::error::SendTimeoutError::Timeout(_)) => { + tracing::error!( + bucket_name, + "KeyValueStoreManager.watch timed out adding existing key to channel" + ); + } } } @@ -362,8 +375,21 @@ impl Manager { None => break, } }; - if let Err(err) = tx.send_timeout(event, WATCH_SEND_TIMEOUT).await { - tracing::error!(bucket_name, %err, "KeyValueStoreManager.watch failed adding new key to channel"); + match tx.send_timeout(event, WATCH_SEND_TIMEOUT).await { + Ok(()) => {} + Err(tokio::sync::mpsc::error::SendTimeoutError::Closed(_)) => { + tracing::debug!( + bucket_name, + "KeyValueStoreManager.watch receiver closed while adding new key; stopping watch task" + ); + break; + } + Err(tokio::sync::mpsc::error::SendTimeoutError::Timeout(_)) => { + tracing::error!( + bucket_name, + "KeyValueStoreManager.watch timed out adding new key to channel" + ); + } } } @@ -595,6 +621,24 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_manager_watch_exits_when_receiver_dropped_during_initial_replay() + -> anyhow::Result<()> { + init(); + + let manager = Arc::new(Manager::memory()); + let bucket = manager.get_or_create_bucket(BUCKET_NAME, None).await?; + bucket.insert(&"test1".into(), "value1".into(), 0).await?; + + let cancel = CancellationToken::new(); + let (watch_task, rx) = manager.watch(BUCKET_NAME, None, cancel); + drop(rx); + + let result = tokio::time::timeout(Duration::from_secs(2), watch_task).await??; + assert!(result.is_ok()); + Ok(()) + } + #[tokio::test] async fn test_broadcast_stream() -> anyhow::Result<()> { init(); diff --git a/lib/runtime/tests/bulwark_cutover.rs b/lib/runtime/tests/bulwark_cutover.rs new file mode 100644 index 000000000000..c82bc59feb45 --- /dev/null +++ b/lib/runtime/tests/bulwark_cutover.rs @@ -0,0 +1,421 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + sync::{ + Arc, + atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use anyhow::{Context, Result, anyhow, bail}; +use async_stream::stream; +use dynamo_runtime::{ + DistributedRuntime, Runtime, + pipeline::{ + AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, PushRouter, ResponseStream, + RouterMode, SingleIn, async_trait, network::Ingress, + }, + protocols::annotated::Annotated, +}; +use futures::StreamExt; +use tokio::time::{Instant, timeout}; + +const COMPONENT: &str = "worker"; +const TOKENS_PER_RESPONSE: usize = 5; +const FIRST_TOKEN_DELAY: Duration = Duration::from_millis(20); +const INTER_TOKEN_DELAY: Duration = Duration::from_millis(8); +const REQUEST_TIMEOUT: Duration = Duration::from_secs(2); +const PHASE_BASELINE: u8 = 0; +const PHASE_TRANSITION: u8 = 1; +const PHASE_POST: u8 = 2; + +static CUTOVER_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +#[derive(Clone)] +struct StreamingHandler { + name: &'static str, + active: Arc, + total: Arc, +} + +impl StreamingHandler { + fn new(name: &'static str) -> Arc { + Arc::new(Self { + name, + active: Arc::new(AtomicUsize::new(0)), + total: Arc::new(AtomicUsize::new(0)), + }) + } + + fn active(&self) -> usize { + self.active.load(Ordering::SeqCst) + } + + fn total(&self) -> usize { + self.total.load(Ordering::SeqCst) + } +} + +struct ActiveGuard(Arc); + +impl Drop for ActiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } +} + +#[async_trait] +impl AsyncEngine, ManyOut>, Error> for StreamingHandler { + async fn generate(&self, input: SingleIn) -> Result>> { + let (_request, ctx) = input.into_parts(); + self.total.fetch_add(1, Ordering::SeqCst); + self.active.fetch_add(1, Ordering::SeqCst); + + let active = self.active.clone(); + let name = self.name; + let output = stream! { + let _active = ActiveGuard(active); + tokio::time::sleep(FIRST_TOKEN_DELAY).await; + for idx in 0..TOKENS_PER_RESPONSE { + if idx > 0 { + tokio::time::sleep(INTER_TOKEN_DELAY).await; + } + yield Annotated::from_data(format!("{name}:{idx}")); + } + }; + + Ok(ResponseStream::new(Box::pin(output), ctx.context())) + } +} + +#[derive(Clone, Debug)] +struct Sample { + phase: u8, + source: String, + ttft: Duration, + avg_itl: Duration, +} + +fn spawn_engine( + runtime: DistributedRuntime, + namespace: &'static str, + endpoint: &'static str, + handler: Arc, +) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + let ingress = Ingress::for_engine(handler)?; + let component = runtime.namespace(namespace)?.component(COMPONENT)?; + component + .endpoint(endpoint) + .endpoint_builder() + .handler(ingress) + .start() + .await + }) +} + +async fn issue_request( + router: Arc>>, + phase: Arc, +) -> Result { + let start = Instant::now(); + let mut response = timeout( + REQUEST_TIMEOUT, + router.round_robin("measure cutover".to_string().into()), + ) + .await + .context("timed out waiting for routed request")??; + + let mut first_token_at = None; + let mut previous_token_at = None; + let mut intervals = Vec::new(); + let mut source = None; + let mut sample_phase = None; + + while let Some(chunk) = timeout(REQUEST_TIMEOUT, response.next()) + .await + .context("timed out waiting for response token")? + { + let now = Instant::now(); + let chunk = chunk.ok().map_err(|err| anyhow!(err))?; + let Some(token) = chunk.data else { + continue; + }; + + if first_token_at.is_none() { + first_token_at = Some(now); + sample_phase = Some(phase.load(Ordering::SeqCst)); + source = Some( + token + .split_once(':') + .map(|(prefix, _)| prefix.to_string()) + .unwrap_or(token), + ); + } else if let Some(previous) = previous_token_at { + intervals.push(now.duration_since(previous)); + } + previous_token_at = Some(now); + } + + let Some(first_token_at) = first_token_at else { + bail!("response stream completed without tokens"); + }; + let Some(source) = source else { + bail!("response stream did not identify a source worker"); + }; + let Some(sample_phase) = sample_phase else { + bail!("response stream did not identify a sample phase"); + }; + + let avg_itl = if intervals.is_empty() { + Duration::ZERO + } else { + let total_nanos: u128 = intervals.iter().map(Duration::as_nanos).sum(); + Duration::from_nanos((total_nanos / intervals.len() as u128) as u64) + }; + + Ok(Sample { + phase: sample_phase, + source, + ttft: first_token_at.duration_since(start), + avg_itl, + }) +} + +async fn wait_for(description: &str, mut condition: F) -> Result<()> +where + F: FnMut() -> bool, +{ + let deadline = Instant::now() + Duration::from_secs(3); + while Instant::now() < deadline { + if condition() { + return Ok(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + bail!("timed out waiting for {description}") +} + +fn p95(values: impl Iterator) -> Result { + let mut values: Vec<_> = values.collect(); + if values.is_empty() { + bail!("cannot compute p95 over an empty sample set"); + } + values.sort_unstable(); + let idx = ((values.len() * 95).div_ceil(100)).saturating_sub(1); + Ok(values[idx]) +} + +#[derive(Debug)] +struct CutoverMetrics { + baseline_samples: usize, + transition_samples: usize, + post_samples: usize, + baseline_ttft_p95: Duration, + post_ttft_p95: Duration, + baseline_itl_p95: Duration, + post_itl_p95: Duration, +} + +async fn run_bulwark_style_shadow_cutover( + namespace: &'static str, + endpoint: &'static str, + client_count: usize, + warmup_duration: Duration, + post_duration: Duration, + latency_budget: Duration, + min_samples: usize, +) -> Result { + dynamo_runtime::logging::init(); + + let runtime = Runtime::from_current()?; + let distributed = DistributedRuntime::new( + runtime.clone(), + dynamo_runtime::distributed::DistributedConfig::process_local(), + ) + .await?; + + let primary = StreamingHandler::new("primary"); + let shadow = StreamingHandler::new("shadow"); + let primary_task = spawn_engine(distributed.clone(), namespace, endpoint, primary.clone()); + + let client = distributed + .namespace(namespace)? + .component(COMPONENT)? + .endpoint(endpoint) + .client() + .await?; + client.wait_for_instances().await?; + assert_eq!( + client.instance_ids().len(), + 1, + "router should see one worker" + ); + + let router = Arc::new( + PushRouter::>::from_client( + client.clone(), + RouterMode::RoundRobin, + ) + .await?, + ); + + let phase = Arc::new(AtomicU8::new(PHASE_BASELINE)); + let stop = Arc::new(AtomicBool::new(false)); + let (sample_tx, mut sample_rx) = + tokio::sync::mpsc::unbounded_channel::>(); + + let mut clients = Vec::new(); + for _ in 0..client_count { + let router = router.clone(); + let phase = phase.clone(); + let stop = stop.clone(); + let sample_tx = sample_tx.clone(); + clients.push(tokio::spawn(async move { + while !stop.load(Ordering::SeqCst) { + let result = issue_request(router.clone(), phase.clone()) + .await + .map_err(|err| err.to_string()); + let _ = sample_tx.send(result); + } + })); + } + drop(sample_tx); + + tokio::time::sleep(warmup_duration).await; + + phase.store(PHASE_TRANSITION, Ordering::SeqCst); + let shadow_task = spawn_engine(distributed.clone(), namespace, endpoint, shadow.clone()); + wait_for("shadow to receive traffic", || shadow.total() > 0).await?; + assert_eq!( + client.instance_ids().len(), + 1, + "router should still see one worker" + ); + + wait_for("primary in-flight requests to drain", || { + primary.active() == 0 + }) + .await?; + phase.store(PHASE_POST, Ordering::SeqCst); + tokio::time::sleep(post_duration).await; + + stop.store(true, Ordering::SeqCst); + for client_task in clients { + client_task.await?; + } + + let mut samples = Vec::new(); + let mut errors = Vec::new(); + while let Some(result) = sample_rx.recv().await { + match result { + Ok(sample) => samples.push(sample), + Err(err) => errors.push(err), + } + } + + distributed.shutdown(); + primary_task.await??; + shadow_task.await??; + drop(runtime); + + if !errors.is_empty() { + bail!("client observed request errors during cutover: {errors:?}"); + } + + let baseline: Vec<_> = samples + .iter() + .filter(|sample| sample.phase == PHASE_BASELINE) + .cloned() + .collect(); + let post: Vec<_> = samples + .iter() + .filter(|sample| sample.phase == PHASE_POST) + .cloned() + .collect(); + let transition: Vec<_> = samples + .iter() + .filter(|sample| sample.phase == PHASE_TRANSITION) + .cloned() + .collect(); + + assert!( + baseline.len() >= min_samples, + "not enough baseline samples: {}", + baseline.len() + ); + assert!( + post.len() >= min_samples, + "not enough post-cutover samples: {}", + post.len() + ); + assert!( + !transition.is_empty(), + "transition phase did not exercise live traffic" + ); + + assert!( + baseline.iter().all(|sample| sample.source == "primary"), + "baseline traffic should be served by primary: {baseline:?}" + ); + assert!( + post.iter().all(|sample| sample.source == "shadow"), + "post-drain traffic should be served by shadow only: {post:?}" + ); + + let baseline_ttft_p95 = p95(baseline.iter().map(|sample| sample.ttft))?; + let post_ttft_p95 = p95(post.iter().map(|sample| sample.ttft))?; + let baseline_itl_p95 = p95(baseline.iter().map(|sample| sample.avg_itl))?; + let post_itl_p95 = p95(post.iter().map(|sample| sample.avg_itl))?; + + assert!( + post_ttft_p95 <= baseline_ttft_p95 + latency_budget, + "post-cutover TTFT p95 regressed: baseline={baseline_ttft_p95:?}, post={post_ttft_p95:?}" + ); + assert!( + post_itl_p95 <= baseline_itl_p95 + latency_budget, + "post-cutover ITL p95 regressed: baseline={baseline_itl_p95:?}, post={post_itl_p95:?}" + ); + + Ok(CutoverMetrics { + baseline_samples: baseline.len(), + transition_samples: transition.len(), + post_samples: post.len(), + baseline_ttft_p95, + post_ttft_p95, + baseline_itl_p95, + post_itl_p95, + }) +} + +fn print_metrics(label: &str, metrics: &CutoverMetrics) { + println!( + "{label}: baseline={}, transition={}, post={}, ttft_p95 {:?}->{:?}, itl_p95 {:?}->{:?}", + metrics.baseline_samples, + metrics.transition_samples, + metrics.post_samples, + metrics.baseline_ttft_p95, + metrics.post_ttft_p95, + metrics.baseline_itl_p95, + metrics.post_itl_p95 + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +async fn bulwark_style_shadow_cutover_under_high_client_load() -> Result<()> { + let _guard = CUTOVER_TEST_LOCK.lock().await; + let metrics = run_bulwark_style_shadow_cutover( + "bulwark_cutover_stress", + "generate_stress", + 64, + Duration::from_millis(1500), + Duration::from_millis(1500), + Duration::from_millis(75), + 64, + ) + .await?; + print_metrics("bulwark cutover stress", &metrics); + Ok(()) +} diff --git a/tests/utils/managed_process.py b/tests/utils/managed_process.py index d2189f69323d..e09277e2804a 100644 --- a/tests/utils/managed_process.py +++ b/tests/utils/managed_process.py @@ -176,6 +176,8 @@ class ManagedProcess: straggler_commands: List[str] = field(default_factory=list) log_dir: str = os.getcwd() display_name: Optional[str] = None + terminate_parent_only_first: bool = False + graceful_shutdown_timeout: float = 8.0 # Ensure attributes exist even if startup fails early proc: Optional[subprocess.Popen] = None @@ -430,6 +432,25 @@ def _force_stop_process( if process.poll() is None: raise ManagedProcessStopTimeoutError(attr_name, process.pid, wait_timeout) + @staticmethod + def _live_process_groups_for_pids(pids: set[int]) -> set[int]: + """Return process groups for snapshotted PIDs that are still non-zombie.""" + live_pgids: set[int] = set() + for pid in pids: + try: + process = psutil.Process(pid) + if process.status() == psutil.STATUS_ZOMBIE: + continue + live_pgids.add(os.getpgid(pid)) + except (ProcessLookupError, psutil.NoSuchProcess, OSError): + pass + except psutil.AccessDenied: + try: + live_pgids.add(os.getpgid(pid)) + except (ProcessLookupError, OSError): + pass + return live_pgids + def _start_process(self): assert self._command_name assert self._log_path @@ -495,7 +516,7 @@ def _start_process(self): ) self._tee_proc = None - def _terminate_process_group(self, timeout: float = 8.0): + def _terminate_process_group(self, timeout: Optional[float] = None): """Terminate the entire process group/session started for the child. Kill Sequence: @@ -529,15 +550,23 @@ def _terminate_process_group(self, timeout: float = 8.0): if self._pgid is None: return + if timeout is None: + timeout = self.graceful_shutdown_timeout + # Step 1: Snapshot all process groups before signaling. Only self.proc runs # in self._pgid (start_new_session=True); _tee_proc/_sed_proc are pipe # helpers in pytest's group with no children. all_pgids: set[int] = {self._pgid} + child_pids: set[int] = set() + pgid_to_pids: dict[int, set[int]] = {self._pgid: set()} if self.proc and self.proc.poll() is None: try: for child in psutil.Process(self.proc.pid).children(recursive=True): + child_pids.add(child.pid) try: - all_pgids.add(os.getpgid(child.pid)) + child_pgid = os.getpgid(child.pid) + all_pgids.add(child_pgid) + pgid_to_pids.setdefault(child_pgid, set()).add(child.pid) except (ProcessLookupError, OSError): pass except (PermissionError, psutil.AccessDenied) as exc: @@ -549,6 +578,49 @@ def _terminate_process_group(self, timeout: float = 8.0): except psutil.NoSuchProcess: pass + if ( + self.terminate_parent_only_first + and self.proc is not None + and self.proc.poll() is None + ): + try: + self._logger.info( + "Sending SIGTERM to parent process: %s", self.proc.pid + ) + self.proc.terminate() + except ProcessLookupError: + return + except Exception as e: + self._logger.warning( + "Error sending SIGTERM to parent process %s: %s", + self.proc.pid, + e, + ) + + poll_interval = 0.1 + elapsed = 0.0 + while elapsed < timeout: + if self.proc.poll() is not None: + live_pgids = self._live_process_groups_for_pids(child_pids) + if not live_pgids: + return + self._logger.info( + "Parent process %s exited; terminating child process groups: %s", + self.proc.pid, + sorted(live_pgids), + ) + all_pgids = live_pgids + break + time.sleep(poll_interval) + elapsed += poll_interval + + if self.proc.poll() is None: + self._logger.warning( + "Parent process %s did not exit within %.1fs; terminating process groups", + self.proc.pid, + timeout, + ) + # Step 2: SIGTERM all snapshotted process groups (graceful shutdown). # Delivers the signal to every group so children that called # setpgid()/setsid() also get a chance to shut down gracefully. @@ -577,9 +649,14 @@ def _terminate_process_group(self, timeout: float = 8.0): while elapsed < timeout: if self.proc is not None: self.proc.poll() - # Check if any signaled group is still alive + # Check if any signaled group is still alive. Zombie-only + # child groups (common with MPI launch helpers such as orted) do not + # hold resources and cannot be killed, so don't burn the timeout. still_alive = False for pgid in sigtermed_pgids: + known_pids = pgid_to_pids.get(pgid, set()) + if known_pids and not self._live_process_groups_for_pids(known_pids): + continue try: os.killpg(pgid, 0) # signal 0 = check existence still_alive = True @@ -595,6 +672,9 @@ def _terminate_process_group(self, timeout: float = 8.0): # survives (e.g. processes that ignored SIGTERM or timed out). # _stop_started_processes handles per-process wait/escalation afterward. for pgid in all_pgids: + known_pids = pgid_to_pids.get(pgid, set()) + if known_pids and not self._live_process_groups_for_pids(known_pids): + continue try: os.killpg(pgid, signal.SIGKILL) except ProcessLookupError: diff --git a/tests/utils/test_managed_process_teardown.py b/tests/utils/test_managed_process_teardown.py index 0e0445cfef93..a0cb30287f3d 100644 --- a/tests/utils/test_managed_process_teardown.py +++ b/tests/utils/test_managed_process_teardown.py @@ -16,6 +16,7 @@ Always use unique markers scoped to the test invocation. """ +import logging import os import signal import subprocess @@ -353,3 +354,104 @@ def test_process_gets_sigterm_grace_before_sigkill(self, tmp_path): assert os.path.exists( marker_file ), "Process was SIGKILLed before SIGTERM handler could run" + + +# --------------------------------------------------------------------------- +# Scenario 7: parent-first graceful teardown +# --------------------------------------------------------------------------- +class TestParentOnlyFirstTeardown: + def test_parent_only_signal_lets_parent_cleanup_child(self, tmp_path): + parent_marker = str(tmp_path / "parent_sigterm") + child_marker = str(tmp_path / "child_cleaned") + ready_file = str(tmp_path / "ready") + child_pid_file = str(tmp_path / "child_pid") + script = "\n".join( + [ + "import pathlib, signal, subprocess, sys, time", + f"parent_marker = pathlib.Path({parent_marker!r})", + f"child_marker = pathlib.Path({child_marker!r})", + f"ready = pathlib.Path({ready_file!r})", + f"child_pid_file = pathlib.Path({child_pid_file!r})", + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(300)'])", + "child_pid_file.write_text(str(child.pid))", + "def on_term(*_):", + " parent_marker.touch()", + " child.terminate()", + " child.wait(timeout=5)", + " child_marker.touch()", + " raise SystemExit(0)", + "signal.signal(signal.SIGTERM, on_term)", + "ready.touch()", + "while True:", + " time.sleep(0.1)", + ] + ) + mp = ManagedProcess( + command=["python3", "-c", script], + timeout=10, + display_output=False, + terminate_all_matching_process_names=False, + terminate_parent_only_first=True, + graceful_shutdown_timeout=5.0, + log_dir=str(tmp_path), + ) + + child_pid = None + with mp: + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline and not os.path.exists(ready_file): + time.sleep(0.05) + assert os.path.exists(ready_file) + child_pid = int(open(child_pid_file, encoding="utf-8").read()) + assert _pid_alive(child_pid) + + assert os.path.exists(parent_marker) + assert os.path.exists(child_marker) + assert child_pid is not None + assert _wait_for_pid_death(child_pid, timeout=5) + + def test_parent_only_ignores_zombie_child_process_group(self, tmp_path, caplog): + ready_file = str(tmp_path / "ready") + child_pid_file = str(tmp_path / "child_pid") + script = "\n".join( + [ + "import os, pathlib, time", + f"ready = pathlib.Path({ready_file!r})", + f"child_pid_file = pathlib.Path({child_pid_file!r})", + "pid = os.fork()", + "if pid == 0:", + " os.setpgid(0, 0)", + " os._exit(0)", + "child_pid_file.write_text(str(pid))", + "ready.touch()", + "while True:", + " time.sleep(0.1)", + ] + ) + mp = ManagedProcess( + command=["python3", "-c", script], + timeout=10, + display_output=False, + terminate_all_matching_process_names=False, + terminate_parent_only_first=True, + graceful_shutdown_timeout=0.5, + log_dir=str(tmp_path), + ) + + with caplog.at_level(logging.WARNING, logger="ManagedProcess"): + with mp: + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline and not os.path.exists(ready_file): + time.sleep(0.05) + assert os.path.exists(ready_file) + child_pid = int(open(child_pid_file, encoding="utf-8").read()) + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline: + try: + if psutil.Process(child_pid).status() == psutil.STATUS_ZOMBIE: + break + except psutil.NoSuchProcess: + break + time.sleep(0.05) + + assert "child process groups are still live" not in caplog.text