[ROCm] Give EngineCore cleanup grace after request abort (#52281)
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>
Andreas Karatzas committed
Aug 19, 2026 at 11:39 UTC
3a386cfaf51f872334c64edd08fbf71dd4bb0120
5 files changed
+221
-3
tests/entrypoints/openai/test_dp_supervisor.py
+34
@@ -242,6 +242,40 @@ def test_handles_shutdown_event():
242
assert supervisor.is_ready is False
243
244
245
+@pytest.mark.asyncio
246
+async def test_shutdown_children_uses_engine_process_timeout(
247
+ monkeypatch: pytest.MonkeyPatch,
248
+):
249
+ supervisor = DPSupervisor(_make_unit_args(shutdown_timeout=0.0))
250
+ supervisor._processes = [
251
+ SimpleNamespace(name="APIServer_DPRank_4", pid=None, is_alive=lambda: False)
252
+ ]
253
+ calls = []
254
+ timeout_calls = []
255
+
256
+ def get_process_timeout(request_timeout, manager_timeout):
257
+ timeout_calls.append((request_timeout, manager_timeout))
258
+ return 15.0
259
+
260
+ monkeypatch.setattr(
261
+ dp_sup,
262
+ "get_engine_process_shutdown_timeout",
263
+ get_process_timeout,
264
+ )
265
+ monkeypatch.setattr(
266
+ dp_sup,
267
+ "_join_processes_with_timeout",
268
+ lambda processes, timeout: calls.append((processes, timeout)),
269
+ )
270
+
271
+ await supervisor._shutdown_children()
272
+
273
+ assert timeout_calls == [(0.0, 0.0)]
274
+ assert calls == [
275
+ (supervisor._processes, 15.0 + CHILD_EXIT_GRACE_S),
276
+ ]
277
+
278
+
279
@pytest.mark.asyncio
280
async def test_handles_child_exit(
281
monkeypatch: pytest.MonkeyPatch,
tests/v1/engine/test_startup_watch_processes.py
+129
@@ -2,14 +2,20 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from multiprocessing import connection
5
+from threading import Event
6
from types import SimpleNamespace
7
8
import pytest
9
import zmq
10
11
+import vllm.platforms as platforms
12
+from vllm.v1.engine import core as core_module
13
+from vllm.v1.engine import utils as engine_utils
14
+from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState
15
from vllm.v1.engine.utils import (
16
CoreEngine,
17
CoreEngineLaunch,
18
+ CoreEngineProcManager,
19
EngineZmqAddresses,
20
wait_for_engine_startup,
21
)
@@ -17,6 +23,129 @@ from vllm.v1.engine.utils import (
23
pytestmark = pytest.mark.skip_global_cleanup
24
25
26
+@pytest.mark.parametrize(
27
+ ("is_rocm", "request_timeout", "manager_timeout", "process_timeout"),
28
+ [
29
+ (True, 0, 0, 15.0),
30
+ (True, 0, 7, 7),
31
+ (True, 0, None, None),
32
+ (False, 0, 0, 0),
33
+ (True, 7, 0, 0),
34
+ ],
35
+)
36
+def test_engine_core_process_shutdown_timeout(
37
+ monkeypatch: pytest.MonkeyPatch,
38
+ is_rocm: bool,
39
+ request_timeout: float | None,
40
+ manager_timeout: float | None,
41
+ process_timeout: float | None,
42
+):
43
+ manager = object.__new__(CoreEngineProcManager)
44
+ manager._request_shutdown_timeout = request_timeout
45
+ manager.manager_stopped = Event()
46
+ manager.processes = [object()]
47
+ detach_results = iter((object(), None))
48
+ manager._finalizer = SimpleNamespace(detach=lambda: next(detach_results))
49
+
50
+ shutdown_calls = []
51
+ monkeypatch.setattr(
52
+ engine_utils,
53
+ "current_platform",
54
+ SimpleNamespace(is_rocm=lambda: is_rocm),
55
+ )
56
+ monkeypatch.setattr(
57
+ engine_utils,
58
+ "shutdown",
59
+ lambda processes, timeout: shutdown_calls.append((processes, timeout)),
60
+ )
61
+
62
+ manager.shutdown(timeout=manager_timeout)
63
+ manager.shutdown(timeout=manager_timeout)
64
+
65
+ assert manager.manager_stopped.is_set()
66
+ assert shutdown_calls == [(manager.processes, process_timeout)]
67
+
68
+
69
+@pytest.mark.parametrize(
70
+ (
71
+ "is_rocm",
72
+ "shutdown_state",
73
+ "has_work",
74
+ "shutdown_timeout",
75
+ "exit_code",
76
+ "expected_calls",
77
+ ),
78
+ [
79
+ (
80
+ True,
81
+ EngineShutdownState.SHUTTING_DOWN,
82
+ False,
83
+ 0,
84
+ None,
85
+ ["shutdown", "freeze"],
86
+ ),
87
+ (False, EngineShutdownState.SHUTTING_DOWN, False, 0, None, ["shutdown"]),
88
+ (True, EngineShutdownState.RUNNING, False, 0, None, ["shutdown"]),
89
+ (True, EngineShutdownState.SHUTTING_DOWN, True, 0, None, ["shutdown"]),
90
+ (True, EngineShutdownState.SHUTTING_DOWN, False, 7, None, ["shutdown"]),
91
+ (True, EngineShutdownState.SHUTTING_DOWN, False, 0, 1, ["shutdown"]),
92
+ ],
93
+)
94
+def test_freeze_gc_after_clean_rocm_engine_core_shutdown(
95
+ monkeypatch: pytest.MonkeyPatch,
96
+ is_rocm: bool,
97
+ shutdown_state: EngineShutdownState,
98
+ has_work: bool,
99
+ shutdown_timeout: int,
100
+ exit_code: int | None,
101
+ expected_calls: list[str],
102
+):
103
+ calls: list[str] = []
104
+ vllm_config = SimpleNamespace(shutdown_timeout=shutdown_timeout)
105
+ proc = SimpleNamespace(
106
+ shutdown_state=EngineShutdownState.RUNNING,
107
+ has_work=lambda: has_work,
108
+ vllm_config=vllm_config,
109
+ )
110
+
111
+ def run_busy_loop():
112
+ proc.shutdown_state = shutdown_state
113
+ raise SystemExit(exit_code)
114
+
115
+ proc.run_busy_loop = run_busy_loop
116
+ proc.shutdown = lambda: calls.append("shutdown")
117
+ parallel_config = SimpleNamespace(
118
+ data_parallel_size=1,
119
+ numa_bind=False,
120
+ reconfigure_for_independent_dp_rank=lambda: None,
121
+ )
122
+ vllm_config.parallel_config = parallel_config
123
+
124
+ for name in (
125
+ "maybe_register_config_serialize_by_value",
126
+ "set_process_title",
127
+ "maybe_init_worker_tracer",
128
+ "decorate_logs",
129
+ ):
130
+ monkeypatch.setattr(core_module, name, lambda *args, **kwargs: None)
131
+ monkeypatch.setattr(core_module, "EngineCoreProc", lambda *args, **kwargs: proc)
132
+ monkeypatch.setattr(
133
+ core_module,
134
+ "SignalCallback",
135
+ lambda callback: SimpleNamespace(trigger=lambda: None, stop=lambda: None),
136
+ )
137
+ monkeypatch.setattr(core_module.signal, "signal", lambda *args: None)
138
+ monkeypatch.setattr(
139
+ platforms, "current_platform", SimpleNamespace(is_rocm=lambda: is_rocm)
140
+ )
141
+ monkeypatch.setattr(core_module.gc, "freeze", lambda: calls.append("freeze"))
142
+
143
+ with pytest.raises(SystemExit):
144
+ EngineCoreProc.run_engine_core(vllm_config=vllm_config)
145
+
146
+ assert calls == expected_calls
147
+
148
+
149
class _FinishedProcess:
150
name = "RustFrontend"
151
vllm/entrypoints/openai/dp_supervisor.py
+6
-1
@@ -30,6 +30,7 @@ from vllm.utils.system_utils import (
30
kill_process_tree,
31
set_process_title,
32
)
33
+from vllm.v1.engine.utils import get_engine_process_shutdown_timeout
34
35
logger = init_logger(__name__)
36
@@ -513,7 +514,11 @@ class DPSupervisor:
514
515
async def _shutdown_children(self) -> None:
516
"""Terminate the vLLM DP servers."""
516
- timeout = self.args.shutdown_timeout + CHILD_EXIT_GRACE_S
517
+ process_timeout = get_engine_process_shutdown_timeout(
518
+ self.args.shutdown_timeout, self.args.shutdown_timeout
519
+ )
520
+ assert process_timeout is not None
521
+ timeout = process_timeout + CHILD_EXIT_GRACE_S
522
523
try:
524
logger.info(
vllm/v1/engine/core.py
+17
-1
@@ -1276,6 +1276,7 @@ class EngineCoreProc(EngineCore):
1276
1277
engine_core: EngineCoreProc | None = None
1278
signal_callback: SignalCallback | None = None
1279
+ clean_shutdown = False
1280
try:
1281
vllm_config: VllmConfig = kwargs["vllm_config"]
1282
parallel_config: ParallelConfig = vllm_config.parallel_config
@@ -1338,8 +1339,15 @@ class EngineCoreProc(EngineCore):
1339
1340
engine_core.run_busy_loop()
1341
1341
- except SystemExit:
1342
+ except SystemExit as e:
1343
logger.info_once("[shutdown] EngineCore: exiting busy loop")
1344
+ clean_shutdown = (
1345
+ e.code in (None, 0)
1346
+ and engine_core is not None
1347
+ and engine_core.shutdown_state == EngineShutdownState.SHUTTING_DOWN
1348
+ and not engine_core.has_work()
1349
+ and engine_core.vllm_config.shutdown_timeout == 0
1350
+ )
1351
raise
1352
except Exception as e:
1353
if engine_core is None:
@@ -1355,6 +1363,14 @@ class EngineCoreProc(EngineCore):
1363
signal_callback.stop()
1364
if engine_core is not None:
1365
engine_core.shutdown()
1366
+ if clean_shutdown:
1367
+ from vllm.platforms import current_platform
1368
+
1369
+ if current_platform.is_rocm():
1370
+ # Cleanup above already unfreezes and collects the heap.
1371
+ # Freeze the surviving graph to skip another slow cyclic-GC
1372
+ # scan during finalization; process exit reclaims it.
1373
+ gc.freeze()
1374
1375
def _init_data_parallel(self, vllm_config: VllmConfig):
1376
pass
vllm/v1/engine/utils.py
+35
-1
@@ -40,6 +40,30 @@ if TYPE_CHECKING:
40
logger = init_logger(__name__)
41
42
STARTUP_POLL_PERIOD_MS = 10000
43
+ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S = 15.0
44
+
45
+
46
+def get_engine_process_shutdown_timeout(
47
+ request_timeout: float | None,
48
+ process_timeout: float | None,
49
+) -> float | None:
50
+ """Return the EngineCore process-manager shutdown timeout.
51
+
52
+ ``VllmConfig.shutdown_timeout`` controls how long in-flight requests may
53
+ drain. A value of zero therefore tells EngineCore to abort requests as soon
54
+ as it receives SIGTERM. The parent process manager still needs a separate
55
+ window in which the EngineCore can release device resources before it is
56
+ force-killed. ROCm teardown can take longer than the generic best-effort
57
+ window, and force-killing during teardown can leave VRAM resident.
58
+
59
+ ``process_timeout`` may be a remaining budget computed by an outer process
60
+ manager. Keep it unchanged unless both values are zero: a zero remaining
61
+ budget for a positive request timeout must not receive a fresh grace period
62
+ because EngineCore relies on that deadline to enforce request draining.
63
+ """
64
+ if request_timeout == 0 and process_timeout == 0 and current_platform.is_rocm():
65
+ return ROCM_ENGINE_PROCESS_SHUTDOWN_TIMEOUT_S
66
+ return process_timeout
67
68
69
class CoreEngineState(Enum):
@@ -136,6 +160,7 @@ class CoreEngineProcManager:
160
client_handshake_address: str | None = None,
161
tensor_queue: Queue | None = None,
162
):
163
+ self._request_shutdown_timeout = vllm_config.shutdown_timeout
164
context = get_mp_context()
165
common_kwargs = {
166
"vllm_config": vllm_config,
@@ -217,7 +242,16 @@ class CoreEngineProcManager:
242
"""Shutdown engine core processes with configurable timeout."""
243
self.manager_stopped.set()
244
if self._finalizer.detach() is not None:
220
- shutdown(self.processes, timeout=timeout)
245
+ process_timeout = get_engine_process_shutdown_timeout(
246
+ self._request_shutdown_timeout, timeout
247
+ )
248
+ if process_timeout != timeout:
249
+ logger.info(
250
+ "[shutdown] EngineCore process manager: using %ss ROCm "
251
+ "cleanup grace after immediate request abort",
252
+ process_timeout,
253
+ )
254
+ shutdown(self.processes, timeout=process_timeout)
255
256
def monitor_engine_liveness(self) -> None:
257
"""Monitor engine core process liveness."""