feat(backend): add worker drain Admin API

docs/fern/index.yml

@@ -1112,6 +1112,8 @@ navigation:
             path: pages/reference/components/tls-configuration.mdx
           - page: Frontend Configuration
             path: pages/reference/components/frontend-configuration.mdx
+          - page: Worker Admin API
+            path: pages/reference/components/worker-admin-api.mdx
           - page: Planner Configuration
             path: pages/reference/components/planner-configuration.mdx
           - page: Profiler Configuration

docs/fern/pages/reference/components/worker-admin-api.mdx

@@ -0,0 +1,54 @@
+---
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+title: Worker Admin API
+subtitle: Worker-local lifecycle endpoints for draining, resuming, and checking worker state
+---
+
+**Experimental.** Unified-backend workers expose lifecycle controls through the worker system-status
+server. Set `DYN_SYSTEM_PORT` to zero or a positive port to enable this server.
+
+<Warning>
+The system-status server does not authenticate Admin API requests. Restrict the system port to the
+operator network.
+</Warning>
+
+## Endpoints
+
+| Method | Path | Behavior |
+| --- | --- | --- |
+| `POST` | `/engine/drain` | Unregister the worker, wait for discovery convergence, then stop new request admission and drain admitted work |
+| `POST` | `/engine/resume` | Resume request admission, then re-register the worker in discovery |
+| `GET` | `/engine/status` | Return the worker lifecycle status |
+
+Send an empty body or an empty JSON object. Each endpoint returns the same status shape:
+
+```json
+{
+  "state": "draining",
+  "inflight_requests": 2,
+  "discovery_registered": false
+}
+```
+
+`state` is one of `serving`, `draining`, `drained`, or `stopping`. Delete a worker only when its state
+is `drained`. An aggregated worker reaches `drained` after all admitted requests finish. A
+disaggregated prefill worker also requires backend quiescence support. Until its backend can confirm
+quiescence, it remains `draining`.
+
+While discovery converges, a draining worker continues to accept requests that a frontend already
+selected. After the configured graceful-shutdown grace period, the worker rejects late admissions
+with a worker-draining signal that makes the frontend reselect another worker independently of the
+configured migration budget, then waits for all admitted work to finish.
+
+The existing `SIGTERM` graceful-shutdown path remains the default fallback and uses the same
+discovery-first convergence ordering. The Admin API does not delete the worker or replace the
+shutdown timeout.
+
+## Example
+
+```bash
+curl --fail-with-body -X POST http://worker-host:8081/engine/drain
+curl --fail-with-body http://worker-host:8081/engine/status
+curl --fail-with-body -X POST http://worker-host:8081/engine/resume
+```

lib/backend-common/src/adapter.rs

@@ -30,6 +30,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 use crate::disagg::DisaggregationMode;
 use crate::engine::{GenerateContext, LLMEngine, RawEngine};
+use crate::lifecycle::RequestTracker;
 
 /// Test-only override count. Compiled out of release builds — tests acquire
 /// an `OtlpExportOverride` RAII guard to force-enable the recording
@@ -135,11 +136,25 @@ impl Drop for CancelMonitorGuard {
 pub(crate) struct EngineAdapter {
     engine: Arc<dyn LLMEngine>,
     mode: DisaggregationMode,
+    request_tracker: Arc<RequestTracker>,
 }
 
 impl EngineAdapter {
+    #[cfg(test)]
     pub(crate) fn new(engine: Arc<dyn LLMEngine>, mode: DisaggregationMode) -> Self {
-        Self { engine, mode }
+        Self::with_request_tracker(engine, mode, RequestTracker::new())
+    }
+
+    pub(crate) fn with_request_tracker(
+        engine: Arc<dyn LLMEngine>,
+        mode: DisaggregationMode,
+        request_tracker: Arc<RequestTracker>,
+    ) -> Self {
+        Self {
+            engine,
+            mode,
+            request_tracker,
+        }
     }
 }
 
@@ -201,6 +216,7 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
         &self,
         input: SingleIn<PreprocessedRequest>,
     ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
+        let request_guard = self.request_tracker.try_acquire()?;
         let (request, handle) = input.into_parts();
         let ctx: Arc<dyn AsyncEngineContext> = handle.context();
 
@@ -369,6 +385,7 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
         let is_handoff_terminal_mode = self.mode.is_prefill() || self.mode.is_encode();
         let finalizer_span = span.clone();
         let mapped = async_stream::stream! {
+            let _request_guard = request_guard;
             let _guard = guard;
             let finalizer = StreamSpanFinalizer::new(finalizer_span);
             let mut inner = chunks;
@@ -494,11 +511,23 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutpu
 /// `JsonProbeAdapter` wrapper needed.
 pub(crate) struct RawEngineAdapter {
     engine: Arc<dyn RawEngine>,
+    request_tracker: Arc<RequestTracker>,
 }
 
 impl RawEngineAdapter {
+    #[cfg(test)]
     pub(crate) fn new(engine: Arc<dyn RawEngine>) -> Self {
-        Self { engine }
+        Self::with_request_tracker(engine, RequestTracker::new())
+    }
+
+    pub(crate) fn with_request_tracker(
+        engine: Arc<dyn RawEngine>,
+        request_tracker: Arc<RequestTracker>,
+    ) -> Self {
+        Self {
+            engine,
+            request_tracker,
+        }
     }
 }
 
@@ -510,6 +539,7 @@ impl AsyncEngine<SingleIn<serde_json::Value>, ManyOut<Annotated<serde_json::Valu
         &self,
         input: SingleIn<serde_json::Value>,
     ) -> Result<ManyOut<Annotated<serde_json::Value>>, Error> {
+        let request_guard = self.request_tracker.try_acquire()?;
         let (request, handle) = input.into_parts();
         let ctx: Arc<dyn AsyncEngineContext> = handle.context();
 
@@ -569,6 +599,7 @@ impl AsyncEngine<SingleIn<serde_json::Value>, ManyOut<Annotated<serde_json::Valu
         let stream_ctx = ctx.clone();
         let finalizer_span = span.clone();
         let mapped = async_stream::stream! {
+            let _request_guard = request_guard;
             let _guard = guard;
             let finalizer = StreamSpanFinalizer::new(finalizer_span);
             let mut inner = chunks;
@@ -713,6 +744,29 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn adapter_tracks_request_until_response_stream_is_dropped() {
+        let (engine, _) = MockEngine::new(vec![chunk::token(11)]);
+        let tracker = RequestTracker::new();
+        let adapter = EngineAdapter::with_request_tracker(
+            engine,
+            DisaggregationMode::Aggregated,
+            Arc::clone(&tracker),
+        );
+
+        let input = Context::new(make_request(vec![1]));
+        let stream = adapter.generate(input).await.unwrap();
+        assert_eq!(tracker.inflight(), 1);
+
+        tracker.stop_accepting();
+        let rejected = adapter.generate(Context::new(make_request(vec![2]))).await;
+        assert!(rejected.is_err());
+        assert_eq!(tracker.inflight(), 1);
+
+        drop(stream);
+        assert_eq!(tracker.inflight(), 0);
+    }
+
     #[tokio::test]
     async fn adapter_cancellation_triggers_engine_abort() {
         let engine = Arc::new(MockEngine {

lib/backend-common/src/lib.rs

@@ -17,6 +17,7 @@ pub mod args;
 pub mod disagg;
 pub mod engine;
 pub mod error;
+mod lifecycle;
 pub mod metrics;
 mod publisher;
 mod rl;