fix: prevent stranded prefill requests after subprocess failure

opt_prime/examples/pd/http_matrix/STATUS.md

@@ -17,3 +17,7 @@ Long hardware follow-up on1e54786:6423–6428token inputs,10normal/reuse HTTP re
 ## Single-block static PD (#236)
 
 Static B1 with one KV block now supports cache-off full-prefix PD. KONI live RNGD->ATOM HTTP14/14 output checks and16remote/0local prefills pass, including cancel/disconnect reuse.43CPU tests cover queued reuse, late cancellation and partial-prefix error recovery; cache/multi-lane guards remain. No new artifact, hardware batch or cache-restoration claim. See single-block-pd236 record and evidence236.
+
+## Subprocess failure lifecycle (#238)
+
+Structured SDK logs and cancelled/malformed replies previously killed the receiver and could strand later requests. Terminal errors now settle all queued Futures and reject subsequent submits. Failed startup and bounded close release the owned process and slab; fresh work succeeds on an explicitly created replacement worker. Nine CPU fault scenarios plus the existing tiny Torch subprocess path pass in the 64-test focused set. Full suite: 460 passed, 15 skipped, 135 subtests passed. Accelerator fault recovery and service-level automatic restart remain separate acceptance items. See prefill-worker-failure238 record.

opt_prime/examples/pd/http_matrix/evidence238/cpu-validation.json

@@ -0,0 +1,17 @@
+{
+  "baseline": "df02223",
+  "baseline_fault_failures": 9,
+  "baseline_cleanup_errors": 1,
+  "focused": {
+    "passed": 64,
+    "warnings": 1
+  },
+  "full": {
+    "passed": 460,
+    "skipped": 15,
+    "subtests_passed": 135,
+    "warnings": 1,
+    "seconds": 134.34
+  },
+  "hardware_fault_injection": false
+}

opt_prime/examples/pd/http_matrix/evidence238/events.jsonl

@@ -0,0 +1,2 @@
+{"id": "35cfd0e51eeb424cad4794ac11fed9f3", "ts": "2026-10-23T23:24:13+09:00", "task": "Prefill subprocess failure recovery (aicomp #238)", "result": "Reproduced 9 failure cases on df02223. Candidate passes64 targeted tests; full suite460passed/15skipped/135subtests. Final full rerun after bounded cleanup adjustment pending.", "next": "Complete final suite and publish scoped PR; inject failure on owned hardware separately after long queue.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {}, "artifacts": [], "experiments": ["experiments/2026-10-23-prefill-worker-failure238.md"], "labels": ["bug"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "xeon-max", "branch": "main", "commit": "351d39d-dirty"}
+{"id": "e0e5e02d0dd140e5884aa9c3c4a80cd3", "ts": "2026-10-23T23:26:56+09:00", "task": "Prefill subprocess failure recovery (aicomp #238)", "result": "Final candidate:64 focused tests; full460passed/15skipped/135subtests, no failures. Nine real subprocess fault cases reproduce on df02223 and pass after cleanup/cancellation/terminal-error fix.", "next": "Publish review PR; hardware failure injection remains separate from CPU lifecycle checks.", "state": "active", "done": false, "blocker": "", "metrics": {}, "params": {"validation": "cpu_process_lifecycle"}, "artifacts": [], "experiments": ["experiments/2026-10-23-prefill-worker-failure238.md"], "labels": ["bug"], "review": true, "todos": [], "note": "", "agent": "codex", "host": "xeon-max", "branch": "main", "commit": "351d39d-dirty"}

opt_prime/examples/pd/http_matrix/experiments/2026-10-23-prefill-worker-failure238.md

@@ -0,0 +1,42 @@
+# Subprocess prefill failure and cancellation recovery
+
+- **Date**: 2026-10-23
+- **Task**: aicomp #238
+- **Host / device**: CPU subprocess fault tests on development host; no accelerator reset or live hardware injection
+- **Commit**: df02223 baseline; fix/238-prefill-worker-failure working tree (final commit recorded in PR)
+
+## Question
+
+Can structured SDK logs, cancellation, malformed replies, worker exit or failed startup strand requests or leak the worker's shared-memory segment?
+
+## Setup
+
+A tiny real subprocess replaces only the expensive model engine. The production worker still owns real stdin/stdout pipes, Futures, receiver thread and shared memory. A file gate ensures cancellation occurs before a response. Faults are startup stall/exit, process exit with queued work, malformed response, ignored SIGTERM and Popen failure. The normal Torch prefill subprocess test continues to exercise actual KV generation and transfer.
+
+Issue: https://github.com/ai-computing/aicomp/issues/238
+
+## Command
+
+Run from the aicomp checkout with a Torch/Transformers test environment:
+
+```bash
+# Baseline: run this test file with PYTHONPATH pointing to df02223/opt_prime.
+PYTHONPATH=opt_prime python -m pytest -q opt_prime/tests/test_subprocess_prefill_failure.py
+# Focused regression set:
+PYTHONPATH=opt_prime python -m pytest -q opt_prime/tests/test_subprocess_prefill_failure.py opt_prime/tests/test_pd_batching_disagg.py opt_prime/tests/test_atom_batching.py opt_prime/tests/test_atom_hybrid.py opt_prime/tests/test_pd_recovery.py opt_prime/tests/test_serve_batched.py
+# Full project suite, from opt_prime/:
+PYTHONPATH=. python -m pytest -q tests
+```
+
+## Results
+
+- Baseline df02223: 9 failure cases fail; 3 unhandled receiver-thread warnings and 1 cleanup error. Structured JSON logs trigger KeyError, a cancelled Future triggers InvalidStateError, malformed span length triggers TypeError. Failed startup leaves shared memory or a live child. A post-EOF submission writes a dead pipe instead of returning a failed Future.
+- Fixed focused set: 64 passed, 1 existing flex_attention warning, 21.32 s. Includes all 9 fault tests plus actual tiny Torch subprocess prefill.
+- Final full suite: 460 passed, 15 skipped, 135 subtests passed, 1 existing flex_attention warning, 134.34 s. No failed tests.
+- No automatic request replay. A healthy replacement worker explicitly processes a fresh request after terminal failure.
+
+Local raw logs: /tmp/worker238-baseline.log, /tmp/worker238-targeted.log, /tmp/worker238-final-full.log.
+
+## Conclusion
+
+The receiver now ignores non-protocol JSON logs, settles cancelled responses safely, fails all outstanding Futures after terminal protocol/pipe failure and rejects later work through failed Futures. Spawn/startup failure and bounded terminate/kill cleanup release the owned child and slab. Request/response and shared-memory layouts remain unchanged. This is CPU process-lifecycle evidence, not accelerator recovery, service-level automatic restart, compute/copy overlap or performance evidence. Long four-model hardware validation continues independently under #234.

opt_prime/opt_prime/accel/pd/pd_batching.py

@@ -263,64 +263,138 @@ def __init__(self, python, engine_args, *, slots, layers, kv_heads, head_dim, ma
         cmd = [python, "-m", "opt_prime.accel.pd.prefill_server", "--shm", self._shm.name, "--slots", str(slots),
                "--layers", str(layers), "--kv-heads", str(kv_heads), "--head-dim", str(head_dim),
                "--max-context", str(max_context), *engine_args]
-        self._proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, bufsize=1,
-                                      env=env)
         self._json = json
         self._futures, self._next, self._lock = {}, 0, threading.Lock()
+        self._write_lock = threading.Lock()
+        self._failed = None
+        self._proc = self._recv = None
+        self._closed = False
         ready = Future()
-        self._recv = threading.Thread(target=self._receive, args=(ready,), daemon=True, name="prefill-subprocess")
-        self._recv.start()
-        ready.result(timeout=startup_timeout)
+        try:
+            self._proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True,
+                                          bufsize=1, env=env)
+            self._recv = threading.Thread(target=self._receive, args=(ready,), daemon=True,
+                                          name="prefill-subprocess")
+            self._recv.start()
+            ready.result(timeout=startup_timeout)
+        except BaseException:
+            self.close()
+            raise
 
     def submit(self, request, send_from=0, slot=None) -> Future:
-        if slot is None:
-            raise ValueError("SubprocessPrefillWorker needs the decode lane (slot)")
+        if slot is None or not 0 <= slot < self.slab.shape[0]:
+            raise ValueError("SubprocessPrefillWorker needs a valid decode lane (slot)")
         f = Future()
-        with self._lock:
-            key, self._next = self._next, self._next + 1
-            self._futures[key] = (f, slot, send_from)
         s = request.sampling
-        line = self._json.dumps(dict(id=key, tokens=list(request.input_ids), send_from=send_from, slot=slot,
-                                     sampling=dict(temperature=s.temperature, top_p=s.top_p, top_k=s.top_k, seed=s.seed,
-                                                   stop_token_ids=list(s.stop_token_ids))))
-        self._proc.stdin.write(line + "\n")
-        self._proc.stdin.flush()
+        write_error = None
+        # Keep pipe writes ordered without holding the response-map lock across blocking I/O.
+        with self._write_lock:
+            with self._lock:
+                error = self._failed
+                key, self._next = self._next, self._next + 1
+                if error is None:
+                    self._futures[key] = (f, slot, send_from)
+            if error is not None:
+                f.set_exception(error)
+                return f
+            try:
+                line = self._json.dumps(dict(id=key, tokens=list(request.input_ids), send_from=send_from, slot=slot,
+                                             sampling=dict(temperature=s.temperature, top_p=s.top_p, top_k=s.top_k,
+                                                           seed=s.seed, stop_token_ids=list(s.stop_token_ids))))
+                self._proc.stdin.write(line + "\n")
+                self._proc.stdin.flush()
+            except (OSError, ValueError, TypeError) as error:
+                write_error = RuntimeError(f"prefill subprocess write failed: {error}")
+        if write_error is not None:
+            self._fail(write_error)
         return f
 
-    def _receive(self, ready):
-        for line in self._proc.stdout:
-            try:
-                msg = self._json.loads(line)
-            except ValueError:
-                continue                                       # SDK log lines on stdout
-            if msg.get("ready"):
-                ready.set_result(True)
-                continue
-            with self._lock:
-                f, slot, start = self._futures.pop(msg["id"])
-            if "error" in msg:
-                f.set_exception(RuntimeError(f"prefill subprocess: {msg['error']}"))
-                continue
-            n = msg["n"]
-            kv = torch.from_numpy(self.slab[slot, :, :, :, :n].view("int16")).view(torch.bfloat16)
-            f.set_result(PrefillOut(msg["token"], start, kv[0], kv[1], msg.get("cached", 0), msg.get("prefill_s", 0.0)))
-        if not ready.done():
-            ready.set_exception(RuntimeError(f"prefill subprocess exited ({self._proc.wait()})"))
+    def _fail(self, error):
         with self._lock:
-            for f, _, _ in self._futures.values():
-                f.set_exception(RuntimeError("prefill subprocess exited"))
-            self._futures.clear()
+            if self._failed is None:
+                self._failed = error
+            pending, self._futures = self._futures, {}
+            error = self._failed
+        # Future callbacks may call back into the worker; never invoke them under its lock.
+        for f, _, _ in pending.values():
+            if f.set_running_or_notify_cancel():
+                f.set_exception(error)
 
-    def close(self):
+    def _receive(self, ready):
         try:
-            self._proc.stdin.close()
-            self._proc.wait(timeout=60)
-        finally:
-            self._shm.close()
+            for line in self._proc.stdout:
+                try:
+                    msg = self._json.loads(line)
+                except ValueError:
+                    continue                                  # SDK log lines on stdout
+                if not isinstance(msg, dict):
+                    continue
+                if msg.get("ready"):
+                    if not ready.done():
+                        ready.set_result(True)
+                    continue
+                if "id" not in msg:
+                    continue                                  # structured SDK logs
+                with self._lock:
+                    entry = self._futures.get(msg["id"])
+                    failed = self._failed
+                if failed is not None:
+                    return
+                if entry is None:
+                    raise ValueError("unknown response ID")
+                f, slot, start = entry
+                error = None
+                if "error" in msg:
+                    error = RuntimeError(f"prefill subprocess: {msg['error']}")
+                else:
+                    n = msg["n"]
+                    if type(n) is not int or not 0 <= n <= self.slab.shape[-2]:
+                        raise ValueError("invalid KV span length")
+                    kv = torch.from_numpy(self.slab[slot, :, :, :, :n].view("int16")).view(torch.bfloat16)
+                    out = PrefillOut(msg["token"], start, kv[0], kv[1], msg.get("cached", 0),
+                                     msg.get("prefill_s", 0.0))
+                # Keep the request registered until validation succeeds so failure settles it too.
+                with self._lock:
+                    entry = self._futures.pop(msg["id"], None)
+                if entry is not None and f.set_running_or_notify_cancel():
+                    if error is not None:
+                        f.set_exception(error)
+                    else:
+                        f.set_result(out)
+            raise RuntimeError("prefill subprocess exited")
+        except Exception as error:
+            failure = RuntimeError(f"prefill subprocess response failed: {error}")
+            self._fail(failure)
+            if not ready.done():
+                ready.set_exception(failure)
+
+    def close(self, timeout=5):
+        import subprocess
+        if self._closed:
+            return
+        self._closed = True
+        self._fail(RuntimeError("prefill subprocess closed"))
+        if self._proc is not None:
+            # Stop the writer before unmapping its slab, including when startup never completed.
+            if self._proc.poll() is None:
+                self._proc.terminate()
+                try:
+                    self._proc.wait(timeout=timeout)
+                except subprocess.TimeoutExpired:
+                    self._proc.kill()
+                    self._proc.wait()
+            if self._recv is not None and self._recv is not threading.current_thread():
+                self._recv.join(timeout=timeout)
             try:
-                self._shm.unlink()
-            except FileNotFoundError:
-                pass
+                self._proc.stdin.close()
+            except OSError:
+                pass                                         # a broken pipe must not leak the slab
+            self._proc.stdout.close()
+        self._shm.close()
+        try:
+            self._shm.unlink()
+        except FileNotFoundError:
+            pass
 
 
 class PrefillWorkerPool: