fix(runtime): propagate MP failures and return unhealthy status

atom/entrypoints/openai/api_server.py

@@ -2750,7 +2750,9 @@ async def list_models():
 
 @app.get("/health")
 async def health():
-    """Health check endpoint."""
+    """Report readiness of all engine workers, not just the HTTP process."""
+    if engine is None or not engine.core_mgr.is_healthy():
+        raise HTTPException(status_code=503, detail="Engine workers unavailable")
     return {"status": "ok"}
 
 

atom/kv_transfer/offload/README.md

@@ -222,12 +222,22 @@ Engine memory under an MP transfer (PAGE sources, restore destinations,
 checkpoint pins, descriptor slots) is released only on a terminal report from
 the server. A timeout cannot prove that a remote DMA stopped, so no clock ever
 frees it. Instead, a transfer that is still not terminal after
-`lmcache.mp.transfer_deadline_s` (default 1200 s) is fatal: the worker, or the
-scheduler for a report that never arrives, raises `LMCacheTransferUnprovable`
-and the engine stops. This covers a submission that raised, a future whose
-poll keeps raising, a restore whose event cannot be queried, and a lost or
-never-aggregated completion. The engine-wide stalled-save reclaim leaves MP
-leases alone for the same reason.
+`lmcache.mp.transfer_deadline_s` (default 30 s) is fatal: the worker raises
+`LMCacheTransferUnprovable` and the engine stops. This is a per-transfer deadline
+from submission, including time queued at the MP server, not a model request
+timeout or a heartbeat timeout. It also catches a live server whose ZMQ handler
+failed without replying; no LMCache transport changes are required. The scheduler
+independently watches dispatched operations with an additional 10 s margin for
+lost or never-aggregated completion reports. These checks run when the respective
+worker/scheduler polls completions; they are not an independent watchdog for a
+blocked process. A future query/result exception stops the engine immediately.
+The engine-wide stalled-save reclaim leaves MP leases alone for the same reason.
+
+Set `"lmcache.mp.transfer_deadline_s": 180` in `kv_connector_extra_config`, for
+example, if measured normal transfer/queue latency on a slower storage tier
+requires a longer deadline. The value must be finite and positive. The default
+is a failure-detection policy, not a guarantee that every overloaded deployment
+can complete transfers within 30 s.
 
 The model namespace includes PAGE/model geometry, TP and speculation settings,
 native layout/image sizes, the Hugging Face commit when available, and the

atom/kv_transfer/offload/mp/transfer.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import logging
+import math
 import time
 from collections import deque
 from collections.abc import Iterable, Sequence
@@ -48,15 +49,15 @@ class _PendingSave:
     started_at: float = field(default_factory=lambda: time.monotonic())
 
 
-_DEFAULT_TRANSFER_DEADLINE_S = 1200.0
+_DEFAULT_TRANSFER_DEADLINE_S = 30.0
 # The scheduler watches the same operations from dispatch, a step or more
 # before a worker submits them. The margin lets the worker, which can name the
 # exact transfer, fail first.
-_SCHEDULER_DEADLINE_MARGIN_S = 60.0
+_SCHEDULER_DEADLINE_MARGIN_S = 10.0
 
 
 class LMCacheTransferUnprovable(RuntimeError):
-    """An MP transfer never reached a terminal state within the deadline."""
+    """An MP transfer failed or exceeded its deadline without proving DMA safety."""
 
 
 def _transfer_deadline_s(config: Any) -> float:
@@ -73,8 +74,8 @@ def _transfer_deadline_s(config: Any) -> float:
     deadline = float(
         extra.get("lmcache.mp.transfer_deadline_s", _DEFAULT_TRANSFER_DEADLINE_S)
     )
-    if not deadline > 0:
-        raise ValueError("lmcache.mp.transfer_deadline_s must be > 0")
+    if not math.isfinite(deadline) or deadline <= 0:
+        raise ValueError("lmcache.mp.transfer_deadline_s must be finite and > 0")
     return deadline
 
 
@@ -121,15 +122,14 @@ def _terminal_future_result(future: Any | None) -> tuple[bool, Any]:
         if not future.query():
             return False, None
         return True, future.result(timeout=0)
-    except Exception:
+    except Exception as exc:
         # A query/result exception is not proof that the remote GPU stream has
-        # quiesced. Keep the operation pending; the transfer deadline bounds a
-        # future that keeps raising.
-        logger.warning(
-            "LMCache MP transfer future could not prove terminal state",
-            exc_info=True,
-        )
-        return False, None
+        # quiesced. Fail the engine immediately, retaining the pending lease:
+        # a known failure must not masquerade as a slow transfer for 20 minutes.
+        raise LMCacheTransferUnprovable(
+            "LMCache MP transfer failed without proof that remote GPU work "
+            f"stopped; retaining its memory and stopping the engine: {exc}"
+        ) from exc
 
 
 def _transfer_operation_id(kind: str, completion: Any) -> str:

atom/model_engine/async_proc.py

@@ -463,8 +463,11 @@ class AsyncIOProcManager:
         {*_PROTOCOL_RPC_NAMES, *UTILITY_MANAGED_RUNNER_METHODS}
     )
 
-    def __init__(self, finalizer, proc_num: int, runner: str, *args):
+    def __init__(
+        self, finalizer, proc_num: int, runner: str, *args, on_worker_failure=None
+    ):
         self.parent_finalizer = finalizer
+        self.on_worker_failure = on_worker_failure
         self.proc_num = proc_num
 
         io_addrs = [get_open_zmq_ipc_path(), get_open_zmq_ipc_path()]
@@ -995,6 +998,14 @@ def monitor_engine_cores():
                 f"{self.label}: [{dead_proc.name}] proc died unexpectedly "
                 f"(exitcode={dead_proc.exitcode}), shutting down.",
             )
+            # Notify the frontend before exit() can block draining GPU work.
+            # A live EngineCore process alone does not imply healthy runners.
+            try:
+                callback = getattr(_self, "on_worker_failure", None)
+                if callback is not None:
+                    callback()
+            except Exception:
+                logger.exception("Failed to report ModelRunner death")
             _self.exit()
 
         Thread(

atom/model_engine/engine_core.py

@@ -133,6 +133,7 @@ def __init__(self, config: Config, input_address: str, output_address: str):
                 config.tp_world_size * config.prefill_context_parallel_size,
                 config.runner_qualname,
                 config,
+                on_worker_failure=self._send_engine_dead,
             )
             self._post_model_load_hook()
             block_info = self.runner_mgr.call_func("get_num_blocks", wait_out=True)

atom/model_engine/engine_core_mgr.py

@@ -15,7 +15,7 @@
 from collections import deque
 from contextlib import contextmanager
 from dataclasses import dataclass
-from threading import Lock, RLock, Thread, local
+from threading import Event, Lock, RLock, Thread, local
 
 import zmq
 import zmq.asyncio
@@ -224,6 +224,7 @@ def _init_shared_state(
         """
         self.label = label
         self._closed = False  # Track whether already closed
+        self._engine_failed = Event()
         self.local_engine_count = local_engine_count
         self.global_engine_count = (
             local_engine_count if global_engine_count is None else global_engine_count
@@ -732,6 +733,7 @@ def process_outputs_socket():
                         )
                         continue
                     if request_type == EngineCoreRequestType.SHUTDOWN:
+                        self._engine_failed.set()
                         logger.debug(
                             f"{self.label} (DP {dp_rank}): output thread receive SHUTDOWN request"
                         )
@@ -2090,6 +2092,16 @@ def is_alive(self):
             proc is not None and proc.is_alive() for proc in self.engine_core_processes
         )
 
+    def is_healthy(self) -> bool:
+        """Include reported nested-worker failures, even during slow teardown."""
+        if self._closed or self._engine_failed.is_set():
+            return False
+        local_procs = [p for p in self.engine_core_processes if p is not None]
+        try:
+            return bool(local_procs) and all(p.is_alive() for p in local_procs)
+        except ValueError:  # A concurrent shutdown closed a process handle.
+            return False
+
     def finished_procs(self):
         return any(
             proc is not None and not proc.is_alive()

tests/entrypoints/test_api_server_helpers.py

@@ -144,6 +144,21 @@ def _boom(*_args, **_kwargs):
     raise RuntimeError("boom")
 
 
+@pytest.mark.parametrize("state,expected", [(None, 503), (False, 503), (True, 200)])
+def test_health_reports_worker_readiness(monkeypatch, state, expected):
+    from fastapi.testclient import TestClient
+
+    value = (
+        None
+        if state is None
+        else SimpleNamespace(core_mgr=SimpleNamespace(is_healthy=lambda: state))
+    )
+    monkeypatch.setattr(api_server, "engine", value)
+    # No lifespan: this test supplies the engine and does not load a model.
+    response = TestClient(api_server.app).get("/health")
+    assert response.status_code == expected
+
+
 class TestMultimodalInputs:
     @pytest.mark.parametrize(
         "part",

tests/test_collective_rpc_transport.py

@@ -46,6 +46,50 @@
 # ── helpers ────────────────────────────────────────────────────────────────
 
 
+def test_worker_death_reports_failure_before_blocking_cleanup(monkeypatch):
+    mgr = object.__new__(AsyncIOProcManager)
+    proc = SimpleNamespace(
+        sentinel=17, name="ModelRunner0", exitcode=1, join=lambda timeout: None
+    )
+    mgr.procs = [proc]
+    mgr._dead_ranks = {}
+    mgr.runner_label = mgr.label = "test"
+    calls = []
+    mgr.on_worker_failure = lambda: calls.append("notify")
+    mgr.exit = lambda: calls.append("cleanup")
+    targets = []
+    monkeypatch.setattr(
+        async_proc_module,
+        "Thread",
+        lambda **kw: SimpleNamespace(start=lambda: targets.append(kw["target"])),
+    )
+    monkeypatch.setattr(
+        async_proc_module.multiprocessing.connection, "wait", lambda sentinels: [17]
+    )
+    mgr.monitor_procs()
+    targets[0]()
+    assert calls == ["notify", "cleanup"]
+    assert mgr._dead_ranks == {0: 1}
+
+
+def test_expected_worker_shutdown_does_not_report_failure(monkeypatch):
+    mgr = object.__new__(AsyncIOProcManager)
+    mgr.procs = [SimpleNamespace(sentinel=17)]
+    mgr.runner_label = "test"
+    targets = []
+    monkeypatch.setattr(
+        async_proc_module,
+        "Thread",
+        lambda **kw: SimpleNamespace(start=lambda: targets.append(kw["target"])),
+    )
+    monkeypatch.setattr(
+        async_proc_module.multiprocessing.connection, "wait", lambda sentinels: [17]
+    )
+    mgr.monitor_procs()
+    mgr.keep_monitoring = False
+    targets[0]()  # Must return without notifying or running failure cleanup.
+
+
 class _Runner:
     """Stand-in for a ModelRunner: one method per outcome under test."""
 

tests/test_engine_core_mgr_shared_state.py

@@ -19,10 +19,15 @@
 
 import ast
 import pathlib
+import pickle
 import unittest
+import uuid
 from types import SimpleNamespace
 
+import zmq
+
 from atom.model_engine.engine_core_mgr import CoreManager, DisaggCoreManager
+from atom.model_engine.engine_core_protocol import EngineCoreRequestType
 
 MODULE_PATH = pathlib.Path(CoreManager.__module__.replace(".", "/") + ".py")
 SOURCE = pathlib.Path(__file__).resolve().parent.parent / MODULE_PATH
@@ -85,6 +90,46 @@ def setUp(self):
         self.managers = [_bare_manager(CoreManager), _bare_manager(DisaggCoreManager)]
         self.addCleanup(self._terminate)
 
+    def test_health_detects_nested_failure_even_when_all_cores_are_alive(self):
+        for mgr in self.managers:
+            mgr.engine_core_processes = [SimpleNamespace(is_alive=lambda: True)] * 2
+            self.assertTrue(mgr.is_healthy())
+            mgr._engine_failed.set()
+            self.assertFalse(mgr.is_healthy())
+
+    def test_health_requires_every_local_core_and_open_manager(self):
+        for mgr in self.managers:
+            mgr.engine_core_processes = []
+            self.assertFalse(mgr.is_healthy())
+            mgr.engine_core_processes = [
+                SimpleNamespace(is_alive=lambda: True),
+                SimpleNamespace(is_alive=lambda: False),
+            ]
+            self.assertFalse(mgr.is_healthy())
+            mgr.engine_core_processes = [SimpleNamespace(is_alive=lambda: True)]
+            mgr._closed = True
+            self.assertFalse(mgr.is_healthy())
+
+    def test_shutdown_frame_marks_unhealthy_before_cleanup(self):
+        mgr = self.managers[0]
+        mgr.engine_core_processes = [SimpleNamespace(is_alive=lambda: True)]
+        url = f"inproc://failure-{uuid.uuid4().hex}"
+        receiver = mgr.ctx.socket(zmq.PULL)
+        receiver.bind(url)
+        sender = mgr.ctx.socket(zmq.PUSH)
+        sender.connect(url)
+        observed = []
+        mgr._shutdown_engine_core_rank = lambda rank: observed.append(mgr.is_healthy())
+        thread = mgr._create_output_thread(0, receiver, url + "-shutdown")
+        thread.start()
+        try:
+            sender.send(pickle.dumps((EngineCoreRequestType.SHUTDOWN, None)))
+            thread.join(timeout=3)
+            self.assertFalse(thread.is_alive())
+            self.assertEqual(observed, [False])
+        finally:
+            sender.close(linger=0)
+
     def _terminate(self):
         for mgr in self.managers:
             mgr.ctx.term()