fix(runtime): bound endpoint drain + route etcd lease loss through Runtime::shutdown()

[nnshah1]

Summary

Fixes the etcd lease-loss "zombie worker" failure mode: when a worker loses its etcd lease while a request is stuck in-flight, it could fail to shut down — staying Running with /health green but unable to serve, until manually killed.

Two root causes:

1. Unbounded per-endpoint drain. On teardown the graceful path waited on while inflight > 0 { notify.notified().await } with no timeout. A single stuck in-flight request (e.g. a request whose engine can no longer make progress and cannot be aborted) keeps inflight > 0 forever, so the drain wedges, the serve future never returns, and Runtime::shutdown() is never reached.
2. Lease loss bypassed Runtime::shutdown(). The keep-alive task called a bare primary_token().cancel() instead of Runtime::shutdown(), skipping the phased shutdown sequence (and the documented etcd::Client::new contract that a lost lease shuts the worker down).

Changes

• push_endpoint.rs — extract the drain into drain_inflight, bounded by the existing graceful_shutdown_timeout() from fix(runtime): bound graceful shutdown drain #10705 (made pub(crate); no new env var or const). Returns the count still inflight if the bound fires.
• etcd/lease.rs — on an unrecoverable keep-alive failure, route through Runtime::shutdown() instead of a bare token cancel, so Phases 1–3 run in order.
• Unit tests (paused-time): the bounded drain returns on a stuck counter; lease-loss teardown is phased (endpoint token cancels first, primary token only after graceful tasks complete).
• tests/fault_tolerance/etcd_ha/test_vllm.py — GPU regression test reproducing the zombie with a frozen vLLM engine (SIGSTOP the engine mid-generation so the in-flight request is non-cancellable, then kill etcd). Verified RED/GREEN on an RTX A6000 / Qwen3-0.6B: without the bound the worker wedges (still running at 60s); with it the worker times out the drain (remaining=1) and exits ~27s.

Notes

• The zombie only reproduces with a non-cancellable in-flight request (modeled by freezing the engine). Cancellable requests drain via other backstops, which is why the existing no-inflight shutdown tests pass either way.
• Reuses fix(runtime): bound graceful shutdown drain #10705's bounded Phase-2; relationship to the open fix: runtime graceful shutdown #9951 (per-endpoint drain + abortable discovery cleanup) called out in code comments — happy to fold these together.
• Complementary to making the worker /health reflect shutdown (so k8s cycles the pod even if a process can't exit); that readiness-side change is tracked separately.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Improved graceful shutdown behavior so in-flight requests are drained with a timeout instead of waiting indefinitely.
  • Endpoints now shut down more reliably when a service lease is lost, helping prevent stuck workers and lingering requests.

• Tests

  • Added end-to-end coverage for lease-loss scenarios with long-running requests to verify phased shutdown and worker termination.

Limitation

The phased shutdown and bounded drain run on the tokio runtime, so they rely on the primary executor staying healthy (the common case: a stuck request on the engine/native side). A fully-stalled tokio executor would defeat any cooperative mechanism here (including the existing exit(911) watchdog); guaranteeing termination in that case needs an OS-thread watchdog or the k8s liveness probe — out of scope for this drain fix.

[devin-ai-integration]

Devin Review found 8 potential issues.

[coderabbitai]

Walkthrough

Routes etcd lease-loss through phased Runtime::shutdown(), adds bounded inflight draining for ingress shutdown paths, and adds unit and e2e coverage for lease-loss and frozen-engine shutdown behavior.

DIS-2295: Phased shutdown on lease loss with bounded inflight drain

Layer / File(s)	Summary
Expose graceful_shutdown_timeout crate-wide lib/runtime/src/runtime.rs	graceful_shutdown_timeout() is now pub(crate) for use by ingress shutdown code.
Lease loss routes through phased Runtime::shutdown lib/runtime/src/transports/etcd.rs, lib/runtime/src/transports/etcd/lease.rs	Client::new clones and passes Runtime into create_lease; create_lease now derives its token from runtime.primary_token() and calls runtime.shutdown() on keep-alive failure. A unit test checks phased token-cancellation ordering during lease loss.
Bounded inflight drain in ingress shutdown lib/runtime/src/pipeline/network/ingress.rs, lib/runtime/src/pipeline/network/ingress/push_endpoint.rs, lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs	Adds drain_inflight with timeout-bound waiting on Notify, and updates push endpoint shutdown plus TCP endpoint unregistering to use it. Unit tests cover timeout and clean-drain behavior.
E2E zombie frozen-engine test tests/fault_tolerance/etcd_ha/test_vllm.py	Adds a test that freezes EngineCore descendants during an in-flight request, stops etcd to trigger lease loss, and asserts worker termination within a deadline.

Changes

DIS-2295: Phased shutdown on lease loss with bounded inflight drain

Layer / File(s)	Summary
Expose graceful_shutdown_timeout crate-wide lib/runtime/src/runtime.rs	graceful_shutdown_timeout() visibility widened to pub(crate).
Lease loss routes through phased Runtime::shutdown lib/runtime/src/transports/etcd.rs, lib/runtime/src/transports/etcd/lease.rs	Client::new clones Runtime for lease creation; create_lease now accepts Runtime, derives the lease token from runtime.primary_token(), and routes unrecoverable keep-alive failure through runtime.shutdown(). A unit test checks cancellation ordering.
Bounded inflight drain in ingress shutdown lib/runtime/src/pipeline/network/ingress.rs, lib/runtime/src/pipeline/network/ingress/push_endpoint.rs, lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs	Adds drain_inflight and switches push endpoint graceful shutdown plus shared TCP endpoint unregistering to the bounded helper. Unit tests cover timeout and successful drain paths.
E2E zombie frozen-engine test tests/fault_tolerance/etcd_ha/test_vllm.py	Adds helper functions and a parametrized end-to-end test that freezes EngineCore descendants, induces lease loss, and waits for worker exit under a bounded deadline.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the fix well, but it does not follow the required template and is missing Overview, reviewer start, and Related Issues sections.	Rewrite the PR description to match the template, adding Overview, Details, Where should reviewer start?, and the required Related Issues section.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main runtime shutdown change and is specific enough for history.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/runtime/src/pipeline/network/ingress/push_endpoint.rs`:
- Line 175: The source comments in push_endpoint’s ingress path still reference
an internal Linear ticket, which must be removed. Update the comment near the
lease-loss worker note to replace DIS-2295 with the corresponding public GitHub
issue reference, or omit the ticket reference entirely. Also scan the nearby
comment(s) in push_endpoint.rs for any other internal ticket IDs and replace
them consistently.

In `@lib/runtime/src/transports/etcd/lease.rs`:
- Line 17: Replace the internal Linear ticket reference in the comment on the
lease-related code with a public GitHub issue reference or remove the ticket
mention entirely. Update the comment near the lease field in the etcd transport
code, and also scan the surrounding lease-related comment(s) for any other
internal ticket IDs so they no longer reference DIS-2295 or similar internal
tracking numbers.

In `@tests/fault_tolerance/etcd_ha/test_vllm.py`:
- Line 454: Replace the internal Linear ticket marker in the added comment with
the matching public GitHub issue reference. Update the source comment/docstring
in the test_vllm-related change so it uses a public format like GH-NNNN or
`#NNNN`, and ensure any other added occurrences in the same diff are updated
consistently.
- Around line 519-530: The _freeze_vllm_engine_processes helper currently uses
pgrep -f EngineCore, which can match unrelated vLLM engines on the host; update
it to scope the search to this test worker’s process tree by starting from the
worker PID and only selecting descendant EngineCore processes. Keep the freezing
logic in _freeze_vllm_engine_processes, and switch the logger.info call to lazy
formatting so pid interpolation is deferred.
- Around line 499-516: The background helper in _run currently swallows all
exceptions and returns a thread that the test never validates, so the test can
pass even if the in-flight request ends early. Update the thread logic in _run
to catch only the expected request/connection-related exception(s), use lazy
logger formatting instead of an f-string, and make the failure visible to the
caller. Then, in the test flow that starts the thread and freezes the engine,
assert the thread is still alive before proceeding so the pinned request is
actually enforced; apply the same fix to the duplicated logic noted by the other
affected locations.
- Around line 576-594: The frozen vLLM engine cleanup is happening too late
because the outer finally around the whole test block only runs after the
worker/frontend/etcd context managers unwind. Move the SIGCONT/SIGKILL handling
for frozen_pids into an inner finally immediately after
_freeze_vllm_engine_processes() and before the teardown-prone
wait_for_processes_to_terminate flow, so the cleanup runs even if context
teardown hangs.
- Around line 534-538: Add the missing vLLM sizing markers to the test decorated
near the existing gpu_1/e2e/nightly marks in test_vllm.py so the scheduler can
account for its resource needs. Update the same test function that already has
FAULT_TOLERANCE_MODEL_NAME, timeout, and the module-level pytest.mark.vllm by
applying both pytest.mark.profiled_vram_gib(...) and
pytest.mark.requested_vllm_kv_cache_bytes(...) with the appropriate values used
elsewhere in this suite.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 122bc114-ab7e-4d6a-a585-07e985509247

📥 Commits

Reviewing files that changed from the base of the PR and between 4cba833 and 2954065.

📒 Files selected for processing (5)

• lib/runtime/src/pipeline/network/ingress/push_endpoint.rs
• lib/runtime/src/runtime.rs
• lib/runtime/src/transports/etcd.rs
• lib/runtime/src/transports/etcd/lease.rs
• tests/fault_tolerance/etcd_ha/test_vllm.py

[datadog-official]

🎯 Code Coverage (details)
• Patch Coverage: 28.17%
• Overall Coverage: 41.98% (-2.67%)

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 3ba5369 | Docs | Give us feedback!}

[nnshah1]

@CodeRabbit full review

[coderabbitai]

✅ Action performed

Full review finished.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/fault_tolerance/etcd_ha/test_vllm.py`:
- Around line 462-471: The parametrized E2E test is still using shared fixed
ports via FRONTEND_PORT and the default NatsServer/stack setup, which can
collide with xdist or other HA tests. Update the test helpers around
_zombie_verify_serving and the related frontend/worker startup paths to accept
injected ports, then switch to dynamic allocation using NatsServer(request,
port=0) and dynamo_dynamic_ports or allocate_port()/allocate_ports(). Thread the
allocated frontend and system/NATS ports through the request helpers and process
managers so the tcp/nats cases run in isolation without hardcoded port values.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 387dbdff-e6bc-483a-a9f8-17cef2472533

📥 Commits

Reviewing files that changed from the base of the PR and between 61865e9 and c272f57.

📒 Files selected for processing (7)

• lib/runtime/src/pipeline/network/ingress.rs
• lib/runtime/src/pipeline/network/ingress/push_endpoint.rs
• lib/runtime/src/pipeline/network/ingress/shared_tcp_endpoint.rs
• lib/runtime/src/runtime.rs
• lib/runtime/src/transports/etcd.rs
• lib/runtime/src/transports/etcd/lease.rs
• tests/fault_tolerance/etcd_ha/test_vllm.py

[kthui]

LGTM!

One limitation worth calling out: this assumes the stuck in-flight request affects only the Python/engine side and that the primary tokio executor remains healthy, so it can run the phased shutdown and enforce the drain timeout. If tokio itself is stalled, this change would not guarantee worker termination.

[nnshah1]

@kthui agreed, and good to call out. This fix targets the common case — a stuck request on the engine/native side while the tokio runtime stays healthy — so the phased shutdown + bounded drain can run. A fully-stalled tokio executor would defeat any cooperative mechanism here (including the existing exit(911) watchdog, which is also tokio-based); guaranteeing termination in that case needs an OS-thread watchdog or the k8s liveness probe, which I'd treat as a separate runtime-level change rather than fold into this drain fix. Noted the assumption in the PR description.

— Neelay + 🤖

[michaelfeil]

seems good on high level! Thanks for bounding the draining of inflight requests.

[michaelfeil]

actually, i think fully stuck tokio runtime is again OK, since in this case the k8s probe would not come though, etcd would loose the lease, and the pod would get cleared by k8s due to missing response, if the health check is e.g. async and requires tokio.