feat(backend-common): drain and exit the worker when its engine is lost

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

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:

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<Vec<KvEventSource>, 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<EngineConfig, DynamoError> {
+                Ok(EngineConfig {
+                    model: "lost-mock".to_string(),
+                    ..EngineConfig::default()
+                })
+            }
+
+            async fn generate(
+                &self,
+                _request: PreprocessedRequest,
+                _ctx: crate::engine::GenerateContext,
+            ) -> Result<
+                BoxStream<'static, Result<crate::engine::LLMEngineOutput, DynamoError>>,
+                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]