diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 0b549f58e..c07762ad5 100644 --- a/tests/engine/test_orchestrator.py +++ b/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] = [] diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index f6683c1e0..844be7194 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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): diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 689378a79..2ab8b37dd 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/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()