From 1fbd6a59e5d480f11fe43b47ec436c6dc97f26e0 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Wed, 5 Aug 2026 10:49:45 +0800 Subject: [PATCH 1/8] feat(backend): add worker drain Admin API Signed-off-by: xianlubird --- docs/fern/index.yml | 2 + .../reference/components/worker-admin-api.mdx | 51 +++ lib/backend-common/src/adapter.rs | 58 ++- lib/backend-common/src/lib.rs | 1 + lib/backend-common/src/lifecycle.rs | 429 ++++++++++++++++++ lib/backend-common/src/worker.rs | 43 +- 6 files changed, 574 insertions(+), 10 deletions(-) create mode 100644 docs/fern/pages/reference/components/worker-admin-api.mdx create mode 100644 lib/backend-common/src/lifecycle.rs diff --git a/docs/fern/index.yml b/docs/fern/index.yml index 49afa031a76a..5dd4444db6f1 100644 --- a/docs/fern/index.yml +++ b/docs/fern/index.yml @@ -997,6 +997,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 diff --git a/docs/fern/pages/reference/components/worker-admin-api.mdx b/docs/fern/pages/reference/components/worker-admin-api.mdx new file mode 100644 index 000000000000..f001f722e515 --- /dev/null +++ b/docs/fern/pages/reference/components/worker-admin-api.mdx @@ -0,0 +1,51 @@ +--- +# 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 deletion safety +--- + +**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. + + +The system-status server does not authenticate Admin API requests. Restrict the system port to the +operator network. + + +## Endpoints + +| Method | Path | Behavior | +| --- | --- | --- | +| `POST` | `/engine/drain` | Stop new request admission, unregister the worker from discovery, and begin draining admitted work | +| `POST` | `/engine/resume` | Re-register the worker in discovery and resume request admission | +| `GET` | `/engine/status` | Return the worker lifecycle and deletion-safety status | + +Send an empty body or an empty JSON object. Each endpoint returns the same status shape: + +```json +{ + "state": "draining", + "inflight_requests": 2, + "kv_transfers": "not_applicable", + "safe_to_delete": false, + "discovery_registered": false +} +``` + +`state` is one of `serving`, `draining`, `drained`, or `stopping`. Delete a worker only when +`safe_to_delete` is `true`. For a disaggregated prefill worker, the worker reports `drained` only when +the backend confirms that both admitted requests and pending KV transfers are complete. If the +backend cannot confirm KV-transfer completion, `kv_transfers` remains `unknown` and +`safe_to_delete` remains `false`. + +The existing `SIGTERM` graceful-shutdown path remains the default fallback. 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 +``` diff --git a/lib/backend-common/src/adapter.rs b/lib/backend-common/src/adapter.rs index 30248ff7acd8..b9fc8795f58e 100644 --- a/lib/backend-common/src/adapter.rs +++ b/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, mode: DisaggregationMode, + request_tracker: Arc, } impl EngineAdapter { + #[cfg(test)] pub(crate) fn new(engine: Arc, mode: DisaggregationMode) -> Self { - Self { engine, mode } + Self::with_request_tracker(engine, mode, RequestTracker::new()) + } + + pub(crate) fn with_request_tracker( + engine: Arc, + mode: DisaggregationMode, + request_tracker: Arc, + ) -> Self { + Self { + engine, + mode, + request_tracker, + } } } @@ -201,6 +216,7 @@ impl AsyncEngine, ManyOut, ) -> Result>, Error> { + let request_guard = self.request_tracker.try_acquire()?; let (request, handle) = input.into_parts(); let ctx: Arc = handle.context(); @@ -369,6 +385,7 @@ impl AsyncEngine, ManyOut, ManyOut, + request_tracker: Arc, } impl RawEngineAdapter { + #[cfg(test)] pub(crate) fn new(engine: Arc) -> Self { - Self { engine } + Self::with_request_tracker(engine, RequestTracker::new()) + } + + pub(crate) fn with_request_tracker( + engine: Arc, + request_tracker: Arc, + ) -> Self { + Self { + engine, + request_tracker, + } } } @@ -510,6 +539,7 @@ impl AsyncEngine, ManyOut, ) -> Result>, Error> { + let request_guard = self.request_tracker.try_acquire()?; let (request, handle) = input.into_parts(); let ctx: Arc = handle.context(); @@ -569,6 +599,7 @@ impl AsyncEngine, ManyOut Self { + match value { + 0 => Self::Serving, + 1 => Self::Draining, + 2 => Self::Drained, + 3 => Self::Stopping, + _ => unreachable!("invalid worker lifecycle state {value}"), + } + } +} + +#[repr(u8)] +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +enum KvTransferState { + NotApplicable = 0, + Unknown = 1, + Pending = 2, + Complete = 3, +} + +impl KvTransferState { + fn from_u8(value: u8) -> Self { + match value { + 0 => Self::NotApplicable, + 1 => Self::Unknown, + 2 => Self::Pending, + 3 => Self::Complete, + _ => unreachable!("invalid KV transfer state {value}"), + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +pub(crate) struct WorkerLifecycleStatus { + state: WorkerLifecycleState, + inflight_requests: u64, + kv_transfers: KvTransferState, + safe_to_delete: bool, + discovery_registered: bool, + #[serde(skip_serializing_if = "Option::is_none")] + last_error: Option, +} + +/// Shared request-admission and in-flight tracker. +/// +/// The double-check around `fetch_add` closes the drain/admission race: a +/// request is either rejected after draining starts or counted so the drain +/// monitor waits for its response stream to be dropped. +#[derive(Debug)] +pub(crate) struct RequestTracker { + accepting: AtomicBool, + inflight: AtomicU64, + changed: Notify, +} + +impl RequestTracker { + pub(crate) fn new() -> Arc { + Arc::new(Self { + accepting: AtomicBool::new(true), + inflight: AtomicU64::new(0), + changed: Notify::new(), + }) + } + + pub(crate) fn try_acquire(self: &Arc) -> Result { + if !self.accepting.load(Ordering::Acquire) { + bail!("worker is not accepting new requests"); + } + + self.inflight.fetch_add(1, Ordering::AcqRel); + if !self.accepting.load(Ordering::Acquire) { + self.release(); + bail!("worker started draining before request admission completed"); + } + + Ok(RequestGuard { + tracker: Arc::clone(self), + }) + } + + pub(crate) fn stop_accepting(&self) { + self.accepting.store(false, Ordering::Release); + self.changed.notify_waiters(); + } + + fn start_accepting(&self) { + self.accepting.store(true, Ordering::Release); + self.changed.notify_waiters(); + } + + pub(crate) fn inflight(&self) -> u64 { + self.inflight.load(Ordering::Acquire) + } + + fn release(&self) { + let previous = self.inflight.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "request tracker underflow"); + if previous == 1 { + self.changed.notify_waiters(); + } + } +} + +pub(crate) struct RequestGuard { + tracker: Arc, +} + +impl Drop for RequestGuard { + fn drop(&mut self) { + self.tracker.release(); + } +} + +/// Coordinates the Admin API, request admission, discovery, and SIGTERM. +pub(crate) struct WorkerLifecycleController { + state: AtomicU8, + kv_transfers: AtomicU8, + discovery_registered: AtomicBool, + last_error: RwLock>, + operation_lock: Mutex<()>, + drain_generation: AtomicU64, + tracker: Arc, + endpoint: Endpoint, + engine: EngineKind, + mode: DisaggregationMode, +} + +impl WorkerLifecycleController { + pub(crate) fn new( + tracker: Arc, + endpoint: Endpoint, + engine: EngineKind, + mode: DisaggregationMode, + ) -> Arc { + Arc::new(Self { + state: AtomicU8::new(WorkerLifecycleState::Serving as u8), + kv_transfers: AtomicU8::new(Self::initial_kv_state(mode) as u8), + discovery_registered: AtomicBool::new(true), + last_error: RwLock::new(None), + operation_lock: Mutex::new(()), + drain_generation: AtomicU64::new(0), + tracker, + endpoint, + engine, + mode, + }) + } + + fn initial_kv_state(mode: DisaggregationMode) -> KvTransferState { + if mode.is_prefill() { + KvTransferState::Unknown + } else { + KvTransferState::NotApplicable + } + } + + fn state(&self) -> WorkerLifecycleState { + WorkerLifecycleState::from_u8(self.state.load(Ordering::Acquire)) + } + + pub(crate) fn status(&self) -> WorkerLifecycleStatus { + let state = self.state(); + WorkerLifecycleStatus { + state, + inflight_requests: self.tracker.inflight(), + kv_transfers: KvTransferState::from_u8(self.kv_transfers.load(Ordering::Acquire)), + safe_to_delete: state == WorkerLifecycleState::Drained, + discovery_registered: self.discovery_registered.load(Ordering::Acquire), + last_error: self.last_error.read().clone(), + } + } + + pub(crate) fn register_admin_routes(self: &Arc, registry: &EngineRouteRegistry) { + registry.register( + "drain", + lifecycle_callback(Arc::clone(self), LifecycleAction::Drain), + ); + registry.register( + "resume", + lifecycle_callback(Arc::clone(self), LifecycleAction::Resume), + ); + registry.register( + "status", + lifecycle_callback(Arc::clone(self), LifecycleAction::Status), + ); + tracing::info!("registered worker Admin API routes under /engine"); + } + + pub(crate) async fn drain(self: &Arc) -> Result { + let operation = self.operation_lock.lock().await; + let generation = match self.state() { + WorkerLifecycleState::Serving => { + self.last_error.write().take(); + self.tracker.stop_accepting(); + self.kv_transfers + .store(Self::initial_kv_state(self.mode) as u8, Ordering::Release); + self.state + .store(WorkerLifecycleState::Draining as u8, Ordering::Release); + Some(self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1) + } + WorkerLifecycleState::Draining => None, + WorkerLifecycleState::Drained => return Ok(self.status()), + WorkerLifecycleState::Stopping => bail!("worker shutdown is already in progress"), + }; + + if self.discovery_registered.load(Ordering::Acquire) { + if let Err(error) = self.endpoint.unregister_endpoint_instance().await { + let message = format!("failed to unregister worker from discovery: {error}"); + *self.last_error.write() = Some(message.clone()); + self.drain_generation.fetch_add(1, Ordering::AcqRel); + self.state + .store(WorkerLifecycleState::Serving as u8, Ordering::Release); + self.tracker.start_accepting(); + return Err(error).context("failed to unregister worker from discovery"); + } + self.discovery_registered.store(false, Ordering::Release); + } + + drop(operation); + if let Some(generation) = generation { + let controller = Arc::clone(self); + tokio::spawn(async move { controller.monitor_drain(generation).await }); + } + Ok(self.status()) + } + + pub(crate) async fn resume(self: &Arc) -> Result { + let _operation = self.operation_lock.lock().await; + let previous_state = match self.state() { + WorkerLifecycleState::Serving => return Ok(self.status()), + WorkerLifecycleState::Stopping => bail!("worker shutdown is already in progress"), + state @ (WorkerLifecycleState::Draining | WorkerLifecycleState::Drained) => state, + }; + + let generation = self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1; + self.state + .store(WorkerLifecycleState::Draining as u8, Ordering::Release); + + if !self.discovery_registered.load(Ordering::Acquire) { + if let Err(error) = self.endpoint.register_endpoint_instance().await { + let message = format!("failed to re-register worker in discovery: {error}"); + *self.last_error.write() = Some(message); + self.state.store(previous_state as u8, Ordering::Release); + if previous_state == WorkerLifecycleState::Draining { + let controller = Arc::clone(self); + tokio::spawn(async move { controller.monitor_drain(generation).await }); + } + return Err(error).context("failed to re-register worker in discovery"); + } + self.discovery_registered.store(true, Ordering::Release); + } + + self.last_error.write().take(); + self.kv_transfers + .store(Self::initial_kv_state(self.mode) as u8, Ordering::Release); + self.state + .store(WorkerLifecycleState::Serving as u8, Ordering::Release); + self.tracker.start_accepting(); + Ok(self.status()) + } + + /// Move into the irreversible SIGTERM path and ensure discovery is down. + pub(crate) async fn begin_shutdown(&self) { + let _operation = self.operation_lock.lock().await; + self.tracker.stop_accepting(); + self.drain_generation.fetch_add(1, Ordering::AcqRel); + self.state + .store(WorkerLifecycleState::Stopping as u8, Ordering::Release); + + if self.discovery_registered.load(Ordering::Acquire) { + if let Err(error) = self.endpoint.unregister_endpoint_instance().await { + tracing::warn!(%error, "discovery unregister failed during shutdown"); + *self.last_error.write() = Some(error.to_string()); + } else { + self.discovery_registered.store(false, Ordering::Release); + } + } + } + + async fn monitor_drain(self: Arc, generation: u64) { + loop { + if self.state() != WorkerLifecycleState::Draining + || self.drain_generation.load(Ordering::Acquire) != generation + { + return; + } + + if self.tracker.inflight() == 0 { + let kv_complete = if self.mode.is_prefill() { + match self.engine.is_quiescent().await { + Ok(Some(true)) => { + self.kv_transfers + .store(KvTransferState::Complete as u8, Ordering::Release); + true + } + Ok(Some(false)) => { + self.kv_transfers + .store(KvTransferState::Pending as u8, Ordering::Release); + false + } + Ok(None) => { + self.kv_transfers + .store(KvTransferState::Unknown as u8, Ordering::Release); + false + } + Err(error) => { + self.kv_transfers + .store(KvTransferState::Unknown as u8, Ordering::Release); + *self.last_error.write() = Some(error.to_string()); + false + } + } + } else { + true + }; + + if kv_complete + && self.drain_generation.load(Ordering::Acquire) == generation + && self + .state + .compare_exchange( + WorkerLifecycleState::Draining as u8, + WorkerLifecycleState::Drained as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + tracing::info!("worker drained and is safe to delete"); + return; + } + } + + tokio::select! { + _ = self.tracker.changed.notified() => {} + _ = tokio::time::sleep(QUIESCENCE_POLL_INTERVAL) => {} + } + } + } +} + +#[derive(Clone, Copy)] +enum LifecycleAction { + Drain, + Resume, + Status, +} + +fn lifecycle_callback( + controller: Arc, + action: LifecycleAction, +) -> EngineRouteCallback { + Arc::new(move |body| { + let controller = Arc::clone(&controller); + Box::pin(async move { + if !body.is_object() { + bail!("worker Admin API request body must be a JSON object"); + } + let status = match action { + LifecycleAction::Drain => controller.drain().await?, + LifecycleAction::Resume => controller.resume().await?, + LifecycleAction::Status => controller.status(), + }; + serde_json::to_value(status).context("serialize worker lifecycle status") + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn request_tracker_rejects_after_drain_starts() { + let tracker = RequestTracker::new(); + let guard = tracker.try_acquire().unwrap(); + assert_eq!(tracker.inflight(), 1); + + tracker.stop_accepting(); + assert!(tracker.try_acquire().is_err()); + assert_eq!(tracker.inflight(), 1); + + drop(guard); + assert_eq!(tracker.inflight(), 0); + } + + #[test] + fn request_tracker_accepts_again_after_resume() { + let tracker = RequestTracker::new(); + tracker.stop_accepting(); + tracker.start_accepting(); + let guard = tracker.try_acquire().unwrap(); + assert_eq!(tracker.inflight(), 1); + drop(guard); + } +} diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 82403f7e219f..fb4b334b38e7 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -34,6 +34,7 @@ use crate::engine::{ EngineConfig, KvEventSource, LLMEngine, MetricsBindings, MetricsCtx, RawEngine, }; use crate::error::{BackendError, DynamoError, ErrorType}; +use crate::lifecycle::{RequestTracker, WorkerLifecycleController}; use crate::publisher::{PublisherHandles, setup_publishers}; /// Default grace-period in seconds between discovery unregister and engine drain. @@ -281,7 +282,7 @@ impl EngineKind { } /// See [`LLMEngine::is_quiescent`]. - async fn is_quiescent(&self) -> Result, DynamoError> { + pub(crate) async fn is_quiescent(&self) -> Result, DynamoError> { match self { EngineKind::Llm(e) => e.is_quiescent().await, EngineKind::Raw(e) => e.is_quiescent().await, @@ -616,7 +617,7 @@ impl Worker { // alive. if shutdown.is_cancelled() { tracing::info!("Shutdown signal observed during engine.start(); running orchestrator"); - self.orchestrator_steps(&endpoint).await; + self.orchestrator_steps(&endpoint, None).await; return Ok(()); } @@ -761,8 +762,14 @@ impl Worker { /// Full graceful-shutdown orchestrator: discovery unregister → /// grace period → engine drain → cleanup. Shared by every shutdown path — /// pre-serve (mid-start signal) and the serve loop's signal arm. - async fn orchestrator_steps(&mut self, endpoint: &dynamo_runtime::component::Endpoint) { - if let Err(e) = endpoint.unregister_endpoint_instance().await { + async fn orchestrator_steps( + &mut self, + endpoint: &dynamo_runtime::component::Endpoint, + lifecycle: Option<&WorkerLifecycleController>, + ) { + if let Some(lifecycle) = lifecycle { + lifecycle.begin_shutdown().await; + } else if let Err(e) = endpoint.unregister_endpoint_instance().await { tracing::warn!(error = %e, "discovery unregister failed"); } else { tracing::info!("Endpoint unregistered from discovery"); @@ -894,14 +901,23 @@ impl Worker { // `serde_json::Value` probe surface; the raw pipeline // (`RawEngineAdapter`) is already JSON-shaped, so it serves as its // own probe. The tuple annotation drives the trait-object coercions. + let request_tracker = RequestTracker::new(); + let lifecycle = WorkerLifecycleController::new( + Arc::clone(&request_tracker), + endpoint.clone(), + self.engine.clone(), + self.config.disaggregation_mode, + ); + let (ingress, probe_engine): ( Arc, dynamo_runtime::local_endpoint_registry::LocalAsyncEngine, ) = match &self.engine { EngineKind::Llm(engine) => { - let engine_adapter = Arc::new(EngineAdapter::new( + let engine_adapter = Arc::new(EngineAdapter::with_request_tracker( engine.clone(), self.config.disaggregation_mode, + Arc::clone(&request_tracker), )); let ingress = Ingress::for_engine(engine_adapter.clone()).map_err(|e| { err( @@ -913,7 +929,10 @@ impl Worker { (ingress, probe) } EngineKind::Raw(engine) => { - let raw_adapter = Arc::new(RawEngineAdapter::new(engine.clone())); + let raw_adapter = Arc::new(RawEngineAdapter::with_request_tracker( + engine.clone(), + Arc::clone(&request_tracker), + )); let ingress = Ingress::for_engine(raw_adapter.clone()).map_err(|e| { err( ErrorType::Backend(BackendError::Unknown), @@ -979,7 +998,15 @@ impl Worker { ) })?; } - let serve_fut = builder.start(); + let started_endpoint = builder.start_with_registration().await.map_err(|e| { + err( + ErrorType::Backend(BackendError::Unknown), + format!("serve registration: {e}"), + ) + })?; + lifecycle.register_admin_routes(endpoint.drt().engine_routes()); + + let serve_fut = started_endpoint.wait(); tokio::pin!(serve_fut); tokio::select! { @@ -1008,7 +1035,7 @@ impl Worker { } } - self.orchestrator_steps(&endpoint).await; + self.orchestrator_steps(&endpoint, Some(&lifecycle)).await; Ok(()) } From cbb9476179d1523a6d3405b35d2a64db28687635 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Tue, 11 Aug 2026 10:25:16 +0800 Subject: [PATCH 2/8] fix(backend): harden worker drain lifecycle Signed-off-by: xianlubird --- .../reference/components/worker-admin-api.mdx | 6 +- lib/backend-common/src/lifecycle.rs | 460 +++++++++++++++--- lib/backend-common/src/worker.rs | 1 + lib/runtime/src/engine_routes.rs | 67 ++- lib/runtime/src/system_status_server.rs | 103 +++- 5 files changed, 545 insertions(+), 92 deletions(-) diff --git a/docs/fern/pages/reference/components/worker-admin-api.mdx b/docs/fern/pages/reference/components/worker-admin-api.mdx index f001f722e515..05e792d2e737 100644 --- a/docs/fern/pages/reference/components/worker-admin-api.mdx +++ b/docs/fern/pages/reference/components/worker-admin-api.mdx @@ -17,7 +17,7 @@ operator network. | Method | Path | Behavior | | --- | --- | --- | -| `POST` | `/engine/drain` | Stop new request admission, unregister the worker from discovery, and begin draining admitted work | +| `POST` | `/engine/drain` | Unregister the worker, wait for discovery convergence, then stop new request admission and drain admitted work | | `POST` | `/engine/resume` | Re-register the worker in discovery and resume request admission | | `GET` | `/engine/status` | Return the worker lifecycle and deletion-safety status | @@ -39,6 +39,10 @@ the backend confirms that both admitted requests and pending KV transfers are co backend cannot confirm KV-transfer completion, `kv_transfers` remains `unknown` and `safe_to_delete` remains `false`. +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 retryable worker-scoped error and waits for all admitted work to finish. + The existing `SIGTERM` graceful-shutdown path remains the default fallback. The Admin API does not delete the worker or replace the shutdown timeout. diff --git a/lib/backend-common/src/lifecycle.rs b/lib/backend-common/src/lifecycle.rs index ec23199d83ed..a0c8fed077d5 100644 --- a/lib/backend-common/src/lifecycle.rs +++ b/lib/backend-common/src/lifecycle.rs @@ -7,12 +7,16 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering}; use std::time::Duration; -use anyhow::{Context, Result, bail}; +use anyhow::{Context, Result, anyhow, bail}; +use async_trait::async_trait; use dynamo_runtime::component::Endpoint; -use dynamo_runtime::engine_routes::{EngineRouteCallback, EngineRouteRegistry}; +use dynamo_runtime::engine_routes::{ + EngineRouteCallback, EngineRouteMethod, EngineRouteRegistry, +}; +use dynamo_runtime::error::{DynamoError, ErrorType}; use parking_lot::RwLock; use serde::Serialize; -use tokio::sync::{Mutex, Notify}; +use tokio::sync::{Mutex, Notify, oneshot}; use crate::disagg::DisaggregationMode; use crate::worker::EngineKind; @@ -74,6 +78,35 @@ pub(crate) struct WorkerLifecycleStatus { last_error: Option, } +#[async_trait] +trait DiscoveryRegistration: Send + Sync { + async fn unregister(&self) -> Result<()>; + async fn register(&self) -> Result<()>; +} + +#[async_trait] +impl DiscoveryRegistration for Endpoint { + async fn unregister(&self) -> Result<()> { + self.unregister_endpoint_instance().await + } + + async fn register(&self) -> Result<()> { + self.register_endpoint_instance().await + } +} + +#[async_trait] +trait QuiescenceCheck: Send + Sync { + async fn is_quiescent(&self) -> Result>; +} + +#[async_trait] +impl QuiescenceCheck for EngineKind { + async fn is_quiescent(&self) -> Result> { + Ok(EngineKind::is_quiescent(self).await?) + } +} + /// Shared request-admission and in-flight tracker. /// /// The double-check around `fetch_add` closes the drain/admission race: a @@ -97,13 +130,15 @@ impl RequestTracker { pub(crate) fn try_acquire(self: &Arc) -> Result { if !self.accepting.load(Ordering::Acquire) { - bail!("worker is not accepting new requests"); + return Err(worker_draining_error("worker is not accepting new requests")); } self.inflight.fetch_add(1, Ordering::AcqRel); if !self.accepting.load(Ordering::Acquire) { self.release(); - bail!("worker started draining before request admission completed"); + return Err(worker_draining_error( + "worker started draining before request admission completed", + )); } Ok(RequestGuard { @@ -134,6 +169,14 @@ impl RequestTracker { } } +fn worker_draining_error(message: &'static str) -> anyhow::Error { + DynamoError::builder() + .error_type(ErrorType::WorkerOverloaded) + .message(message) + .build() + .into() +} + pub(crate) struct RequestGuard { tracker: Arc, } @@ -153,9 +196,10 @@ pub(crate) struct WorkerLifecycleController { operation_lock: Mutex<()>, drain_generation: AtomicU64, tracker: Arc, - endpoint: Endpoint, - engine: EngineKind, + discovery: Arc, + quiescence: Arc, mode: DisaggregationMode, + discovery_grace_period: Duration, } impl WorkerLifecycleController { @@ -164,6 +208,23 @@ impl WorkerLifecycleController { endpoint: Endpoint, engine: EngineKind, mode: DisaggregationMode, + discovery_grace_period: Duration, + ) -> Arc { + Self::with_dependencies( + tracker, + Arc::new(endpoint), + Arc::new(engine), + mode, + discovery_grace_period, + ) + } + + fn with_dependencies( + tracker: Arc, + discovery: Arc, + quiescence: Arc, + mode: DisaggregationMode, + discovery_grace_period: Duration, ) -> Arc { Arc::new(Self { state: AtomicU8::new(WorkerLifecycleState::Serving as u8), @@ -173,9 +234,10 @@ impl WorkerLifecycleController { operation_lock: Mutex::new(()), drain_generation: AtomicU64::new(0), tracker, - endpoint, - engine, + discovery, + quiescence, mode, + discovery_grace_period, }) } @@ -204,16 +266,19 @@ impl WorkerLifecycleController { } pub(crate) fn register_admin_routes(self: &Arc, registry: &EngineRouteRegistry) { - registry.register( + registry.register_method( "drain", + EngineRouteMethod::Post, lifecycle_callback(Arc::clone(self), LifecycleAction::Drain), ); - registry.register( + registry.register_method( "resume", + EngineRouteMethod::Post, lifecycle_callback(Arc::clone(self), LifecycleAction::Resume), ); - registry.register( + registry.register_method( "status", + EngineRouteMethod::Get, lifecycle_callback(Arc::clone(self), LifecycleAction::Status), ); tracing::info!("registered worker Admin API routes under /engine"); @@ -221,40 +286,95 @@ impl WorkerLifecycleController { pub(crate) async fn drain(self: &Arc) -> Result { let operation = self.operation_lock.lock().await; - let generation = match self.state() { + let started = match self.state() { WorkerLifecycleState::Serving => { self.last_error.write().take(); - self.tracker.stop_accepting(); self.kv_transfers .store(Self::initial_kv_state(self.mode) as u8, Ordering::Release); self.state .store(WorkerLifecycleState::Draining as u8, Ordering::Release); - Some(self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1) + let generation = self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1; + let (started_tx, started_rx) = oneshot::channel(); + let controller = Arc::clone(self); + tokio::spawn(async move { + controller + .run_drain_operation(generation, started_tx) + .await; + }); + Some(started_rx) } WorkerLifecycleState::Draining => None, WorkerLifecycleState::Drained => return Ok(self.status()), WorkerLifecycleState::Stopping => bail!("worker shutdown is already in progress"), }; + drop(operation); - if self.discovery_registered.load(Ordering::Acquire) { - if let Err(error) = self.endpoint.unregister_endpoint_instance().await { - let message = format!("failed to unregister worker from discovery: {error}"); - *self.last_error.write() = Some(message.clone()); - self.drain_generation.fetch_add(1, Ordering::AcqRel); - self.state - .store(WorkerLifecycleState::Serving as u8, Ordering::Release); - self.tracker.start_accepting(); - return Err(error).context("failed to unregister worker from discovery"); + if let Some(started) = started { + started + .await + .context("drain operation ended before discovery was updated")??; + } + Ok(self.status()) + } + + /// Own the drain after the HTTP handler starts it. The task is detached + /// from the request so cancelling the client cannot strand `Draining`. + async fn run_drain_operation( + self: Arc, + generation: u64, + started: oneshot::Sender>, + ) { + let mut started = Some(started); + { + let _operation = self.operation_lock.lock().await; + if self.state() != WorkerLifecycleState::Draining + || self.drain_generation.load(Ordering::Acquire) != generation + { + if let Some(started) = started.take() { + let _ = started.send(Err(anyhow!( + "drain operation was superseded before discovery was updated" + ))); + } + return; + } + + if self.discovery_registered.load(Ordering::Acquire) { + if let Err(error) = self.discovery.unregister().await { + let message = format!("failed to unregister worker from discovery: {error}"); + *self.last_error.write() = Some(message); + self.drain_generation.fetch_add(1, Ordering::AcqRel); + self.state + .store(WorkerLifecycleState::Serving as u8, Ordering::Release); + self.tracker.start_accepting(); + if let Some(started) = started.take() { + let _ = started.send( + Err(error).context("failed to unregister worker from discovery"), + ); + } + return; + } + self.discovery_registered.store(false, Ordering::Release); + } + + if let Some(started) = started.take() { + let _ = started.send(Ok(())); } - self.discovery_registered.store(false, Ordering::Release); } - drop(operation); - if let Some(generation) = generation { - let controller = Arc::clone(self); - tokio::spawn(async move { controller.monitor_drain(generation).await }); + // Keep accepting requests already selected by frontends until their + // discovery views have had time to observe the unregister. + tokio::time::sleep(self.discovery_grace_period).await; + { + let _operation = self.operation_lock.lock().await; + if self.state() != WorkerLifecycleState::Draining + || self.drain_generation.load(Ordering::Acquire) != generation + { + return; + } + self.tracker.stop_accepting(); } - Ok(self.status()) + + self.monitor_drain(generation).await; } pub(crate) async fn resume(self: &Arc) -> Result { @@ -270,11 +390,15 @@ impl WorkerLifecycleController { .store(WorkerLifecycleState::Draining as u8, Ordering::Release); if !self.discovery_registered.load(Ordering::Acquire) { - if let Err(error) = self.endpoint.register_endpoint_instance().await { + if let Err(error) = self.discovery.register().await { let message = format!("failed to re-register worker in discovery: {error}"); *self.last_error.write() = Some(message); self.state.store(previous_state as u8, Ordering::Release); if previous_state == WorkerLifecycleState::Draining { + // The old operation was fenced by the generation bump. + // Continue toward Drained rather than leaving the worker + // unregistered with no controller-owned monitor. + self.tracker.stop_accepting(); let controller = Arc::clone(self); tokio::spawn(async move { controller.monitor_drain(generation).await }); } @@ -301,7 +425,7 @@ impl WorkerLifecycleController { .store(WorkerLifecycleState::Stopping as u8, Ordering::Release); if self.discovery_registered.load(Ordering::Acquire) { - if let Err(error) = self.endpoint.unregister_endpoint_instance().await { + if let Err(error) = self.discovery.unregister().await { tracing::warn!(%error, "discovery unregister failed during shutdown"); *self.last_error.write() = Some(error.to_string()); } else { @@ -319,49 +443,45 @@ impl WorkerLifecycleController { } if self.tracker.inflight() == 0 { - let kv_complete = if self.mode.is_prefill() { - match self.engine.is_quiescent().await { - Ok(Some(true)) => { - self.kv_transfers - .store(KvTransferState::Complete as u8, Ordering::Release); - true - } - Ok(Some(false)) => { - self.kv_transfers - .store(KvTransferState::Pending as u8, Ordering::Release); - false - } - Ok(None) => { - self.kv_transfers - .store(KvTransferState::Unknown as u8, Ordering::Release); - false - } - Err(error) => { - self.kv_transfers - .store(KvTransferState::Unknown as u8, Ordering::Release); - *self.last_error.write() = Some(error.to_string()); - false - } - } + let quiescence = if self.mode.is_prefill() { + self.quiescence.is_quiescent().await } else { - true + Ok(Some(true)) }; - if kv_complete - && self.drain_generation.load(Ordering::Acquire) == generation - && self - .state - .compare_exchange( - WorkerLifecycleState::Draining as u8, - WorkerLifecycleState::Drained as u8, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() + // Publish both KV state and Drained under the same operation + // fence used by resume. A stale monitor may finish an engine + // check, but it can never publish after its generation ends. + let _operation = self.operation_lock.lock().await; + if self.state() != WorkerLifecycleState::Draining + || self.drain_generation.load(Ordering::Acquire) != generation { - tracing::info!("worker drained and is safe to delete"); return; } + + match quiescence { + Ok(Some(true)) => { + if self.mode.is_prefill() { + self.kv_transfers + .store(KvTransferState::Complete as u8, Ordering::Release); + } + self.state + .store(WorkerLifecycleState::Drained as u8, Ordering::Release); + tracing::info!("worker drained and is safe to delete"); + return; + } + Ok(Some(false)) => self + .kv_transfers + .store(KvTransferState::Pending as u8, Ordering::Release), + Ok(None) => self + .kv_transfers + .store(KvTransferState::Unknown as u8, Ordering::Release), + Err(error) => { + self.kv_transfers + .store(KvTransferState::Unknown as u8, Ordering::Release); + *self.last_error.write() = Some(error.to_string()); + } + } } tokio::select! { @@ -402,6 +522,106 @@ fn lifecycle_callback( #[cfg(test)] mod tests { use super::*; + use tokio::sync::Semaphore; + + struct MockDiscovery { + unregister_started: Semaphore, + unregister_gate: Semaphore, + } + + impl MockDiscovery { + fn new(unregister_permits: usize) -> Arc { + Arc::new(Self { + unregister_started: Semaphore::new(0), + unregister_gate: Semaphore::new(unregister_permits), + }) + } + } + + #[async_trait] + impl DiscoveryRegistration for MockDiscovery { + async fn unregister(&self) -> Result<()> { + self.unregister_started.add_permits(1); + self.unregister_gate + .acquire() + .await + .expect("unregister gate should stay open") + .forget(); + Ok(()) + } + + async fn register(&self) -> Result<()> { + Ok(()) + } + } + + struct MockQuiescence { + result: Option, + check_started: Semaphore, + check_gate: Semaphore, + } + + impl MockQuiescence { + fn new(result: Option, check_permits: usize) -> Arc { + Arc::new(Self { + result, + check_started: Semaphore::new(0), + check_gate: Semaphore::new(check_permits), + }) + } + } + + #[async_trait] + impl QuiescenceCheck for MockQuiescence { + async fn is_quiescent(&self) -> Result> { + self.check_started.add_permits(1); + self.check_gate + .acquire() + .await + .expect("quiescence gate should stay open") + .forget(); + Ok(self.result) + } + } + + fn test_controller( + discovery: Arc, + quiescence: Arc, + mode: DisaggregationMode, + grace_period: Duration, + ) -> (Arc, Arc) { + let tracker = RequestTracker::new(); + let controller = WorkerLifecycleController::with_dependencies( + Arc::clone(&tracker), + discovery, + quiescence, + mode, + grace_period, + ); + (controller, tracker) + } + + async fn yield_to_background_tasks() { + for _ in 0..10 { + tokio::task::yield_now().await; + } + } + + async fn wait_for_state( + controller: &WorkerLifecycleController, + expected: WorkerLifecycleState, + ) { + for _ in 0..20 { + if controller.status().state == expected { + return; + } + tokio::task::yield_now().await; + } + panic!( + "worker did not reach {expected:?}; current state is {:?}", + controller.status().state + ); + } #[test] fn request_tracker_rejects_after_drain_starts() { @@ -426,4 +646,106 @@ mod tests { assert_eq!(tracker.inflight(), 1); drop(guard); } + + #[test] + fn request_tracker_uses_migratable_error_while_draining() { + let tracker = RequestTracker::new(); + tracker.stop_accepting(); + + let error = match tracker.try_acquire() { + Ok(_) => panic!("draining worker must reject new requests"), + Err(error) => error, + }; + let error = error + .downcast_ref::() + .expect("drain rejection should preserve its Dynamo error type"); + assert_eq!(error.error_type(), ErrorType::WorkerOverloaded); + } + + #[tokio::test(start_paused = true)] + async fn drain_keeps_admission_open_during_discovery_grace_period() { + let discovery = MockDiscovery::new(1); + let quiescence = MockQuiescence::new(Some(true), 1); + let grace_period = Duration::from_secs(5); + let (controller, tracker) = test_controller( + discovery, + quiescence, + DisaggregationMode::Aggregated, + grace_period, + ); + + let status = controller.drain().await.unwrap(); + assert_eq!(status.state, WorkerLifecycleState::Draining); + let admitted_during_convergence = tracker + .try_acquire() + .expect("late frontend selections should be admitted during convergence"); + + yield_to_background_tasks().await; + tokio::time::advance(grace_period).await; + yield_to_background_tasks().await; + assert!(tracker.try_acquire().is_err()); + + drop(admitted_during_convergence); + tokio::time::advance(QUIESCENCE_POLL_INTERVAL).await; + wait_for_state(&controller, WorkerLifecycleState::Drained).await; + } + + #[tokio::test(start_paused = true)] + async fn cancelled_drain_caller_does_not_strand_the_worker() { + let discovery = MockDiscovery::new(0); + let quiescence = MockQuiescence::new(Some(true), 1); + let (controller, _) = test_controller( + Arc::clone(&discovery), + quiescence, + DisaggregationMode::Aggregated, + Duration::ZERO, + ); + + let first_controller = Arc::clone(&controller); + let first_call = tokio::spawn(async move { first_controller.drain().await }); + discovery + .unregister_started + .acquire() + .await + .unwrap() + .forget(); + first_call.abort(); + let _ = first_call.await; + + let retry_controller = Arc::clone(&controller); + let retry = tokio::spawn(async move { retry_controller.drain().await }); + discovery.unregister_gate.add_permits(1); + retry.await.unwrap().unwrap(); + wait_for_state(&controller, WorkerLifecycleState::Drained).await; + } + + #[tokio::test(start_paused = true)] + async fn stale_monitor_cannot_publish_after_resume() { + let discovery = MockDiscovery::new(1); + let quiescence = MockQuiescence::new(Some(true), 0); + let (controller, _) = test_controller( + discovery, + Arc::clone(&quiescence), + DisaggregationMode::Prefill, + Duration::ZERO, + ); + + controller.drain().await.unwrap(); + quiescence + .check_started + .acquire() + .await + .unwrap() + .forget(); + + let resumed = controller.resume().await.unwrap(); + assert_eq!(resumed.state, WorkerLifecycleState::Serving); + quiescence.check_gate.add_permits(1); + yield_to_background_tasks().await; + + let status = controller.status(); + assert_eq!(status.state, WorkerLifecycleState::Serving); + assert_eq!(status.kv_transfers, KvTransferState::Unknown); + assert!(!status.safe_to_delete); + } } diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index fb4b334b38e7..bd64ad8254cf 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -907,6 +907,7 @@ impl Worker { endpoint.clone(), self.engine.clone(), self.config.disaggregation_mode, + Duration::from_secs_f64(grace_period_secs()), ); let (ingress, probe_engine): ( diff --git a/lib/runtime/src/engine_routes.rs b/lib/runtime/src/engine_routes.rs index eacb7a7b6d0b..cd2adf87c839 100644 --- a/lib/runtime/src/engine_routes.rs +++ b/lib/runtime/src/engine_routes.rs @@ -16,13 +16,37 @@ pub type EngineRouteCallback = Arc< + Sync, >; +/// HTTP method accepted by an engine route. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum EngineRouteMethod { + Get, + Post, +} + +/// Registered engine route, including an optional method restriction. +#[derive(Clone)] +pub struct EngineRoute { + callback: EngineRouteCallback, + method: Option, +} + +impl EngineRoute { + pub fn callback(&self) -> EngineRouteCallback { + Arc::clone(&self.callback) + } + + pub fn method(&self) -> Option { + self.method + } +} + /// Registry for engine route callbacks /// /// This registry stores callbacks that handle requests to `/engine/*` routes. /// Routes are registered from Python via `runtime.register_engine_route()`. #[derive(Clone, Default)] pub struct EngineRouteRegistry { - routes: Arc>>, + routes: Arc>>, } impl EngineRouteRegistry { @@ -40,8 +64,28 @@ impl EngineRouteRegistry { /// it usually signals two registration mechanisms colliding rather than an /// intentional replacement. pub fn register(&self, route: &str, callback: EngineRouteCallback) { + self.register_inner(route, None, callback); + } + + /// Register a callback that accepts only `method`. + pub fn register_method( + &self, + route: &str, + method: EngineRouteMethod, + callback: EngineRouteCallback, + ) { + self.register_inner(route, Some(method), callback); + } + + fn register_inner( + &self, + route: &str, + method: Option, + callback: EngineRouteCallback, + ) { let mut routes = self.routes.write().unwrap(); - if routes.insert(route.to_string(), callback).is_some() { + let entry = EngineRoute { callback, method }; + if routes.insert(route.to_string(), entry).is_some() { tracing::warn!("Overwriting already-registered engine route: /engine/{route}"); } else { tracing::debug!("Registered engine route: /engine/{route}"); @@ -50,6 +94,12 @@ impl EngineRouteRegistry { /// Get callback for a route pub fn get(&self, route: &str) -> Option { + let routes = self.routes.read().unwrap(); + routes.get(route).map(EngineRoute::callback) + } + + /// Get a route together with its method restriction. + pub fn get_route(&self, route: &str) -> Option { let routes = self.routes.read().unwrap(); routes.get(route).cloned() } @@ -130,4 +180,17 @@ mod tests { // Original should also see it (they share the Arc) assert!(registry.get("test2").is_some()); } + + #[test] + fn method_registration_preserves_method_metadata() { + let registry = EngineRouteRegistry::new(); + let callback: EngineRouteCallback = + Arc::new(|_| Box::pin(async { Ok(serde_json::json!({})) })); + + registry.register_method("drain", EngineRouteMethod::Post, callback); + + let route = registry.get_route("drain").unwrap(); + assert_eq!(route.method(), Some(EngineRouteMethod::Post)); + assert!(registry.get("drain").is_some()); + } } diff --git a/lib/runtime/src/system_status_server.rs b/lib/runtime/src/system_status_server.rs index 1a5d69d7749e..14f2ec6284d4 100644 --- a/lib/runtime/src/system_status_server.rs +++ b/lib/runtime/src/system_status_server.rs @@ -8,6 +8,7 @@ use crate::config::HealthStatus; use crate::config::environment_names::logging as env_logging; use crate::config::environment_names::runtime::canary as env_canary; use crate::config::environment_names::runtime::system as env_system; +use crate::engine_routes::EngineRouteMethod; use crate::logging::make_system_request_span; use crate::metrics::MetricsHierarchy; use crate::traits::DistributedRuntimeProvider; @@ -15,7 +16,7 @@ use axum::{ Router, body::Bytes, extract::{Json, Path, State}, - http::StatusCode, + http::{Method, StatusCode}, response::IntoResponse, routing::{any, delete, get, post}, }; @@ -188,7 +189,7 @@ pub async fn spawn_system_status_server( "/engine/{*path}", any({ let state = Arc::clone(&server_state); - move |path, body| engine_route_handler(state, path, body) + move |method, path, body| engine_route_handler(state, method, path, body) }), ); @@ -642,12 +643,51 @@ fn parse_lora_response(response_data: &serde_json::Value) -> LoraResponse { #[tracing::instrument(skip_all, level = "trace", fields(path = %path))] async fn engine_route_handler( state: Arc, + method: Method, Path(path): Path, body: Bytes, ) -> impl IntoResponse { tracing::trace!("Engine route request to /engine/{path}"); - // Parse body as JSON (empty object for GET/empty body) + // Look up callback + let route = match state.drt().engine_routes().get_route(&path) { + Some(route) => route, + None => { + tracing::debug!("Route /engine/{path} not found"); + return ( + StatusCode::NOT_FOUND, + json!({ + "error": "Route not found", + "message": format!("Route /engine/{} not found", path) + }) + .to_string(), + ) + .into_response(); + } + }; + + let method_allowed = match route.method() { + None => true, + Some(EngineRouteMethod::Get) => method == Method::GET, + Some(EngineRouteMethod::Post) => method == Method::POST, + }; + if !method_allowed { + tracing::debug!(%method, "method not allowed for /engine/{path}"); + return ( + StatusCode::METHOD_NOT_ALLOWED, + json!({ + "error": "Method not allowed", + "message": format!("{} is not allowed for /engine/{}", method, path) + }) + .to_string(), + ) + .into_response(); + } + + let callback = route.callback(); + + // Parse body as JSON (empty object for GET/empty body) only after the + // method check, so a wrong method consistently returns 405. let body_json: serde_json::Value = if body.is_empty() { serde_json::json!({}) } else { @@ -668,23 +708,6 @@ async fn engine_route_handler( } }; - // Look up callback - let callback = match state.drt().engine_routes().get(&path) { - Some(cb) => cb, - None => { - tracing::debug!("Route /engine/{path} not found"); - return ( - StatusCode::NOT_FOUND, - json!({ - "error": "Route not found", - "message": format!("Route /engine/{} not found", path) - }) - .to_string(), - ) - .into_response(); - } - }; - // Call callback (it's async, so await it) match callback(body_json).await { Ok(response) => { @@ -753,6 +776,7 @@ mod integration_tests { use crate::metrics::MetricsHierarchy; use anyhow::Result; use rstest::rstest; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use tokio::time::Duration; @@ -1180,6 +1204,45 @@ mod integration_tests { .await; } + #[tokio::test] + async fn test_engine_route_rejects_wrong_method_without_invoking_callback() { + temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, Some("0"))], async { + let drt = Arc::new(create_test_drt_async().await); + let invocations = Arc::new(AtomicUsize::new(0)); + let callback_invocations = Arc::clone(&invocations); + let callback: crate::engine_routes::EngineRouteCallback = Arc::new(move |_| { + let callback_invocations = Arc::clone(&callback_invocations); + Box::pin(async move { + callback_invocations.fetch_add(1, Ordering::Relaxed); + Ok(serde_json::json!({"state": "draining"})) + }) + }); + drt.engine_routes() + .register_method("drain", EngineRouteMethod::Post, callback); + + let addr = drt + .system_status_server_info() + .expect("system status server should be started") + .socket_addr; + let client = reqwest::Client::new(); + let url = format!("http://{addr}/engine/drain"); + + let response = client.get(&url).send().await.unwrap(); + assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED); + assert_eq!(invocations.load(Ordering::Relaxed), 0); + + let response = client + .post(&url) + .json(&serde_json::json!({})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(invocations.load(Ordering::Relaxed), 1); + }) + .await; + } + #[cfg(feature = "integration")] #[tokio::test] async fn test_health_check_with_payload_and_timeout() { From bc0e097731d6b262251fe3b8dda50350d8df6c2e Mon Sep 17 00:00:00 2001 From: xianlubird Date: Tue, 11 Aug 2026 10:47:28 +0800 Subject: [PATCH 3/8] fix(backend): address worker drain CI failures Signed-off-by: xianlubird --- lib/backend-common/src/lifecycle.rs | 28 +++++++++---------------- lib/runtime/src/system_status_server.rs | 2 +- 2 files changed, 11 insertions(+), 19 deletions(-) diff --git a/lib/backend-common/src/lifecycle.rs b/lib/backend-common/src/lifecycle.rs index a0c8fed077d5..528c02efcb18 100644 --- a/lib/backend-common/src/lifecycle.rs +++ b/lib/backend-common/src/lifecycle.rs @@ -10,9 +10,7 @@ use std::time::Duration; use anyhow::{Context, Result, anyhow, bail}; use async_trait::async_trait; use dynamo_runtime::component::Endpoint; -use dynamo_runtime::engine_routes::{ - EngineRouteCallback, EngineRouteMethod, EngineRouteRegistry, -}; +use dynamo_runtime::engine_routes::{EngineRouteCallback, EngineRouteMethod, EngineRouteRegistry}; use dynamo_runtime::error::{DynamoError, ErrorType}; use parking_lot::RwLock; use serde::Serialize; @@ -130,7 +128,9 @@ impl RequestTracker { pub(crate) fn try_acquire(self: &Arc) -> Result { if !self.accepting.load(Ordering::Acquire) { - return Err(worker_draining_error("worker is not accepting new requests")); + return Err(worker_draining_error( + "worker is not accepting new requests", + )); } self.inflight.fetch_add(1, Ordering::AcqRel); @@ -171,7 +171,7 @@ impl RequestTracker { fn worker_draining_error(message: &'static str) -> anyhow::Error { DynamoError::builder() - .error_type(ErrorType::WorkerOverloaded) + .error_type(ErrorType::Unavailable) .message(message) .build() .into() @@ -297,9 +297,7 @@ impl WorkerLifecycleController { let (started_tx, started_rx) = oneshot::channel(); let controller = Arc::clone(self); tokio::spawn(async move { - controller - .run_drain_operation(generation, started_tx) - .await; + controller.run_drain_operation(generation, started_tx).await; }); Some(started_rx) } @@ -347,9 +345,8 @@ impl WorkerLifecycleController { .store(WorkerLifecycleState::Serving as u8, Ordering::Release); self.tracker.start_accepting(); if let Some(started) = started.take() { - let _ = started.send( - Err(error).context("failed to unregister worker from discovery"), - ); + let _ = started + .send(Err(error).context("failed to unregister worker from discovery")); } return; } @@ -659,7 +656,7 @@ mod tests { let error = error .downcast_ref::() .expect("drain rejection should preserve its Dynamo error type"); - assert_eq!(error.error_type(), ErrorType::WorkerOverloaded); + assert_eq!(error.error_type(), ErrorType::Unavailable); } #[tokio::test(start_paused = true)] @@ -731,12 +728,7 @@ mod tests { ); controller.drain().await.unwrap(); - quiescence - .check_started - .acquire() - .await - .unwrap() - .forget(); + quiescence.check_started.acquire().await.unwrap().forget(); let resumed = controller.resume().await.unwrap(); assert_eq!(resumed.state, WorkerLifecycleState::Serving); diff --git a/lib/runtime/src/system_status_server.rs b/lib/runtime/src/system_status_server.rs index 14f2ec6284d4..74835edaebde 100644 --- a/lib/runtime/src/system_status_server.rs +++ b/lib/runtime/src/system_status_server.rs @@ -776,8 +776,8 @@ mod integration_tests { use crate::metrics::MetricsHierarchy; use anyhow::Result; use rstest::rstest; - use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::time::Duration; #[tokio::test] From bce4014152f502af166f2b03782d8487d000ac30 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Wed, 12 Aug 2026 10:00:04 +0800 Subject: [PATCH 4/8] fix(backend): harden worker drain lifecycle Signed-off-by: xianlubird --- .../reference/components/worker-admin-api.mdx | 10 +- lib/backend-common/src/lifecycle.rs | 324 +++++++++++++++--- lib/backend-common/src/worker.rs | 13 +- lib/llm/src/http/service/generate.rs | 15 +- lib/llm/src/http/service/metrics.rs | 5 +- lib/llm/src/migration.rs | 186 ++++++++-- lib/runtime/src/engine_routes.rs | 108 +++++- lib/runtime/src/error.rs | 5 + .../network/egress/addressed_router.rs | 2 +- .../pipeline/network/egress/push_router.rs | 51 ++- lib/runtime/src/system_status_server.rs | 5 +- 11 files changed, 617 insertions(+), 107 deletions(-) diff --git a/docs/fern/pages/reference/components/worker-admin-api.mdx b/docs/fern/pages/reference/components/worker-admin-api.mdx index 05e792d2e737..e287739d0de6 100644 --- a/docs/fern/pages/reference/components/worker-admin-api.mdx +++ b/docs/fern/pages/reference/components/worker-admin-api.mdx @@ -18,7 +18,7 @@ operator network. | 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` | Re-register the worker in discovery and resume request admission | +| `POST` | `/engine/resume` | Resume request admission, then re-register the worker in discovery | | `GET` | `/engine/status` | Return the worker lifecycle and deletion-safety status | Send an empty body or an empty JSON object. Each endpoint returns the same status shape: @@ -41,10 +41,12 @@ backend cannot confirm KV-transfer completion, `kv_transfers` remains `unknown` 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 retryable worker-scoped error and waits for all admitted work to finish. +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. The Admin API does not -delete the worker or replace the shutdown timeout. +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 diff --git a/lib/backend-common/src/lifecycle.rs b/lib/backend-common/src/lifecycle.rs index 528c02efcb18..ebe7a001c047 100644 --- a/lib/backend-common/src/lifecycle.rs +++ b/lib/backend-common/src/lifecycle.rs @@ -10,11 +10,15 @@ use std::time::Duration; use anyhow::{Context, Result, anyhow, bail}; use async_trait::async_trait; use dynamo_runtime::component::Endpoint; -use dynamo_runtime::engine_routes::{EngineRouteCallback, EngineRouteMethod, EngineRouteRegistry}; +use dynamo_runtime::engine_routes::{ + EngineRouteCallback, EngineRouteMethod, EngineRouteRegistration, EngineRouteRegistry, +}; use dynamo_runtime::error::{DynamoError, ErrorType}; use parking_lot::RwLock; use serde::Serialize; use tokio::sync::{Mutex, Notify, oneshot}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; use crate::disagg::DisaggregationMode; use crate::worker::EngineKind; @@ -156,6 +160,10 @@ impl RequestTracker { self.changed.notify_waiters(); } + fn is_accepting(&self) -> bool { + self.accepting.load(Ordering::Acquire) + } + pub(crate) fn inflight(&self) -> u64 { self.inflight.load(Ordering::Acquire) } @@ -171,7 +179,7 @@ impl RequestTracker { fn worker_draining_error(message: &'static str) -> anyhow::Error { DynamoError::builder() - .error_type(ErrorType::Unavailable) + .error_type(ErrorType::WorkerDraining) .message(message) .build() .into() @@ -181,6 +189,15 @@ pub(crate) struct RequestGuard { tracker: Arc, } +struct DrainTask { + cancel: CancellationToken, + handle: JoinHandle<()>, +} + +pub(crate) struct AdminRouteRegistration { + _routes: Vec, +} + impl Drop for RequestGuard { fn drop(&mut self) { self.tracker.release(); @@ -194,6 +211,7 @@ pub(crate) struct WorkerLifecycleController { discovery_registered: AtomicBool, last_error: RwLock>, operation_lock: Mutex<()>, + drain_task: Mutex>, drain_generation: AtomicU64, tracker: Arc, discovery: Arc, @@ -232,6 +250,7 @@ impl WorkerLifecycleController { discovery_registered: AtomicBool::new(true), last_error: RwLock::new(None), operation_lock: Mutex::new(()), + drain_task: Mutex::new(None), drain_generation: AtomicU64::new(0), tracker, discovery, @@ -265,23 +284,29 @@ impl WorkerLifecycleController { } } - pub(crate) fn register_admin_routes(self: &Arc, registry: &EngineRouteRegistry) { - registry.register_method( - "drain", - EngineRouteMethod::Post, - lifecycle_callback(Arc::clone(self), LifecycleAction::Drain), - ); - registry.register_method( - "resume", - EngineRouteMethod::Post, - lifecycle_callback(Arc::clone(self), LifecycleAction::Resume), - ); - registry.register_method( - "status", - EngineRouteMethod::Get, - lifecycle_callback(Arc::clone(self), LifecycleAction::Status), - ); + pub(crate) fn register_admin_routes( + self: &Arc, + registry: &EngineRouteRegistry, + ) -> AdminRouteRegistration { + let routes = vec![ + registry.register_scoped_method( + "drain", + EngineRouteMethod::Post, + lifecycle_callback(Arc::clone(self), LifecycleAction::Drain), + ), + registry.register_scoped_method( + "resume", + EngineRouteMethod::Post, + lifecycle_callback(Arc::clone(self), LifecycleAction::Resume), + ), + registry.register_scoped_method( + "status", + EngineRouteMethod::Get, + lifecycle_callback(Arc::clone(self), LifecycleAction::Status), + ), + ]; tracing::info!("registered worker Admin API routes under /engine"); + AdminRouteRegistration { _routes: routes } } pub(crate) async fn drain(self: &Arc) -> Result { @@ -295,10 +320,23 @@ impl WorkerLifecycleController { .store(WorkerLifecycleState::Draining as u8, Ordering::Release); let generation = self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1; let (started_tx, started_rx) = oneshot::channel(); + let cancel = CancellationToken::new(); let controller = Arc::clone(self); - tokio::spawn(async move { - controller.run_drain_operation(generation, started_tx).await; + let task_cancel = cancel.clone(); + let handle = tokio::spawn(async move { + controller + .run_drain_operation(generation, started_tx, task_cancel) + .await; }); + let previous = self + .drain_task + .lock() + .await + .replace(DrainTask { cancel, handle }); + debug_assert!( + previous.as_ref().is_none_or(|task| task.handle.is_finished()), + "serving worker must not have active drain work" + ); Some(started_rx) } WorkerLifecycleState::Draining => None, @@ -321,10 +359,14 @@ impl WorkerLifecycleController { self: Arc, generation: u64, started: oneshot::Sender>, + cancel: CancellationToken, ) { let mut started = Some(started); { - let _operation = self.operation_lock.lock().await; + let _operation = tokio::select! { + _ = cancel.cancelled() => return, + operation = self.operation_lock.lock() => operation, + }; if self.state() != WorkerLifecycleState::Draining || self.drain_generation.load(Ordering::Acquire) != generation { @@ -360,9 +402,15 @@ impl WorkerLifecycleController { // Keep accepting requests already selected by frontends until their // discovery views have had time to observe the unregister. - tokio::time::sleep(self.discovery_grace_period).await; + tokio::select! { + _ = cancel.cancelled() => return, + _ = tokio::time::sleep(self.discovery_grace_period) => {} + } { - let _operation = self.operation_lock.lock().await; + let _operation = tokio::select! { + _ = cancel.cancelled() => return, + operation = self.operation_lock.lock() => operation, + }; if self.state() != WorkerLifecycleState::Draining || self.drain_generation.load(Ordering::Acquire) != generation { @@ -371,10 +419,24 @@ impl WorkerLifecycleController { self.tracker.stop_accepting(); } - self.monitor_drain(generation).await; + self.monitor_drain(generation, cancel).await; } pub(crate) async fn resume(self: &Arc) -> Result { + let (completed_tx, completed_rx) = oneshot::channel(); + let controller = Arc::clone(self); + tokio::spawn(async move { + let result = controller.run_resume_operation().await; + let _ = completed_tx.send(result); + }); + completed_rx + .await + .context("resume operation ended before publishing its result")? + } + + /// Resume is controller-owned so cancellation of the Admin request cannot + /// leave admission and discovery in different states. + async fn run_resume_operation(self: Arc) -> Result { let _operation = self.operation_lock.lock().await; let previous_state = match self.state() { WorkerLifecycleState::Serving => return Ok(self.status()), @@ -385,19 +447,20 @@ impl WorkerLifecycleController { let generation = self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1; self.state .store(WorkerLifecycleState::Draining as u8, Ordering::Release); + self.cancel_and_join_drain_task().await; + + // A worker must be able to accept requests before discovery can make + // it visible. Roll this back if registration fails. + self.tracker.start_accepting(); if !self.discovery_registered.load(Ordering::Acquire) { if let Err(error) = self.discovery.register().await { let message = format!("failed to re-register worker in discovery: {error}"); *self.last_error.write() = Some(message); + self.tracker.stop_accepting(); self.state.store(previous_state as u8, Ordering::Release); if previous_state == WorkerLifecycleState::Draining { - // The old operation was fenced by the generation bump. - // Continue toward Drained rather than leaving the worker - // unregistered with no controller-owned monitor. - self.tracker.stop_accepting(); - let controller = Arc::clone(self); - tokio::spawn(async move { controller.monitor_drain(generation).await }); + self.start_monitor_task(generation).await; } return Err(error).context("failed to re-register worker in discovery"); } @@ -409,30 +472,68 @@ impl WorkerLifecycleController { .store(Self::initial_kv_state(self.mode) as u8, Ordering::Release); self.state .store(WorkerLifecycleState::Serving as u8, Ordering::Release); - self.tracker.start_accepting(); Ok(self.status()) } - /// Move into the irreversible SIGTERM path and ensure discovery is down. + /// Move into the irreversible SIGTERM path using the same discovery-first + /// ordering as Admin drain. pub(crate) async fn begin_shutdown(&self) { - let _operation = self.operation_lock.lock().await; - self.tracker.stop_accepting(); - self.drain_generation.fetch_add(1, Ordering::AcqRel); - self.state - .store(WorkerLifecycleState::Stopping as u8, Ordering::Release); + let needs_convergence_grace = { + let _operation = self.operation_lock.lock().await; + self.drain_generation.fetch_add(1, Ordering::AcqRel); + self.state + .store(WorkerLifecycleState::Stopping as u8, Ordering::Release); + self.cancel_and_join_drain_task().await; - if self.discovery_registered.load(Ordering::Acquire) { - if let Err(error) = self.discovery.unregister().await { - tracing::warn!(%error, "discovery unregister failed during shutdown"); - *self.last_error.write() = Some(error.to_string()); - } else { - self.discovery_registered.store(false, Ordering::Release); + if self.discovery_registered.load(Ordering::Acquire) { + if let Err(error) = self.discovery.unregister().await { + tracing::warn!(%error, "discovery unregister failed during shutdown"); + *self.last_error.write() = Some(error.to_string()); + } else { + self.discovery_registered.store(false, Ordering::Release); + } } + self.tracker.is_accepting() + }; + + if needs_convergence_grace { + tokio::time::sleep(self.discovery_grace_period).await; + } + self.tracker.stop_accepting(); + } + + async fn cancel_and_join_drain_task(&self) { + let Some(task) = self.drain_task.lock().await.take() else { + return; + }; + task.cancel.cancel(); + if let Err(error) = task.handle.await + && !error.is_cancelled() + { + tracing::warn!(%error, "drain task failed while being joined"); } } - async fn monitor_drain(self: Arc, generation: u64) { + async fn start_monitor_task(self: &Arc, generation: u64) { + let cancel = CancellationToken::new(); + let task_cancel = cancel.clone(); + let controller = Arc::clone(self); + let handle = tokio::spawn(async move { + controller.monitor_drain(generation, task_cancel).await; + }); + let previous = self + .drain_task + .lock() + .await + .replace(DrainTask { cancel, handle }); + debug_assert!(previous.is_none(), "drain task slot must be empty"); + } + + async fn monitor_drain(self: Arc, generation: u64, cancel: CancellationToken) { loop { + if cancel.is_cancelled() { + return; + } if self.state() != WorkerLifecycleState::Draining || self.drain_generation.load(Ordering::Acquire) != generation { @@ -441,7 +542,10 @@ impl WorkerLifecycleController { if self.tracker.inflight() == 0 { let quiescence = if self.mode.is_prefill() { - self.quiescence.is_quiescent().await + tokio::select! { + _ = cancel.cancelled() => return, + result = self.quiescence.is_quiescent() => result, + } } else { Ok(Some(true)) }; @@ -449,7 +553,10 @@ impl WorkerLifecycleController { // Publish both KV state and Drained under the same operation // fence used by resume. A stale monitor may finish an engine // check, but it can never publish after its generation ends. - let _operation = self.operation_lock.lock().await; + let _operation = tokio::select! { + _ = cancel.cancelled() => return, + operation = self.operation_lock.lock() => operation, + }; if self.state() != WorkerLifecycleState::Draining || self.drain_generation.load(Ordering::Acquire) != generation { @@ -482,6 +589,7 @@ impl WorkerLifecycleController { } tokio::select! { + _ = cancel.cancelled() => return, _ = self.tracker.changed.notified() => {} _ = tokio::time::sleep(QUIESCENCE_POLL_INTERVAL) => {} } @@ -524,13 +632,26 @@ mod tests { struct MockDiscovery { unregister_started: Semaphore, unregister_gate: Semaphore, + register_started: Semaphore, + register_gate: Semaphore, + register_fails: AtomicBool, } impl MockDiscovery { fn new(unregister_permits: usize) -> Arc { + Self::with_register_permits(unregister_permits, 1) + } + + fn with_register_permits( + unregister_permits: usize, + register_permits: usize, + ) -> Arc { Arc::new(Self { unregister_started: Semaphore::new(0), unregister_gate: Semaphore::new(unregister_permits), + register_started: Semaphore::new(0), + register_gate: Semaphore::new(register_permits), + register_fails: AtomicBool::new(false), }) } } @@ -548,6 +669,15 @@ mod tests { } async fn register(&self) -> Result<()> { + self.register_started.add_permits(1); + self.register_gate + .acquire() + .await + .expect("register gate should stay open") + .forget(); + if self.register_fails.load(Ordering::Acquire) { + bail!("injected register failure"); + } Ok(()) } } @@ -645,7 +775,7 @@ mod tests { } #[test] - fn request_tracker_uses_migratable_error_while_draining() { + fn request_tracker_uses_worker_draining_error() { let tracker = RequestTracker::new(); tracker.stop_accepting(); @@ -656,7 +786,7 @@ mod tests { let error = error .downcast_ref::() .expect("drain rejection should preserve its Dynamo error type"); - assert_eq!(error.error_type(), ErrorType::Unavailable); + assert_eq!(error.error_type(), ErrorType::WorkerDraining); } #[tokio::test(start_paused = true)] @@ -734,10 +864,108 @@ mod tests { assert_eq!(resumed.state, WorkerLifecycleState::Serving); quiescence.check_gate.add_permits(1); yield_to_background_tasks().await; + assert_eq!( + quiescence.check_gate.available_permits(), + 1, + "resume must cancel and join the blocked quiescence check" + ); let status = controller.status(); assert_eq!(status.state, WorkerLifecycleState::Serving); assert_eq!(status.kv_transfers, KvTransferState::Unknown); assert!(!status.safe_to_delete); } + + #[tokio::test(start_paused = true)] + async fn resume_opens_admission_before_registration_and_survives_caller_cancel() { + let discovery = MockDiscovery::with_register_permits(1, 0); + let quiescence = MockQuiescence::new(Some(true), 0); + let (controller, tracker) = test_controller( + Arc::clone(&discovery), + Arc::clone(&quiescence), + DisaggregationMode::Prefill, + Duration::ZERO, + ); + + controller.drain().await.unwrap(); + quiescence.check_started.acquire().await.unwrap().forget(); + + let resume_controller = Arc::clone(&controller); + let resume_call = tokio::spawn(async move { resume_controller.resume().await }); + discovery + .register_started + .acquire() + .await + .unwrap() + .forget(); + + let admitted = tracker + .try_acquire() + .expect("admission must be open before discovery registration publishes"); + drop(admitted); + + resume_call.abort(); + let _ = resume_call.await; + discovery.register_gate.add_permits(1); + wait_for_state(&controller, WorkerLifecycleState::Serving).await; + assert!(controller.status().discovery_registered); + } + + #[tokio::test(start_paused = true)] + async fn failed_resume_registration_rolls_back_admission_and_monitoring() { + let discovery = MockDiscovery::new(1); + discovery.register_fails.store(true, Ordering::Release); + let quiescence = MockQuiescence::new(Some(true), 0); + let (controller, tracker) = test_controller( + discovery, + Arc::clone(&quiescence), + DisaggregationMode::Prefill, + Duration::ZERO, + ); + + controller.drain().await.unwrap(); + quiescence.check_started.acquire().await.unwrap().forget(); + assert!(controller.resume().await.is_err()); + + assert_eq!(controller.status().state, WorkerLifecycleState::Draining); + assert!(tracker.try_acquire().is_err()); + quiescence.check_started.acquire().await.unwrap().forget(); + } + + #[tokio::test(start_paused = true)] + async fn shutdown_unregisters_before_closing_admission() { + let discovery = MockDiscovery::new(0); + let quiescence = MockQuiescence::new(Some(true), 1); + let grace_period = Duration::from_secs(5); + let (controller, tracker) = test_controller( + Arc::clone(&discovery), + quiescence, + DisaggregationMode::Aggregated, + grace_period, + ); + + let shutdown_controller = Arc::clone(&controller); + let shutdown = tokio::spawn(async move { shutdown_controller.begin_shutdown().await }); + discovery + .unregister_started + .acquire() + .await + .unwrap() + .forget(); + let during_unregister = tracker + .try_acquire() + .expect("SIGTERM must keep admission open while unregistering"); + drop(during_unregister); + + discovery.unregister_gate.add_permits(1); + yield_to_background_tasks().await; + let during_convergence = tracker + .try_acquire() + .expect("SIGTERM must keep admission open during discovery convergence"); + drop(during_convergence); + + tokio::time::advance(grace_period).await; + shutdown.await.unwrap(); + assert!(tracker.try_acquire().is_err()); + } } diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 7150eee5b483..8ecb90e62db5 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -774,7 +774,13 @@ impl Worker { } else { tracing::info!("Endpoint unregistered from discovery"); } - self.run_engine_shutdown_steps().await; + if lifecycle.is_some() { + // The lifecycle controller already kept admission open for the + // discovery convergence grace period before closing it. + self.run_engine_shutdown_steps_with_grace(0.0).await; + } else { + self.run_engine_shutdown_steps().await; + } } /// Start the engine exactly once. `Worker::run` consumes `self`, so all @@ -1004,7 +1010,7 @@ impl Worker { format!("serve registration: {e}"), ) })?; - lifecycle.register_admin_routes(endpoint.drt().engine_routes()); + let admin_routes = lifecycle.register_admin_routes(endpoint.drt().engine_routes()); let serve_fut = started_endpoint.wait(); tokio::pin!(serve_fut); @@ -1035,6 +1041,9 @@ impl Worker { } } + // The DRT can outlive this Python worker. Remove callbacks that retain + // this worker's controller and engine before engine cleanup begins. + drop(admin_routes); self.orchestrator_steps(&endpoint, Some(&lifecycle)).await; Ok(()) } diff --git a/lib/llm/src/http/service/generate.rs b/lib/llm/src/http/service/generate.rs index 17f4bc8e1789..53b91cb7ce40 100644 --- a/lib/llm/src/http/service/generate.rs +++ b/lib/llm/src/http/service/generate.rs @@ -426,9 +426,10 @@ async fn generate_dispatch( let was_cancelled = request_context.is_killed() || super::metrics::request_was_cancelled(error.as_ref()); let was_rejected = super::metrics::request_was_rejected(error.as_ref()); + let was_unavailable = super::metrics::request_was_unavailable(error.as_ref()); inflight_guard.mark_error(if was_cancelled { ErrorType::Cancelled - } else if was_rejected { + } else if was_rejected || was_unavailable { ErrorType::Unavailable } else { ErrorType::Internal @@ -447,6 +448,18 @@ async fn generate_dispatch( "engine rejected the request".to_string(), ); } + if was_unavailable { + tracing::warn!( + %request_id, + error = %format!("{error:#}"), + "no backend worker is available" + ); + return generate_error_response( + StatusCode::SERVICE_UNAVAILABLE, + "service_unavailable", + "no backend worker is available".to_string(), + ); + } tracing::error!(%request_id, error = %format!("{error:#}"), "engine generate call failed"); return generate_internal_error_response(); } diff --git a/lib/llm/src/http/service/metrics.rs b/lib/llm/src/http/service/metrics.rs index 8cc3042ba8ac..0865016a23b9 100644 --- a/lib/llm/src/http/service/metrics.rs +++ b/lib/llm/src/http/service/metrics.rs @@ -48,7 +48,10 @@ pub fn request_was_rejected(err: &(dyn std::error::Error + 'static)) -> bool { /// Check whether an error chain indicates that no backend worker is available. pub fn request_was_unavailable(err: &(dyn std::error::Error + 'static)) -> bool { - const UNAVAILABLE: &[DynamoErrorType] = &[DynamoErrorType::Unavailable]; + const UNAVAILABLE: &[DynamoErrorType] = &[ + DynamoErrorType::Unavailable, + DynamoErrorType::WorkerDraining, + ]; const AVAILABLE: &[DynamoErrorType] = &[]; dynamo_runtime::error::match_error_chain(err, UNAVAILABLE, AVAILABLE) } diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs index 9c72627d5c1b..6585677081ff 100644 --- a/lib/llm/src/migration.rs +++ b/lib/llm/src/migration.rs @@ -30,6 +30,11 @@ use dynamo_runtime::pipeline::{ }; use dynamo_runtime::protocols::annotated::Annotated; +/// Reselections caused by stale discovery are operational retries, not user +/// migrations. Bound them separately so a broken discovery view cannot loop +/// forever while the default `migration_limit = 0` still remains safe. +const WORKER_DRAINING_RESELECTION_LIMIT: u32 = 3; + /// Accessors the migration RetryManager needs from a response chunk. /// `token_ids` lets it replay already-delivered tokens; `worker_trace_link` /// lets it stamp the failed worker's span onto the next attempt's @@ -78,6 +83,10 @@ fn is_migratable(err: &(dyn StdError + 'static)) -> bool { error::match_error_chain(err, MIGRATABLE, NON_MIGRATABLE) } +fn is_worker_draining(err: &(dyn StdError + 'static)) -> bool { + error::match_error_chain(err, &[ErrorType::WorkerDraining], &[]) +} + /// Whether a worker-scoped failure can be retried without violating an explicit route. /// /// The phase is read after the failed attempt because disaggregated routing updates the @@ -125,6 +134,43 @@ fn is_migratable_for_request( allows_phase(phase) } +fn can_reselect_worker_for_request(request: &PreprocessedRequest) -> bool { + let allows_phase = |phase| match explicit_target(request, phase) { + Ok(None) => true, + Ok(Some(target)) => { + tracing::debug!( + ?phase, + worker_id = target.worker_id, + dp_rank = ?target.dp_rank, + "Worker reselection disabled for explicitly pinned worker" + ); + false + } + Err(error) => { + tracing::warn!(?phase, %error, "Worker reselection disabled for invalid target"); + false + } + }; + + let Some(tracker) = request.tracker.as_ref() else { + return [ + RequestPhase::Prefill, + RequestPhase::Decode, + RequestPhase::Aggregated, + ] + .into_iter() + .all(allows_phase); + }; + allows_phase(tracker.phase()) +} + +fn is_worker_draining_for_request( + request: &PreprocessedRequest, + err: &(dyn StdError + 'static), +) -> bool { + is_worker_draining(err) && can_reselect_worker_for_request(request) +} + pub struct Migration { migration_limit: u32, max_seq_len: Option, @@ -245,6 +291,7 @@ where next_generate: ServerStreamingEngine>, next_stream: Option>>, retries_left: u32, + worker_draining_reselections_left: u32, max_seq_len: Option, model_name: Arc, metrics: Arc, @@ -297,9 +344,10 @@ where } retries_left = 0; } - if retries_left > 0 { - preprocessed_request.migration_state = Some(Default::default()); - } + // The routing layer also uses this state to exclude a worker that + // rejected admission because it is draining. That operational + // reselection remains enabled even when user migration is disabled. + preprocessed_request.migration_state = Some(Default::default()); let mut slf = Self { context, metadata, @@ -307,13 +355,14 @@ where session_affinity, next_generate: next, next_stream: None, - retries_left: retries_left + 1, // +1 to account for the initial attempt + retries_left, + worker_draining_reselections_left: WORKER_DRAINING_RESELECTION_LIMIT, max_seq_len, model_name, metrics, last_worker_link: None, }; - slf.new_stream().await?; + slf.new_stream(RetryReason::Initial).await?; slf.exceed_max_seq_len(0); // disable migration if prompt len > max_seq_len Ok(slf) } @@ -329,15 +378,22 @@ where }; if let Some(response) = response_stream.next().await { // Check if this is a migratable error that should trigger stream recreation. - if let Some(err) = response.error.as_ref() - && is_migratable_for_request(&self.request, err) - { - tracing::warn!(error = %err, "Stream disconnected, recreating stream"); - self.metrics.inc_migration_ongoing_request(&self.model_name); - if let Err(err) = self.new_stream().await { - tracing::warn!(error = ?err, "Cannot recreate stream"); + if let Some(err) = response.error.as_ref() { + let retry_reason = if is_worker_draining_for_request(&self.request, err) { + Some(RetryReason::WorkerDraining) + } else if is_migratable_for_request(&self.request, err) { + self.metrics.inc_migration_ongoing_request(&self.model_name); + Some(RetryReason::Migration) } else { - continue; + None + }; + if let Some(retry_reason) = retry_reason { + tracing::warn!(error = %err, ?retry_reason, "Recreating request stream"); + if let Err(err) = self.new_stream(retry_reason).await { + tracing::warn!(error = ?err, "Cannot recreate stream"); + } else { + continue; + } } } self.track_response(&response); @@ -347,10 +403,24 @@ where } } - async fn new_stream(&mut self) -> Result<()> { + async fn new_stream(&mut self, mut retry_reason: RetryReason) -> Result<()> { let mut response_stream: Option>>> = None; - while self.retries_left > 0 { - self.retries_left -= 1; + loop { + match retry_reason { + RetryReason::Initial => {} + RetryReason::Migration if self.retries_left > 0 => { + self.retries_left -= 1; + } + RetryReason::WorkerDraining + if self.worker_draining_reselections_left > 0 => + { + self.worker_draining_reselections_left -= 1; + } + RetryReason::Migration => return Err(Error::msg("Migration limit exhausted")), + RetryReason::WorkerDraining => { + return Err(Error::msg("Worker-draining reselection limit exhausted")); + } + } // Once any chunks have arrived from a previous attempt, stamp // that worker's span as `migration_link` so the next worker's // span renders an OTel Link back to it. Guarded so the initial @@ -379,12 +449,22 @@ where .into()); } response_stream = Some(self.next_generate.generate(request).await); - if let Some(err) = response_stream.as_ref().unwrap().as_ref().err() - && is_migratable_for_request(&self.request, err.as_ref()) - { - tracing::warn!(error = %err, "Creating new stream, retrying"); - self.metrics.inc_migration_new_request(&self.model_name); - continue; + if let Some(err) = response_stream.as_ref().unwrap().as_ref().err() { + if is_worker_draining_for_request(&self.request, err.as_ref()) + && self.worker_draining_reselections_left > 0 + { + tracing::warn!(error = %err, "Selected worker is draining, reselecting"); + retry_reason = RetryReason::WorkerDraining; + continue; + } + if is_migratable_for_request(&self.request, err.as_ref()) { + tracing::warn!(error = %err, "Creating new stream, retrying"); + self.metrics.inc_migration_new_request(&self.model_name); + if self.retries_left > 0 { + retry_reason = RetryReason::Migration; + continue; + } + } } break; } @@ -394,9 +474,7 @@ where Ok(()) } Some(Err(err)) => Err(err), // should propagate original error if any - None => Err(Error::msg( - "Migration limit exhausted", // should propagate original error if any - )), + None => Err(Error::msg("Request stream was not created")), } } @@ -450,6 +528,13 @@ where } } +#[derive(Clone, Copy, Debug)] +enum RetryReason { + Initial, + Migration, + WorkerDraining, +} + #[cfg(test)] mod tests { use super::*; @@ -596,6 +681,21 @@ mod tests { assert!(is_migratable_for_request(&unpinned, &error)); } + #[test] + fn explicit_worker_pin_blocks_draining_reselection() { + let error = DynamoError::builder() + .error_type(ErrorType::WorkerDraining) + .message("worker is draining") + .build(); + let mut request = create_mock_request(1); + request.routing = Some(RoutingHints { + backend_instance_id: Some(7), + ..Default::default() + }); + + assert!(!is_worker_draining_for_request(&request, &error)); + } + // Helper to create a mock preprocessed request fn create_mock_request(max_tokens: u32) -> PreprocessedRequest { PreprocessedRequest::builder() @@ -641,6 +741,8 @@ mod tests { /// Fails on first call with NoResponders error, then succeeds on subsequent calls FailThenSuccess, FailThenSuccessWithAffinity, + /// A stale discovery view selects a draining worker before a healthy replacement. + WorkerDrainingThenSuccess, /// One addressed worker rejects admission, then a replacement succeeds. WorkerOverloadSequence { worker_ids: Vec, @@ -753,6 +855,18 @@ mod tests { .await } } + MockBehavior::WorkerDrainingThenSuccess => { + if call_num == 0 { + return Err(anyhow::anyhow!( + DynamoError::builder() + .error_type(ErrorType::WorkerDraining) + .message("selected worker is draining") + .build() + )); + } + self.send_responses(responses_already_generated, self.num_responses) + .await + } MockBehavior::WorkerOverloadSequence { worker_ids } => { let excluded = preprocessed_request .migration_state @@ -1036,6 +1150,28 @@ mod tests { assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); } + #[tokio::test] + async fn draining_worker_is_reselected_when_migration_is_disabled() { + let context_id = uuid::Uuid::new_v4().to_string(); + let mock_engine = Arc::new(MockEngine::new( + MockBehavior::WorkerDrainingThenSuccess, + 1, + 100, + context_id.clone(), + )); + let calls = Arc::clone(&mock_engine.call_count); + let request = + Context::with_id_and_metadata(create_mock_request(1), context_id, BTreeMap::new()); + + let migration = Migration::new(0, None, TEST_MODEL.to_string(), Arc::new(Metrics::new())); + let mut stream = migration.generate(request, mock_engine).await.unwrap(); + let responses = stream.by_ref().collect::>().await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!(responses.len(), 1); + assert!(responses[0].error.is_none()); + } + #[tokio::test] async fn test_migration_preserves_session_affinity_across_retry() { let context_id = uuid::Uuid::new_v4().to_string(); diff --git a/lib/runtime/src/engine_routes.rs b/lib/runtime/src/engine_routes.rs index cd2adf87c839..7380d8fc2b63 100644 --- a/lib/runtime/src/engine_routes.rs +++ b/lib/runtime/src/engine_routes.rs @@ -4,7 +4,10 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, RwLock}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use parking_lot::RwLock; /// Callback type for engine routes (async) /// Takes JSON body, returns JSON response (or error) wrapped in a Future @@ -28,6 +31,7 @@ pub enum EngineRouteMethod { pub struct EngineRoute { callback: EngineRouteCallback, method: Option, + registration_id: u64, } impl EngineRoute { @@ -44,9 +48,32 @@ impl EngineRoute { /// /// This registry stores callbacks that handle requests to `/engine/*` routes. /// Routes are registered from Python via `runtime.register_engine_route()`. -#[derive(Clone, Default)] +#[derive(Clone)] pub struct EngineRouteRegistry { routes: Arc>>, + next_registration_id: Arc, +} + +impl Default for EngineRouteRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Removes a scoped route when its owner goes away without deleting a newer +/// registration that reused the same path. +#[must_use = "dropping the registration removes the engine route"] +pub struct EngineRouteRegistration { + registry: EngineRouteRegistry, + route: String, + registration_id: u64, +} + +impl Drop for EngineRouteRegistration { + fn drop(&mut self) { + self.registry + .remove_if_current(&self.route, self.registration_id); + } } impl EngineRouteRegistry { @@ -54,6 +81,7 @@ impl EngineRouteRegistry { pub fn new() -> Self { Self { routes: Arc::new(RwLock::new(HashMap::new())), + next_registration_id: Arc::new(AtomicU64::new(1)), } } @@ -77,36 +105,69 @@ impl EngineRouteRegistry { self.register_inner(route, Some(method), callback); } + /// Register a method-restricted route whose lifetime is owned by the + /// returned guard. + pub fn register_scoped_method( + &self, + route: &str, + method: EngineRouteMethod, + callback: EngineRouteCallback, + ) -> EngineRouteRegistration { + let registration_id = self.register_inner(route, Some(method), callback); + EngineRouteRegistration { + registry: self.clone(), + route: route.to_string(), + registration_id, + } + } + fn register_inner( &self, route: &str, method: Option, callback: EngineRouteCallback, - ) { - let mut routes = self.routes.write().unwrap(); - let entry = EngineRoute { callback, method }; + ) -> u64 { + let registration_id = self.next_registration_id.fetch_add(1, Ordering::Relaxed); + let mut routes = self.routes.write(); + let entry = EngineRoute { + callback, + method, + registration_id, + }; if routes.insert(route.to_string(), entry).is_some() { tracing::warn!("Overwriting already-registered engine route: /engine/{route}"); } else { tracing::debug!("Registered engine route: /engine/{route}"); } + registration_id + } + + fn remove_if_current(&self, route: &str, registration_id: u64) { + let mut routes = self.routes.write(); + if routes + .get(route) + .is_some_and(|entry| entry.registration_id == registration_id) + { + routes.remove(route); + tracing::debug!("Unregistered engine route: /engine/{route}"); + } } /// Get callback for a route pub fn get(&self, route: &str) -> Option { - let routes = self.routes.read().unwrap(); + let routes = self.routes.read(); routes.get(route).map(EngineRoute::callback) } /// Get a route together with its method restriction. pub fn get_route(&self, route: &str) -> Option { - let routes = self.routes.read().unwrap(); + let routes = self.routes.read(); routes.get(route).cloned() } /// List all registered routes pub fn routes(&self) -> Vec { - let routes = self.routes.read().unwrap(); + let routes = self.routes.read(); routes.keys().cloned().collect() } } @@ -193,4 +254,35 @@ mod tests { assert_eq!(route.method(), Some(EngineRouteMethod::Post)); assert!(registry.get("drain").is_some()); } + + #[test] + fn scoped_registration_is_removed_on_drop() { + let registry = EngineRouteRegistry::new(); + let callback: EngineRouteCallback = + Arc::new(|_| Box::pin(async { Ok(serde_json::json!({})) })); + + let registration = + registry.register_scoped_method("drain", EngineRouteMethod::Post, callback); + assert!(registry.get("drain").is_some()); + + drop(registration); + assert!(registry.get("drain").is_none()); + } + + #[test] + fn stale_scoped_registration_does_not_remove_replacement() { + let registry = EngineRouteRegistry::new(); + let first: EngineRouteCallback = + Arc::new(|_| Box::pin(async { Ok(serde_json::json!({"version": 1})) })); + let second: EngineRouteCallback = + Arc::new(|_| Box::pin(async { Ok(serde_json::json!({"version": 2})) })); + + let stale = registry.register_scoped_method("status", EngineRouteMethod::Get, first); + let current = registry.register_scoped_method("status", EngineRouteMethod::Get, second); + drop(stale); + + assert!(registry.get("status").is_some()); + drop(current); + assert!(registry.get("status").is_none()); + } } diff --git a/lib/runtime/src/error.rs b/lib/runtime/src/error.rs index 8adcefc6b37f..b58260fa9110 100644 --- a/lib/runtime/src/error.rs +++ b/lib/runtime/src/error.rs @@ -62,6 +62,9 @@ pub enum ErrorType { /// Distinct from [`Self::ResourceExhausted`] so a request whose routing /// constraints permit reassignment can migrate; both surface as HTTP 529. WorkerOverloaded, + /// The selected worker is intentionally draining and cannot admit this request. + /// Frontends may reselect another worker without consuming the user migration budget. + WorkerDraining, /// No backend worker is currently available to handle the request. Unavailable, /// Error originating from a backend engine. @@ -80,6 +83,7 @@ impl fmt::Display for ErrorType { ErrorType::Cancelled => write!(f, "Cancelled"), ErrorType::ResourceExhausted => write!(f, "ResourceExhausted"), ErrorType::WorkerOverloaded => write!(f, "WorkerOverloaded"), + ErrorType::WorkerDraining => write!(f, "WorkerDraining"), ErrorType::Unavailable => write!(f, "Unavailable"), ErrorType::Backend(sub) => write!(f, "Backend{sub}"), } @@ -483,6 +487,7 @@ mod tests { "ResourceExhausted" ); assert_eq!(ErrorType::WorkerOverloaded.to_string(), "WorkerOverloaded"); + assert_eq!(ErrorType::WorkerDraining.to_string(), "WorkerDraining"); assert_eq!(ErrorType::Unavailable.to_string(), "Unavailable"); assert_eq!( ErrorType::Backend(BackendError::Unknown).to_string(), diff --git a/lib/runtime/src/pipeline/network/egress/addressed_router.rs b/lib/runtime/src/pipeline/network/egress/addressed_router.rs index d1da8218af63..338c7bd2f9e2 100644 --- a/lib/runtime/src/pipeline/network/egress/addressed_router.rs +++ b/lib/runtime/src/pipeline/network/egress/addressed_router.rs @@ -784,7 +784,7 @@ where /// /// Impls MUST surface faults as top-level [`crate::error::ErrorType`] variants /// (`CannotConnect` / `Disconnected` / `ConnectionTimeout` / `ResponseTimeout` / -/// `WorkerOverloaded` / `ResourceExhausted` / `Cancelled`), or +/// `WorkerOverloaded` / `WorkerDraining` / `ResourceExhausted` / `Cancelled`), or /// `wrap_with_fault_detection`'s /// report-down / overload / migration won't fire. /// diff --git a/lib/runtime/src/pipeline/network/egress/push_router.rs b/lib/runtime/src/pipeline/network/egress/push_router.rs index 21940f3a53b2..ea9e13961155 100644 --- a/lib/runtime/src/pipeline/network/egress/push_router.rs +++ b/lib/runtime/src/pipeline/network/egress/push_router.rs @@ -53,6 +53,14 @@ fn is_inhibited(err: &(dyn std::error::Error + 'static)) -> bool { match_error_chain(err, INHIBITED, &[]) } +fn should_skip_for_reselection(err: &(dyn std::error::Error + 'static)) -> bool { + match_error_chain( + err, + &[ErrorType::WorkerOverloaded, ErrorType::WorkerDraining], + &[], + ) +} + /// Read the backend response inactivity timeout from the environment. /// Reuses `DYN_HTTP_BACKEND_STREAM_TIMEOUT_SECS` — the same env var /// as the HTTP-layer safety net in `disconnect.rs`. @@ -1545,14 +1553,13 @@ where "Reporting instance {instance_id} down due to error: {err}" ); self.client.report_instance_down(instance_id); - } else if match_error_chain(err.as_ref(), &[ErrorType::WorkerOverloaded], &[]) { - // Backpressure: worker said "my queue is full, - // retry later". Mark overloaded so this FE skips it on - // the next selection; the next ActiveLoad event from the - // worker monitor overwrites the overloaded set from fresh - // metrics. This is NOT report_instance_down (fault path). + } else if should_skip_for_reselection(err.as_ref()) { + // Backpressure or intentional draining: skip this worker + // on the next selection. Discovery or the next fresh load + // event will converge the local view. This is not the + // report-instance-down fault path. tracing::debug!( - "Marking instance {instance_id} overloaded due to backpressure: {err}" + "Temporarily skipping instance {instance_id} after rejection: {err}" ); self.client.mark_overloaded_immediate(instance_id); } @@ -1569,13 +1576,18 @@ where let client = self.client.clone(); let client_for_timeout = self.client.clone(); let stream = stream.map(move |res| { - if let Some(err) = res.err() - && is_inhibited(&err) - { - tracing::debug!( - "Reporting instance {instance_id} down due to migratable error: {err}" - ); - client.report_instance_down(instance_id); + if let Some(err) = res.err() { + if is_inhibited(&err) { + tracing::debug!( + "Reporting instance {instance_id} down due to migratable error: {err}" + ); + client.report_instance_down(instance_id); + } else if should_skip_for_reselection(&err) { + tracing::debug!( + "Temporarily skipping instance {instance_id} after stream rejection: {err}" + ); + client.mark_overloaded_immediate(instance_id); + } } res }); @@ -1963,6 +1975,17 @@ mod tests { ); } + #[test] + fn draining_worker_is_skipped_without_being_quarantined() { + let err = DynamoError::builder() + .error_type(ErrorType::WorkerDraining) + .message("worker is draining") + .build(); + + assert!(should_skip_for_reselection(&err)); + assert!(!is_inhibited(&err)); + } + #[test] fn p2c_lifecycle_tracks_inflight_counts_with_shared_tracker() { let state = Arc::new(RoutingOccupancyState::default()); diff --git a/lib/runtime/src/system_status_server.rs b/lib/runtime/src/system_status_server.rs index a9b2663c6cc2..5667cd695eaa 100644 --- a/lib/runtime/src/system_status_server.rs +++ b/lib/runtime/src/system_status_server.rs @@ -666,12 +666,12 @@ async fn engine_route_handler( } }; - let method_allowed = match route.method() { + let is_method_allowed = match route.method() { None => true, Some(EngineRouteMethod::Get) => method == Method::GET, Some(EngineRouteMethod::Post) => method == Method::POST, }; - if !method_allowed { + if !is_method_allowed { tracing::debug!(%method, "method not allowed for /engine/{path}"); return ( StatusCode::METHOD_NOT_ALLOWED, @@ -708,7 +708,6 @@ async fn engine_route_handler( } }; - // Call callback (it's async, so await it) match callback(body_json).await { Ok(response) => { tracing::trace!("Engine route handler succeeded for /engine/{path}"); From e8fedabd1e98ab422ac98f4aa7439be05163c240 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Wed, 12 Aug 2026 10:04:49 +0800 Subject: [PATCH 5/8] style(rust): apply worker lifecycle formatting Signed-off-by: xianlubird --- lib/backend-common/src/lifecycle.rs | 16 +++++----------- lib/llm/src/migration.rs | 4 +--- 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/lib/backend-common/src/lifecycle.rs b/lib/backend-common/src/lifecycle.rs index ebe7a001c047..36ecdf3f7478 100644 --- a/lib/backend-common/src/lifecycle.rs +++ b/lib/backend-common/src/lifecycle.rs @@ -334,7 +334,9 @@ impl WorkerLifecycleController { .await .replace(DrainTask { cancel, handle }); debug_assert!( - previous.as_ref().is_none_or(|task| task.handle.is_finished()), + previous + .as_ref() + .is_none_or(|task| task.handle.is_finished()), "serving worker must not have active drain work" ); Some(started_rx) @@ -642,10 +644,7 @@ mod tests { Self::with_register_permits(unregister_permits, 1) } - fn with_register_permits( - unregister_permits: usize, - register_permits: usize, - ) -> Arc { + fn with_register_permits(unregister_permits: usize, register_permits: usize) -> Arc { Arc::new(Self { unregister_started: Semaphore::new(0), unregister_gate: Semaphore::new(unregister_permits), @@ -892,12 +891,7 @@ mod tests { let resume_controller = Arc::clone(&controller); let resume_call = tokio::spawn(async move { resume_controller.resume().await }); - discovery - .register_started - .acquire() - .await - .unwrap() - .forget(); + discovery.register_started.acquire().await.unwrap().forget(); let admitted = tracker .try_acquire() diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs index 6585677081ff..234a583c3308 100644 --- a/lib/llm/src/migration.rs +++ b/lib/llm/src/migration.rs @@ -411,9 +411,7 @@ where RetryReason::Migration if self.retries_left > 0 => { self.retries_left -= 1; } - RetryReason::WorkerDraining - if self.worker_draining_reselections_left > 0 => - { + RetryReason::WorkerDraining if self.worker_draining_reselections_left > 0 => { self.worker_draining_reselections_left -= 1; } RetryReason::Migration => return Err(Error::msg("Migration limit exhausted")), From 0362d7af61b771427e32da0c8df8d7031a96aed2 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Wed, 12 Aug 2026 10:15:18 +0800 Subject: [PATCH 6/8] fix(llm): remove redundant stream initialization Signed-off-by: xianlubird --- lib/llm/src/migration.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs index 234a583c3308..3ae5eeaa2a1e 100644 --- a/lib/llm/src/migration.rs +++ b/lib/llm/src/migration.rs @@ -404,8 +404,7 @@ where } async fn new_stream(&mut self, mut retry_reason: RetryReason) -> Result<()> { - let mut response_stream: Option>>> = None; - loop { + let response_stream = loop { match retry_reason { RetryReason::Initial => {} RetryReason::Migration if self.retries_left > 0 => { @@ -446,8 +445,8 @@ where .build() .into()); } - response_stream = Some(self.next_generate.generate(request).await); - if let Some(err) = response_stream.as_ref().unwrap().as_ref().err() { + let response_stream = self.next_generate.generate(request).await; + if let Some(err) = response_stream.as_ref().err() { if is_worker_draining_for_request(&self.request, err.as_ref()) && self.worker_draining_reselections_left > 0 { @@ -464,15 +463,14 @@ where } } } - break; - } + break response_stream; + }; match response_stream { - Some(Ok(next_stream)) => { + Ok(next_stream) => { self.next_stream = Some(next_stream); Ok(()) } - Some(Err(err)) => Err(err), // should propagate original error if any - None => Err(Error::msg("Request stream was not created")), + Err(err) => Err(err), // should propagate original error if any } } From 644c419b7b5f8d2db7e01e4f7fe237099a19dbd2 Mon Sep 17 00:00:00 2001 From: xianlubird Date: Thu, 13 Aug 2026 16:44:52 +0800 Subject: [PATCH 7/8] fix(backend): address worker drain review feedback Signed-off-by: xianlubird --- .../reference/components/worker-admin-api.mdx | 15 +- lib/backend-common/src/lifecycle.rs | 198 ++++++++++-------- lib/backend-common/src/worker.rs | 9 +- lib/llm/src/http/service/generate.rs | 65 ++++++ lib/runtime/src/engine_routes.rs | 69 +++++- 5 files changed, 256 insertions(+), 100 deletions(-) diff --git a/docs/fern/pages/reference/components/worker-admin-api.mdx b/docs/fern/pages/reference/components/worker-admin-api.mdx index e287739d0de6..8d5c5d26413f 100644 --- a/docs/fern/pages/reference/components/worker-admin-api.mdx +++ b/docs/fern/pages/reference/components/worker-admin-api.mdx @@ -2,7 +2,7 @@ # 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 deletion safety +subtitle: Worker-local lifecycle endpoints for draining, resuming, and checking worker state --- **Experimental.** Unified-backend workers expose lifecycle controls through the worker system-status @@ -19,7 +19,7 @@ operator network. | --- | --- | --- | | `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 and deletion-safety status | +| `GET` | `/engine/status` | Return the worker lifecycle status | Send an empty body or an empty JSON object. Each endpoint returns the same status shape: @@ -27,17 +27,14 @@ Send an empty body or an empty JSON object. Each endpoint returns the same statu { "state": "draining", "inflight_requests": 2, - "kv_transfers": "not_applicable", - "safe_to_delete": false, "discovery_registered": false } ``` -`state` is one of `serving`, `draining`, `drained`, or `stopping`. Delete a worker only when -`safe_to_delete` is `true`. For a disaggregated prefill worker, the worker reports `drained` only when -the backend confirms that both admitted requests and pending KV transfers are complete. If the -backend cannot confirm KV-transfer completion, `kv_transfers` remains `unknown` and -`safe_to_delete` remains `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 diff --git a/lib/backend-common/src/lifecycle.rs b/lib/backend-common/src/lifecycle.rs index 36ecdf3f7478..6612bdb2c7e2 100644 --- a/lib/backend-common/src/lifecycle.rs +++ b/lib/backend-common/src/lifecycle.rs @@ -47,34 +47,10 @@ impl WorkerLifecycleState { } } -#[repr(u8)] -#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] -#[serde(rename_all = "snake_case")] -enum KvTransferState { - NotApplicable = 0, - Unknown = 1, - Pending = 2, - Complete = 3, -} - -impl KvTransferState { - fn from_u8(value: u8) -> Self { - match value { - 0 => Self::NotApplicable, - 1 => Self::Unknown, - 2 => Self::Pending, - 3 => Self::Complete, - _ => unreachable!("invalid KV transfer state {value}"), - } - } -} - #[derive(Clone, Debug, Eq, PartialEq, Serialize)] pub(crate) struct WorkerLifecycleStatus { state: WorkerLifecycleState, inflight_requests: u64, - kv_transfers: KvTransferState, - safe_to_delete: bool, discovery_registered: bool, #[serde(skip_serializing_if = "Option::is_none")] last_error: Option, @@ -111,38 +87,46 @@ impl QuiescenceCheck for EngineKind { /// Shared request-admission and in-flight tracker. /// -/// The double-check around `fetch_add` closes the drain/admission race: a -/// request is either rejected after draining starts or counted so the drain -/// monitor waits for its response stream to be dropped. +/// Admission and the in-flight count share one atomic word, so closing +/// admission is linearizable with request acquisition. #[derive(Debug)] pub(crate) struct RequestTracker { - accepting: AtomicBool, - inflight: AtomicU64, + state: AtomicU64, changed: Notify, } +const ACCEPTING_BIT: u64 = 1 << 63; +const INFLIGHT_MASK: u64 = !ACCEPTING_BIT; + impl RequestTracker { pub(crate) fn new() -> Arc { Arc::new(Self { - accepting: AtomicBool::new(true), - inflight: AtomicU64::new(0), + state: AtomicU64::new(ACCEPTING_BIT), changed: Notify::new(), }) } pub(crate) fn try_acquire(self: &Arc) -> Result { - if !self.accepting.load(Ordering::Acquire) { - return Err(worker_draining_error( - "worker is not accepting new requests", - )); - } - - self.inflight.fetch_add(1, Ordering::AcqRel); - if !self.accepting.load(Ordering::Acquire) { - self.release(); - return Err(worker_draining_error( - "worker started draining before request admission completed", - )); + let mut current = self.state.load(Ordering::Acquire); + loop { + if current & ACCEPTING_BIT == 0 { + return Err(worker_draining_error( + "worker is not accepting new requests", + )); + } + assert!( + current & INFLIGHT_MASK < INFLIGHT_MASK, + "request tracker overflow" + ); + match self.state.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => current = observed, + } } Ok(RequestGuard { @@ -151,27 +135,28 @@ impl RequestTracker { } pub(crate) fn stop_accepting(&self) { - self.accepting.store(false, Ordering::Release); + self.state.fetch_and(INFLIGHT_MASK, Ordering::AcqRel); self.changed.notify_waiters(); } fn start_accepting(&self) { - self.accepting.store(true, Ordering::Release); + self.state.fetch_or(ACCEPTING_BIT, Ordering::AcqRel); self.changed.notify_waiters(); } fn is_accepting(&self) -> bool { - self.accepting.load(Ordering::Acquire) + self.state.load(Ordering::Acquire) & ACCEPTING_BIT != 0 } pub(crate) fn inflight(&self) -> u64 { - self.inflight.load(Ordering::Acquire) + self.state.load(Ordering::Acquire) & INFLIGHT_MASK } fn release(&self) { - let previous = self.inflight.fetch_sub(1, Ordering::AcqRel); - debug_assert!(previous > 0, "request tracker underflow"); - if previous == 1 { + let previous = self.state.fetch_sub(1, Ordering::AcqRel); + let previous_inflight = previous & INFLIGHT_MASK; + debug_assert!(previous_inflight > 0, "request tracker underflow"); + if previous_inflight == 1 { self.changed.notify_waiters(); } } @@ -207,7 +192,6 @@ impl Drop for RequestGuard { /// Coordinates the Admin API, request admission, discovery, and SIGTERM. pub(crate) struct WorkerLifecycleController { state: AtomicU8, - kv_transfers: AtomicU8, discovery_registered: AtomicBool, last_error: RwLock>, operation_lock: Mutex<()>, @@ -246,7 +230,6 @@ impl WorkerLifecycleController { ) -> Arc { Arc::new(Self { state: AtomicU8::new(WorkerLifecycleState::Serving as u8), - kv_transfers: AtomicU8::new(Self::initial_kv_state(mode) as u8), discovery_registered: AtomicBool::new(true), last_error: RwLock::new(None), operation_lock: Mutex::new(()), @@ -260,14 +243,6 @@ impl WorkerLifecycleController { }) } - fn initial_kv_state(mode: DisaggregationMode) -> KvTransferState { - if mode.is_prefill() { - KvTransferState::Unknown - } else { - KvTransferState::NotApplicable - } - } - fn state(&self) -> WorkerLifecycleState { WorkerLifecycleState::from_u8(self.state.load(Ordering::Acquire)) } @@ -277,8 +252,6 @@ impl WorkerLifecycleController { WorkerLifecycleStatus { state, inflight_requests: self.tracker.inflight(), - kv_transfers: KvTransferState::from_u8(self.kv_transfers.load(Ordering::Acquire)), - safe_to_delete: state == WorkerLifecycleState::Drained, discovery_registered: self.discovery_registered.load(Ordering::Acquire), last_error: self.last_error.read().clone(), } @@ -287,26 +260,26 @@ impl WorkerLifecycleController { pub(crate) fn register_admin_routes( self: &Arc, registry: &EngineRouteRegistry, - ) -> AdminRouteRegistration { - let routes = vec![ - registry.register_scoped_method( + ) -> Result { + let routes = registry.try_register_scoped_methods(vec![ + ( "drain", EngineRouteMethod::Post, lifecycle_callback(Arc::clone(self), LifecycleAction::Drain), ), - registry.register_scoped_method( + ( "resume", EngineRouteMethod::Post, lifecycle_callback(Arc::clone(self), LifecycleAction::Resume), ), - registry.register_scoped_method( + ( "status", EngineRouteMethod::Get, lifecycle_callback(Arc::clone(self), LifecycleAction::Status), ), - ]; + ])?; tracing::info!("registered worker Admin API routes under /engine"); - AdminRouteRegistration { _routes: routes } + Ok(AdminRouteRegistration { _routes: routes }) } pub(crate) async fn drain(self: &Arc) -> Result { @@ -314,8 +287,6 @@ impl WorkerLifecycleController { let started = match self.state() { WorkerLifecycleState::Serving => { self.last_error.write().take(); - self.kv_transfers - .store(Self::initial_kv_state(self.mode) as u8, Ordering::Release); self.state .store(WorkerLifecycleState::Draining as u8, Ordering::Release); let generation = self.drain_generation.fetch_add(1, Ordering::AcqRel) + 1; @@ -470,8 +441,6 @@ impl WorkerLifecycleController { } self.last_error.write().take(); - self.kv_transfers - .store(Self::initial_kv_state(self.mode) as u8, Ordering::Release); self.state .store(WorkerLifecycleState::Serving as u8, Ordering::Release); Ok(self.status()) @@ -552,8 +521,8 @@ impl WorkerLifecycleController { Ok(Some(true)) }; - // Publish both KV state and Drained under the same operation - // fence used by resume. A stale monitor may finish an engine + // Publish Drained under the same operation fence used by + // resume. A stale monitor may finish an engine // check, but it can never publish after its generation ends. let _operation = tokio::select! { _ = cancel.cancelled() => return, @@ -567,24 +536,13 @@ impl WorkerLifecycleController { match quiescence { Ok(Some(true)) => { - if self.mode.is_prefill() { - self.kv_transfers - .store(KvTransferState::Complete as u8, Ordering::Release); - } self.state .store(WorkerLifecycleState::Drained as u8, Ordering::Release); tracing::info!("worker drained and is safe to delete"); return; } - Ok(Some(false)) => self - .kv_transfers - .store(KvTransferState::Pending as u8, Ordering::Release), - Ok(None) => self - .kv_transfers - .store(KvTransferState::Unknown as u8, Ordering::Release), + Ok(Some(false)) | Ok(None) => {} Err(error) => { - self.kv_transfers - .store(KvTransferState::Unknown as u8, Ordering::Release); *self.last_error.write() = Some(error.to_string()); } } @@ -788,6 +746,70 @@ mod tests { assert_eq!(error.error_type(), ErrorType::WorkerDraining); } + #[test] + fn concurrent_admission_is_counted_or_rejected_when_drain_closes() { + const CONTENDERS: usize = 16; + + let tracker = RequestTracker::new(); + let start = Arc::new(std::sync::Barrier::new(CONTENDERS + 1)); + let admitted = Arc::new(std::sync::Mutex::new(Vec::new())); + + std::thread::scope(|scope| { + for _ in 0..CONTENDERS { + let tracker = Arc::clone(&tracker); + let start = Arc::clone(&start); + let admitted = Arc::clone(&admitted); + scope.spawn(move || { + start.wait(); + if let Ok(guard) = tracker.try_acquire() { + admitted.lock().unwrap().push(guard); + } + }); + } + + start.wait(); + tracker.stop_accepting(); + }); + + let guards = admitted.lock().unwrap(); + assert_eq!(tracker.inflight(), guards.len() as u64); + assert!(tracker.try_acquire().is_err()); + drop(guards); + admitted.lock().unwrap().clear(); + assert_eq!(tracker.inflight(), 0); + } + + #[test] + fn second_lifecycle_controller_cannot_replace_admin_routes() { + let registry = EngineRouteRegistry::new(); + let (first, _) = test_controller( + MockDiscovery::new(1), + MockQuiescence::new(Some(true), 1), + DisaggregationMode::Aggregated, + Duration::ZERO, + ); + let (second, _) = test_controller( + MockDiscovery::new(1), + MockQuiescence::new(Some(true), 1), + DisaggregationMode::Aggregated, + Duration::ZERO, + ); + + let first_routes = first.register_admin_routes(®istry).unwrap(); + let error = match second.register_admin_routes(®istry) { + Ok(_) => panic!("a second lifecycle controller must not replace Admin API routes"), + Err(error) => error, + }; + + assert!(error.to_string().contains("already registered")); + assert!(registry.get("drain").is_some()); + assert!(registry.get("resume").is_some()); + assert!(registry.get("status").is_some()); + + drop(first_routes); + assert!(registry.routes().is_empty()); + } + #[tokio::test(start_paused = true)] async fn drain_keeps_admission_open_during_discovery_grace_period() { let discovery = MockDiscovery::new(1); @@ -871,8 +893,6 @@ mod tests { let status = controller.status(); assert_eq!(status.state, WorkerLifecycleState::Serving); - assert_eq!(status.kv_transfers, KvTransferState::Unknown); - assert!(!status.safe_to_delete); } #[tokio::test(start_paused = true)] diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 4a00a13b3d27..b723d3027aa0 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -1034,7 +1034,14 @@ impl Worker { format!("serve registration: {e}"), ) })?; - let admin_routes = lifecycle.register_admin_routes(endpoint.drt().engine_routes()); + let admin_routes = lifecycle + .register_admin_routes(endpoint.drt().engine_routes()) + .map_err(|error| { + err( + ErrorType::Backend(BackendError::Unknown), + format!("register worker Admin API routes: {error}"), + ) + })?; let serve_fut = started_endpoint.wait(); tokio::pin!(serve_fut); diff --git a/lib/llm/src/http/service/generate.rs b/lib/llm/src/http/service/generate.rs index 53b91cb7ce40..4a8c185acd99 100644 --- a/lib/llm/src/http/service/generate.rs +++ b/lib/llm/src/http/service/generate.rs @@ -504,6 +504,19 @@ async fn generate_dispatch( inflight_guard.mark_error(ErrorType::Cancelled); return generate_cancelled_response(); } + if super::metrics::request_was_unavailable(error.as_ref()) { + inflight_guard.mark_error(ErrorType::Unavailable); + tracing::warn!( + %request_id, + error = %format!("{error:#}"), + "no backend worker is available while folding generate stream" + ); + return generate_error_response( + StatusCode::SERVICE_UNAVAILABLE, + "service_unavailable", + "no backend worker is available".to_string(), + ); + } inflight_guard.mark_error(ErrorType::Internal); tracing::error!(%request_id, %error, "failed to fold generate stream"); generate_internal_error_response() @@ -529,7 +542,9 @@ mod tests { use crate::protocols::{Annotated, common::llm_backend::LLMEngineOutput}; use dynamo_runtime::{ engine::{AsyncEngine, ResponseStream}, + error::{DynamoError, ErrorType as DynamoErrorType}, pipeline::{Error, ManyOut, SingleIn}, + protocols::maybe_error::MaybeError, }; use futures::Stream; use tokio::sync::Notify; @@ -602,6 +617,8 @@ mod tests { struct CancelledEngine; + struct WorkerDrainingStreamEngine; + #[async_trait::async_trait] impl AsyncEngine, ManyOut>, Error> for CancelledEngine @@ -618,6 +635,23 @@ mod tests { } } + #[async_trait::async_trait] + impl AsyncEngine, ManyOut>, Error> + for WorkerDrainingStreamEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let error = DynamoError::builder() + .error_type(DynamoErrorType::WorkerDraining) + .message("worker started draining") + .build(); + let stream = futures::stream::iter([Annotated::from_err(error)]); + Ok(ResponseStream::new(Box::pin(stream), request.context())) + } + } + #[async_trait::async_trait] impl AsyncEngine, ManyOut>, Error> for TerminalEngine @@ -1144,6 +1178,37 @@ mod tests { assert_cancelled_dispatch_metrics(state.as_ref()); } + #[tokio::test] + async fn worker_draining_stream_error_returns_503() { + let engine: crate::types::openai::generate::GenerateStreamingEngine = + Arc::new(WorkerDrainingStreamEngine); + let service = HttpService::builder().build().unwrap(); + let state = service.state_clone(); + + let response = generate_dispatch( + engine, + dispatch_test_context(), + "req-worker-draining-stream".to_string(), + "test-model".to_string(), + state.clone(), + GenerateResponseOptions::default(), + ) + .await; + + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let metric_model = state.manager().metric_model_for("test-model"); + assert_eq!( + state.metrics_clone().get_request_counter( + metric_model, + &Endpoint::Generate, + &RequestType::Unary, + &Status::Error, + &ErrorType::Unavailable, + ), + 1 + ); + } + #[test] fn generate_control_shadow_carries_dp_rank_and_inverted_priority() { let request: GenerateRequest = serde_json::from_value(serde_json::json!({ diff --git a/lib/runtime/src/engine_routes.rs b/lib/runtime/src/engine_routes.rs index 7380d8fc2b63..9fc0db53199a 100644 --- a/lib/runtime/src/engine_routes.rs +++ b/lib/runtime/src/engine_routes.rs @@ -1,13 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use parking_lot::RwLock; +use thiserror::Error; /// Callback type for engine routes (async) /// Takes JSON body, returns JSON response (or error) wrapped in a Future @@ -69,6 +70,13 @@ pub struct EngineRouteRegistration { registration_id: u64, } +/// Returned when a set of scoped routes cannot be registered atomically. +#[derive(Debug, Error)] +#[error("engine route is already registered: /engine/{route}")] +pub struct EngineRouteConflict { + route: String, +} + impl Drop for EngineRouteRegistration { fn drop(&mut self) { self.registry @@ -121,6 +129,45 @@ impl EngineRouteRegistry { } } + /// Register method-restricted routes as one all-or-nothing operation. + /// + /// This is intended for route groups owned by one component. If any name + /// is already registered, no route in the group is changed. + pub fn try_register_scoped_methods( + &self, + registrations: Vec<(&str, EngineRouteMethod, EngineRouteCallback)>, + ) -> Result, EngineRouteConflict> { + let mut routes = self.routes.write(); + let mut names = HashSet::with_capacity(registrations.len()); + for (route, _, _) in ®istrations { + if routes.contains_key(*route) || !names.insert((*route).to_string()) { + return Err(EngineRouteConflict { + route: (*route).to_string(), + }); + } + } + + let mut guards = Vec::with_capacity(registrations.len()); + for (route, method, callback) in registrations { + let registration_id = self.next_registration_id.fetch_add(1, Ordering::Relaxed); + routes.insert( + route.to_string(), + EngineRoute { + callback, + method: Some(method), + registration_id, + }, + ); + tracing::debug!("Registered engine route: /engine/{route}"); + guards.push(EngineRouteRegistration { + registry: self.clone(), + route: route.to_string(), + registration_id, + }); + } + Ok(guards) + } + fn register_inner( &self, route: &str, @@ -285,4 +332,24 @@ mod tests { drop(current); assert!(registry.get("status").is_none()); } + + #[test] + fn scoped_method_group_rejects_conflicts_without_partial_registration() { + let registry = EngineRouteRegistry::new(); + let existing: EngineRouteCallback = + Arc::new(|_| Box::pin(async { Ok(serde_json::json!({"owner": "first"})) })); + let _existing = + registry.register_scoped_method("status", EngineRouteMethod::Get, existing); + + let callback: EngineRouteCallback = + Arc::new(|_| Box::pin(async { Ok(serde_json::json!({"owner": "second"})) })); + let result = registry.try_register_scoped_methods(vec![ + ("drain", EngineRouteMethod::Post, callback.clone()), + ("status", EngineRouteMethod::Get, callback), + ]); + + assert!(result.is_err()); + assert!(registry.get("drain").is_none()); + assert!(registry.get("status").is_some()); + } } From e8a7b94a2165c573af620659bb26da32d506cf2c Mon Sep 17 00:00:00 2001 From: xianlubird Date: Thu, 13 Aug 2026 16:51:25 +0800 Subject: [PATCH 8/8] style(runtime): apply engine route formatting Signed-off-by: xianlubird --- lib/runtime/src/engine_routes.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lib/runtime/src/engine_routes.rs b/lib/runtime/src/engine_routes.rs index 9fc0db53199a..512a8cd2858f 100644 --- a/lib/runtime/src/engine_routes.rs +++ b/lib/runtime/src/engine_routes.rs @@ -338,8 +338,7 @@ mod tests { let registry = EngineRouteRegistry::new(); let existing: EngineRouteCallback = Arc::new(|_| Box::pin(async { Ok(serde_json::json!({"owner": "first"})) })); - let _existing = - registry.register_scoped_method("status", EngineRouteMethod::Get, existing); + let _existing = registry.register_scoped_method("status", EngineRouteMethod::Get, existing); let callback: EngineRouteCallback = Arc::new(|_| Box::pin(async { Ok(serde_json::json!({"owner": "second"})) }));