From 57c102a15e900f47cb17587dcf2badabd904290c Mon Sep 17 00:00:00 2001 From: Julien Darve Date: Fri, 30 Oct 2026 14:07:56 -0700 Subject: [PATCH 1/5] feat(backend-common): withdraw and recover workers on engine loss Add optional availability(), recover() and on_recovered() engine hooks. The worker withdraws its endpoint while the engine is NotReady, also withdraws the model card while it is Lost, calls recover(), re-registers fresh metadata with the caller's taints, and calls on_recovered() before the endpoint is routable again. Operator pause survives recovery. Pause, resume, taint, recovery and shutdown discovery writes share one lock, record an unknown state before each write so a dropped HTTP request cannot leave a stale record, and skip endpoint writes that discovery already matches. A prefill worker whose engine is lost skips the shutdown drain. Signed-off-by: Julien Darve --- lib/backend-common/AGENTS.md | 16 +- lib/backend-common/src/engine.rs | 42 ++- lib/backend-common/src/lib.rs | 12 +- lib/backend-common/src/worker.rs | 601 +++++++++++++++++++++++++++++-- 4 files changed, 623 insertions(+), 48 deletions(-) diff --git a/lib/backend-common/AGENTS.md b/lib/backend-common/AGENTS.md index 93e88d7b65e6..22ee52ce4299 100644 --- a/lib/backend-common/AGENTS.md +++ b/lib/backend-common/AGENTS.md @@ -29,7 +29,7 @@ parse args, start engine, wire Prometheus serve requests drain-poll sh return engine return metadata (optional) (concurrent) predicate release ``` -The trait has twelve methods. `from_args` is NOT on the trait — each +The trait has seventeen methods. `from_args` is NOT on the trait — each backend exposes a backend-specific constructor (typically a sync `from_args(argv) -> Result<(Self, WorkerConfig)>` inherent method). This keeps the trait fully object-safe without a `where Self: Sized` @@ -123,6 +123,20 @@ opt-out and lets `run.rs` stay non-generic. startup**. Engines that publish their own discovery records stash it (e.g. vLLM dynamic LoRA via `register_model` / `unregister_model`). Mirrors the `on_publisher_ready` handoff idiom. +- `availability()` / `recover()` / `on_recovered()` — optional, for + engines that can fail while the worker lives (e.g. a sidecar's separate + engine process). `availability()` returns a `watch` of `Ready` / + `NotReady` / `Lost`; the default `None` keeps today's behavior. On + `NotReady` the Worker withdraws the endpoint; on `Lost` it also + withdraws the model card, calls `recover()`, registers the returned + metadata with the caller's taints, then calls `on_recovered()` so the + engine can restore records that need the card (e.g. LoRA adapters). It + re-registers the endpoint only while the engine is `Ready` and no + operator pause is in effect. In-flight requests are + never cancelled, and a `recover()` error stops the worker. Discovery + writes for pause, resume, taints, engine state and shutdown go through + one lock and are never cancelled mid-write; endpoint writes are skipped + when discovery already matches. ## Contract for `generate` diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index 35fb5ff7ccf4..00b2de329a2d 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -19,7 +19,7 @@ use async_trait::async_trait; use futures::stream::BoxStream; use tokio::sync::watch; -use crate::error::DynamoError; +use crate::error::{BackendError, DynamoError, ErrorType}; pub use dynamo_llm::first_token::FirstTokenNotifier; pub use dynamo_llm::kv_router::publisher::KvEventPublisher; @@ -184,6 +184,17 @@ pub struct EngineConfig { pub llm: Option, } +/// Whether the engine behind an [`LLMEngine`] can serve, as reported by +/// [`LLMEngine::availability`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EngineAvailability { + Ready, + /// Alive but not accepting work, for example while paused or asleep. + NotReady, + /// Gone or unreachable until [`LLMEngine::recover`] reconnects. + Lost, +} + /// Inference engine trait. /// /// Lifecycle: @@ -422,6 +433,35 @@ pub trait LLMEngine: Send + Sync + 'static { ) -> Result<(), DynamoError> { Ok(()) } + + /// Live availability of an engine that can fail independently of the + /// worker. `None` (default) keeps the worker registered until shutdown. + /// On `NotReady` the worker withdraws its endpoint; on `Lost` it also + /// withdraws the model card and calls [`recover`](LLMEngine::recover). + /// `generate` must reject new requests while `Lost`. + fn availability(&self) -> Option> { + None + } + + /// Reconnect after [`EngineAvailability::Lost`] and return fresh + /// registration metadata once the engine serves again, with availability + /// no longer `Lost`. Runs while the worker is out of discovery and is + /// dropped if the worker shuts down; in-flight requests are not cancelled. + /// An error stops the worker, so the engine bounds its own retries. KV event + /// sources and metrics bindings from `start` are kept. + async fn recover(&self) -> Result { + Err(DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::EngineShutdown)) + .message("engine does not support recovery") + .build()) + } + + /// Called after recovery re-registers the model card and before the + /// endpoint is routable again, to restore discovery records that depend on + /// the card (for example LoRA adapters). An error is logged. + async fn on_recovered(&self) -> Result<(), DynamoError> { + Ok(()) + } } /// Raw media-generation engine trait — the non-token sibling of [`LLMEngine`]. diff --git a/lib/backend-common/src/lib.rs b/lib/backend-common/src/lib.rs index fd749de7e1ab..3321be891102 100644 --- a/lib/backend-common/src/lib.rs +++ b/lib/backend-common/src/lib.rs @@ -34,12 +34,12 @@ pub use args::CommonArgs; pub use disagg::DisaggregationMode; pub use dynamo_llm::model_type::ModelInput; pub use engine::{ - AsyncEngineContext, BootstrapInfo, CompletionUsage, ComponentSnapshot, EngineConfig, - FinishReason, FirstTokenNotifier, GenerateContext, GuidedDecodingOptions, HEALTH_CHECK_KEY, - KV_HINT_TRANSFER_CAPABILITY_KEY, KvEventPublisher, KvEventSource, KvHint, KvHintAction, - KvSourceLocationsPayload, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, LlmRegistration, - LogProbs, Metrics, MetricsBindings, MetricsCtx, MultimodalData, OnPublisherReady, - OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, + AsyncEngineContext, BootstrapInfo, CompletionUsage, ComponentSnapshot, EngineAvailability, + EngineConfig, FinishReason, FirstTokenNotifier, GenerateContext, GuidedDecodingOptions, + HEALTH_CHECK_KEY, KV_HINT_TRANSFER_CAPABILITY_KEY, KvEventPublisher, KvEventSource, KvHint, + KvHintAction, KvSourceLocationsPayload, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, + LlmRegistration, LogProbs, Metrics, MetricsBindings, MetricsCtx, MultimodalData, + OnPublisherReady, OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, PromptTokensDetails, RawEngine, SamplingOptions, StopConditions, StopReason, TopLogprob, TopLogprobs, chunk, usage, }; diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 342236f7fcf3..6d21288ecbd8 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -35,7 +35,8 @@ use tokio_util::sync::CancellationToken; use crate::adapter::{EngineAdapter, RawEngineAdapter}; use crate::disagg::DisaggregationMode; use crate::engine::{ - EngineConfig, KvEventSource, LLMEngine, MetricsBindings, MetricsCtx, RawEngine, + EngineAvailability, EngineConfig, KvEventSource, LLMEngine, MetricsBindings, MetricsCtx, + RawEngine, }; use crate::error::{BackendError, DynamoError, ErrorType}; use crate::publisher::{PublisherHandles, setup_publishers}; @@ -296,6 +297,69 @@ enum EngineRouteLifecycle { ShuttingDown, } +/// The worker's discovery registration, shared by serving, operator pause and +/// engine recovery. The lock is held across discovery writes but never across +/// engine calls, and those writes are not cancelled, so the recorded state +/// matches discovery. +struct Registration { + /// Whether the endpoint instance is in discovery; `None` after a failed write. + endpoint: Option, + /// Whether the base model card is in discovery; `None` after a failed write. + card: Option, + is_paused: bool, + engine: EngineAvailability, + is_shutting_down: bool, + /// Caller-managed taints, restored when recovery re-registers the model card. + taints: Option>, +} + +impl Default for Registration { + fn default() -> Self { + Self { + endpoint: None, + card: None, + is_paused: false, + engine: EngineAvailability::Ready, + is_shutting_down: false, + taints: None, + } + } +} + +impl Registration { + fn is_routable(&self) -> bool { + !self.is_shutting_down && !self.is_paused && self.engine == EngineAvailability::Ready + } + + /// Register or withdraw the endpoint to match `is_routable`, skipping the + /// write when discovery already matches, and publish readiness to match. + async fn sync_endpoint( + &mut self, + endpoint: &dynamo_runtime::component::Endpoint, + ) -> anyhow::Result<()> { + let is_routable = self.is_routable(); + let mut result = Ok(()); + if self.endpoint != Some(is_routable) { + // Unknown until the write returns: an HTTP route's future can be + // dropped after discovery commits. + self.endpoint = None; + result = if is_routable { + endpoint.register_endpoint_instance().await + } else { + endpoint.unregister_endpoint_instance().await + }; + self.endpoint = result.is_ok().then_some(is_routable); + } + let status = if self.endpoint == Some(true) { + HealthStatus::Ready + } else { + HealthStatus::NotReady + }; + set_worker_health(endpoint, status); + result + } +} + /// The engine a [`Worker`] drives, tagged by request modality. Both variants /// share the lifecycle (driven via the forwarders below); they differ only in /// the serve-loop adapter: `Llm` → token pipeline ([`EngineAdapter`]), `Raw` → @@ -418,6 +482,30 @@ impl EngineKind { } } + fn availability(&self) -> Option> { + match self { + EngineKind::Llm(e) => e.availability(), + EngineKind::Raw(_) => None, + } + } + + async fn recover(&self) -> Result { + match self { + EngineKind::Llm(e) => e.recover().await, + EngineKind::Raw(_) => Err(err( + ErrorType::Backend(BackendError::EngineShutdown), + "raw engines do not support recovery", + )), + } + } + + async fn on_recovered(&self) -> Result<(), DynamoError> { + match self { + EngineKind::Llm(e) => e.on_recovered().await, + EngineKind::Raw(_) => Ok(()), + } + } + /// Raw media engines (image/video/audio) register name-only — the engine /// loads the model itself and the model has no LLM artifacts (tokenizer / /// chat template / config.json) for Dynamo to fetch. @@ -443,6 +531,7 @@ pub struct Worker { /// Serializes controls that mutate discovery registration and shutdown's /// final transition, preventing a resume from re-registering a stale worker. engine_route_mutation: Arc>, + registration: Arc>, /// Signals in-flight Rust administrative route futures to stop. Engine /// adapters that detach work (such as a separately scheduled language /// runtime task) remain responsible for cancelling that work themselves. @@ -481,6 +570,7 @@ impl Worker { EngineRouteLifecycle::Starting, )), engine_route_mutation: Arc::new(tokio::sync::Mutex::new(())), + registration: Arc::default(), engine_route_shutdown: CancellationToken::new(), publishers: None, lifecycle: None, @@ -855,6 +945,7 @@ impl Worker { endpoint.clone(), self.engine_route_lifecycle.clone(), self.engine_route_mutation.clone(), + self.registration.clone(), self.engine_route_shutdown.clone(), ); // Namespace control routes under `/engine/control/` so they @@ -915,6 +1006,7 @@ impl Worker { async fn begin_engine_route_shutdown(&self) { self.engine_route_shutdown.cancel(); let _mutation = self.engine_route_mutation.lock().await; + self.registration.lock().await.is_shutting_down = true; let mut lifecycle = self.engine_route_lifecycle.write().await; *lifecycle = EngineRouteLifecycle::ShuttingDown; } @@ -929,6 +1021,7 @@ impl Worker { model_taint_update_callback( endpoint.clone(), self.engine_route_lifecycle.clone(), + self.registration.clone(), self.engine_route_shutdown.clone(), ), ); @@ -939,11 +1032,14 @@ impl Worker { /// pre-serve (mid-start signal) and the serve loop's signal arm. async fn orchestrator_steps(&mut self, endpoint: &dynamo_runtime::component::Endpoint) { endpoint.drt().runtime().mark_shutting_down(); - if let Err(e) = endpoint.unregister_endpoint_instance().await { + let mut registration = self.registration.lock().await; + registration.is_shutting_down = true; + if let Err(e) = registration.sync_endpoint(endpoint).await { tracing::warn!(error = %e, "discovery unregister failed"); } else { tracing::info!("Endpoint unregistered from discovery"); } + drop(registration); self.run_engine_shutdown_steps().await; } @@ -1052,7 +1148,7 @@ impl Worker { self.config.model_input, None, Some(worker_type), - needs, + needs.clone(), ) .await .map_err(|e| { @@ -1061,6 +1157,7 @@ impl Worker { format!("model attach: {e}"), ) })?; + self.registration.lock().await.card = Some(true); tracing::debug!("model registered with discovery"); self.register_engine_controls(&endpoint).await?; @@ -1217,6 +1314,7 @@ impl Worker { return Ok(()); } }; + self.registration.lock().await.endpoint = Some(true); // A signal can arrive while primary registration is in flight. Keep // routes closed and tear the endpoint back down rather than briefly @@ -1282,37 +1380,73 @@ impl Worker { // ready before this point; drop it here, because the write below // publishes readiness through the very signal it suppresses. drop(readiness_hold); - set_worker_health(&endpoint, HealthStatus::Ready); + // A pause control may already have withdrawn the endpoint. + if let Err(error) = self + .registration + .lock() + .await + .sync_endpoint(&endpoint) + .await + { + tracing::warn!(%error, "failed to sync endpoint registration"); + } let serve_fut = primary_endpoint.wait(); tokio::pin!(serve_fut); - let serve_result = tokio::select! { - biased; - result = &mut serve_fut => { - match result { - // Endpoint exited cleanly (e.g. DRT primary token - // cancelled it) — run the orchestrator so drain/cleanup - // don't race transport teardown. - Ok(()) => { - tracing::info!( - "Endpoint completed gracefully; running shutdown orchestration" - ); - Ok(()) + // Apply the engine's current availability first, then each change. + let mut availability = self.engine.availability(); + if let Some(receiver) = availability.as_mut() { + receiver.mark_changed(); + } + let serve_result = loop { + tokio::select! { + biased; + result = &mut serve_fut => { + break match result { + // Endpoint exited cleanly (e.g. DRT primary token + // cancelled it) — run the orchestrator so drain/cleanup + // don't race transport teardown. + Ok(()) => { + tracing::info!( + "Endpoint completed gracefully; running shutdown orchestration" + ); + Ok(()) + } + // Serve errored; cleanup_once in run() is the safety net. + Err(e) => { + Err(err( + ErrorType::Backend(BackendError::Unknown), + format!("serve: {e}"), + )) + } + }; + } + _ = shutdown.cancelled() => { + tracing::info!("Received shutdown signal; running graceful orchestration"); + break Ok(()); + } + state = next_availability(&mut availability) => { + if state != EngineAvailability::Lost { + let mut registration = self.registration.lock().await; + registration.engine = state; + if let Err(error) = registration.sync_endpoint(&endpoint).await { + tracing::warn!(%error, ?state, "failed to sync endpoint registration"); + } + continue; + } + let recovered = self + .recover_engine(&endpoint, model_type, worker_type, &needs, &shutdown) + .await; + if let Err(error) = recovered { + break Err(error); } - // Serve errored; cleanup_once in run() is the safety net. - Err(e) => { - Err(err( - ErrorType::Backend(BackendError::Unknown), - format!("serve: {e}"), - )) + // Re-apply whatever the engine reports now. + if let Some(receiver) = availability.as_mut() { + receiver.mark_changed(); } } } - _ = shutdown.cancelled() => { - tracing::info!("Received shutdown signal; running graceful orchestration"); - Ok(()) - } }; // Cancel accepted Rust route futures, wait for their shared lifecycle @@ -1334,6 +1468,81 @@ impl Worker { serve_result } + /// Withdraw the endpoint and model card, reconnect through + /// [`LLMEngine::recover`], and register a fresh model card with the caller's + /// taints. The endpoint stays withdrawn until the serve loop sees the + /// engine `Ready`; in-flight requests keep running throughout. + async fn recover_engine( + &mut self, + endpoint: &dynamo_runtime::component::Endpoint, + model_type: ModelType, + worker_type: WorkerType, + needs: &[Vec], + shutdown: &CancellationToken, + ) -> Result<(), DynamoError> { + tracing::warn!("Engine lost; withdrawing worker from discovery until it recovers"); + { + let mut registration = self.registration.lock().await; + registration.engine = EngineAvailability::Lost; + if let Err(error) = registration.sync_endpoint(endpoint).await { + tracing::warn!(%error, "failed to withdraw endpoint after engine loss"); + } + if registration.card != Some(false) { + let result = LocalModel::detach_from_endpoint(endpoint, None).await; + registration.card = result.is_ok().then_some(false); + if let Err(error) = result { + tracing::warn!(%error, "failed to withdraw model card after engine loss"); + } + } + } + + let engine_config = tokio::select! { + biased; + _ = shutdown.cancelled() => return Ok(()), + result = self.engine.recover() => result?, + }; + let mut local_model = + build_local_model(&self.config, &engine_config, self.engine.is_raw()).await?; + + let mut registration = self.registration.lock().await; + local_model + .attach( + endpoint, + model_type, + self.config.model_input, + None, + Some(worker_type), + needs.to_vec(), + ) + .await + .map_err(|e| { + err( + ErrorType::Backend(BackendError::Unknown), + format!("model attach after engine recovery: {e}"), + ) + })?; + registration.card = Some(true); + if let Some(taints) = registration.taints.clone() { + update_model_taints(endpoint, taints).await.map_err(|e| { + err( + ErrorType::Backend(BackendError::Unknown), + format!("restore model taints after engine recovery: {e}"), + ) + })?; + } + drop(registration); + if matches!(self.engine, EngineKind::Llm(_)) { + self.admission_gate + .get_or_create(endpoint.component()) + .record_model_card(local_model.card()); + } + if let Err(error) = self.engine.on_recovered().await { + tracing::warn!(%error, "engine on_recovered failed"); + } + tracing::info!(model = %engine_config.model, "Engine recovered; model card re-registered"); + Ok(()) + } + /// Engine-facing shutdown sequence: grace period sleep → drain loop on /// `engine.is_quiescent()` → `cleanup_once()`. Each engine step swallows /// non-fatal failures so a misbehaving engine can't block the worker @@ -1372,7 +1581,10 @@ impl Worker { /// expires. Budget = `DYN_PREFILL_DRAIN_TIMEOUT_S` capped at /// `graceful_shutdown_timeout - CLEANUP_RESERVE_S`. async fn drain_until_idle_or_deadline(&self) { - if !self.config.disaggregation_mode.is_prefill() { + // A lost engine has no KV transfers left to finish. + if !self.config.disaggregation_mode.is_prefill() + || self.registration.lock().await.engine == EngineAvailability::Lost + { return; } let configured = drain_timeout_secs(); @@ -1427,6 +1639,18 @@ impl Worker { } } +/// The engine's next availability change; pending forever without one. +async fn next_availability( + availability: &mut Option>, +) -> EngineAvailability { + if let Some(receiver) = availability + && receiver.changed().await.is_ok() + { + return *receiver.borrow_and_update(); + } + std::future::pending().await +} + /// Publish worker readiness on both layers the runtime's health route reads. /// /// `SystemHealth::get_health_status` consults `use_endpoint_health_status`, then @@ -1750,11 +1974,13 @@ fn parse_model_taint_update_request(body: serde_json::Value) -> anyhow::Result>, + registration: Arc>, route_shutdown: CancellationToken, ) -> EngineRouteCallback { Arc::new(move |body| { let endpoint = endpoint.clone(); let route_lifecycle = route_lifecycle.clone(); + let registration = registration.clone(); let route_shutdown = route_shutdown.clone(); Box::pin(async move { let taints = parse_model_taint_update_request(body)?; @@ -1763,13 +1989,20 @@ fn model_taint_update_callback( Ok(lifecycle) => lifecycle, Err(response) => return Ok(response), }; - tokio::select! { + let mut registration = tokio::select! { biased; _ = route_shutdown.cancelled() => { return Ok(engine_route_lifecycle_error(EngineRouteLifecycle::ShuttingDown)); } - result = update_model_taints(&endpoint, taints.clone()) => result?, + registration = registration.lock() => registration, + }; + // Recorded first so a dropped request cannot leave an older set to + // restore; while recovery has the card withdrawn, they apply on re-registration. + registration.taints = Some(taints.clone()); + if registration.card != Some(false) { + update_model_taints(&endpoint, taints.clone()).await?; } + drop(registration); let mut response_taints: Vec<_> = taints.into_iter().collect(); response_taints.sort(); @@ -1827,6 +2060,7 @@ fn engine_update_callback( }) } +#[allow(clippy::too_many_arguments)] fn wrap_engine_control_callback( control_name: String, callback: EngineRouteCallback, @@ -1834,6 +2068,7 @@ fn wrap_engine_control_callback( endpoint: dynamo_runtime::component::Endpoint, route_lifecycle: Arc>, route_mutation: Arc>, + registration: Arc>, route_shutdown: CancellationToken, ) -> EngineRouteCallback { let policy = engine_control_policy(&control_name); @@ -1844,6 +2079,7 @@ fn wrap_engine_control_callback( let control_name = control_name.clone(); let route_lifecycle = route_lifecycle.clone(); let route_mutation = route_mutation.clone(); + let registration = registration.clone(); let route_shutdown = route_shutdown.clone(); Box::pin(async move { if let Some(response) = control_request_body_error(&body) { @@ -1881,23 +2117,27 @@ fn wrap_engine_control_callback( Ok(lifecycle) => lifecycle, Err(response) => return Ok(response), }; - let unregister_result = tokio::select! { + let mut registration = tokio::select! { biased; _ = route_shutdown.cancelled() => { return Ok(engine_route_lifecycle_error(EngineRouteLifecycle::ShuttingDown)); } - result = endpoint.unregister_endpoint_instance() => result, + registration = registration.lock() => registration, }; - if let Err(e) = unregister_result { + registration.is_paused = true; + if let Err(e) = registration.sync_endpoint(&endpoint).await { + registration.is_paused = false; + if let Err(error) = registration.sync_endpoint(&endpoint).await { + tracing::warn!(%error, "failed to restore endpoint after a failed pause"); + } return Ok(control_error_response(format!( "failed to unregister endpoint before /engine/control/{control_name}: {e}" ))); } // Out of discovery, so no longer routable. Whether the - // control itself then succeeds or fails, the endpoint is - // left unregistered, so readiness stays withdrawn until a - // resume control re-registers it. - set_worker_health(&endpoint, HealthStatus::NotReady); + // control itself then succeeds or fails, the endpoint and + // readiness stay withdrawn until a resume control succeeds. + drop(registration); let callback_result = tokio::select! { biased; @@ -1955,14 +2195,16 @@ fn wrap_engine_control_callback( } return Ok(response); } - let register_result = tokio::select! { + let mut registration = tokio::select! { biased; _ = route_shutdown.cancelled() => { return Ok(engine_route_lifecycle_error(EngineRouteLifecycle::ShuttingDown)); } - result = endpoint.register_endpoint_instance() => result, + registration = registration.lock() => registration, }; - if let Err(e) = register_result { + registration.is_paused = false; + // A lost or recovering engine registers once it is ready again. + if let Err(e) = registration.sync_endpoint(&endpoint).await { // The engine is serving-safe but absent from discovery. The // operation is idempotent: retrying /engine/control/{control_name} // re-registers without repeating the wake/resume work (the @@ -1972,7 +2214,6 @@ fn wrap_engine_control_callback( "engine resumed but re-registration failed after /engine/control/{control_name}: {e}; retry /engine/control/{control_name} to rejoin discovery" ))); } - set_worker_health(&endpoint, HealthStatus::Ready); Ok(response) } } @@ -4039,6 +4280,7 @@ mod handoff_and_lifecycle_tests { endpoint.clone(), worker.engine_route_lifecycle.clone(), worker.engine_route_mutation.clone(), + worker.registration.clone(), worker.engine_route_shutdown.clone(), ); let request = tokio::spawn(async move { callback(serde_json::json!({})).await.unwrap() }); @@ -4097,6 +4339,7 @@ mod handoff_and_lifecycle_tests { endpoint.clone(), worker.engine_route_lifecycle.clone(), worker.engine_route_mutation.clone(), + worker.registration.clone(), worker.engine_route_shutdown.clone(), ); let resume_request = @@ -4112,6 +4355,7 @@ mod handoff_and_lifecycle_tests { endpoint, worker.engine_route_lifecycle.clone(), worker.engine_route_mutation.clone(), + worker.registration.clone(), worker.engine_route_shutdown.clone(), ); let response = tokio::time::timeout( @@ -4577,4 +4821,281 @@ mod handoff_and_lifecycle_tests { "model taint updates must not be registered after a fatal handoff" ); } + + /// Engine whose availability the test drives. `recover` waits for the test + /// and reports a new context length, so re-registration must carry it. + struct RecoveringEngine { + availability: tokio::sync::watch::Sender, + release_recovery: Notify, + is_recovered: std::sync::atomic::AtomicBool, + } + + fn recovering_config(context_length: u32) -> EngineConfig { + EngineConfig { + model: "recovering-mock".to_string(), + llm: Some(crate::engine::LlmRegistration { + context_length: Some(context_length), + ..Default::default() + }), + ..EngineConfig::default() + } + } + + #[async_trait] + impl LLMEngine for RecoveringEngine { + async fn start(&self, _worker_id: u64) -> Result { + Ok(recovering_config(1024)) + } + + async fn generate( + &self, + _request: PreprocessedRequest, + _ctx: crate::engine::GenerateContext, + ) -> Result< + BoxStream<'static, Result>, + DynamoError, + > { + unreachable!("not used in recovery tests") + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + Ok(()) + } + + async fn supported_controls(&self) -> Result, DynamoError> { + Ok(vec![ + "pause_generation".to_string(), + "resume_generation".to_string(), + ]) + } + + async fn engine_control( + &self, + _control: String, + _body: serde_json::Value, + ) -> Result { + Ok(serde_json::json!({"status": "ok"})) + } + + fn availability(&self) -> Option> { + Some(self.availability.subscribe()) + } + + async fn recover(&self) -> Result { + self.release_recovery.notified().await; + self.availability.send_replace(EngineAvailability::Ready); + Ok(recovering_config(2048)) + } + + async fn on_recovered(&self) -> Result<(), DynamoError> { + self.is_recovered + .store(true, std::sync::atomic::Ordering::SeqCst); + Ok(()) + } + } + + struct RecoveryHarness { + engine: Arc, + drt: DistributedRuntime, + shutdown: CancellationToken, + run: tokio::task::JoinHandle>, + _discovery_dir: tempfile::TempDir, + } + + impl RecoveryHarness { + async fn start() -> Self { + use dynamo_runtime::distributed::{DiscoveryBackend, DistributedConfig}; + + let discovery_dir = tempfile::tempdir().unwrap(); + let config = DistributedConfig { + discovery_backend: DiscoveryBackend::KvStore( + dynamo_runtime::storage::kv::Selector::File(discovery_dir.path().to_owned()), + ), + ..DistributedConfig::process_local() + }; + let drt = DistributedRuntime::new(Runtime::from_current().unwrap(), config) + .await + .unwrap(); + let engine = Arc::new(RecoveringEngine { + availability: tokio::sync::watch::Sender::new(EngineAvailability::Ready), + release_recovery: Notify::new(), + is_recovered: Default::default(), + }); + let worker = Worker::new( + engine.clone(), + WorkerConfig { + enable_kv_routing: false, + ..WorkerConfig::default() + }, + ); + let shutdown = CancellationToken::new(); + let run = tokio::spawn(worker.run_with_drt(drt.clone(), shutdown.clone())); + let harness = Self { + engine, + drt, + shutdown, + run, + _discovery_dir: discovery_dir, + }; + harness.wait_for(1, Some(1024)).await; + harness + } + + async fn discovered(&self) -> (usize, Option) { + use dynamo_runtime::discovery::DiscoveryInstance; + + let discovery = self.drt.discovery(); + let config = WorkerConfig::default(); + let endpoints = discovery + .list(DiscoveryQuery::Endpoint { + namespace: config.namespace.clone(), + component: config.component.clone(), + endpoint: config.endpoint.clone(), + }) + .await + .unwrap(); + let models = discovery + .list(DiscoveryQuery::EndpointModels { + namespace: config.namespace, + component: config.component, + endpoint: config.endpoint, + }) + .await + .unwrap(); + let card = models.into_iter().find_map(|model| match model { + DiscoveryInstance::Model { card_json, .. } => Some(card_json), + _ => None, + }); + (endpoints.len(), card) + } + + /// `context_length: None` waits for no model card. + async fn wait_for( + &self, + endpoints: usize, + context_length: Option, + ) -> Option { + let mut last = None; + for _ in 0..200 { + let (count, card) = self.discovered().await; + let card_context = card.as_ref().map(|card| { + card["runtime_config"]["context_length"] + .as_u64() + .unwrap_or_default() + }); + if count == endpoints && card_context == context_length.map(u64::from) { + return card; + } + last = Some((count, card_context)); + tokio::time::sleep(Duration::from_millis(25)).await; + } + panic!("expected {endpoints} endpoints and card {context_length:?}; last saw {last:?}"); + } + + async fn route(&self, name: &str, body: serde_json::Value) -> serde_json::Value { + let callback = self.drt.engine_routes().get(name).unwrap(); + let response = callback(body).await.unwrap(); + assert!(!control_response_is_error(&response), "{name}: {response}"); + response + } + + async fn stop(self) { + self.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(30), self.run) + .await + .expect("worker must shut down") + .unwrap() + .expect("worker must exit cleanly"); + } + } + + /// Regression: a worker kept advertising an engine that had stopped + /// serving or gone away. `NotReady` must withdraw only the endpoint; `Lost` + /// must also withdraw the model card until recovery registers fresh metadata. + #[tokio::test] + async fn engine_loss_withdraws_until_recovery_registers_fresh_metadata() { + let harness = RecoveryHarness::start().await; + + harness + .engine + .availability + .send_replace(EngineAvailability::NotReady); + harness.wait_for(0, Some(1024)).await; + harness + .engine + .availability + .send_replace(EngineAvailability::Ready); + harness.wait_for(1, Some(1024)).await; + + harness + .engine + .availability + .send_replace(EngineAvailability::Lost); + harness.wait_for(0, None).await; + harness.engine.release_recovery.notify_one(); + harness.wait_for(1, Some(2048)).await; + assert!( + harness + .engine + .is_recovered + .load(std::sync::atomic::Ordering::SeqCst) + ); + + harness.stop().await; + } + + /// Regression: recovery could re-register a worker the operator had paused, + /// or drop the caller's taints when it re-registered the model card. + #[tokio::test] + async fn operator_pause_and_taints_survive_engine_recovery() { + let harness = RecoveryHarness::start().await; + harness + .route( + MODEL_TAINT_UPDATE_ROUTE, + serde_json::json!({"taints": ["capacity/fast"]}), + ) + .await; + harness + .route("control/pause_generation", serde_json::json!({})) + .await; + harness.wait_for(0, Some(1024)).await; + + harness + .engine + .availability + .send_replace(EngineAvailability::Lost); + harness.wait_for(0, None).await; + harness.engine.release_recovery.notify_one(); + harness.wait_for(0, Some(2048)).await; + let has_taint = |card: &serde_json::Value| { + card["runtime_config"]["taints"] + .as_array() + .is_some_and(|taints| taints.contains(&serde_json::json!("capacity/fast"))) + }; + let mut card = None; + for _ in 0..200 { + card = harness.discovered().await.1; + if card.as_ref().is_some_and(has_taint) { + break; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + assert!( + card.as_ref().is_some_and(has_taint), + "recovery must restore caller taints: {card:?}" + ); + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!( + harness.discovered().await.0, + 0, + "a paused worker must stay withdrawn after recovery" + ); + + harness + .route("control/resume_generation", serde_json::json!({})) + .await; + harness.wait_for(1, Some(2048)).await; + + harness.stop().await; + } } From e929de3f1771ec4be0fff05bb312aa9b37c564d8 Mon Sep 17 00:00:00 2001 From: Julien Darve Date: Fri, 30 Oct 2026 17:32:36 -0700 Subject: [PATCH 2/5] fix(backend-common): stop engine recovery when the endpoint stops The serve loop waited on LLMEngine::recover() without watching its endpoint, so a runtime shutdown during recovery left the worker waiting and could re-register the model card after shutdown started. Race recover() against the endpoint and the shutdown token; the discovery writes before and after it still run to completion. Signed-off-by: Julien Darve --- lib/backend-common/src/worker.rs | 127 +++++++++++++++++++------------ 1 file changed, 79 insertions(+), 48 deletions(-) diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 6d21288ecbd8..96a08c39a770 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -1402,26 +1402,7 @@ impl Worker { let serve_result = loop { tokio::select! { biased; - result = &mut serve_fut => { - break match result { - // Endpoint exited cleanly (e.g. DRT primary token - // cancelled it) — run the orchestrator so drain/cleanup - // don't race transport teardown. - Ok(()) => { - tracing::info!( - "Endpoint completed gracefully; running shutdown orchestration" - ); - Ok(()) - } - // Serve errored; cleanup_once in run() is the safety net. - Err(e) => { - Err(err( - ErrorType::Backend(BackendError::Unknown), - format!("serve: {e}"), - )) - } - }; - } + result = &mut serve_fut => break endpoint_exit(result), _ = shutdown.cancelled() => { tracing::info!("Received shutdown signal; running graceful orchestration"); break Ok(()); @@ -1435,8 +1416,27 @@ impl Worker { } continue; } + self.withdraw_lost_engine(&endpoint).await; + let engine_config = tokio::select! { + biased; + result = &mut serve_fut => break endpoint_exit(result), + _ = shutdown.cancelled() => { + tracing::info!("Received shutdown signal; running graceful orchestration"); + break Ok(()); + } + result = self.engine.recover() => match result { + Ok(engine_config) => engine_config, + Err(error) => break Err(error), + }, + }; let recovered = self - .recover_engine(&endpoint, model_type, worker_type, &needs, &shutdown) + .register_recovered_engine( + &endpoint, + engine_config, + model_type, + worker_type, + &needs, + ) .await; if let Err(error) = recovered { break Err(error); @@ -1468,39 +1468,35 @@ impl Worker { serve_result } - /// Withdraw the endpoint and model card, reconnect through - /// [`LLMEngine::recover`], and register a fresh model card with the caller's - /// taints. The endpoint stays withdrawn until the serve loop sees the - /// engine `Ready`; in-flight requests keep running throughout. - async fn recover_engine( + /// Withdraw the endpoint and model card after engine loss; in-flight + /// requests keep running. + async fn withdraw_lost_engine(&self, endpoint: &dynamo_runtime::component::Endpoint) { + tracing::warn!("Engine lost; withdrawing worker from discovery until it recovers"); + let mut registration = self.registration.lock().await; + registration.engine = EngineAvailability::Lost; + if let Err(error) = registration.sync_endpoint(endpoint).await { + tracing::warn!(%error, "failed to withdraw endpoint after engine loss"); + } + if registration.card != Some(false) { + let result = LocalModel::detach_from_endpoint(endpoint, None).await; + registration.card = result.is_ok().then_some(false); + if let Err(error) = result { + tracing::warn!(%error, "failed to withdraw model card after engine loss"); + } + } + } + + /// Register a fresh model card from [`LLMEngine::recover`] with the + /// caller's taints. The endpoint stays withdrawn until the serve loop sees + /// the engine `Ready`. + async fn register_recovered_engine( &mut self, endpoint: &dynamo_runtime::component::Endpoint, + engine_config: EngineConfig, model_type: ModelType, worker_type: WorkerType, needs: &[Vec], - shutdown: &CancellationToken, ) -> Result<(), DynamoError> { - tracing::warn!("Engine lost; withdrawing worker from discovery until it recovers"); - { - let mut registration = self.registration.lock().await; - registration.engine = EngineAvailability::Lost; - if let Err(error) = registration.sync_endpoint(endpoint).await { - tracing::warn!(%error, "failed to withdraw endpoint after engine loss"); - } - if registration.card != Some(false) { - let result = LocalModel::detach_from_endpoint(endpoint, None).await; - registration.card = result.is_ok().then_some(false); - if let Err(error) = result { - tracing::warn!(%error, "failed to withdraw model card after engine loss"); - } - } - } - - let engine_config = tokio::select! { - biased; - _ = shutdown.cancelled() => return Ok(()), - result = self.engine.recover() => result?, - }; let mut local_model = build_local_model(&self.config, &engine_config, self.engine.is_raw()).await?; @@ -1639,6 +1635,22 @@ impl Worker { } } +fn endpoint_exit(result: anyhow::Result<()>) -> Result<(), DynamoError> { + match result { + // Endpoint exited cleanly (e.g. DRT primary token cancelled it) — run + // the orchestrator so drain/cleanup don't race transport teardown. + Ok(()) => { + tracing::info!("Endpoint completed gracefully; running shutdown orchestration"); + Ok(()) + } + // Serve errored; cleanup_once in run() is the safety net. + Err(e) => Err(err( + ErrorType::Backend(BackendError::Unknown), + format!("serve: {e}"), + )), + } +} + /// The engine's next availability change; pending forever without one. async fn next_availability( availability: &mut Option>, @@ -5044,6 +5056,25 @@ mod handoff_and_lifecycle_tests { harness.stop().await; } + /// Regression: a runtime shutdown during engine recovery left the worker + /// waiting on `recover()` instead of shutting down. + #[tokio::test] + async fn runtime_shutdown_stops_engine_recovery() { + let harness = RecoveryHarness::start().await; + harness + .engine + .availability + .send_replace(EngineAvailability::Lost); + harness.wait_for(0, None).await; + + harness.drt.shutdown(); + tokio::time::timeout(Duration::from_secs(30), harness.run) + .await + .expect("worker must shut down while recovery is pending") + .unwrap() + .expect("worker must exit cleanly"); + } + /// Regression: recovery could re-register a worker the operator had paused, /// or drop the caller's taints when it re-registered the model card. #[tokio::test] From 892dc3ac76e9f1b8ad1274ca2a089e1247dd65ab Mon Sep 17 00:00:00 2001 From: Julien Darve Date: Mon, 2 Nov 2026 13:20:26 -0800 Subject: [PATCH 3/5] fix(backend-common): retry failed endpoint syncs from the serve loop A discovery write that failed when the engine's availability changed was only logged, so one transient error could leave a healthy worker out of discovery until its next change. The serve loop now retries the sync every second until it succeeds. Also correct the registration docs, which said these writes are never cancelled, and the comment on skipping the drain for a lost engine. Signed-off-by: Julien Darve --- lib/backend-common/AGENTS.md | 5 ++-- lib/backend-common/src/worker.rs | 46 +++++++++++++++++++++----------- 2 files changed, 34 insertions(+), 17 deletions(-) diff --git a/lib/backend-common/AGENTS.md b/lib/backend-common/AGENTS.md index 22ee52ce4299..c27795a7065d 100644 --- a/lib/backend-common/AGENTS.md +++ b/lib/backend-common/AGENTS.md @@ -135,8 +135,9 @@ opt-out and lets `run.rs` stay non-generic. operator pause is in effect. In-flight requests are never cancelled, and a `recover()` error stops the worker. Discovery writes for pause, resume, taints, engine state and shutdown go through - one lock and are never cancelled mid-write; endpoint writes are skipped - when discovery already matches. + one lock. A write dropped with its HTTP request leaves its state unknown + until the next sync. Endpoint writes are skipped when discovery already + matches, and the serve loop retries its own failed writes. ## Contract for `generate` diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 50d2b1e5eea0..51e1e54e2f5a 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -60,6 +60,8 @@ const DRAIN_HEARTBEAT_INTERVAL_S: f64 = 5.0; /// Budget reserved for `cleanup()` so the drain loop can't consume the whole /// graceful-shutdown deadline and trip the hard-exit that skips cleanup. const CLEANUP_RESERVE_S: f64 = 5.0; +/// Delay before the serve loop retries a failed endpoint sync. +const SYNC_RETRY_INTERVAL: Duration = Duration::from_secs(1); /// Operator override for the health-check canary, mirrors the Python helper /// in `lib/bindings/python/src/dynamo/health_check.py`. @@ -299,8 +301,8 @@ enum EngineRouteLifecycle { /// The worker's discovery registration, shared by serving, operator pause and /// engine recovery. The lock is held across discovery writes but never across -/// engine calls, and those writes are not cancelled, so the recorded state -/// matches discovery. +/// engine calls. A write can be dropped with its HTTP request, so each write +/// records its state as unknown until it returns. struct Registration { /// Whether the endpoint instance is in discovery; `None` after a failed write. endpoint: Option, @@ -1389,15 +1391,13 @@ impl Worker { // publishes readiness through the very signal it suppresses. drop(readiness_hold); // A pause control may already have withdrawn the endpoint. - if let Err(error) = self - .registration - .lock() - .await - .sync_endpoint(&endpoint) - .await - { - tracing::warn!(%error, "failed to sync endpoint registration"); - } + let mut resync = sync_failed( + self.registration + .lock() + .await + .sync_endpoint(&endpoint) + .await, + ); let serve_fut = primary_endpoint.wait(); tokio::pin!(serve_fut); @@ -1419,9 +1419,7 @@ impl Worker { if state != EngineAvailability::Lost { let mut registration = self.registration.lock().await; registration.engine = state; - if let Err(error) = registration.sync_endpoint(&endpoint).await { - tracing::warn!(%error, ?state, "failed to sync endpoint registration"); - } + resync = sync_failed(registration.sync_endpoint(&endpoint).await); continue; } self.withdraw_lost_engine(&endpoint).await; @@ -1454,6 +1452,15 @@ impl Worker { receiver.mark_changed(); } } + _ = tokio::time::sleep(SYNC_RETRY_INTERVAL), if resync => { + resync = sync_failed( + self.registration + .lock() + .await + .sync_endpoint(&endpoint) + .await, + ); + } } }; @@ -1585,7 +1592,8 @@ impl Worker { /// expires. Budget = `DYN_PREFILL_DRAIN_TIMEOUT_S` capped at /// `graceful_shutdown_timeout - CLEANUP_RESERVE_S`. async fn drain_until_idle_or_deadline(&self) { - // A lost engine has no KV transfers left to finish. + // A lost engine can't report quiescence, and its KV memory lives in its + // own process, which this worker's cleanup doesn't free. if !self.config.disaggregation_mode.is_prefill() || self.registration.lock().await.engine == EngineAvailability::Lost { @@ -1659,6 +1667,14 @@ fn endpoint_exit(result: anyhow::Result<()>) -> Result<(), DynamoError> { } } +/// Log a failed serve-loop endpoint sync; true when it needs a retry. +fn sync_failed(result: anyhow::Result<()>) -> bool { + if let Err(error) = &result { + tracing::warn!(%error, "failed to sync endpoint registration; retrying"); + } + result.is_err() +} + /// The engine's next availability change; pending forever without one. async fn next_availability( availability: &mut Option>, From a68fccb848b815dcd46f516e11017bc07d1a4073 Mon Sep 17 00:00:00 2001 From: Julien Darve Date: Tue, 3 Nov 2026 11:39:20 -0800 Subject: [PATCH 4/5] feat(backend-common): drain and exit the worker when its engine is lost Replace the availability, recover and on_recovered hooks with one optional LLMEngine hook, engine_lost(). When it resolves, the Worker shuts down as on SIGTERM (withdraw from discovery, grace period, drain, cleanup) under the same shutdown deadline, and returns the engine's error so the process exits and its supervisor restarts it. Engines that don't implement it never resolve it, so their behavior is unchanged. Signed-off-by: Julien Darve --- lib/backend-common/AGENTS.md | 21 +- lib/backend-common/src/engine.rs | 50 +- lib/backend-common/src/lib.rs | 12 +- lib/backend-common/src/worker.rs | 789 +++++++------------------------ 4 files changed, 196 insertions(+), 676 deletions(-) diff --git a/lib/backend-common/AGENTS.md b/lib/backend-common/AGENTS.md index c27795a7065d..63f08743534c 100644 --- a/lib/backend-common/AGENTS.md +++ b/lib/backend-common/AGENTS.md @@ -29,7 +29,7 @@ parse args, start engine, wire Prometheus serve requests drain-poll sh return engine return metadata (optional) (concurrent) predicate release ``` -The trait has seventeen methods. `from_args` is NOT on the trait — each +The trait has twelve methods. `from_args` is NOT on the trait — each backend exposes a backend-specific constructor (typically a sync `from_args(argv) -> Result<(Self, WorkerConfig)>` inherent method). This keeps the trait fully object-safe without a `where Self: Sized` @@ -83,6 +83,10 @@ opt-out and lets `run.rs` stay non-generic. reliable idle signal); SGLang overrides it by counting in-flight prefill streams. The mode-gate, poll loop, and SIGTERM/SIGINT ownership live in the `Worker`. +- `engine_lost(&self) -> DynamoError` — optional, default never resolves. + Resolves once the engine can no longer serve. The `Worker` then shuts down + as on SIGTERM (withdraw, grace period, drain, `cleanup`) and returns this + error, so the process exits and its supervisor restarts it. - `cleanup(&self) -> Result<(), DynamoError>` — called exactly once. Runs after `start()` returns Ok on shutdown (even if registration / serve fails), **and** after `start()` raises — so implementations @@ -123,21 +127,6 @@ opt-out and lets `run.rs` stay non-generic. startup**. Engines that publish their own discovery records stash it (e.g. vLLM dynamic LoRA via `register_model` / `unregister_model`). Mirrors the `on_publisher_ready` handoff idiom. -- `availability()` / `recover()` / `on_recovered()` — optional, for - engines that can fail while the worker lives (e.g. a sidecar's separate - engine process). `availability()` returns a `watch` of `Ready` / - `NotReady` / `Lost`; the default `None` keeps today's behavior. On - `NotReady` the Worker withdraws the endpoint; on `Lost` it also - withdraws the model card, calls `recover()`, registers the returned - metadata with the caller's taints, then calls `on_recovered()` so the - engine can restore records that need the card (e.g. LoRA adapters). It - re-registers the endpoint only while the engine is `Ready` and no - operator pause is in effect. In-flight requests are - never cancelled, and a `recover()` error stops the worker. Discovery - writes for pause, resume, taints, engine state and shutdown go through - one lock. A write dropped with its HTTP request leaves its state unknown - until the next sync. Endpoint writes are skipped when discovery already - matches, and the serve loop retries its own failed writes. ## Contract for `generate` diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index 00b2de329a2d..c20abe18fcef 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -19,7 +19,7 @@ use async_trait::async_trait; use futures::stream::BoxStream; use tokio::sync::watch; -use crate::error::{BackendError, DynamoError, ErrorType}; +use crate::error::DynamoError; pub use dynamo_llm::first_token::FirstTokenNotifier; pub use dynamo_llm::kv_router::publisher::KvEventPublisher; @@ -184,17 +184,6 @@ pub struct EngineConfig { pub llm: Option, } -/// Whether the engine behind an [`LLMEngine`] can serve, as reported by -/// [`LLMEngine::availability`]. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum EngineAvailability { - Ready, - /// Alive but not accepting work, for example while paused or asleep. - NotReady, - /// Gone or unreachable until [`LLMEngine::recover`] reconnects. - Lost, -} - /// Inference engine trait. /// /// Lifecycle: @@ -291,6 +280,14 @@ pub trait LLMEngine: Send + Sync + 'static { Ok(None) } + /// Resolves once the engine can no longer serve, with the error the + /// `Worker` returns. The `Worker` then shuts down as on SIGTERM, withdrawing + /// from discovery, draining and running `cleanup`, so the process exits + /// and its supervisor can restart it. Default: never resolves. + async fn engine_lost(&self) -> DynamoError { + std::future::pending().await + } + /// Release all engine resources. Called exactly once. /// /// `Worker` guarantees: @@ -433,35 +430,6 @@ pub trait LLMEngine: Send + Sync + 'static { ) -> Result<(), DynamoError> { Ok(()) } - - /// Live availability of an engine that can fail independently of the - /// worker. `None` (default) keeps the worker registered until shutdown. - /// On `NotReady` the worker withdraws its endpoint; on `Lost` it also - /// withdraws the model card and calls [`recover`](LLMEngine::recover). - /// `generate` must reject new requests while `Lost`. - fn availability(&self) -> Option> { - None - } - - /// Reconnect after [`EngineAvailability::Lost`] and return fresh - /// registration metadata once the engine serves again, with availability - /// no longer `Lost`. Runs while the worker is out of discovery and is - /// dropped if the worker shuts down; in-flight requests are not cancelled. - /// An error stops the worker, so the engine bounds its own retries. KV event - /// sources and metrics bindings from `start` are kept. - async fn recover(&self) -> Result { - Err(DynamoError::builder() - .error_type(ErrorType::Backend(BackendError::EngineShutdown)) - .message("engine does not support recovery") - .build()) - } - - /// Called after recovery re-registers the model card and before the - /// endpoint is routable again, to restore discovery records that depend on - /// the card (for example LoRA adapters). An error is logged. - async fn on_recovered(&self) -> Result<(), DynamoError> { - Ok(()) - } } /// Raw media-generation engine trait — the non-token sibling of [`LLMEngine`]. diff --git a/lib/backend-common/src/lib.rs b/lib/backend-common/src/lib.rs index 3321be891102..fd749de7e1ab 100644 --- a/lib/backend-common/src/lib.rs +++ b/lib/backend-common/src/lib.rs @@ -34,12 +34,12 @@ pub use args::CommonArgs; pub use disagg::DisaggregationMode; pub use dynamo_llm::model_type::ModelInput; pub use engine::{ - AsyncEngineContext, BootstrapInfo, CompletionUsage, ComponentSnapshot, EngineAvailability, - EngineConfig, FinishReason, FirstTokenNotifier, GenerateContext, GuidedDecodingOptions, - HEALTH_CHECK_KEY, KV_HINT_TRANSFER_CAPABILITY_KEY, KvEventPublisher, KvEventSource, KvHint, - KvHintAction, KvSourceLocationsPayload, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, - LlmRegistration, LogProbs, Metrics, MetricsBindings, MetricsCtx, MultimodalData, - OnPublisherReady, OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, + AsyncEngineContext, BootstrapInfo, CompletionUsage, ComponentSnapshot, EngineConfig, + FinishReason, FirstTokenNotifier, GenerateContext, GuidedDecodingOptions, HEALTH_CHECK_KEY, + KV_HINT_TRANSFER_CAPABILITY_KEY, KvEventPublisher, KvEventSource, KvHint, KvHintAction, + KvSourceLocationsPayload, LLMEngine, LLMEngineOutput, LLMEngineOutputExt, LlmRegistration, + LogProbs, Metrics, MetricsBindings, MetricsCtx, MultimodalData, OnPublisherReady, + OnSnapshotPublisherReady, OutputOptions, PrefillResult, PreprocessedRequest, PromptTokensDetails, RawEngine, SamplingOptions, StopConditions, StopReason, TopLogprob, TopLogprobs, chunk, usage, }; diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 51e1e54e2f5a..b7906f1e2d98 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -35,8 +35,7 @@ use tokio_util::sync::CancellationToken; use crate::adapter::{EngineAdapter, RawEngineAdapter}; use crate::disagg::DisaggregationMode; use crate::engine::{ - EngineAvailability, EngineConfig, KvEventSource, LLMEngine, MetricsBindings, MetricsCtx, - RawEngine, + EngineConfig, KvEventSource, LLMEngine, MetricsBindings, MetricsCtx, RawEngine, }; use crate::error::{BackendError, DynamoError, ErrorType}; use crate::publisher::{PublisherHandles, setup_publishers}; @@ -60,8 +59,6 @@ const DRAIN_HEARTBEAT_INTERVAL_S: f64 = 5.0; /// Budget reserved for `cleanup()` so the drain loop can't consume the whole /// graceful-shutdown deadline and trip the hard-exit that skips cleanup. const CLEANUP_RESERVE_S: f64 = 5.0; -/// Delay before the serve loop retries a failed endpoint sync. -const SYNC_RETRY_INTERVAL: Duration = Duration::from_secs(1); /// Operator override for the health-check canary, mirrors the Python helper /// in `lib/bindings/python/src/dynamo/health_check.py`. @@ -299,69 +296,6 @@ enum EngineRouteLifecycle { ShuttingDown, } -/// The worker's discovery registration, shared by serving, operator pause and -/// engine recovery. The lock is held across discovery writes but never across -/// engine calls. A write can be dropped with its HTTP request, so each write -/// records its state as unknown until it returns. -struct Registration { - /// Whether the endpoint instance is in discovery; `None` after a failed write. - endpoint: Option, - /// Whether the base model card is in discovery; `None` after a failed write. - card: Option, - is_paused: bool, - engine: EngineAvailability, - is_shutting_down: bool, - /// Caller-managed taints, restored when recovery re-registers the model card. - taints: Option>, -} - -impl Default for Registration { - fn default() -> Self { - Self { - endpoint: None, - card: None, - is_paused: false, - engine: EngineAvailability::Ready, - is_shutting_down: false, - taints: None, - } - } -} - -impl Registration { - fn is_routable(&self) -> bool { - !self.is_shutting_down && !self.is_paused && self.engine == EngineAvailability::Ready - } - - /// Register or withdraw the endpoint to match `is_routable`, skipping the - /// write when discovery already matches, and publish readiness to match. - async fn sync_endpoint( - &mut self, - endpoint: &dynamo_runtime::component::Endpoint, - ) -> anyhow::Result<()> { - let is_routable = self.is_routable(); - let mut result = Ok(()); - if self.endpoint != Some(is_routable) { - // Unknown until the write returns: an HTTP route's future can be - // dropped after discovery commits. - self.endpoint = None; - result = if is_routable { - endpoint.register_endpoint_instance().await - } else { - endpoint.unregister_endpoint_instance().await - }; - self.endpoint = result.is_ok().then_some(is_routable); - } - let status = if self.endpoint == Some(true) { - HealthStatus::Ready - } else { - HealthStatus::NotReady - }; - set_worker_health(endpoint, status); - result - } -} - /// The engine a [`Worker`] drives, tagged by request modality. Both variants /// share the lifecycle (driven via the forwarders below); they differ only in /// the serve-loop adapter: `Llm` → token pipeline ([`EngineAdapter`]), `Raw` → @@ -403,6 +337,14 @@ impl EngineKind { } } + /// See [`LLMEngine::engine_lost`]. + async fn engine_lost(&self) -> DynamoError { + match self { + EngineKind::Llm(e) => e.engine_lost().await, + EngineKind::Raw(_) => std::future::pending().await, + } + } + async fn kv_event_sources(&self) -> Result, DynamoError> { match self { EngineKind::Llm(e) => e.kv_event_sources().await, @@ -484,30 +426,6 @@ impl EngineKind { } } - fn availability(&self) -> Option> { - match self { - EngineKind::Llm(e) => e.availability(), - EngineKind::Raw(_) => None, - } - } - - async fn recover(&self) -> Result { - match self { - EngineKind::Llm(e) => e.recover().await, - EngineKind::Raw(_) => Err(err( - ErrorType::Backend(BackendError::EngineShutdown), - "raw engines do not support recovery", - )), - } - } - - async fn on_recovered(&self) -> Result<(), DynamoError> { - match self { - EngineKind::Llm(e) => e.on_recovered().await, - EngineKind::Raw(_) => Ok(()), - } - } - /// Raw media engines (image/video/audio) register name-only — the engine /// loads the model itself and the model has no LLM artifacts (tokenizer / /// chat template / config.json) for Dynamo to fetch. @@ -533,7 +451,6 @@ pub struct Worker { /// Serializes controls that mutate discovery registration and shutdown's /// final transition, preventing a resume from re-registering a stale worker. engine_route_mutation: Arc>, - registration: Arc>, /// Signals in-flight Rust administrative route futures to stop. Engine /// adapters that detach work (such as a separately scheduled language /// runtime task) remain responsible for cancelling that work themselves. @@ -572,7 +489,6 @@ impl Worker { EngineRouteLifecycle::Starting, )), engine_route_mutation: Arc::new(tokio::sync::Mutex::new(())), - registration: Arc::default(), engine_route_shutdown: CancellationToken::new(), publishers: None, lifecycle: None, @@ -691,14 +607,25 @@ impl Worker { // grace-period sleep, which is a fixed wait rather than a hang // risk), otherwise we exit(911). Healthy long-running workers // never hit this — the timer only starts after `shutdown_token` - // is cancelled. + // is cancelled. A lost engine starts the same shutdown. + let engine = self.engine.clone(); let outcome = { let inner_fut = self.run_inner(runtime.clone(), drt, &shutdown_token); tokio::pin!(inner_fut); - tokio::select! { - result = &mut inner_fut => result, - _ = shutdown_token.cancelled() => { + let (finished, lost) = tokio::select! { + biased; + result = &mut inner_fut => (Some(result), None), + _ = shutdown_token.cancelled() => (None, None), + error = engine.engine_lost() => (None, Some(error)), + }; + match finished { + Some(result) => result, + None => { + if let Some(error) = &lost { + tracing::error!(%error, "Engine lost; shutting down the worker"); + shutdown_token.cancel(); + } runtime.mark_shutting_down(); let timeout = graceful_shutdown_timeout(); let grace = grace_period_secs(); @@ -709,7 +636,7 @@ impl Worker { timeout.as_secs(), grace, ); - match tokio::time::timeout(deadline, &mut inner_fut).await { + let result = match tokio::time::timeout(deadline, &mut inner_fut).await { Ok(result) => result, Err(_) => { tracing::error!( @@ -719,6 +646,15 @@ impl Worker { ); std::process::exit(911); } + }; + match lost { + Some(error) => { + if let Err(stopped) = result { + tracing::warn!(error = %stopped, "Worker shutdown failed"); + } + Err(error) + } + None => result, } } } @@ -947,7 +883,6 @@ impl Worker { endpoint.clone(), self.engine_route_lifecycle.clone(), self.engine_route_mutation.clone(), - self.registration.clone(), self.engine_route_shutdown.clone(), ); // Namespace control routes under `/engine/control/` so they @@ -1008,7 +943,6 @@ impl Worker { async fn begin_engine_route_shutdown(&self) { self.engine_route_shutdown.cancel(); let _mutation = self.engine_route_mutation.lock().await; - self.registration.lock().await.is_shutting_down = true; let mut lifecycle = self.engine_route_lifecycle.write().await; *lifecycle = EngineRouteLifecycle::ShuttingDown; } @@ -1023,7 +957,6 @@ impl Worker { model_taint_update_callback( endpoint.clone(), self.engine_route_lifecycle.clone(), - self.registration.clone(), self.engine_route_shutdown.clone(), ), ); @@ -1034,14 +967,11 @@ impl Worker { /// pre-serve (mid-start signal) and the serve loop's signal arm. async fn orchestrator_steps(&mut self, endpoint: &dynamo_runtime::component::Endpoint) { endpoint.drt().runtime().mark_shutting_down(); - let mut registration = self.registration.lock().await; - registration.is_shutting_down = true; - if let Err(e) = registration.sync_endpoint(endpoint).await { + if let Err(e) = endpoint.unregister_endpoint_instance().await { tracing::warn!(error = %e, "discovery unregister failed"); } else { tracing::info!("Endpoint unregistered from discovery"); } - drop(registration); self.run_engine_shutdown_steps().await; } @@ -1158,7 +1088,7 @@ impl Worker { self.config.model_input, None, Some(worker_type), - needs.clone(), + needs, ) .await .map_err(|e| { @@ -1167,7 +1097,6 @@ impl Worker { format!("model attach: {e}"), ) })?; - self.registration.lock().await.card = Some(true); tracing::debug!("model registered with discovery"); self.register_engine_controls(&endpoint).await?; @@ -1324,7 +1253,6 @@ impl Worker { return Ok(()); } }; - self.registration.lock().await.endpoint = Some(true); // A signal can arrive while primary registration is in flight. Keep // routes closed and tear the endpoint back down rather than briefly @@ -1390,77 +1318,36 @@ impl Worker { // ready before this point; drop it here, because the write below // publishes readiness through the very signal it suppresses. drop(readiness_hold); - // A pause control may already have withdrawn the endpoint. - let mut resync = sync_failed( - self.registration - .lock() - .await - .sync_endpoint(&endpoint) - .await, - ); + set_worker_health(&endpoint, HealthStatus::Ready); let serve_fut = primary_endpoint.wait(); tokio::pin!(serve_fut); - // Apply the engine's current availability first, then each change. - let mut availability = self.engine.availability(); - if let Some(receiver) = availability.as_mut() { - receiver.mark_changed(); - } - let serve_result = loop { - tokio::select! { - biased; - result = &mut serve_fut => break endpoint_exit(result), - _ = shutdown.cancelled() => { - tracing::info!("Received shutdown signal; running graceful orchestration"); - break Ok(()); - } - state = next_availability(&mut availability) => { - if state != EngineAvailability::Lost { - let mut registration = self.registration.lock().await; - registration.engine = state; - resync = sync_failed(registration.sync_endpoint(&endpoint).await); - continue; - } - self.withdraw_lost_engine(&endpoint).await; - let engine_config = tokio::select! { - biased; - result = &mut serve_fut => break endpoint_exit(result), - _ = shutdown.cancelled() => { - tracing::info!("Received shutdown signal; running graceful orchestration"); - break Ok(()); - } - result = self.engine.recover() => match result { - Ok(engine_config) => engine_config, - Err(error) => break Err(error), - }, - }; - let recovered = self - .register_recovered_engine( - &endpoint, - engine_config, - model_type, - worker_type, - &needs, - ) - .await; - if let Err(error) = recovered { - break Err(error); + let serve_result = tokio::select! { + biased; + result = &mut serve_fut => { + match result { + // Endpoint exited cleanly (e.g. DRT primary token + // cancelled it) — run the orchestrator so drain/cleanup + // don't race transport teardown. + Ok(()) => { + tracing::info!( + "Endpoint completed gracefully; running shutdown orchestration" + ); + Ok(()) } - // Re-apply whatever the engine reports now. - if let Some(receiver) = availability.as_mut() { - receiver.mark_changed(); + // Serve errored; cleanup_once in run() is the safety net. + Err(e) => { + Err(err( + ErrorType::Backend(BackendError::Unknown), + format!("serve: {e}"), + )) } } - _ = tokio::time::sleep(SYNC_RETRY_INTERVAL), if resync => { - resync = sync_failed( - self.registration - .lock() - .await - .sync_endpoint(&endpoint) - .await, - ); - } + } + _ = shutdown.cancelled() => { + tracing::info!("Received shutdown signal; running graceful orchestration"); + Ok(()) } }; @@ -1483,77 +1370,6 @@ impl Worker { serve_result } - /// Withdraw the endpoint and model card after engine loss; in-flight - /// requests keep running. - async fn withdraw_lost_engine(&self, endpoint: &dynamo_runtime::component::Endpoint) { - tracing::warn!("Engine lost; withdrawing worker from discovery until it recovers"); - let mut registration = self.registration.lock().await; - registration.engine = EngineAvailability::Lost; - if let Err(error) = registration.sync_endpoint(endpoint).await { - tracing::warn!(%error, "failed to withdraw endpoint after engine loss"); - } - if registration.card != Some(false) { - let result = LocalModel::detach_from_endpoint(endpoint, None).await; - registration.card = result.is_ok().then_some(false); - if let Err(error) = result { - tracing::warn!(%error, "failed to withdraw model card after engine loss"); - } - } - } - - /// Register a fresh model card from [`LLMEngine::recover`] with the - /// caller's taints. The endpoint stays withdrawn until the serve loop sees - /// the engine `Ready`. - async fn register_recovered_engine( - &mut self, - endpoint: &dynamo_runtime::component::Endpoint, - engine_config: EngineConfig, - model_type: ModelType, - worker_type: WorkerType, - needs: &[Vec], - ) -> Result<(), DynamoError> { - let mut local_model = - build_local_model(&self.config, &engine_config, self.engine.is_raw()).await?; - - let mut registration = self.registration.lock().await; - local_model - .attach( - endpoint, - model_type, - self.config.model_input, - None, - Some(worker_type), - needs.to_vec(), - ) - .await - .map_err(|e| { - err( - ErrorType::Backend(BackendError::Unknown), - format!("model attach after engine recovery: {e}"), - ) - })?; - registration.card = Some(true); - if let Some(taints) = registration.taints.clone() { - update_model_taints(endpoint, taints).await.map_err(|e| { - err( - ErrorType::Backend(BackendError::Unknown), - format!("restore model taints after engine recovery: {e}"), - ) - })?; - } - drop(registration); - if matches!(self.engine, EngineKind::Llm(_)) { - self.admission_gate - .get()? - .record_model_card(local_model.card()); - } - if let Err(error) = self.engine.on_recovered().await { - tracing::warn!(%error, "engine on_recovered failed"); - } - tracing::info!(model = %engine_config.model, "Engine recovered; model card re-registered"); - Ok(()) - } - /// Engine-facing shutdown sequence: grace period sleep → drain loop on /// `engine.is_quiescent()` → `cleanup_once()`. Each engine step swallows /// non-fatal failures so a misbehaving engine can't block the worker @@ -1592,11 +1408,7 @@ impl Worker { /// expires. Budget = `DYN_PREFILL_DRAIN_TIMEOUT_S` capped at /// `graceful_shutdown_timeout - CLEANUP_RESERVE_S`. async fn drain_until_idle_or_deadline(&self) { - // A lost engine can't report quiescence, and its KV memory lives in its - // own process, which this worker's cleanup doesn't free. - if !self.config.disaggregation_mode.is_prefill() - || self.registration.lock().await.engine == EngineAvailability::Lost - { + if !self.config.disaggregation_mode.is_prefill() { return; } let configured = drain_timeout_secs(); @@ -1651,42 +1463,6 @@ impl Worker { } } -fn endpoint_exit(result: anyhow::Result<()>) -> Result<(), DynamoError> { - match result { - // Endpoint exited cleanly (e.g. DRT primary token cancelled it) — run - // the orchestrator so drain/cleanup don't race transport teardown. - Ok(()) => { - tracing::info!("Endpoint completed gracefully; running shutdown orchestration"); - Ok(()) - } - // Serve errored; cleanup_once in run() is the safety net. - Err(e) => Err(err( - ErrorType::Backend(BackendError::Unknown), - format!("serve: {e}"), - )), - } -} - -/// Log a failed serve-loop endpoint sync; true when it needs a retry. -fn sync_failed(result: anyhow::Result<()>) -> bool { - if let Err(error) = &result { - tracing::warn!(%error, "failed to sync endpoint registration; retrying"); - } - result.is_err() -} - -/// The engine's next availability change; pending forever without one. -async fn next_availability( - availability: &mut Option>, -) -> EngineAvailability { - if let Some(receiver) = availability - && receiver.changed().await.is_ok() - { - return *receiver.borrow_and_update(); - } - std::future::pending().await -} - /// Publish worker readiness on both layers the runtime's health route reads. /// /// `SystemHealth::get_health_status` consults `use_endpoint_health_status`, then @@ -2010,13 +1786,11 @@ fn parse_model_taint_update_request(body: serde_json::Value) -> anyhow::Result>, - registration: Arc>, route_shutdown: CancellationToken, ) -> EngineRouteCallback { Arc::new(move |body| { let endpoint = endpoint.clone(); let route_lifecycle = route_lifecycle.clone(); - let registration = registration.clone(); let route_shutdown = route_shutdown.clone(); Box::pin(async move { let taints = parse_model_taint_update_request(body)?; @@ -2025,20 +1799,13 @@ fn model_taint_update_callback( Ok(lifecycle) => lifecycle, Err(response) => return Ok(response), }; - let mut registration = tokio::select! { + tokio::select! { biased; _ = route_shutdown.cancelled() => { return Ok(engine_route_lifecycle_error(EngineRouteLifecycle::ShuttingDown)); } - registration = registration.lock() => registration, - }; - // Recorded first so a dropped request cannot leave an older set to - // restore; while recovery has the card withdrawn, they apply on re-registration. - registration.taints = Some(taints.clone()); - if registration.card != Some(false) { - update_model_taints(&endpoint, taints.clone()).await?; + result = update_model_taints(&endpoint, taints.clone()) => result?, } - drop(registration); let mut response_taints: Vec<_> = taints.into_iter().collect(); response_taints.sort(); @@ -2096,7 +1863,6 @@ fn engine_update_callback( }) } -#[allow(clippy::too_many_arguments)] fn wrap_engine_control_callback( control_name: String, callback: EngineRouteCallback, @@ -2104,7 +1870,6 @@ fn wrap_engine_control_callback( endpoint: dynamo_runtime::component::Endpoint, route_lifecycle: Arc>, route_mutation: Arc>, - registration: Arc>, route_shutdown: CancellationToken, ) -> EngineRouteCallback { let policy = engine_control_policy(&control_name); @@ -2115,7 +1880,6 @@ fn wrap_engine_control_callback( let control_name = control_name.clone(); let route_lifecycle = route_lifecycle.clone(); let route_mutation = route_mutation.clone(); - let registration = registration.clone(); let route_shutdown = route_shutdown.clone(); Box::pin(async move { if let Some(response) = control_request_body_error(&body) { @@ -2153,27 +1917,23 @@ fn wrap_engine_control_callback( Ok(lifecycle) => lifecycle, Err(response) => return Ok(response), }; - let mut registration = tokio::select! { + let unregister_result = tokio::select! { biased; _ = route_shutdown.cancelled() => { return Ok(engine_route_lifecycle_error(EngineRouteLifecycle::ShuttingDown)); } - registration = registration.lock() => registration, + result = endpoint.unregister_endpoint_instance() => result, }; - registration.is_paused = true; - if let Err(e) = registration.sync_endpoint(&endpoint).await { - registration.is_paused = false; - if let Err(error) = registration.sync_endpoint(&endpoint).await { - tracing::warn!(%error, "failed to restore endpoint after a failed pause"); - } + if let Err(e) = unregister_result { return Ok(control_error_response(format!( "failed to unregister endpoint before /engine/control/{control_name}: {e}" ))); } // Out of discovery, so no longer routable. Whether the - // control itself then succeeds or fails, the endpoint and - // readiness stay withdrawn until a resume control succeeds. - drop(registration); + // control itself then succeeds or fails, the endpoint is + // left unregistered, so readiness stays withdrawn until a + // resume control re-registers it. + set_worker_health(&endpoint, HealthStatus::NotReady); let callback_result = tokio::select! { biased; @@ -2231,16 +1991,14 @@ fn wrap_engine_control_callback( } return Ok(response); } - let mut registration = tokio::select! { + let register_result = tokio::select! { biased; _ = route_shutdown.cancelled() => { return Ok(engine_route_lifecycle_error(EngineRouteLifecycle::ShuttingDown)); } - registration = registration.lock() => registration, + result = endpoint.register_endpoint_instance() => result, }; - registration.is_paused = false; - // A lost or recovering engine registers once it is ready again. - if let Err(e) = registration.sync_endpoint(&endpoint).await { + if let Err(e) = register_result { // The engine is serving-safe but absent from discovery. The // operation is idempotent: retrying /engine/control/{control_name} // re-registers without repeating the wake/resume work (the @@ -2250,6 +2008,7 @@ fn wrap_engine_control_callback( "engine resumed but re-registration failed after /engine/control/{control_name}: {e}; retry /engine/control/{control_name} to rejoin discovery" ))); } + set_worker_health(&endpoint, HealthStatus::Ready); Ok(response) } } @@ -4316,7 +4075,6 @@ mod handoff_and_lifecycle_tests { endpoint.clone(), worker.engine_route_lifecycle.clone(), worker.engine_route_mutation.clone(), - worker.registration.clone(), worker.engine_route_shutdown.clone(), ); let request = tokio::spawn(async move { callback(serde_json::json!({})).await.unwrap() }); @@ -4375,7 +4133,6 @@ mod handoff_and_lifecycle_tests { endpoint.clone(), worker.engine_route_lifecycle.clone(), worker.engine_route_mutation.clone(), - worker.registration.clone(), worker.engine_route_shutdown.clone(), ); let resume_request = @@ -4391,7 +4148,6 @@ mod handoff_and_lifecycle_tests { endpoint, worker.engine_route_lifecycle.clone(), worker.engine_route_mutation.clone(), - worker.registration.clone(), worker.engine_route_shutdown.clone(), ); let response = tokio::time::timeout( @@ -4625,6 +4381,109 @@ mod handoff_and_lifecycle_tests { worker.begin_engine_route_shutdown().await; } + /// A lost engine withdraws the serving worker, runs cleanup, and makes the + /// worker return the engine's error. + #[tokio::test] + #[serial_test::serial] + async fn lost_engine_withdraws_the_worker_and_returns_its_error() { + struct LostEngine { + lost: CancellationToken, + cleanup_calls: std::sync::atomic::AtomicUsize, + } + + #[async_trait] + impl LLMEngine for LostEngine { + async fn start(&self, _worker_id: u64) -> Result { + Ok(EngineConfig { + model: "lost-mock".to_string(), + ..EngineConfig::default() + }) + } + + async fn generate( + &self, + _request: PreprocessedRequest, + _ctx: crate::engine::GenerateContext, + ) -> Result< + BoxStream<'static, Result>, + DynamoError, + > { + unreachable!("not used in engine-loss tests") + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + self.cleanup_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(()) + } + + async fn engine_lost(&self) -> DynamoError { + self.lost.cancelled().await; + DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::EngineShutdown)) + .message("synthetic engine loss") + .build() + } + } + + let runtime = dynamo_runtime::Runtime::from_current().unwrap(); + let config = dynamo_runtime::distributed::DistributedConfig::process_local(); + let drt = DistributedRuntime::new(runtime, config).await.unwrap(); + let engine = Arc::new(LostEngine { + lost: CancellationToken::new(), + cleanup_calls: Default::default(), + }); + let config = WorkerConfig::default(); + let query = DiscoveryQuery::Endpoint { + namespace: config.namespace.clone(), + component: config.component.clone(), + endpoint: config.endpoint.clone(), + }; + let shutdown = CancellationToken::new(); + let run = tokio::spawn( + Worker::new(engine.clone(), config).run_with_drt(drt.clone(), shutdown.clone()), + ); + tokio::time::timeout(Duration::from_secs(30), async { + while drt + .discovery() + .list(query.clone()) + .await + .unwrap() + .is_empty() + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("worker registers its endpoint"); + + engine.lost.cancel(); + let error = tokio::time::timeout(Duration::from_secs(30), run) + .await + .expect("a lost engine stops the worker") + .unwrap() + .expect_err("a lost engine fails the worker"); + assert_eq!( + error.error_type(), + ErrorType::Backend(BackendError::EngineShutdown) + ); + assert!( + error.to_string().contains("synthetic engine loss"), + "{error}" + ); + assert!( + drt.discovery().list(query).await.unwrap().is_empty(), + "a lost engine must leave no endpoint" + ); + assert_eq!( + engine + .cleanup_calls + .load(std::sync::atomic::Ordering::SeqCst), + 1 + ); + assert!(shutdown.is_cancelled(), "engine loss uses the SIGTERM path"); + } + /// Ensures a payload-free Rust backend publishes readiness while it is /// serviceable and withdraws it again on shutdown. #[tokio::test] @@ -4857,300 +4716,4 @@ mod handoff_and_lifecycle_tests { "model taint updates must not be registered after a fatal handoff" ); } - - /// Engine whose availability the test drives. `recover` waits for the test - /// and reports a new context length, so re-registration must carry it. - struct RecoveringEngine { - availability: tokio::sync::watch::Sender, - release_recovery: Notify, - is_recovered: std::sync::atomic::AtomicBool, - } - - fn recovering_config(context_length: u32) -> EngineConfig { - EngineConfig { - model: "recovering-mock".to_string(), - llm: Some(crate::engine::LlmRegistration { - context_length: Some(context_length), - ..Default::default() - }), - ..EngineConfig::default() - } - } - - #[async_trait] - impl LLMEngine for RecoveringEngine { - async fn start(&self, _worker_id: u64) -> Result { - Ok(recovering_config(1024)) - } - - async fn generate( - &self, - _request: PreprocessedRequest, - _ctx: crate::engine::GenerateContext, - ) -> Result< - BoxStream<'static, Result>, - DynamoError, - > { - unreachable!("not used in recovery tests") - } - - async fn cleanup(&self) -> Result<(), DynamoError> { - Ok(()) - } - - async fn supported_controls(&self) -> Result, DynamoError> { - Ok(vec![ - "pause_generation".to_string(), - "resume_generation".to_string(), - ]) - } - - async fn engine_control( - &self, - _control: String, - _body: serde_json::Value, - ) -> Result { - Ok(serde_json::json!({"status": "ok"})) - } - - fn availability(&self) -> Option> { - Some(self.availability.subscribe()) - } - - async fn recover(&self) -> Result { - self.release_recovery.notified().await; - self.availability.send_replace(EngineAvailability::Ready); - Ok(recovering_config(2048)) - } - - async fn on_recovered(&self) -> Result<(), DynamoError> { - self.is_recovered - .store(true, std::sync::atomic::Ordering::SeqCst); - Ok(()) - } - } - - struct RecoveryHarness { - engine: Arc, - drt: DistributedRuntime, - shutdown: CancellationToken, - run: tokio::task::JoinHandle>, - _discovery_dir: tempfile::TempDir, - } - - impl RecoveryHarness { - async fn start() -> Self { - use dynamo_runtime::distributed::{DiscoveryBackend, DistributedConfig}; - - let discovery_dir = tempfile::tempdir().unwrap(); - let config = DistributedConfig { - discovery_backend: DiscoveryBackend::KvStore( - dynamo_runtime::storage::kv::Selector::File(discovery_dir.path().to_owned()), - ), - ..DistributedConfig::process_local() - }; - let drt = DistributedRuntime::new(Runtime::from_current().unwrap(), config) - .await - .unwrap(); - let engine = Arc::new(RecoveringEngine { - availability: tokio::sync::watch::Sender::new(EngineAvailability::Ready), - release_recovery: Notify::new(), - is_recovered: Default::default(), - }); - let worker = Worker::new( - engine.clone(), - WorkerConfig { - enable_kv_routing: false, - ..WorkerConfig::default() - }, - ); - let shutdown = CancellationToken::new(); - let run = tokio::spawn(worker.run_with_drt(drt.clone(), shutdown.clone())); - let harness = Self { - engine, - drt, - shutdown, - run, - _discovery_dir: discovery_dir, - }; - harness.wait_for(1, Some(1024)).await; - harness - } - - async fn discovered(&self) -> (usize, Option) { - use dynamo_runtime::discovery::DiscoveryInstance; - - let discovery = self.drt.discovery(); - let config = WorkerConfig::default(); - let endpoints = discovery - .list(DiscoveryQuery::Endpoint { - namespace: config.namespace.clone(), - component: config.component.clone(), - endpoint: config.endpoint.clone(), - }) - .await - .unwrap(); - let models = discovery - .list(DiscoveryQuery::EndpointModels { - namespace: config.namespace, - component: config.component, - endpoint: config.endpoint, - }) - .await - .unwrap(); - let card = models.into_iter().find_map(|model| match model { - DiscoveryInstance::Model { card_json, .. } => Some(card_json), - _ => None, - }); - (endpoints.len(), card) - } - - /// `context_length: None` waits for no model card. - async fn wait_for( - &self, - endpoints: usize, - context_length: Option, - ) -> Option { - let mut last = None; - for _ in 0..200 { - let (count, card) = self.discovered().await; - let card_context = card.as_ref().map(|card| { - card["runtime_config"]["context_length"] - .as_u64() - .unwrap_or_default() - }); - if count == endpoints && card_context == context_length.map(u64::from) { - return card; - } - last = Some((count, card_context)); - tokio::time::sleep(Duration::from_millis(25)).await; - } - panic!("expected {endpoints} endpoints and card {context_length:?}; last saw {last:?}"); - } - - async fn route(&self, name: &str, body: serde_json::Value) -> serde_json::Value { - let callback = self.drt.engine_routes().get(name).unwrap(); - let response = callback(body).await.unwrap(); - assert!(!control_response_is_error(&response), "{name}: {response}"); - response - } - - async fn stop(self) { - self.shutdown.cancel(); - tokio::time::timeout(Duration::from_secs(30), self.run) - .await - .expect("worker must shut down") - .unwrap() - .expect("worker must exit cleanly"); - } - } - - /// Regression: a worker kept advertising an engine that had stopped - /// serving or gone away. `NotReady` must withdraw only the endpoint; `Lost` - /// must also withdraw the model card until recovery registers fresh metadata. - #[tokio::test] - async fn engine_loss_withdraws_until_recovery_registers_fresh_metadata() { - let harness = RecoveryHarness::start().await; - - harness - .engine - .availability - .send_replace(EngineAvailability::NotReady); - harness.wait_for(0, Some(1024)).await; - harness - .engine - .availability - .send_replace(EngineAvailability::Ready); - harness.wait_for(1, Some(1024)).await; - - harness - .engine - .availability - .send_replace(EngineAvailability::Lost); - harness.wait_for(0, None).await; - harness.engine.release_recovery.notify_one(); - harness.wait_for(1, Some(2048)).await; - assert!( - harness - .engine - .is_recovered - .load(std::sync::atomic::Ordering::SeqCst) - ); - - harness.stop().await; - } - - /// Regression: a runtime shutdown during engine recovery left the worker - /// waiting on `recover()` instead of shutting down. - #[tokio::test] - async fn runtime_shutdown_stops_engine_recovery() { - let harness = RecoveryHarness::start().await; - harness - .engine - .availability - .send_replace(EngineAvailability::Lost); - harness.wait_for(0, None).await; - - harness.drt.shutdown(); - tokio::time::timeout(Duration::from_secs(30), harness.run) - .await - .expect("worker must shut down while recovery is pending") - .unwrap() - .expect("worker must exit cleanly"); - } - - /// Regression: recovery could re-register a worker the operator had paused, - /// or drop the caller's taints when it re-registered the model card. - #[tokio::test] - async fn operator_pause_and_taints_survive_engine_recovery() { - let harness = RecoveryHarness::start().await; - harness - .route( - MODEL_TAINT_UPDATE_ROUTE, - serde_json::json!({"taints": ["capacity/fast"]}), - ) - .await; - harness - .route("control/pause_generation", serde_json::json!({})) - .await; - harness.wait_for(0, Some(1024)).await; - - harness - .engine - .availability - .send_replace(EngineAvailability::Lost); - harness.wait_for(0, None).await; - harness.engine.release_recovery.notify_one(); - harness.wait_for(0, Some(2048)).await; - let has_taint = |card: &serde_json::Value| { - card["runtime_config"]["taints"] - .as_array() - .is_some_and(|taints| taints.contains(&serde_json::json!("capacity/fast"))) - }; - let mut card = None; - for _ in 0..200 { - card = harness.discovered().await.1; - if card.as_ref().is_some_and(has_taint) { - break; - } - tokio::time::sleep(Duration::from_millis(25)).await; - } - assert!( - card.as_ref().is_some_and(has_taint), - "recovery must restore caller taints: {card:?}" - ); - tokio::time::sleep(Duration::from_millis(200)).await; - assert_eq!( - harness.discovered().await.0, - 0, - "a paused worker must stay withdrawn after recovery" - ); - - harness - .route("control/resume_generation", serde_json::json!({})) - .await; - harness.wait_for(1, Some(2048)).await; - - harness.stop().await; - } } From 7629691df9e01252470b86f57bc26ad075b23755 Mon Sep 17 00:00:00 2001 From: Julien Darve Date: Tue, 3 Nov 2026 12:38:28 -0800 Subject: [PATCH 5/5] docs(backend-common): state when the Worker polls and drops engine_lost Signed-off-by: Julien Darve --- lib/backend-common/src/engine.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index c20abe18fcef..48637acd9b63 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -284,6 +284,11 @@ pub trait LLMEngine: Send + Sync + 'static { /// `Worker` returns. The `Worker` then shuts down as on SIGTERM, withdrawing /// from discovery, draining and running `cleanup`, so the process exits /// and its supervisor can restart it. Default: never resolves. + /// + /// The `Worker` may poll this before `start` completes, and drops the + /// future once it finishes or starts shutting down another way. + /// Implementations must allow polling during startup and must not use this + /// future to clean up. async fn engine_lost(&self) -> DynamoError { std::future::pending().await }