fix(workers): serialize reaping so a killed worker never looks alive

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 ### Fixed
 - **Predicts against wide or long contexts ran as one-row calls** (#47). Predict chunks were sized with the fit memory model, which grows with the training table; past ~3,500 x 500 no chunk "fit", chunking fell to one row, and a 100-row predict became 100 sequential model calls: 14.4 s (TabICLv2) and 6.3 s (TabPFN-3.5) instead of ~0.1 s on an A100, ~24 s and ~11 s on prod's A40s. Chunks are now sized by `PredictCostModel` (`memory/predict.py`), built from predict peaks measured under the serving configuration (`benchmarks/calibrate_predict.py`, A100, compiled TabICL, int8 kv; generated `memory/calibration_predict_a100.py`): a cost per test row and a fixed cost per call, both by width, plus a cost per training row. It bounds every measurement it was built from and assumes at most linear growth beyond them (affine in width past the widest, where it scales in proportion to width), and it applies only to the calibrated configuration (kv cache, int8, compiled TabICL; TabPFN). When a call's fixed cost alone exceeds the headroom, only the rows above it are sized against the budget instead of splitting to single rows, which only multiplied that fixed cost; with no headroom a predict still runs as the smallest chunks, and a real OOM still halves a chunk. `LIGHT_PREDICT_ROWS` is now `chunking.MAX_CHUNK_ROWS` (2,048), so a light predict is one call.
+- **A killed worker could look alive to `close()`**, which then kept its host-memory reservations (and failed `test_workers_close_together_within_one_grace_period` intermittently in CI). The reader thread and `close()` both reaped a dying worker; when two threads call waitpid at once, the loser finds no child and CPython reports the process still running until the winner records its exit code. Every wait, liveness check and kill of a worker is now serialized under one lock (waits on its process sentinel happen outside it, so health checks never block), so every caller sees the exit, and its code, once one does; a sentinel a forked grandchild still holds falls back to a non-blocking check.
 ### Added
 - `benchmarks/bench_stalls.py` (fixed-cost stalls across table sizes, with compiled-graph counts per call), `benchmarks/profile_predict.py` (cProfile of one predict), `benchmarks/calibrate_predict.py` and `benchmarks/gen_predict_grid.py` (predict calibration), `benchmarks/sweep_gateway.py` (the table-size sweep through the Vector gateway).
 

src/tabctx/serve/workers.py

@@ -49,6 +49,7 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from concurrent.futures import TimeoutError as FutureTimeoutError
 from dataclasses import dataclass
+from multiprocessing.connection import wait as wait_for_handles
 from pathlib import Path
 from typing import Any
 
@@ -310,6 +311,7 @@ def __init__(
         _LIVE_SHARDS.add(self)
         self._send_lock = threading.Lock()
         self._state_lock = threading.Lock()
+        self._reap_lock = threading.Lock()
         self._pending: dict[int, Future] = {}
         self._ids = itertools.count(1)
         self._dead: str | None = None
@@ -329,12 +331,13 @@ def wait_ready(self, timeout_s: float) -> None:
         try:
             first = self._conn.recv()
         except (EOFError, OSError) as e:
+            self._exited(5)  # reap it, for a meaningful exit code
             raise WorkerDiedError(
                 f"worker for {self.device} exited during startup "
-                f"(exit code {self._proc.exitcode})"
+                f"(exit code {self._exit_code})"
             ) from e
         if isinstance(first, _Failure):
-            self._proc.join(5)
+            self._exited(5)
             raise first.error
         self._backend_names = first.backend_names
         self._default_backend = first.default_backend
@@ -404,27 +407,57 @@ def _request_exit(self) -> None:
                 pass
 
     def _finish_close(self, grace_s: float) -> None:
-        self._proc.join(grace_s)
-        if self._proc.is_alive():
-            self._proc.kill()
-            self._proc.join(5)
+        if not self._exited(grace_s):
+            self._kill()
+            self._exited(5)
         self._conn.close()
         self._release_host_memory()
         self._fail_pending(f"worker for {self.device} was shut down")
 
+    def _exited(self, wait_s: float = 0.0) -> bool:
+        """Whether the worker process has exited (waiting up to wait_s for
+        it to), reaping it if so. Every wait, liveness check and kill goes
+        through _reap_lock, because asking the process itself races:
+        is_alive() and join() call waitpid, and when two threads do at once
+        (the reader thread and close() both wait on a dying worker) the
+        loser finds no child and reports it still running (CPython's
+        Popen.poll) -- so close() could see a killed worker as alive and
+        keep its host reservations. The wait is on the process sentinel,
+        outside the lock, so a health check never waits behind it. Once the
+        sentinel is ready, a blocking reap returns at once (a process's fds
+        close just before its exit status is available); a sentinel a forked
+        grandchild still holds never becomes ready, so the exit is then read
+        without blocking."""
+        ready = wait_for_handles([self._proc.sentinel], wait_s)
+        with self._reap_lock:
+            if ready:
+                self._proc.join(5)
+            return self._proc.exitcode is not None
+
+    def _kill(self) -> None:
+        """SIGKILL the worker unless it was reaped (its pid may be reused)."""
+        with self._reap_lock:
+            if self._proc.exitcode is None:
+                self._proc.kill()
+
+    @property
+    def _exit_code(self) -> int | None:
+        with self._reap_lock:
+            return self._proc.exitcode
+
     def _release_host_memory(self) -> None:
         """Return host memory the (now exited) worker had reserved."""
-        if self._host_guard is not None and not self._proc.is_alive():
+        if self._host_guard is not None and self._exited():
             self._host_guard.release_owner(self._index)
 
     @property
     def alive(self) -> bool:
-        return self._dead is None and self._proc.is_alive()
+        return self._dead is None and not self._exited()
 
     @property
     def death_reason(self) -> str | None:
-        if self._dead is None and not self._proc.is_alive():
-            return f"worker for {self.device} exited (code {self._proc.exitcode})"
+        if self._dead is None and self._exited():
+            return f"worker for {self.device} exited (code {self._exit_code})"
         return self._dead
 
     # ---- transport ---------------------------------------------------
@@ -440,7 +473,7 @@ def _read_loop(self) -> None:
                 # trustworthy: retire the worker (its calls fail, health
                 # checks fail, the serving framework restarts the replica).
                 logger.exception("unreadable reply from worker for %s", self.device)
-                self._proc.kill()
+                self._kill()
                 break
             self._apply(snap)
             with self._state_lock:
@@ -451,13 +484,13 @@ def _read_loop(self) -> None:
                 fut.set_result(payload)
             else:
                 fut.set_exception(payload)
-        self._proc.join(5)  # reap it, for a meaningful exit code
-        if self._dead is None and self._proc.is_alive():
-            self._proc.kill()  # its pipe is gone: it can't serve, and must
-            self._proc.join(5)  # exit before its reservations are returned
+        exited = self._exited(5)  # reap it, for a meaningful exit code
+        if self._dead is None and not exited:
+            self._kill()  # its pipe is gone: it can't serve, and must
+            self._exited(5)  # exit before its reservations are returned
         self._release_host_memory()
         reason = (
-            f"worker for {self.device} exited (code {self._proc.exitcode}); "
+            f"worker for {self.device} exited (code {self._exit_code}); "
             "its contexts are lost"
         )
         if self._dead is None:

tests/unit/test_workers.py

@@ -190,6 +190,58 @@ def test_workers_close_together_within_one_grace_period(monkeypatch):
     assert not any(s._proc.is_alive() for s in shards)
 
 
+def test_concurrent_waits_all_see_a_dead_worker_exit():
+    """The reader thread and close() both wait on a dying worker. Two
+    waitpid calls race, and the loser used to report the worker still
+    running (so close() kept its host reservations); every wait and
+    liveness check now reaps under one lock, so whoever sees the exit also
+    sees its exit code."""
+    b = build_engine(ServeSettings(backends=("fake",), devices=("cuda:0", "cuda:1")))
+    shards = list(b.worker_shards)
+    shard = shards[0]
+    try:
+        barrier = threading.Barrier(8)
+        seen: list[bool] = []
+
+        def wait() -> None:
+            barrier.wait()
+            seen.append(shard._exited(5) and shard._exit_code is not None)
+
+        threads = [threading.Thread(target=wait) for _ in range(7)]
+        for t in threads:
+            t.start()
+        shard._proc.kill()
+        barrier.wait()
+        for t in threads:
+            t.join()
+        assert seen == [True] * 7
+        assert not shard.alive
+    finally:
+        ProcessShard.close_all(shards)
+
+
+def test_a_dead_worker_exits_while_a_grandchild_holds_its_sentinel(monkeypatch):
+    """A process the worker forked inherits its sentinel's write end and
+    keeps it open after the worker dies: the exit is then read without the
+    sentinel, and close() still returns its host reservations."""
+    import tabctx.serve.workers as workers
+
+    b = build_engine(ServeSettings(backends=("fake",), devices=("cuda:0", "cuda:1")))
+    shards = list(b.worker_shards)
+    try:
+        monkeypatch.setattr(workers, "wait_for_handles", lambda handles, timeout: [])
+        shard = shards[0]
+        shard._proc.kill()
+        for _ in range(500):
+            if shard._exited():
+                break
+            threading.Event().wait(0.01)
+        assert shard._exited()
+        assert shard._exit_code is not None
+    finally:
+        ProcessShard.close_all(shards)
+
+
 def test_mp_exit_handler_is_registered_before_ours():
     """Our atexit handler must run first (atexit is LIFO), however the
     module was imported: workers.py imports multiprocessing.util itself."""