From 3a386cfaf51f872334c64edd08fbf71dd4bb0120 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 11:39:59 -0500 Subject: [PATCH] [ROCm] Give EngineCore cleanup grace after request abort (#52281) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../entrypoints/openai/test_dp_supervisor.py | 34 +++++ .../v1/engine/test_startup_watch_processes.py | 129 ++++++++++++++++++ vllm/entrypoints/openai/dp_supervisor.py | 7 +- vllm/v1/engine/core.py | 18 ++- vllm/v1/engine/utils.py | 36 ++++- 5 files changed, 221 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 980f98712f..4bd0b4f7c3 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -242,6 +242,40 @@ def test_handles_shutdown_event(): assert supervisor.is_ready is False +@pytest.mark.asyncio +async def test_shutdown_children_uses_engine_process_timeout( + monkeypatch: pytest.MonkeyPatch, +): + supervisor = DPSupervisor(_make_unit_args(shutdown_timeout=0.0)) + supervisor._processes = [ + SimpleNamespace(name="APIServer_DPRank_4", pid=None, is_alive=lambda: False) + ] + calls = [] + timeout_calls = [] + + def get_process_timeout(request_timeout, manager_timeout): + timeout_calls.append((request_timeout, manager_timeout)) + return 15.0 + + monkeypatch.setattr( + dp_sup, + "get_engine_process_shutdown_timeout", + get_process_timeout, + ) + monkeypatch.setattr( + dp_sup, + "_join_processes_with_timeout", + lambda processes, timeout: calls.append((processes, timeout)), + ) + + await supervisor._shutdown_children() + + assert timeout_calls == [(0.0, 0.0)] + assert calls == [ + (supervisor._processes, 15.0 + CHILD_EXIT_GRACE_S), + ] + + @pytest.mark.asyncio async def test_handles_child_exit( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/v1/engine/test_startup_watch_processes.py b/tests/v1/engine/test_startup_watch_processes.py index b7b4b8a5a6..0b346efb65 100644 --- a/tests/v1/engine/test_startup_watch_processes.py +++ b/tests/v1/engine/test_startup_watch_processes.py @@ -2,14 +2,20 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from multiprocessing import connection +from threading import Event from types import SimpleNamespace import pytest import zmq +import vllm.platforms as platforms +from vllm.v1.engine import core as core_module +from vllm.v1.engine import utils as engine_utils +from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState from vllm.v1.engine.utils import ( CoreEngine, CoreEngineLaunch, + CoreEngineProcManager, EngineZmqAddresses, wait_for_engine_startup, ) @@ -17,6 +23,129 @@ from vllm.v1.engine.utils import ( pytestmark = pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + ("is_rocm", "request_timeout", "manager_timeout", "process_timeout"), + [ + (True, 0, 0, 15.0), + (True, 0, 7, 7), + (True, 0, None, None), + (False, 0, 0, 0), + (True, 7, 0, 0), + ], +) +def test_engine_core_process_shutdown_timeout( + monkeypatch: pytest.MonkeyPatch, + is_rocm: bool, + request_timeout: float | None, + manager_timeout: float | None, + process_timeout: float | None, +): + manager = object.__new__(CoreEngineProcManager) + manager._request_shutdown_timeout = request_timeout + manager.manager_stopped = Event() + manager.processes = [object()] + detach_results = iter((object(), None)) + manager._finalizer = SimpleNamespace(detach=lambda: next(detach_results)) + + shutdown_calls = [] + monkeypatch.setattr( + engine_utils, + "current_platform", + SimpleNamespace(is_rocm=lambda: is_rocm), + ) + monkeypatch.setattr( + engine_utils, + "shutdown", + lambda processes, timeout: shutdown_calls.append((processes, timeout)), + ) + + manager.shutdown(timeout=manager_timeout) + manager.shutdown(timeout=manager_timeout) + + assert manager.manager_stopped.is_set() + assert shutdown_calls == [(manager.processes, process_timeout)] + + +@pytest.mark.parametrize( + ( + "is_rocm", + "shutdown_state", + "has_work", + "shutdown_timeout", + "exit_code", + "expected_calls", + ), + [ + ( + True, + EngineShutdownState.SHUTTING_DOWN, + False, + 0, + None, + ["shutdown", "freeze"], + ), + (False, EngineShutdownState.SHUTTING_DOWN, False, 0, None, ["shutdown"]), + (True, EngineShutdownState.RUNNING, False, 0, None, ["shutdown"]), + (True, EngineShutdownState.SHUTTING_DOWN, True, 0, None, ["shutdown"]), + (True, EngineShutdownState.SHUTTING_DOWN, False, 7, None, ["shutdown"]), + (True, EngineShutdownState.SHUTTING_DOWN, False, 0, 1, ["shutdown"]), + ], +) +def test_freeze_gc_after_clean_rocm_engine_core_shutdown( + monkeypatch: pytest.MonkeyPatch, + is_rocm: bool, + shutdown_state: EngineShutdownState, + has_work: bool, + shutdown_timeout: int, + exit_code: int | None, + expected_calls: list[str], +): + calls: list[str] = [] + vllm_config = SimpleNamespace(shutdown_timeout=shutdown_timeout) + proc = SimpleNamespace( + shutdown_state=EngineShutdownState.RUNNING, + has_work=lambda: has_work, + vllm_config=vllm_config, + ) + + def run_busy_loop(): + proc.shutdown_state = shutdown_state + raise SystemExit(exit_code) + + proc.run_busy_loop = run_busy_loop + proc.shutdown = lambda: calls.append("shutdown") + parallel_config = SimpleNamespace( + data_parallel_size=1, + numa_bind=False, + reconfigure_for_independent_dp_rank=lambda: None, + ) + vllm_config.parallel_config = parallel_config + + for name in ( + "maybe_register_config_serialize_by_value", + "set_process_title", + "maybe_init_worker_tracer", + "decorate_logs", + ): + monkeypatch.setattr(core_module, name, lambda *args, **kwargs: None) + monkeypatch.setattr(core_module, "EngineCoreProc", lambda *args, **kwargs: proc) + monkeypatch.setattr( + core_module, + "SignalCallback", + lambda callback: SimpleNamespace(trigger=lambda: None, stop=lambda: None), + ) + monkeypatch.setattr(core_module.signal, "signal", lambda *args: None) + monkeypatch.setattr( + platforms, "current_platform", SimpleNamespace(is_rocm=lambda: is_rocm) + ) + monkeypatch.setattr(core_module.gc, "freeze", lambda: calls.append("freeze")) + + with pytest.raises(SystemExit): + EngineCoreProc.run_engine_core(vllm_config=vllm_config) + + assert calls == expected_calls + + class _FinishedProcess: name = "RustFrontend" diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index c5efe916c0..da757264f3 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -30,6 +30,7 @@ from vllm.utils.system_utils import ( kill_process_tree, set_process_title, ) +from vllm.v1.engine.utils import get_engine_process_shutdown_timeout logger = init_logger(__name__) @@ -513,7 +514,11 @@ class DPSupervisor: async def _shutdown_children(self) -> None: """Terminate the vLLM DP servers.""" - timeout = self.args.shutdown_timeout + CHILD_EXIT_GRACE_S + process_timeout = get_engine_process_shutdown_timeout( + self.args.shutdown_timeout, self.args.shutdown_timeout + ) + assert process_timeout is not None + timeout = process_timeout + CHILD_EXIT_GRACE_S try: logger.info( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 71a791d721..8946271aaf 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1276,6 +1276,7 @@ class EngineCoreProc(EngineCore): engine_core: EngineCoreProc | None = None signal_callback: SignalCallback | None = None + clean_shutdown = False try: vllm_config: VllmConfig = kwargs["vllm_config"] parallel_config: ParallelConfig = vllm_config.parallel_config @@ -1338,8 +1339,15 @@ class EngineCoreProc(EngineCore): engine_core.run_busy_loop() - except SystemExit: + except SystemExit as e: logger.info_once("[shutdown] EngineCore: exiting busy loop") + clean_shutdown = ( + e.code in (None, 0) + and engine_core is not None + and engine_core.shutdown_state == EngineShutdownState.SHUTTING_DOWN + and not engine_core.has_work() + and engine_core.vllm_config.shutdown_timeout == 0 + ) raise except Exception as e: if engine_core is None: @@ -1355,6 +1363,14 @@ class EngineCoreProc(EngineCore): signal_callback.stop() if engine_core is not None: engine_core.shutdown() + if clean_shutdown: + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + # Cleanup above already unfreezes and collects the heap. + # Freeze the surviving graph to skip another slow cyclic-GC + # scan during finalization; process exit reclaims it. + gc.freeze() def _init_data_parallel(self, vllm_config: VllmConfig): pass diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 7736062d99..9bce2b7687 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -40,6 +40,30 @@ if TYPE_CHECKING: logger = init_logger(__name__) STARTUP_POLL_PERIOD_MS = 10000 +ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S = 15.0 + + +def get_engine_process_shutdown_timeout( + request_timeout: float | None, + process_timeout: float | None, +) -> float | None: + """Return the EngineCore process-manager shutdown timeout. + + ``VllmConfig.shutdown_timeout`` controls how long in-flight requests may + drain. A value of zero therefore tells EngineCore to abort requests as soon + as it receives SIGTERM. The parent process manager still needs a separate + window in which the EngineCore can release device resources before it is + force-killed. ROCm teardown can take longer than the generic best-effort + window, and force-killing during teardown can leave VRAM resident. + + ``process_timeout`` may be a remaining budget computed by an outer process + manager. Keep it unchanged unless both values are zero: a zero remaining + budget for a positive request timeout must not receive a fresh grace period + because EngineCore relies on that deadline to enforce request draining. + """ + if request_timeout == 0 and process_timeout == 0 and current_platform.is_rocm(): + return ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S + return process_timeout class CoreEngineState(Enum): @@ -136,6 +160,7 @@ class CoreEngineProcManager: client_handshake_address: str | None = None, tensor_queue: Queue | None = None, ): + self._request_shutdown_timeout = vllm_config.shutdown_timeout context = get_mp_context() common_kwargs = { "vllm_config": vllm_config, @@ -217,7 +242,16 @@ class CoreEngineProcManager: """Shutdown engine core processes with configurable timeout.""" self.manager_stopped.set() if self._finalizer.detach() is not None: - shutdown(self.processes, timeout=timeout) + process_timeout = get_engine_process_shutdown_timeout( + self._request_shutdown_timeout, timeout + ) + if process_timeout != timeout: + logger.info( + "[shutdown] EngineCore process manager: using %ss ROCm " + "cleanup grace after immediate request abort", + process_timeout, + ) + shutdown(self.processes, timeout=process_timeout) def monitor_engine_liveness(self) -> None: """Monitor engine core process liveness."""