[Bugfix] graceful shutdown for multi-stage engine processes (#3001)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

tests/engine/test_orchestrator.py

@@ -10,11 +10,14 @@ from types import SimpleNamespace
 from typing import Any
 
 import janus
+import psutil
 import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.core_client import AsyncMPClient
 
 from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -104,6 +107,42 @@ class FakeOutputProcessor:
         return None
 
 
+class _FakeProc:
+    pid = 1234
+
+    def __init__(self):
+        self.terminated = False
+        self.killed = False
+        self.join_calls = []
+
+    def is_alive(self):
+        return not self.terminated and not self.killed
+
+    def terminate(self):
+        self.terminated = True
+
+    def kill(self):
+        self.killed = True
+
+    def join(self, timeout=None):
+        self.join_calls.append(timeout)
+
+
+class _FakeChildProc:
+    def __init__(self):
+        self.terminated = False
+        self.killed = False
+
+    def is_running(self):
+        return not self.terminated and not self.killed
+
+    def terminate(self):
+        self.terminated = True
+
+    def kill(self):
+        self.killed = True
+
+
 def _sampling_params(max_tokens: int = 4) -> SamplingParams:
     return SamplingParams(max_tokens=max_tokens)
 
@@ -269,6 +308,61 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ
     )
 
 
+def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch):
+    fake_proc = _FakeProc()
+    fake_child = _FakeChildProc()
+
+    class FakePsutilProcess:
+        def __init__(self, pid):
+            assert pid == fake_proc.pid
+
+        def children(self, recursive=True):
+            assert recursive
+            return [fake_child]
+
+    def fail_base_shutdown(self):
+        raise RuntimeError("base shutdown failed")
+
+    monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
+    monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), []))
+    monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown)
+
+    client = object.__new__(StageEngineCoreClient)
+    client._proc = fake_proc
+
+    with pytest.raises(RuntimeError, match="base shutdown failed"):
+        client.shutdown()
+
+    assert fake_proc.terminated
+    assert fake_proc.join_calls == [5]
+    assert fake_child.terminated
+
+
+def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch):
+    fake_proc = _FakeProc()
+    fake_child = _FakeChildProc()
+
+    class FakePsutilProcess:
+        def __init__(self, pid):
+            assert pid == fake_proc.pid
+
+        def children(self, recursive=True):
+            assert recursive
+            return [fake_child]
+
+    monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
+    monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs)))
+    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self: None)
+
+    client = object.__new__(StageEngineCoreClient)
+    client._proc = fake_proc
+
+    client.shutdown()
+
+    assert fake_child.terminated
+    assert fake_child.killed
+
+
 @pytest.fixture
 def orchestrator_factory():
     fixtures: list[OrchestratorFixture] = []

vllm_omni/engine/stage_engine_core_client.py

@@ -13,6 +13,7 @@ import weakref
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
+import psutil
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
@@ -30,6 +31,8 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+SHUTDOWN_TIMEOUT_S = 5
+
 
 class StageEngineCoreClientBase:
     """Shared stage-aware behavior for async EngineCore clients.
@@ -420,13 +423,37 @@ class StageEngineCoreClientBase:
 
     def shutdown(self) -> None:
         """Shutdown managed resources and any externally spawned subprocess."""
-        super().shutdown()
-        if self._proc is not None and self._proc.is_alive():
-            self._proc.terminate()
-            self._proc.join(timeout=5)
-            if self._proc.is_alive():
-                self._proc.kill()
-        self._proc = None
+        child_procs: list[psutil.Process] = []
+        if self._proc is not None and self._proc.pid is not None:
+            try:
+                child_procs = psutil.Process(self._proc.pid).children(recursive=True)
+            except psutil.Error:
+                child_procs = []
+
+        try:
+            super().shutdown()
+        finally:
+            if self._proc is not None and self._proc.is_alive():
+                self._proc.terminate()
+                self._proc.join(timeout=SHUTDOWN_TIMEOUT_S)
+                if self._proc.is_alive():
+                    self._proc.kill()
+                    self._proc.join(timeout=SHUTDOWN_TIMEOUT_S)
+
+            alive_children = [proc for proc in child_procs if proc.is_running()]
+            for proc in alive_children:
+                try:
+                    proc.terminate()
+                except psutil.Error:
+                    pass
+            _, still_alive = psutil.wait_procs(alive_children, timeout=SHUTDOWN_TIMEOUT_S)
+            for proc in still_alive:
+                try:
+                    proc.kill()
+                except psutil.Error:
+                    pass
+            # The process handle is no longer reliable after best-effort cleanup.
+            self._proc = None
 
 
 class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient):

vllm_omni/engine/stage_engine_core_proc.py

@@ -23,10 +23,12 @@ from vllm.utils.system_utils import (
     get_mp_context,
     set_process_title,
 )
-from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.engine import EngineCoreRequestType
+from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState
 from vllm.v1.engine.utils import (
     EngineHandshakeMetadata,
     EngineZmqAddresses,
+    SignalCallback,
     get_engine_zmq_addresses,
 )
 from vllm.v1.utils import shutdown
@@ -54,18 +56,9 @@ class StageEngineCoreProc(EngineCoreProc):
         **kwargs: Any,
     ) -> None:
         """Launch StageEngineCoreProc busy loop in background process."""
-        shutdown_requested = False
+        signal_callback: SignalCallback | None = None
         maybe_register_config_serialize_by_value()
 
-        def signal_handler(signum: int, frame: Any) -> None:
-            nonlocal shutdown_requested
-            if not shutdown_requested:
-                shutdown_requested = True
-                raise SystemExit()
-
-        signal.signal(signal.SIGTERM, signal_handler)
-        signal.signal(signal.SIGINT, signal_handler)
-
         engine_core: StageEngineCoreProc | None = None
         try:
             vllm_config: VllmConfig = kwargs["vllm_config"]
@@ -88,6 +81,19 @@ class StageEngineCoreProc(EngineCoreProc):
                 engine_index=dp_rank,
                 **kwargs,
             )
+
+            def wakeup_engine() -> None:
+                engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None))
+
+            signal_callback = SignalCallback(wakeup_engine)
+
+            def signal_handler(signum: int, frame: Any) -> None:
+                engine_core.shutdown_state = EngineShutdownState.REQUESTED
+                signal_callback.trigger()
+
+            signal.signal(signal.SIGTERM, signal_handler)
+            signal.signal(signal.SIGINT, signal_handler)
+
             engine_core.run_busy_loop()
 
         except SystemExit:
@@ -101,6 +107,10 @@ class StageEngineCoreProc(EngineCoreProc):
                 engine_core._send_engine_dead()
             raise
         finally:
+            signal.signal(signal.SIGTERM, signal.SIG_DFL)
+            signal.signal(signal.SIGINT, signal.SIG_DFL)
+            if signal_callback is not None:
+                signal_callback.stop()
             if engine_core is not None:
                 engine_core.shutdown()