feat(backend): add worker drain Admin API

[xianlubird]

Background

This is PR 1 of the implementation plan in #11173. The design discussion concluded that lifecycle state must be owned by the worker rather than stored in each frontend replica.

This PR therefore removes the previous frontend-local prototype and introduces a worker-local Admin API. The existing SIGTERM graceful-shutdown flow remains the default fallback.

Summary

• Add worker-local Admin API endpoints on the system-status server:
  • POST /engine/drain
  • POST /engine/resume
  • GET /engine/status
• Add the lifecycle state machine serving -> draining -> drained.
• Stop new admissions before unregistering the worker from discovery.
• Track admitted requests until their response streams finish or are dropped.
• Report safe_to_delete: true only after the worker reaches drained.
• Require prefill workers to confirm that pending KV transfers are complete. If the backend cannot confirm this, the status remains fail-closed and never reports a false safe-to-delete state.
• Re-register the worker and reopen admission on resume.
• Coordinate Admin API operations with the existing SIGTERM shutdown orchestrator.
• Document the experimental API and status response.

Out of Scope

The following work remains in separate follow-up PRs:

• Kubernetes self-Pod annotation trigger
• Drain-completion support for additional backends
• Operator/Grove drain-wait-delete automation
• Frontend aggregation and CLI tooling

Validation

• cargo fmt --all -- --check
• cargo check -p dynamo-backend-common
• cargo test -p dynamo-backend-common --lib — 137 passed
• cargo clippy -p dynamo-backend-common --all-targets -- -D warnings
• python3 docs/fern/scripts/check_asset_paths.py
• git diff origin/main...HEAD --check

[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]

🌿 Fern Docs Preview: https://nvidia-preview-01a01784-401c-726d-8442-3d483bec501d.docs.buildwithfern.com/dynamo/dev

[devin-ai-integration]

Devin Review found 1 potential issue.

[coderabbitai]

Walkthrough

This PR adds manual worker drain/resume support end-to-end: the runtime routing client tracks a drained-id set alongside overloaded ids, WorkerSet stores a routing client, Model/ModelManager expose worker status and drain APIs, and a new admin HTTP router exposes GET /workers and POST drain/resume endpoints, with documentation updates.

Changes

Worker drain/resume feature

Layer / File(s)	Summary
Routing client drained-state support lib/runtime/src/component/client.rs	Adds drained_ids tracking to RoutingInstanceCounts/RoutingInstanceSnapshot/RoutingInstances, updates reconciliation and free-id derivation to exclude drained workers, adds Client::drain_instance, resume_instance, drained_instance_ids, routing_instance_snapshot, and new tests.
WorkerSet routing client wiring lib/llm/src/discovery/worker_set.rs, lib/llm/src/discovery/watcher.rs	WorkerSet gains an optional routing_client field and accessor; watcher.rs wires this field from the appropriate Client in each serving-engine branch (prefill, KV routing, embeddings, chat, completions, tensor, realtime).
Model and ModelManager worker status/drain APIs lib/llm/src/discovery.rs, lib/llm/src/discovery/model.rs, lib/llm/src/discovery/model_manager.rs	Adds WorkerServingState/WorkerStatus types, Model::worker_statuses/set_worker_drained using each WorkerSet's routing client snapshot, ModelManager::worker_statuses/set_worker_drained wrappers, and re-exports.
Worker admin HTTP router and docs lib/llm/src/http/service.rs, lib/llm/src/http/service/service_v2.rs, lib/llm/src/http/service/worker_admin.rs, docs/components/frontend/configuration.md	New worker_admin module registers GET /workers and POST /workers/{worker_id}/drain/resume, wired into service_v2 in place of busy_threshold; docs describe the new admin-gated endpoints.

Estimated code review effort: 3 (Moderate) | ~30 minutes

Related PRs: None identified.

Suggested labels: enhancement, http-api, runtime, discovery

Suggested reviewers: None identified.

🐰 A worker grows weary, needs a rest,
Drain the routes, let it decompress,
Resume the flow when it's fit once more,
Through Client, Model, and admin door,
Hop by hop, the state's addressed.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description is largely off-topic because it describes worker-local /engine lifecycle APIs, while this PR adds frontend /workers drain and resume APIs.	Rewrite the description to cover the frontend worker Admin API, worker status reporting, runtime drain state, documentation, and the required Related Issues section.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title clearly identifies the primary change: adding a backend worker-drain Admin API.

---

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

[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/llm/src/http/service/service_v2.rs (1)

994-999: 🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift

Protect the admin router with an auth boundary.

When DYN_ENABLE_FRONTEND_ADMIN_API is true, these drain/resume routes are merged into the same externally reachable frontend router as health/models and no auth or private-network boundary is added here. Any client that can reach the service can take workers out of rotation. Please wrap the admin routes in auth middleware or expose them on a separate private listener.

🤖 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/llm/src/http/service/service_v2.rs` around lines 994 - 999, The admin
routes added in the router assembly are exposed without any protection, so guard
the `busy_threshold_router` and `worker_admin_router` paths with an
auth/private-network boundary before pushing them into `system_routes`. Update
the `admin_api_enabled` branch in `service_v2.rs` to either wrap those routers
with the existing auth middleware used elsewhere in the service or move them
onto a separate private listener, and make sure the fix is applied around the
`busy_threshold::busy_threshold_router` and `worker_admin::worker_admin_router`
integration points.

🤖 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/component/client.rs`:
- Around line 264-270: `Client::overloaded_instance_ids()` is no longer
overload-only because it now includes manual drains via `excluded_ids()`, which
changes the public contract used by `routing_instance_counts().overloaded` and
`routing_instance_snapshot().overloaded_ids`. Keep `overloaded_ids()` limited to
true overload signals only, and introduce a separate accessor for
excluded/scheduler-facing IDs so drained workers are not reported as overloaded.
Update the relevant helper(s) in `Client` and any callers around
`routing_instance_counts`/`routing_instance_snapshot` to use the new accessor
where appropriate.

---

Outside diff comments:
In `@lib/llm/src/http/service/service_v2.rs`:
- Around line 994-999: The admin routes added in the router assembly are exposed
without any protection, so guard the `busy_threshold_router` and
`worker_admin_router` paths with an auth/private-network boundary before pushing
them into `system_routes`. Update the `admin_api_enabled` branch in
`service_v2.rs` to either wrap those routers with the existing auth middleware
used elsewhere in the service or move them onto a separate private listener, and
make sure the fix is applied around the `busy_threshold::busy_threshold_router`
and `worker_admin::worker_admin_router` integration points.

🪄 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: 220dc0d7-4b25-45b1-87fd-5c045e015c2b

📥 Commits

Reviewing files that changed from the base of the PR and between b4999c4 and 9c14b87.

📒 Files selected for processing (10)

• docs/components/frontend/configuration.md
• lib/llm/src/discovery.rs
• lib/llm/src/discovery/model.rs
• lib/llm/src/discovery/model_manager.rs
• lib/llm/src/discovery/watcher.rs
• lib/llm/src/discovery/worker_set.rs
• lib/llm/src/http/service.rs
• lib/llm/src/http/service/service_v2.rs
• lib/llm/src/http/service/worker_admin.rs
• lib/runtime/src/component/client.rs

[rmccorm4]

@tmonty12 @sttts can you help take a look here?

[tedzhouhk]

It looks like traffic removal before runtime shutdown is already implemented by the existing graceful-shutdown flow. On SIGTERM, a worker first unregisters its endpoints from discovery so routers stop assigning new requests, waits through the routing-convergence grace period, and then shuts down/drains in-flight work (see #6093 and #10705). With Grove OnDelete, manually deleting the Pod should still trigger that same SIGTERM sequence.

Could you clarify which remaining failure mode requires this additional frontend-side drain API? If the intended distinction is administrator-controlled quiescing before initiating Pod deletion, it would be helpful to state that explicitly and explain why the existing discovery unregister path is insufficient, especially since this PR's drain state is frontend-local and in-memory.

[sttts]

Re-review of the worker drain lifecycle at head 4ba7fb8. The earlier method, cancellation-ownership, and generation-fence findings are fixed; these three availability issues remain.

[jthomson04]

Source-only re-review at head 4ba7fb8. These comments are distinct from the existing availability findings and are submitted as a non-blocking review.

[dynamo-ops]

/ok to test bce4014

[dynamo-ops]

/ok to test e8fedab

[dynamo-ops]

/ok to test 4611297

[jthomson04]

Source-only review of the current PR head.

[tmonty12]

Can we simplify the status response by removing both kv_transfers and safe_to_delete?

There is currently no uniform Dynamo-facing engine API that reports KV transfers independently from request activity:

• vLLM internally knows about pending connector work, but its available work/quiescence signal combines requests and pending KV work.
• SGLang internally tracks explicit prefill and decode transfer queues, but Dynamo does not currently wire that state into is_quiescent().
• TensorRT-LLM tracks outgoing transfers in its internal AsyncTransferManager, but its exposed statistics only report active and queued requests. That transfer-manager state is not exposed through the normal LLM API.

The engines are aware of transfer progress, but Dynamo cannot currently consume a consistent transfer-only state. Consequently, no production engine implements is_quiescent(), prefill workers report kv_transfers: unknown, and Admin drain never transitions them to drained.

kv_transfers also does not make sense for aggregated workers, where it is always not_applicable. This leaks a prefill-specific implementation detail into the common worker lifecycle schema without giving consumers actionable information.

safe_to_delete is redundant because it is defined exactly as state == drained. Consumers can derive deletion safety directly from state, avoiding duplicated state that could drift.

As a follow-up, each engine adapter should implement is_quiescent() using the strongest engine-owned work signal available. The probe should determine whether all work is complete—not merely whether Dynamo’s request count is zero—and must include requests retained for active or pending KV transfers:

• vLLM: scheduler/connector pending-work state.
• SGLang: is_fully_idle() or its request and disaggregation-transfer queues.
• TensorRT-LLM: active and queued requests plus an exposed AsyncTransferManager pending-transfer count.

The worker can then publish drained only when:

dynamo_inflight_requests == 0
AND engine_pending_work_including_kv_transfers == 0

Until those engine integrations exist, the public status should stick to state and inflight_requests rather than exposing transfer state that cannot yet be populated reliably.

[dynamo-ops]

/ok to test 644c419

[dynamo-ops]

/ok to test e8a7b94

[xianlubird]

Thanks, agreed. I simplified the public status in 644c419 by removing both kv_transfers and safe_to_delete, and updated the docs accordingly. state == drained is now the canonical deletion-safety signal. A disaggregated prefill worker stays draining until its backend can confirm quiescence; the backend-specific adapters can add those stronger checks in follow-up PRs.

[dynamo-ops]

/ok to test c69da7e

[dynamo-ops]

/ok to test b39e335