fix: keep a PD stage accepting work after a refused stop (#174)

opt_prime/opt_prime/accel/serve/llmd.py

@@ -473,6 +473,8 @@ def run(self):
                         if not pending:
                             del self._pending[args[1]]
             if error is not None:
+                if stopping:
+                    self.accepting = True       # the stop was refused: reset and a retried stop must still get in
                 future.set_exception(error)
             else:
                 future.set_result(result)

opt_prime/tests/test_pd_recovery.py

@@ -4,9 +4,11 @@
 import urllib.error
 import urllib.request
 
+import pytest
+
 from test_pd_llmd import service, body
 from opt_prime.accel.pd import RetryableError
-from opt_prime.accel.serve.llmd import OwnerExecutor, serve_stage
+from opt_prime.accel.serve.llmd import OwnerExecutor, StageError, serve_stage
 
 
 def _http(port, path, rid=None, payload=None, method="POST"):
@@ -108,6 +110,31 @@ def transfer(*args):
         _close(owner, worker, server, http)
 
 
+def test_refused_stop_still_allows_reset_and_stop():
+    """#174: a stop refused on a poisoned stage leaves the owner accepting, so reset and a second stop work"""
+    engine, owner, worker, server, http = _stage()
+    engine.transport.transfer = lambda *a: (_ for _ in ()).throw(RuntimeError("DMA completion unknown"))
+    port = server.server_port
+    try:
+        owner.submit("execute", "prefill", "p", body(1))
+        assert _http(port, "/v1/completions", "p", body())[0] == 500 and engine.state == "poisoned"
+        with pytest.raises(StageError) as refused:
+            owner.submit("stop")
+        assert refused.value.status == 409
+        status, payload, _ = _http(port, "/v1/admin/reset")
+        assert status == 200 and payload == {"state": "idle"}
+        assert owner.submit("stop") is None
+        worker.join(5)
+        assert not worker.is_alive()
+    finally:
+        if worker.is_alive():
+            _close(owner, worker, server, http)
+        else:
+            server.shutdown()
+            server.server_close()
+            http.join(5)
+
+
 def test_reset_refused_without_backend_recovery():
     engine, owner, worker, server, http = _stage(recoverable=False)
     engine.transport.transfer = lambda *a: (_ for _ in ()).throw(RuntimeError("uncertain"))