From 369c8ffa50017aef411383e294d5c00fffbecac6 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 23:27:06 +0900 Subject: [PATCH] fix: settle prefill futures and clean up failed subprocess workers (#238) --- opt_prime/examples/pd/http_matrix/STATUS.md | 4 + .../evidence238/cpu-validation.json | 17 ++ .../pd/http_matrix/evidence238/events.jsonl | 2 + .../2026-10-23-prefill-worker-failure238.md | 42 ++++ opt_prime/opt_prime/accel/pd/pd_batching.py | 164 +++++++++++----- .../tests/test_subprocess_prefill_failure.py | 179 ++++++++++++++++++ 6 files changed, 363 insertions(+), 45 deletions(-) create mode 100644 opt_prime/examples/pd/http_matrix/evidence238/cpu-validation.json create mode 100644 opt_prime/examples/pd/http_matrix/evidence238/events.jsonl create mode 100644 opt_prime/examples/pd/http_matrix/experiments/2026-10-23-prefill-worker-failure238.md create mode 100644 opt_prime/tests/test_subprocess_prefill_failure.py diff --git a/opt_prime/examples/pd/http_matrix/STATUS.md b/opt_prime/examples/pd/http_matrix/STATUS.md index 9f065c3a..f93eb16c 100644 --- a/opt_prime/examples/pd/http_matrix/STATUS.md +++ b/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. diff --git a/opt_prime/examples/pd/http_matrix/evidence238/cpu-validation.json b/opt_prime/examples/pd/http_matrix/evidence238/cpu-validation.json new file mode 100644 index 00000000..551a9e3a --- /dev/null +++ b/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 +} diff --git a/opt_prime/examples/pd/http_matrix/evidence238/events.jsonl b/opt_prime/examples/pd/http_matrix/evidence238/events.jsonl new file mode 100644 index 00000000..bcae6e30 --- /dev/null +++ b/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"} diff --git a/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-prefill-worker-failure238.md b/opt_prime/examples/pd/http_matrix/experiments/2026-10-23-prefill-worker-failure238.md new file mode 100644 index 00000000..b09107b8 --- /dev/null +++ b/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. diff --git a/opt_prime/opt_prime/accel/pd/pd_batching.py b/opt_prime/opt_prime/accel/pd/pd_batching.py index 022c7b35..ed56608c 100644 --- a/opt_prime/opt_prime/accel/pd/pd_batching.py +++ b/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: diff --git a/opt_prime/tests/test_subprocess_prefill_failure.py b/opt_prime/tests/test_subprocess_prefill_failure.py new file mode 100644 index 00000000..d5895c0f --- /dev/null +++ b/opt_prime/tests/test_subprocess_prefill_failure.py @@ -0,0 +1,179 @@ +"""Real pipe/process failures; the tiny child replaces only the expensive engine.""" +import os +import subprocess +import sys +import time +from concurrent.futures import TimeoutError +from multiprocessing import shared_memory + +import pytest +from opt_prime.accel.pd.contracts import PDRequest +from opt_prime.accel.pd.pd_batching import SubprocessPrefillWorker + +CHILD = ''' +import json, os, sys, time +from pathlib import Path +mode = os.environ.get("PREFILL_TEST_MODE", "ok") +if mode == "stubborn": + import signal + signal.signal(signal.SIGTERM, signal.SIG_IGN) +gate = Path(os.environ["PREFILL_TEST_GATE"]) +if mode == "startup_stall": + time.sleep(120) +if mode == "startup_exit": + sys.exit(7) +if mode == "logs": + print('{"compile_start": true}', flush=True) + print('[]', flush=True) +print('{"ready": true}', flush=True) +for line in sys.stdin: + msg = json.loads(line) + if msg["id"] == 0 and mode in ("cancel", "exit", "malformed", "stubborn"): + gate.with_suffix(".seen").touch() + while not gate.exists(): + time.sleep(.005) + if mode == "exit": + sys.exit(8) + if mode == "malformed": + print(json.dumps({"id": msg["id"], "n": "bad"}), flush=True) + continue + print(json.dumps({"id": msg["id"], "n": 1, "token": 17}), flush=True) +''' + + +@pytest.fixture +def children(tmp_path, monkeypatch): + """Observe real allocated resources, and clean up even against the buggy baseline.""" + procs, names = [], [] + popen, shm = subprocess.Popen, shared_memory.SharedMemory + + def spawn(*args, **kwargs): + p = popen(*args, **kwargs) + procs.append(p) + return p + + def allocate(*args, **kwargs): + s = shm(*args, **kwargs) + if kwargs.get("create"): + names.append(s.name) + return s + + monkeypatch.setattr(subprocess, "Popen", spawn) + monkeypatch.setattr(shared_memory, "SharedMemory", allocate) + script = tmp_path / "child" + script.write_text(f"#!{sys.executable}\n" + CHILD) + script.chmod(0o755) + workers = [] + + def make(mode="ok", **kwargs): + env = dict(os.environ, PREFILL_TEST_MODE=mode, PREFILL_TEST_GATE=str(tmp_path / "gate")) + w = SubprocessPrefillWorker(str(script), [], slots=2, layers=1, kv_heads=1, + head_dim=2, max_context=8, env=env, **kwargs) + workers.append(w) + return w + + yield make, procs, names, tmp_path / "gate" + for p in procs: + if p.poll() is None: + p.kill() + p.wait(timeout=5) + for w in workers: + w.close() + for name in names: + try: + s = shm(name=name) + s.close() + s.unlink() + except FileNotFoundError: + pass + + +def request(): + return PDRequest("test", (1, 2), 1) + + +def release(gate): + deadline = time.monotonic() + 3 + while not gate.with_suffix(".seen").exists(): + assert time.monotonic() < deadline, "child did not receive request" + time.sleep(.005) + gate.touch() + + +def test_json_sdk_logs_do_not_break_startup(children): + make, _, _, _ = children + w = make("logs", startup_timeout=1) + assert w.submit(request(), slot=0).result(timeout=2).token == 17 + + +def test_cancelled_response_does_not_strand_next_request(children): + make, _, _, gate = children + w = make("cancel") + first = w.submit(request(), slot=0) + assert first.cancel() + second = w.submit(request(), slot=1) + release(gate) + assert second.result(timeout=2).token == 17 + assert first.cancelled() + + +@pytest.mark.parametrize("mode", ["exit", "malformed"]) +def test_terminal_failure_settles_all_requests_and_rejects_new_ones(children, mode): + make, _, _, gate = children + w = make(mode) + futures = [w.submit(request(), slot=i) for i in range(2)] + release(gate) + for f in futures: + with pytest.raises(RuntimeError): + f.result(timeout=2) + with pytest.raises(RuntimeError): + w.submit(request(), slot=0).result(timeout=2) + replacement = make() + assert replacement.submit(request(), slot=0).result(timeout=2).token == 17 + + +@pytest.mark.parametrize("mode", ["startup_stall", "startup_exit"]) +def test_failed_startup_releases_child_and_shared_memory(children, mode): + make, procs, names, _ = children + with pytest.raises((TimeoutError, RuntimeError)): + make(mode, startup_timeout=.2) + assert procs[-1].poll() is not None + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=names[-1]) + + +def test_close_is_idempotent_and_rejects_new_work(children): + make, _, names, _ = children + w = make() + w.close() + w.close() + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=names[-1]) + with pytest.raises(RuntimeError): + w.submit(request(), slot=0).result(timeout=2) + + +def test_close_kills_unresponsive_worker_and_settles_pending(children): + make, procs, names, _ = children + w = make("stubborn") + pending = w.submit(request(), slot=0) + w.close(timeout=.1) + with pytest.raises(RuntimeError): + pending.result(timeout=2) + assert procs[-1].poll() is not None + assert not w._recv.is_alive() + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=names[-1]) + + +def test_spawn_failure_does_not_leak_shared_memory(children, monkeypatch): + make, _, names, _ = children + + def fail(*args, **kwargs): + raise OSError("cannot start child") + + monkeypatch.setattr(subprocess, "Popen", fail) + with pytest.raises(OSError, match="cannot start child"): + make() + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=names[-1])