From edc4cf0baa5c73c6c35c4c78932e11d75449d049 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:56:41 +0900 Subject: [PATCH] fix(serve): keep a PD stage accepting work after a refused stop (#174) submit("stop") clears `accepting` before the owner checks quiescence. A stop is refused (409) only on a non-quiescent stage, e.g. a poisoned one, and afterwards /v1/admin/reset and a second stop were rejected with 503 "owner not accepting work" while the owner thread kept running. The owner now restores `accepting` when the stop fails. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/accel/serve/llmd.py | 2 ++ opt_prime/tests/test_pd_recovery.py | 29 ++++++++++++++++++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/opt_prime/opt_prime/accel/serve/llmd.py b/opt_prime/opt_prime/accel/serve/llmd.py index 85b4da28..fad81b23 100644 --- a/opt_prime/opt_prime/accel/serve/llmd.py +++ b/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) diff --git a/opt_prime/tests/test_pd_recovery.py b/opt_prime/tests/test_pd_recovery.py index 90dbb8d3..a34f4ec2 100644 --- a/opt_prime/tests/test_pd_recovery.py +++ b/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"))