fix: fail pending prefills when the ProcessPrefillWorker child dies (#168)

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

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)