diff --git a/lib/backend-common/AGENTS.md b/lib/backend-common/AGENTS.md index 93e88d7b65e6..63f08743534c 100644 --- a/lib/backend-common/AGENTS.md +++ b/lib/backend-common/AGENTS.md @@ -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 diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index 35fb5ff7ccf4..48637acd9b63 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -280,6 +280,19 @@ 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. + /// + /// 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 + } + /// Release all engine resources. Called exactly once. /// /// `Worker` guarantees: diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 0f9a04118ea7..b7906f1e2d98 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -337,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, @@ -599,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(); @@ -617,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!( @@ -627,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, } } } @@ -4353,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]