fix: runtime graceful shutdown

[michaelfeil]

Signed-off-by: Michael Feil 63565275+michaelfeil@users.noreply.github.com

Improve Graceful Shutdown

Summary

Improve graceful shutdown to remove workers from routing immediately while still completing accepted work within a bounded budget. This prevents routers from selecting draining workers and ensures active streaming requests complete before service shutdown.

Problem

Runtime graceful shutdown currently drains inflight work too late in two critical places:

1. Endpoint discovery stays published until the runtime tears down the main lease, so routers can still select a worker that has already entered drain
2. PushEndpoint stops the NATS service endpoint as soon as endpoint shutdown starts, before inflight streaming requests finish, causing active streams to receive ControlMessage::Stop and surface as cancelled or 500 responses

During rollout restart or pod termination, a worker should remove itself from routing quickly while still finishing work it has already accepted. We also have issue with orphaned pods. Deleting it from k8s cluster with force would leave the pod running. If something goes wrong in the worker (e.g. drain takes forever due to the casual deadlock of inference engine), the worker would hog gpu memory forever / until the cluster decides to delete worker (5h).

Changes

1. Unpublish Discovery at Drain Start

• Spawn a cleanup task tied to the endpoint shutdown token that immediately unregisters the endpoint from discovery
• Discovery removal happens as soon as endpoint shutdown starts, preventing routers from selecting the draining worker
• Cleanup is idempotent and uses the discovery abstraction (works for etcd, Kubernetes, file, and memory backends)
• Added comprehensive tests for discovery cleanup behavior

2. Drain Before Stopping PushEndpoint

• Changed shutdown ordering to wait for inflight requests to complete before stopping the NATS service
• Prevents active streaming contexts from receiving service stop while they are still healthy
• Respects the graceful_shutdown flag - non-graceful shutdowns maintain immediate stop behavior
• Added timeout support to prevent indefinite blocking
• Added tests verifying correct sequencing and timeout behavior

3. Add Global Runtime Shutdown Budget

• Wrapped Runtime Phase 2 in a configurable timeout (default: 870 seconds)
• Prevents one stuck stream, leaked guard, or broken endpoint from blocking termination indefinitely
• Timeout applies to the whole graceful shutdown window, not per-endpoint
• Configurable via DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS environment variable
• Added test for shutdown phases with timeout

Testing

• Discovery cleanup: Tests verify cleanup waits for shutdown token and properly unregisters endpoints
• PushEndpoint drain: Tests verify service stop happens only after inflight requests drain or timeout expires
• Runtime shutdown phases: Tests verify proper token cancellation sequencing with timeout

All tests use fast deterministic timeouts (milliseconds) and avoid external services.

Configuration

New environment variable:

• DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS: Maximum time in seconds the runtime waits for graceful endpoint drain before tearing down shared transports (default: 870)

Kubernetes Alignment

Recommended configuration:

terminationGracePeriodSeconds: 900
env:
  - name: DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS
    value: "870"

This leaves roughly 30 seconds for final cleanup, transport teardown, and process exit after the runtime stops waiting gracefully.

Related

• Implements the graceful shutdown improvements from plan_graceful.md
• Follows the upstreaming strategy of small, focused PRs for easier review and backports

Summary by CodeRabbit

• New Features

  • Configurable graceful shutdown timeout via env var
  • Runtime uses phased shutdown with a bounded graceful window
  • New utility to parse env vars with defaults

• Bug Fixes

  • Endpoints now drain inflight requests before stopping
  • Improved registration/unregistration lifecycle to ensure discovery cleanup during shutdown

• Documentation

  • Added graceful shutdown design and operational guidance

• Tests

  • Added unit tests covering shutdown, drain, and cleanup behaviors

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

👋 Hi michaelfeil! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

[coderabbitai]

Walkthrough

This PR implements runtime graceful shutdown: adds configurable graceful timeout, defers PushEndpoint stop until inflight requests drain, introduces discovery unregister lifecycle tied to endpoint shutdown, and orchestrates shutdown phases with a timeout.

Changes

Graceful Shutdown with Configurable Timeout

Layer / File(s)	Summary
Configuration and env parsing lib/runtime/src/config.rs, lib/runtime/src/config/environment_names.rs	Adds env_parse_or_default helper and tests; introduces DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS constant and updates related tests.
Runtime shutdown phases and timeout lib/runtime/src/runtime.rs	Extracts shutdown orchestration into run_shutdown_phases, adds runtime_graceful_shutdown_timeout() and default timeout constant, and wraps Phase 2 (graceful wait) in a timeout with warning on expiry.
Push Endpoint inflight drain and shutdown helpers lib/runtime/src/pipeline/network/ingress/push_endpoint.rs	Adds wait_for_inflight_requests and finish_shutdown_after_inflight_drain, changes cancellation handling to defer endpoint.stop() until after inflight drain when graceful, and adds tests for drain and timeout behavior.
Endpoint discovery registration & cleanup lib/runtime/src/component/endpoint.rs	Captures registered_instance from discovery.register, spawns cleanup task that waits for endpoint shutdown token then unregisters the instance via wait_for_endpoint_shutdown_and_cleanup, aborts/awaits cleanup on request-plane exit, and performs a final best-effort unregister; includes tests.
Design and plan document plan_graceful.md	Adds a design document describing the graceful shutdown contract, discovery ordering, drain semantics, global timeout, test strategy, and upstreaming plan.

---

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'fix: runtime graceful shutdown' directly aligns with the main objective to improve graceful shutdown behavior, capturing the core change without unnecessary details.
Description check	✅ Passed	The description comprehensively covers all required sections including overview, detailed problem/solution, specific changes, testing strategy, configuration, and related items, meeting template expectations.
Docstring Coverage	✅ Passed	Docstring coverage is 90.91% 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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

lib/runtime/src/pipeline/network/ingress/push_endpoint.rs (1)

63-77: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Prioritize cancellation before endpoint.next() in tokio::select!

With tokio::select! { biased; ... }, branches are polled top-to-bottom and the first ready branch in that order is selected—so if cancellation and endpoint.next() are both ready, the request branch can run for the iteration. Reorder the branches so shutdown always wins when cancellation is ready.

💡 Suggested change

         let req = tokio::select! {
             biased;
 
-            // await on service request
-            req = endpoint.next() => {
-                req
-            }
-
             // process shutdown
             _ = self.cancellation_token.cancelled() => {
                 tracing::info!(
                     "PushEndpoint received cancellation signal, stopping service after inflight requests drain"
                 );
                 stop_service_after_drain = true;
                 break;
             }
+
+            // await on service request
+            req = endpoint.next() => {
+                req
+            }
         };

🤖 Prompt for 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.

In `@lib/runtime/src/pipeline/network/ingress/push_endpoint.rs` around lines 63 -
77, The tokio::select! in PushEndpoint currently lists the endpoint.next()
branch before self.cancellation_token.cancelled(), so with biased selection a
ready cancellation can be lost; to fix, reorder the branches inside the select!
so the cancellation branch (awaiting self.cancellation_token.cancelled())
appears before the endpoint.next() branch (the req = endpoint.next() => { ... })
while keeping biased; and preserve the existing behavior of setting
stop_service_after_drain and breaking when cancellation fires.

🤖 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/runtime.rs`:
- Around line 327-335: The runtime_graceful_shutdown_timeout function silently
falls back to DEFAULT_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS when the env var
DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS is missing or unparsable; change it
to log invalid values: read the env var first, if missing keep current default
but emit a debug/info log, and if present but s.parse::<u64>() fails emit a
warning (including the variable name and the bad value) that you are falling
back to DEFAULT_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS; update
runtime_graceful_shutdown_timeout to perform these checks and use the crate's
logging/tracing facility so callers can see misconfigurations.

---

Outside diff comments:
In `@lib/runtime/src/pipeline/network/ingress/push_endpoint.rs`:
- Around line 63-77: The tokio::select! in PushEndpoint currently lists the
endpoint.next() branch before self.cancellation_token.cancelled(), so with
biased selection a ready cancellation can be lost; to fix, reorder the branches
inside the select! so the cancellation branch (awaiting
self.cancellation_token.cancelled()) appears before the endpoint.next() branch
(the req = endpoint.next() => { ... }) while keeping biased; and preserve the
existing behavior of setting stop_service_after_drain and breaking when
cancellation fires.

🪄 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: Pro

Run ID: ea7cc821-16ab-4c29-b597-4573755daef0

📥 Commits

Reviewing files that changed from the base of the PR and between c5ca451 and 5671771.

📒 Files selected for processing (4)

• lib/runtime/src/component/endpoint.rs
• lib/runtime/src/config/environment_names.rs
• lib/runtime/src/pipeline/network/ingress/push_endpoint.rs
• lib/runtime/src/runtime.rs

[michaelfeil]

@CodeRabbit can you perform a stricter review and find issues with it across the codebase?

[coderabbitai]

@michaelfeil: Absolutely — I’ll run a stricter full review across the PR and look for cross-codebase issues, especially around shutdown ordering, cancellation/token lifetimes, discovery cleanup idempotency, timeout behavior, and regression risks in active request handling.

✅ Actions performed

Full review triggered.

[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 `@plan_graceful.md`:
- Line 24: The plan document currently hardcodes the global timeout as "840
seconds" which is inconsistent with the PR's runtime contract; update every
occurrence (including the line with "Default global budget: 840 seconds." and
the other instances around the referenced sections) to the new default "870
seconds" and ensure any explanatory text about Kubernetes grace/teardown
references the 900s cluster grace and ~30s teardown accordingly.

🪄 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: Pro

Run ID: 1d95399d-641a-4c53-8d58-3b7cf33a43f8

📥 Commits

Reviewing files that changed from the base of the PR and between 5ec8eee and 1adf149.

📒 Files selected for processing (6)

• lib/runtime/src/component/endpoint.rs
• lib/runtime/src/config.rs
• lib/runtime/src/config/environment_names.rs
• lib/runtime/src/pipeline/network/ingress/push_endpoint.rs
• lib/runtime/src/runtime.rs
• plan_graceful.md

[kthui]

/ok to test 106a8f9

[kthui]

Thanks for implementing graceful shutdown logic at the Rust level! This will help making graceful shutdown behavior unified across backends.

LGTM overall!

Since both discovery and endpoint unregistration happens in parallel at the same time upon shutdown, it creates a ms interval where some request(s) may still be routed from the frontend to the worker. Would it be better if we wait for the full graceful shutdown timeout interval, before stopping the endpoint?

For instance:

1. unregister from discovery
2. wait for the full graceful shutdown timeout interval, while allowing existing/new requests to be served.
3. unregister the endpoint
4. drop all ongoing requests, if any
5. shutdown the worker

For requests that are sent after shutdown begins but before the frontend acknowledges the worker unregistration, they can still be processed normally as long as they complete within the graceful shutdown timeout interval, so we avoid failing any request due to shutdown.

The tradeoff is each shutdown will take longer, but still bounded by the graceful shutdown timeout interval.

[nnshah1]

Picking this up as #11068 — rebased onto post-#10705 main and extended to cover the etcd-lease-loss path.

Context: a production deployment hit a lease-loss "zombie worker" — a worker loses its etcd lease while a request is stuck in-flight (an engine that can't make progress and can't be aborted), the unbounded per-endpoint drain wedges, Runtime::shutdown() is never reached, and the process never exits. This PR's push_endpoint.rs bounded-drain is the right fix for the wedge, so I'd like to land it rather than start fresh.

Two adjustments vs this branch:

1. Drop the runtime.rs Phase-2 timeout + DYN_RUNTIME_GRACEFUL_SHUTDOWN_TIMEOUT_SECS — fix(runtime): bound graceful shutdown drain #10705 (merged) already shipped that and is the source of the current conflict. Keep only the push_endpoint.rs changes on top of it and reuse fix(runtime): bound graceful shutdown drain #10705's graceful_shutdown_timeout() (I made it pub(crate)).
2. Route lease loss through Runtime::shutdown() — we were calling a bare primary_token().cancel(), so the phased shutdown never ran.

I also added a GPU regression test that reproduces the zombie by freezing the vLLM engine mid-generation (a genuinely non-cancellable inflight) and asserting the worker still exits: without the bound it wedges (still running at 60s); with it, the drain times out and the worker exits (~27s).

Does this line up with the internal fix you mentioned — anything from your version that should fold in here?

— Neelay + 🤖