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)