iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm] Give EngineCore cleanup grace after request abort (#52281)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Andreas Karatzas
2026-08-19 11:39:59 -05:00
committed by GitHub
co-authored by OpenAI Codex mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent e9e1630e93
commit 3a386cfaf5
5 changed files with 221 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/entrypoints/openai/test_dp_supervisor.py
+34
View File
@@ -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,
tests/v1/engine/test_startup_watch_processes.py
+129
View File
@@ -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"
vllm/entrypoints/openai/dp_supervisor.py
+6 -1
View File
@@ -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(
vllm/v1/engine/core.py
+17 -1
View File
@@ -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
vllm/v1/engine/utils.py
+35 -1
View File
@@ -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."""