"""Production worker-process isolation for the Core ML / ANE backend. DecisionEngine / 05_serve --backend coreml | JSON lines over stdin/stdout v python coreml_worker.py | v CoreMLEngine -> Core ML / Apple Neural Engine Why a subprocess: on this machine the ANE runtime can corrupt the host heap when a bridged predict output is deallocated on a background dispatch thread (objc_destructInstance -> _PyObject_Free, then unrelated native calls fault; see ANE_REPORT.md "failed experiments"). A disposable child bounds the blast radius - a worker crash is an EOF on a pipe, not a SIGSEGV inside a serving process, and CoreMLEngine's _keepalive list (which must retain every predict output to avoid that dealloc) is bounded by the worker's lifetime. Protocol - one JSON object per line, request in / response out: {"op": "info"} -> engine metadata (compute_units, buckets, temps, ...) {"op": "decide", "state": ..., "questions": ..., "calibrated": bool} -> {"result": {...}, "label_calls": int} {"op": "stats"} -> {"calls": int, "uptime_s": float, "rss_bytes": int} {"op": "ping"} -> {"pong": true} {"op": "quit"} -> worker exits Every response carries "ok": true or {"ok": false, "error": "..."}. WorkerCoreMLEngine is the parent-side client. It exposes the same practical surface as CoreMLEngine (decide + metadata attributes), spawns the worker eagerly at init, respawns once and retries a request once on worker death or read timeout, and can recycle the worker every `recycle_every` decisions so the retained predict outputs stay bounded in a long-running service. """ from __future__ import annotations import json import os import select import subprocess import sys import threading import time from collections import deque READ_TIMEOUT_S = 120.0 # a hung worker is as dead as a crashed one RECYCLE_EVERY = 10000 # decisions per worker lifetime before recycling _CLOSE_TIMEOUT_S = 10.0 _STDERR_LINES = 200 _WORKER_PY = os.path.abspath(__file__) def _self_rss_bytes() -> int: """Current RSS - /proc on Linux, `ps` elsewhere (macOS).""" try: with open("/proc/self/status") as f: for line in f: if line.startswith("VmRSS:"): return int(line.split()[1]) * 1024 except OSError: pass try: out = subprocess.run(["ps", "-o", "rss=", "-p", str(os.getpid())], capture_output=True, text=True, timeout=5) return int(out.stdout.strip()) * 1024 except (OSError, ValueError, subprocess.SubprocessError): return 0 def worker_main(argv) -> int: """Child-process entrypoint: CoreMLEngine behind the JSON-lines protocol.""" pkg_dir, ckpt, cu = argv[0], argv[1], argv[2] seq_lens = None if len(argv) > 3: seq_lens = [int(x) for x in argv[3].split(",") if x] from coreml_backend import CoreMLEngine eng = CoreMLEngine(pkg_dir, ckpt=None if ckpt == "-" else ckpt, compute_units=cu, seq_lens=seq_lens) calls = 0 t_start = time.monotonic() for line in sys.stdin: line = line.strip() if not line: continue try: req = json.loads(line) except json.JSONDecodeError as e: out = {"ok": False, "error": f"bad request json: {e}"} sys.stdout.write(json.dumps(out) + "\n") sys.stdout.flush() continue try: if req["op"] == "info": res = {"compute_units": eng.compute_units, "encoder": eng.encoder_name, "max_seq": eng.max_seq, "max_len": eng.max_len, "encode_len": eng.encode_len, "temperatures": eng.temperatures, "slots": eng.slots, "buckets": sorted(eng._models)} elif req["op"] == "ping": res = {"pong": True} elif req["op"] == "stats": res = {"calls": calls, "uptime_s": time.monotonic() - t_start, "rss_bytes": _self_rss_bytes(), "keepalive_len": len(eng._keepalive)} elif req["op"] == "decide": calls_before = calls orig = eng.label_logits def counting(enc): nonlocal calls calls += 1 return orig(enc) eng.label_logits = counting try: res = {"result": eng.decide( req["state"], req["questions"], req.get("calibrated", True)), "label_calls": calls - calls_before} finally: eng.label_logits = orig elif req["op"] == "quit": return 0 else: raise ValueError(f"unknown op {req['op']!r}") out = {"ok": True, **res} except Exception as e: # noqa: BLE001 - report it out = {"ok": False, "error": f"{type(e).__name__}: {e}"} sys.stdout.write(json.dumps(out, ensure_ascii=False) + "\n") sys.stdout.flush() return 0 class WorkerCoreMLEngine: """decide() proxy for a CoreMLEngine living in a worker subprocess. Same practical interface as CoreMLEngine - construction takes the same (packages, ckpt, compute_units, seq_lens) arguments and metadata arrives via the worker's info response. Lifecycle: - worker dies or a read times out -> kill/reap, respawn, retry once; a second failure raises RuntimeError (no respawn loop). - `recycle_every` decisions per worker bounds CoreMLEngine._keepalive, which must retain every predict output (releasing them is what trips the ANE heap bug). Recycling is transparent and lazy - it happens at the start of the next decide(), so a quiet service pays nothing. - close() asks the worker to quit, then kills if it does not. `worker_cmd` overrides the spawn command - tests use it to point at a fake worker that needs no Core ML runtime. """ def __init__(self, packages="artifacts/coreml", ckpt=None, compute_units="cpu_and_ne", seq_lens=None, read_timeout=READ_TIMEOUT_S, recycle_every=RECYCLE_EVERY, worker_cmd=None): self._args = (packages, ckpt, compute_units, seq_lens) self.read_timeout = read_timeout self.recycle_every = recycle_every self._worker_cmd = list(worker_cmd) if worker_cmd else None self.calls = 0 # successful decide() calls, all workers self.restarts = 0 # crash/timeout respawns (unplanned) self.recycles = 0 # planned respawns (keepalive bound) self.last_label_calls = 0 self._calls_this_worker = 0 self._stderr = deque(maxlen=_STDERR_LINES) self._spawn() info = self._call({"op": "info"}) self.compute_units = info["compute_units"] self.encoder_name = info.get("encoder") self.max_seq = info["max_seq"] self.max_len = info["max_len"] self.encode_len = info["encode_len"] self.temperatures = info["temperatures"] self.slots = info.get("slots", "lead") self.buckets = info["buckets"] # -- process management ------------------------------------------------- def _spawn(self): old = getattr(self, "_p", None) if old is not None and old.poll() is None: old.kill() # hung predecessor, not a dead one try: old.wait(timeout=_CLOSE_TIMEOUT_S) except subprocess.TimeoutExpired: pass pkg_dir, ckpt, cu, seq_lens = self._args cmd = self._worker_cmd if cmd is None: cmd = [sys.executable, _WORKER_PY, pkg_dir, ckpt if ckpt else "-", cu] if seq_lens: cmd.append(",".join(str(s) for s in seq_lens)) self._p = subprocess.Popen( cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=os.path.dirname(_WORKER_PY)) self._rbuf = b"" self._calls_this_worker = 0 # stderr is drained continuously into a bounded buffer: it never # blocks the worker on a full pipe, and it is there for diagnostics # if the worker dies. self._stderr.clear() self._stderr_thread = threading.Thread( target=self._drain_stderr, args=(self._p.stderr,), daemon=True) self._stderr_thread.start() def _drain_stderr(self, stream): try: for line in stream: self._stderr.append( line.decode(errors="replace").rstrip()) except (ValueError, OSError): pass def _stderr_tail(self, n=5): return "; ".join(list(self._stderr)[-n:]) or "(no stderr)" def _read_response_line(self): """One '\n'-terminated line, bounded by read_timeout. None on timeout/EOF. Uses os.read + an own buffer: buffered text IO can hold a second response that select() then never reports.""" while b"\n" not in self._rbuf: ready = select.select([self._p.stdout], [], [], self.read_timeout)[0] if not ready: return None chunk = os.read(self._p.stdout.fileno(), 1 << 16) if not chunk: return None self._rbuf += chunk line, self._rbuf = self._rbuf.split(b"\n", 1) return line def _call(self, req, _retries=1): assert self._p.stdin and self._p.stdout try: self._p.stdin.write( (json.dumps(req, ensure_ascii=False) + "\n").encode()) self._p.stdin.flush() line = self._read_response_line() except (BrokenPipeError, ValueError, OSError): line = None if line is None: if req["op"] == "quit": # already dead - nothing to quit return {"ok": True} if _retries <= 0: raise RuntimeError( f"coreml worker died or timed out on {req['op']!r} " f"(exit {self._p.poll()}, stderr: {self._stderr_tail()})") self.restarts += 1 self._spawn() # respawn and retry once return self._call(req, _retries - 1) res = json.loads(line) if not res["ok"]: raise RuntimeError(f"coreml worker: {res['error']}") return res # -- engine surface ------------------------------------------------------ def decide(self, state, questions, calibrated=True): """Same shape as CoreMLEngine.decide() / 05_serve.decide().""" if self.recycle_every and self._calls_this_worker >= self.recycle_every: # Retained predict outputs bound the ANE heap bug; recycling the # worker frees them with the process instead of deallocating in # place, which is the unsafe operation. self.recycles += 1 self._recycle() res = self._call({"op": "decide", "state": state, "questions": questions, "calibrated": calibrated}) self.calls += 1 self._calls_this_worker += 1 self.last_label_calls = res["label_calls"] return res["result"] def ping(self): return self._call({"op": "ping"})["pong"] def worker_stats(self): """Worker-side counters incl. RSS - for benchmarks and leak checks.""" return self._call({"op": "stats"}) def _recycle(self): try: self._call({"op": "quit"}) self._p.wait(timeout=_CLOSE_TIMEOUT_S) except Exception: try: self._p.kill() except Exception: pass self._spawn() def close(self): try: self._call({"op": "quit"}) except Exception: pass try: self._p.wait(timeout=_CLOSE_TIMEOUT_S) except subprocess.TimeoutExpired: self._p.kill() try: self._p.wait(timeout=_CLOSE_TIMEOUT_S) except subprocess.TimeoutExpired: pass def __enter__(self): return self def __exit__(self, *exc): self.close() def __del__(self): try: self.close() except Exception: pass if __name__ == "__main__": raise SystemExit(worker_main(sys.argv[1:]))