From 6dafd6ac50759dee38d7de6dfa47a5e7680dcc85 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:56:42 +0900 Subject: [PATCH] fix: fail pending prefills when the ProcessPrefillWorker child dies (#168) The receiver thread blocked on the reply queue and never checked the child, so an OOM-killed or crashed prefill process left every outstanding future pending forever. The receiver now polls with a timeout; once the child is gone it fails the outstanding futures, and later submits raise, as SubprocessPrefillWorker already does. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/accel/pd/pd_batching.py | 13 ++++++++++++- opt_prime/tests/test_pd_batching_disagg.py | 14 ++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/opt_prime/opt_prime/accel/pd/pd_batching.py b/opt_prime/opt_prime/accel/pd/pd_batching.py index cc572c57..022c7b35 100644 --- a/opt_prime/opt_prime/accel/pd/pd_batching.py +++ b/opt_prime/opt_prime/accel/pd/pd_batching.py @@ -197,6 +197,8 @@ def __init__(self, factory, out_device): def submit(self, request, send_from=0, slot=None) -> Future: f = Future() with self._lock: + if self._futures is None: + raise RuntimeError("prefill process exited") key, self._next = self._next, self._next + 1 self._futures[key] = f self._req.put((key, request, send_from)) @@ -205,7 +207,16 @@ def submit(self, request, send_from=0, slot=None) -> Future: def _receive(self): stream = torch.cuda.Stream(self.out_device) if self.out_device.type == "cuda" else None while True: - item = self._rep.get() + try: + item = self._rep.get(timeout=0.5) + except queue.Empty: + if self._proc.is_alive(): + continue + with self._lock: # the child died: fail what it will never answer + for f in self._futures.values(): + f.set_exception(RuntimeError(f"prefill process exited ({self._proc.exitcode})")) + self._futures = None + return if item is None: return key, out, error = item diff --git a/opt_prime/tests/test_pd_batching_disagg.py b/opt_prime/tests/test_pd_batching_disagg.py index 0e7ee8a2..037f6859 100644 --- a/opt_prime/tests/test_pd_batching_disagg.py +++ b/opt_prime/tests/test_pd_batching_disagg.py @@ -114,6 +114,20 @@ def test_process_worker(): worker.close() +def test_process_worker_child_death_fails_outstanding_prefills(): + """#168: a killed prefill process fails its pending futures (and later submits) instead of leaving them pending""" + import time + worker = ProcessPrefillWorker(functools.partial(time.sleep, 600), out_device="cpu") # never serves + try: + f = worker.submit(PDRequest("dead", (1, 2, 3), 2)) + worker._proc.kill() + assert isinstance(f.exception(timeout=5), RuntimeError) + with pytest.raises(RuntimeError): + worker.submit(PDRequest("late", (1, 2, 3), 2)) + finally: + worker.close() + + def test_served_by_batched_service(): from opt_prime.accel.serve.batched import BatchedService b = _pd(_llama(), slots=2)