From c7cc95efccfeb29d278c675d39e3bef4093f1575 Mon Sep 17 00:00:00 2001 From: nnshah1 Date: Mon, 29 Jun 2026 20:12:53 -0700 Subject: [PATCH 1/3] fix(runtime): bound the per-endpoint inflight drain on teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On endpoint teardown the graceful-shutdown path waited on `while inflight > 0 { notify.notified().await }` with no timeout, on BOTH request planes (NATS `PushEndpoint` and the default TCP `SharedTcpServer`). A single stuck inflight request (e.g. one whose engine can no longer make progress and cannot be aborted) keeps inflight > 0, so the drain wedges, the serve future never returns, and `Runtime::shutdown()` is never reached — the worker zombies (Running, /health green, unable to serve). Add a shared `drain_inflight` helper in the ingress module, bounded by the existing #10705 `graceful_shutdown_timeout()` (made pub(crate); no new env/const), and call it from both `PushEndpoint::start` and `SharedTcpServer::unregister_endpoint`. Tested with paused time: the bounded wait returns instead of hanging when a request never completes, and still drains cleanly to zero when it does. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: nnshah1 --- lib/runtime/src/pipeline/network/ingress.rs | 93 +++++++++++++++++++ .../pipeline/network/ingress/push_endpoint.rs | 22 ++--- .../network/ingress/shared_tcp_endpoint.rs | 22 ++--- lib/runtime/src/runtime.rs | 2 +- 4 files changed, 108 insertions(+), 31 deletions(-) diff --git a/lib/runtime/src/pipeline/network/ingress.rs b/lib/runtime/src/pipeline/network/ingress.rs index b0bc5172e72b..906e4b4a8481 100644 --- a/lib/runtime/src/pipeline/network/ingress.rs +++ b/lib/runtime/src/pipeline/network/ingress.rs @@ -8,3 +8,96 @@ pub mod shared_tcp_endpoint; pub mod unified_server; use super::*; + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use tokio::sync::Notify; + +/// Wait for inflight requests to drain, bounded by `timeout`. Returns the count +/// still inflight when the wait ends: `0` on a clean drain, `>0` if the timeout +/// fired (a stuck request must not wedge teardown). Shared by the NATS push and +/// TCP request planes. +async fn drain_inflight( + inflight: Arc, + notify: Arc, + endpoint_name: &str, + timeout: Duration, +) -> u64 { + let inflight_count = inflight.load(Ordering::SeqCst); + if inflight_count == 0 { + return 0; + } + + tracing::info!( + endpoint_name, + inflight_count, + "Waiting for inflight requests to complete" + ); + + let wait = async { + while inflight.load(Ordering::SeqCst) > 0 { + notify.notified().await; + } + }; + + match tokio::time::timeout(timeout, wait).await { + Ok(()) => { + tracing::info!(endpoint_name, "All inflight requests completed"); + 0 + } + Err(_) => { + let remaining = inflight.load(Ordering::SeqCst); + tracing::warn!( + endpoint_name, + timeout_secs = timeout.as_secs(), + remaining, + "Timed out waiting for inflight requests to drain; proceeding with shutdown" + ); + remaining + } + } +} + +#[cfg(test)] +mod drain_tests { + use super::*; + + /// Drain returns within the bound even if the request never completes. + #[tokio::test(start_paused = true)] + async fn drain_inflight_is_bounded_when_request_never_completes() { + let inflight = Arc::new(AtomicU64::new(1)); + let notify = Arc::new(Notify::new()); + + let remaining = tokio::time::timeout( + Duration::from_secs(3600), + drain_inflight(inflight, notify, "test-endpoint", Duration::from_secs(5)), + ) + .await + .expect("drain_inflight must be bounded; it hung past the outer guard"); + + assert_eq!( + remaining, 1, + "the stuck inflight request should still be counted as remaining" + ); + } + + /// A clean drain (request completes) returns 0 promptly. + #[tokio::test(start_paused = true)] + async fn drain_inflight_returns_zero_when_requests_complete() { + let inflight = Arc::new(AtomicU64::new(1)); + let notify = Arc::new(Notify::new()); + + let inflight_clone = inflight.clone(); + let notify_clone = notify.clone(); + tokio::spawn(async move { + inflight_clone.fetch_sub(1, Ordering::SeqCst); + notify_clone.notify_one(); + }); + + let remaining = + drain_inflight(inflight, notify, "test-endpoint", Duration::from_secs(5)).await; + + assert_eq!(remaining, 0, "a completed request should drain cleanly"); + } +} diff --git a/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs b/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs index 8472cfb92c12..943c60dd37cd 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_endpoint.rs @@ -148,21 +148,13 @@ impl PushEndpoint { // await for all inflight requests to complete if graceful shutdown if self.graceful_shutdown { - let inflight_count = inflight.load(Ordering::SeqCst); - if inflight_count > 0 { - tracing::info!( - endpoint_name = endpoint_name_local.as_str(), - inflight_count = inflight_count, - "Waiting for inflight NATS requests to complete" - ); - while inflight.load(Ordering::SeqCst) > 0 { - notify.notified().await; - } - tracing::info!( - endpoint_name = endpoint_name_local.as_str(), - "All inflight NATS requests completed" - ); - } + super::drain_inflight( + inflight, + notify, + endpoint_name_local.as_str(), + crate::runtime::graceful_shutdown_timeout(), + ) + .await; } else { tracing::info!( endpoint_name = endpoint_name_local.as_str(), diff --git a/lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs b/lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs index 4975e45b3b78..d16974f1fd97 100644 --- a/lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs +++ b/lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs @@ -463,21 +463,13 @@ impl SharedTcpServer { "Unregistered TCP endpoint handler" ); - let inflight_count = handler.inflight.load(Ordering::SeqCst); - if inflight_count > 0 { - tracing::info!( - endpoint_name = %endpoint_name, - inflight_count = inflight_count, - "Waiting for inflight TCP requests to complete" - ); - while handler.inflight.load(Ordering::SeqCst) > 0 { - handler.notify.notified().await; - } - tracing::info!( - endpoint_name = %endpoint_name, - "All inflight TCP requests completed" - ); - } + super::drain_inflight( + handler.inflight.clone(), + handler.notify.clone(), + endpoint_name, + crate::runtime::graceful_shutdown_timeout(), + ) + .await; } } diff --git a/lib/runtime/src/runtime.rs b/lib/runtime/src/runtime.rs index aaee61ea9939..50edcd11170f 100644 --- a/lib/runtime/src/runtime.rs +++ b/lib/runtime/src/runtime.rs @@ -32,7 +32,7 @@ pub use tokio_util::sync::CancellationToken; const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_SECS: u64 = 15 * 60; -fn graceful_shutdown_timeout() -> Duration { +pub(crate) fn graceful_shutdown_timeout() -> Duration { let timeout_secs = std::env::var( config::environment_names::runtime::DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS, ) From 866f2e904ddec8f520bdc22cd7d2089406b1b123 Mon Sep 17 00:00:00 2001 From: nnshah1 Date: Mon, 29 Jun 2026 20:12:56 -0700 Subject: [PATCH 2/3] fix(runtime): route etcd lease loss through Runtime::shutdown() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On lease loss the keep-alive task called a bare `primary_token().cancel()`, tearing the primary token down at once and skipping the phased shutdown sequence (Phase 1 endpoint-token cancel -> Phase 2 bounded graceful drain -> Phase 3 backend teardown). Combined with the previously unbounded endpoint drain, a stuck inflight request left lease-loss workers wedged. Pass the Runtime into `create_lease` and, on an unrecoverable keep-alive error, call `Runtime::shutdown()` instead of a bare token cancel — honoring the documented `etcd::Client::new` contract that a lost lease shuts the worker down. Unit-tested: lease-loss teardown is phased (endpoint token cancels first, primary token only after graceful tasks complete). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: nnshah1 --- lib/runtime/src/transports/etcd.rs | 14 +++--- lib/runtime/src/transports/etcd/lease.rs | 56 +++++++++++++++++++++--- 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/lib/runtime/src/transports/etcd.rs b/lib/runtime/src/transports/etcd.rs index 65764d7f40b8..8178a8344119 100644 --- a/lib/runtime/src/transports/etcd.rs +++ b/lib/runtime/src/transports/etcd.rs @@ -69,10 +69,10 @@ impl Client { /// If the lease expires, the [`Runtime`] will be shutdown. /// If the [`Runtime`] is shutdown, the lease will be revoked. pub async fn new(config: ClientOptions, runtime: Runtime) -> Result { - let token = runtime.primary_token(); + let runtime_for_lease = runtime.clone(); let ((connector, lease_id), rt) = build_in_runtime( - async move { Self::connect_with_startup_retry(&config, token).await }, + async move { Self::connect_with_startup_retry(&config, runtime_for_lease).await }, 1, ) .await?; @@ -88,8 +88,9 @@ impl Client { /// Connect to etcd during startup, retrying with exponential backoff for up to 2 minutes. async fn connect_with_startup_retry( config: &ClientOptions, - token: CancellationToken, + runtime: Runtime, ) -> Result<(Arc, u64)> { + let token = runtime.primary_token(); let deadline = Instant::now() + STARTUP_CONNECT_TIMEOUT; let mut backoff = STARTUP_CONNECT_INITIAL_BACKOFF; @@ -98,7 +99,7 @@ impl Client { anyhow::bail!("etcd startup connection cancelled"); } - let attempt = Self::connect_startup_attempt(config, &token).await; + let attempt = Self::connect_startup_attempt(config, &runtime).await; match attempt { Ok(connection) => return Ok(connection), @@ -134,13 +135,14 @@ impl Client { async fn connect_startup_attempt( config: &ClientOptions, - token: &CancellationToken, + runtime: &Runtime, ) -> Result<(Arc, u64)> { + let token = runtime.primary_token(); let connector = Connector::new(config.etcd_url.clone(), config.etcd_connect_options.clone()).await?; let lease_id = if config.attach_lease { - create_lease(connector.clone(), config.lease_ttl, token.clone()) + create_lease(connector.clone(), config.lease_ttl, runtime.clone()) .await .with_context(|| { format!( diff --git a/lib/runtime/src/transports/etcd/lease.rs b/lib/runtime/src/transports/etcd/lease.rs index a34127c9e9a2..387a8cc835e6 100644 --- a/lib/runtime/src/transports/etcd/lease.rs +++ b/lib/runtime/src/transports/etcd/lease.rs @@ -2,21 +2,24 @@ // SPDX-License-Identifier: Apache-2.0 use super::connector::Connector; +use crate::runtime::Runtime; use etcd_client::{LeaseKeepAliveStream, LeaseKeeper}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio_util::sync::CancellationToken; -/// Create an etcd lease with the given TTL, attach it to the provided cancellation token, +/// Create an etcd lease with the given TTL, tie its lifetime to the [`Runtime`], /// spawn a keep-alive task, and return the lease id (u64). /// -/// Note: this function spawns a background task that maintains the lease until the token is -/// cancelled or an unrecoverable error occurs. +/// Note: this function spawns a background task that maintains the lease until the runtime is +/// shut down or an unrecoverable error occurs. On an unrecoverable error the runtime is shut +/// down, honoring the contract that a lost lease shuts the worker down. pub async fn create_lease( connector: Arc, ttl: u64, - token: CancellationToken, + runtime: Runtime, ) -> anyhow::Result { + let token = runtime.primary_token(); if token.is_cancelled() { anyhow::bail!("lease creation cancelled"); } @@ -47,7 +50,9 @@ pub async fn create_lease( error = %e, "Unable to maintain lease. Check etcd server status" ); - token.cancel(); + // Phased shutdown (endpoint drain -> backend teardown), not a + // bare primary-token cancel, so teardown is ordered. + runtime.shutdown(); } } }); @@ -204,3 +209,44 @@ async fn keep_alive_with_stream( } } } + +#[cfg(test)] +mod tests { + use super::*; + + /// The lease-loss teardown path (`runtime.shutdown()`) must be phased: + /// Phase 1 cancels the endpoint token promptly, but Phase 3 (primary token + /// + backend teardown) waits for outstanding graceful tasks. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn lease_loss_shutdown_is_phased() { + let runtime = Runtime::from_current().unwrap(); + + // Stands in for a serving endpoint; Phase 1 cancels the endpoint + // shutdown token, cascading to this child. + let endpoint_token = runtime.child_token(); + let primary_token = runtime.primary_token(); + + // Hold a graceful task so Phase 2 cannot advance to Phase 3 yet. + let guard = runtime.graceful_shutdown_tracker().register_task(); + + runtime.shutdown(); + + // Phase 1 cancels the endpoint token promptly. + tokio::time::timeout(Duration::from_secs(5), endpoint_token.cancelled()) + .await + .expect("Phase 1 must cancel the endpoint token on lease loss"); + + // While the graceful task is outstanding, Phase 3 must not have torn + // down the primary token. A bare-cancel regression fails this assert. + assert!( + !primary_token.is_cancelled(), + "primary token must not be cancelled while a graceful task is outstanding" + ); + + // Releasing the task lets Phase 3 proceed and cancel the primary token. + drop(guard); + tokio::time::timeout(Duration::from_secs(5), primary_token.cancelled()) + .await + .expect("Phase 3 must cancel the primary token once graceful tasks complete"); + } +} From 3ba53696d9b733444b95f9de3db376bf8b1e0e7a Mon Sep 17 00:00:00 2001 From: nnshah1 Date: Mon, 29 Jun 2026 20:12:59 -0700 Subject: [PATCH 3/3] test(ft): reproduce lease-loss zombie with a frozen vLLM engine The existing non-HA shutdown tests kill etcd with no request in flight, so the endpoint drain is instant and they pass even on the unbounded-drain code. The zombie requires an in-flight request that can neither complete nor be aborted (a stuck engine / transfer). This test SIGSTOPs the vLLM engine (rank) process mid-generation so the request stays pinned in the endpoint inflight counter, then kills etcd. With the bounded drain the worker times out the drain (remaining=1) and exits (~26s); without it the worker wedges (still running at 60s). Parametrized over both request planes (default `tcp` SharedTcpServer + `nats` PushEndpoint) since the bound must cover both. Verified RED/GREEN on an RTX A6000 with Qwen3-0.6B. Holds the frontend drain open so only the worker-side behavior is measured. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: nnshah1 --- tests/fault_tolerance/etcd_ha/test_vllm.py | 165 +++++++++++++++++++++ 1 file changed, 165 insertions(+) diff --git a/tests/fault_tolerance/etcd_ha/test_vllm.py b/tests/fault_tolerance/etcd_ha/test_vllm.py index 7f4b597b215b..3ca04b50531c 100644 --- a/tests/fault_tolerance/etcd_ha/test_vllm.py +++ b/tests/fault_tolerance/etcd_ha/test_vllm.py @@ -5,9 +5,13 @@ import logging import os import shutil +import threading +import time from enum import Enum +import psutil import pytest +import requests from tests.conftest import NatsServer from tests.fault_tolerance.etcd_ha.utils import ( @@ -443,3 +447,164 @@ def test_etcd_non_ha_shutdown_vllm_disaggregated( "Frontend": frontend, } ) + + +# Lease-loss zombie: a non-cancellable in-flight request (modeled by +# SIGSTOP-ing the vLLM engine) must not wedge worker teardown. +ZOMBIE_GRACEFUL_SHUTDOWN_TIMEOUT_SECS = 10 # worker-side drain bound +# Hold the frontend's drain open so it does not abort the request and mask the +# worker-side behavior under test. +ZOMBIE_FRONTEND_DRAIN_TIMEOUT_SECS = 300 +ZOMBIE_WORKER_EXIT_DEADLINE_SECS = ZOMBIE_GRACEFUL_SHUTDOWN_TIMEOUT_SECS + 50 +ZOMBIE_INFLIGHT_MAX_TOKENS = 8000 + + +def _zombie_verify_serving(): + r = requests.post( + f"http://localhost:{FRONTEND_PORT}/v1/completions", + json={ + "model": FAULT_TOLERANCE_MODEL_NAME, + "prompt": "The capital of France is", + "max_tokens": 5, + "temperature": 0.0, + }, + timeout=120, + ) + assert ( + r.status_code == 200 + ), f"pre-fault completion failed: {r.status_code} {r.text}" + + +def _zombie_start_inflight_request(): + """Fire a long completion in the background so a request is in flight. + + Non-streaming on purpose: the worker's handle_payload runs the full + generation before returning, holding the push-endpoint inflight counter. + """ + + errors: list = [] + + def _run(): + try: + requests.post( + f"http://localhost:{FRONTEND_PORT}/v1/completions", + json={ + "model": FAULT_TOLERANCE_MODEL_NAME, + "prompt": "Tell me a very long story.", + "max_tokens": ZOMBIE_INFLIGHT_MAX_TOKENS, + "temperature": 0.0, + }, + timeout=600, + ) + except requests.RequestException as exc: + errors.append(exc) + logger.info("in-flight request ended: %s", exc) + + t = threading.Thread(target=_run, daemon=True) + t.start() + return t, errors + + +def _freeze_vllm_engine_descendants(worker_pid: int) -> list: + """SIGSTOP the vLLM engine (rank) subprocess(es) under this worker. + + Scoped to the worker's process tree so concurrent tests/workers on the same + host are untouched. A frozen engine cannot finish or abort the in-flight + request, so it stays pinned in the endpoint inflight counter -- a + non-cancellable inflight. + """ + frozen = [] + for child in psutil.Process(worker_pid).children(recursive=True): + try: + if "EngineCore" in child.name() or "EngineCore" in " ".join( + child.cmdline() + ): + child.suspend() + frozen.append(child.pid) + logger.info("SIGSTOP vLLM engine pid=%s", child.pid) + except psutil.Error: + continue + assert frozen, "no vLLM EngineCore descendant of the worker to freeze" + return frozen + + +def _resume_kill(pids: list) -> None: + """Resume then kill frozen engines so they can't hang teardown or hold a GPU.""" + for pid in pids: + try: + proc = psutil.Process(pid) + proc.resume() + proc.kill() + except psutil.NoSuchProcess: + pass + + +@pytest.mark.gpu_1 +@pytest.mark.xpu_1 +@pytest.mark.e2e +@pytest.mark.nightly +@pytest.mark.model(FAULT_TOLERANCE_MODEL_NAME) +@pytest.mark.timeout(420) +@pytest.mark.parametrize("request_plane", ["tcp", "nats"]) +def test_etcd_lease_loss_zombie_vllm_frozen_engine( + request, monkeypatch, request_plane, predownload_models +): + """A worker that loses its etcd lease with a non-cancellable in-flight + request must still exit. + + Repro: freeze the vLLM engine mid-generation so the request can neither + complete nor be aborted, then kill etcd. With the bounded endpoint drain the + worker times out the drain and exits; the unbounded drain wedges (zombie). + + Parametrized over both request planes: the bounded drain must cover the + default ``tcp`` plane (SharedTcpServer) as well as ``nats`` (PushEndpoint). + """ + # Hold the frontend's drain open so only the worker behavior is measured. + monkeypatch.setenv("DYN_REQUEST_PLANE", request_plane) + monkeypatch.setenv( + "DYN_HTTP_GRACEFUL_SHUTDOWN_TIMEOUT_SECS", + str(ZOMBIE_FRONTEND_DRAIN_TIMEOUT_SECS), + ) + + with NatsServer(request): + with EtcdCluster(request, num_replicas=1) as etcd_cluster: + etcd_endpoints = etcd_cluster.get_client_endpoints() + with DynamoFrontendProcess(request, etcd_endpoints): + worker = DynamoWorkerProcess( + request, etcd_endpoints, mode=WorkerMode.AGGREGATED + ) + worker.env["DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS"] = str( + ZOMBIE_GRACEFUL_SHUTDOWN_TIMEOUT_SECS + ) + with worker: + _zombie_verify_serving() + + logger.info("Starting long in-flight request") + ( + inflight_thread, + inflight_errors, + ) = _zombie_start_inflight_request() + time.sleep(5) # let it reach decode + assert inflight_thread.is_alive(), ( + "in-flight request ended before engine freeze: " + f"{inflight_errors!r}" + ) + + # Freeze the engine: the in-flight request is now + # non-cancellable (cannot complete or be aborted). Resume/kill + # in finally before the worker context unwinds, so a SIGSTOP-ed + # engine can't hang teardown or strand a GPU-holding process. + frozen_pids = _freeze_vllm_engine_descendants(worker.proc.pid) + try: + time.sleep(1) + + logger.info("Terminating ETCD to induce lease loss") + etcd_cluster.stop() + + # Bounded drain -> worker exits; unbounded -> zombie. + wait_for_processes_to_terminate( + {"Worker": worker}, + timeout=ZOMBIE_WORKER_EXIT_DEADLINE_SECS, + ) + finally: + _resume_kill(frozen_pids)