| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | from contextlib import nullcontext |
| 6 | from multiprocessing import connection |
| 7 | from threading import Event |
| 8 | from types import SimpleNamespace |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | import zmq |
| 13 | |
| 14 | import vllm.platforms as platforms |
| 15 | from vllm.utils import torch_utils |
| 16 | from vllm.v1.engine import core as core_module |
| 17 | from vllm.v1.engine import utils as engine_utils |
| 18 | from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState |
| 19 | from vllm.v1.engine.utils import ( |
| 20 | CoreEngine, |
| 21 | CoreEngineLaunch, |
| 22 | CoreEngineProcManager, |
| 23 | EngineZmqAddresses, |
| 24 | wait_for_engine_startup, |
| 25 | ) |
| 26 | from vllm.v1.executor import UniProcExecutor, multiproc_executor |
| 27 | |
| 28 | pytestmark = pytest.mark.skip_global_cleanup |
| 29 | |
| 30 | |
| 31 | @pytest.mark.parametrize( |
| 32 | ( |
| 33 | "executor_class", |
| 34 | "local_engine_count", |
| 35 | "initial_threads", |
| 36 | "user_threads", |
| 37 | "fail_start", |
| 38 | "expected", |
| 39 | ), |
| 40 | [ |
| 41 | (UniProcExecutor, 1, 192, None, False, (12, "12", "1")), |
| 42 | (UniProcExecutor, 2, 192, None, False, (6, "6", "1")), |
| 43 | (UniProcExecutor, 2, 192, None, True, (6, "6", "1")), |
| 44 | (UniProcExecutor, 2, 5, None, False, (5, "5", "1")), |
| 45 | (UniProcExecutor, 2, 1, None, False, (1, "1", "1")), |
| 46 | (UniProcExecutor, 2, 5, "5", False, (5, "5", None)), |
| 47 | ( |
| 48 | multiproc_executor.MultiprocExecutor, |
| 49 | 2, |
| 50 | 192, |
| 51 | None, |
| 52 | False, |
| 53 | (192, None, None), |
| 54 | ), |
| 55 | ], |
| 56 | ) |
| 57 | def test_engine_core_startup_threads_are_scoped_to_launch( |
| 58 | monkeypatch: pytest.MonkeyPatch, |
| 59 | executor_class, |
| 60 | local_engine_count: int, |
| 61 | initial_threads: int, |
| 62 | user_threads: str | None, |
| 63 | fail_start: bool, |
| 64 | expected: tuple[int, str | None, str | None], |
| 65 | ): |
| 66 | """UniProc workers inherit their CPU share and restore the parent.""" |
| 67 | marker = torch_utils.OMP_NUM_THREADS_SET_BY_VLLM |
| 68 | monkeypatch.delenv(marker, raising=False) |
| 69 | monkeypatch.delenv("OMP_NUM_THREADS", raising=False) |
| 70 | if user_threads is not None: |
| 71 | monkeypatch.setenv("OMP_NUM_THREADS", user_threads) |
| 72 | monkeypatch.setattr(torch_utils, "available_cpu_count", lambda: 12) |
| 73 | platform = SimpleNamespace( |
| 74 | is_cpu=lambda: False, is_cuda_alike=lambda: True, is_xpu=lambda: False |
| 75 | ) |
| 76 | monkeypatch.setattr(engine_utils, "current_platform", platform) |
| 77 | threads = SimpleNamespace(count=initial_threads) |
| 78 | setter_calls = [] |
| 79 | |
| 80 | def set_threads(count): |
| 81 | setter_calls.append(count) |
| 82 | threads.count = count |
| 83 | |
| 84 | monkeypatch.setattr(torch, "get_num_threads", lambda: threads.count) |
| 85 | monkeypatch.setattr(torch, "set_num_threads", set_threads) |
| 86 | start_states = [] |
| 87 | |
| 88 | def thread_state(): |
| 89 | return ( |
| 90 | threads.count, |
| 91 | os.environ.get("OMP_NUM_THREADS"), |
| 92 | os.environ.get(marker), |
| 93 | ) |
| 94 | |
| 95 | def start_process(): |
| 96 | start_states.append(thread_state()) |
| 97 | if fail_start: |
| 98 | raise RuntimeError("start failed") |
| 99 | |
| 100 | def get_context(): |
| 101 | assert thread_state() == (initial_threads, user_threads, None) |
| 102 | return SimpleNamespace( |
| 103 | Process=lambda **kwargs: SimpleNamespace( |
| 104 | name=kwargs["name"], |
| 105 | exitcode=1 if fail_start else None, |
| 106 | start=start_process, |
| 107 | ) |
| 108 | ) |
| 109 | |
| 110 | monkeypatch.setattr(engine_utils, "get_mp_context", get_context) |
| 111 | monkeypatch.setattr(engine_utils, "shutdown", lambda *_args, **_kwargs: None) |
| 112 | monkeypatch.setattr( |
| 113 | engine_utils.numa_utils, |
| 114 | "configure_subprocess", |
| 115 | lambda *args, **kwargs: nullcontext(), |
| 116 | ) |
| 117 | with ( |
| 118 | pytest.raises(RuntimeError, match="start failed") |
| 119 | if fail_start |
| 120 | else nullcontext() |
| 121 | ): |
| 122 | manager = CoreEngineProcManager( |
| 123 | local_engine_count=local_engine_count, |
| 124 | start_index=0, |
| 125 | local_start_index=0, |
| 126 | vllm_config=SimpleNamespace( |
| 127 | shutdown_timeout=0, |
| 128 | parallel_config=SimpleNamespace( |
| 129 | data_parallel_size=local_engine_count, |
| 130 | assigned_physical_gpu_ids=None, |
| 131 | use_ray=False, |
| 132 | ), |
| 133 | ), |
| 134 | local_client=True, |
| 135 | handshake_address="unused", |
| 136 | executor_class=executor_class, |
| 137 | log_stats=False, |
| 138 | ) |
| 139 | manager._finalizer.detach() |
| 140 | assert start_states == [expected] * (1 if fail_start else local_engine_count) |
| 141 | assert thread_state() == (initial_threads, user_threads, None) |
| 142 | assert setter_calls == ( |
| 143 | [expected[0], initial_threads] |
| 144 | if executor_class is UniProcExecutor and expected[0] < initial_threads |
| 145 | else [] |
| 146 | ) |
| 147 | |
| 148 | |
| 149 | @pytest.mark.parametrize( |
| 150 | ("is_rocm", "request_timeout", "manager_timeout", "process_timeout"), |
| 151 | [ |
| 152 | (True, 0, 0, 15.0), |
| 153 | (True, 0, 7, 7), |
| 154 | (True, 0, None, None), |
| 155 | (False, 0, 0, 0), |
| 156 | (True, 7, 0, 0), |
| 157 | ], |
| 158 | ) |
| 159 | def test_engine_core_process_shutdown_timeout( |
| 160 | monkeypatch: pytest.MonkeyPatch, |
| 161 | is_rocm: bool, |
| 162 | request_timeout: float | None, |
| 163 | manager_timeout: float | None, |
| 164 | process_timeout: float | None, |
| 165 | ): |
| 166 | manager = object.__new__(CoreEngineProcManager) |
| 167 | manager._request_shutdown_timeout = request_timeout |
| 168 | manager.manager_stopped = Event() |
| 169 | manager.processes = [object()] |
| 170 | detach_results = iter((object(), None)) |
| 171 | manager._finalizer = SimpleNamespace(detach=lambda: next(detach_results)) |
| 172 | |
| 173 | shutdown_calls = [] |
| 174 | monkeypatch.setattr( |
| 175 | engine_utils, |
| 176 | "current_platform", |
| 177 | SimpleNamespace(is_rocm=lambda: is_rocm), |
| 178 | ) |
| 179 | monkeypatch.setattr( |
| 180 | engine_utils, |
| 181 | "shutdown", |
| 182 | lambda processes, timeout: shutdown_calls.append((processes, timeout)), |
| 183 | ) |
| 184 | |
| 185 | manager.shutdown(timeout=manager_timeout) |
| 186 | manager.shutdown(timeout=manager_timeout) |
| 187 | |
| 188 | assert manager.manager_stopped.is_set() |
| 189 | assert shutdown_calls == [(manager.processes, process_timeout)] |
| 190 | |
| 191 | |
| 192 | @pytest.mark.parametrize( |
| 193 | ( |
| 194 | "is_rocm", |
| 195 | "shutdown_state", |
| 196 | "has_work", |
| 197 | "shutdown_timeout", |
| 198 | "exit_code", |
| 199 | "expected_calls", |
| 200 | ), |
| 201 | [ |
| 202 | ( |
| 203 | True, |
| 204 | EngineShutdownState.SHUTTING_DOWN, |
| 205 | False, |
| 206 | 0, |
| 207 | None, |
| 208 | ["shutdown", "freeze"], |
| 209 | ), |
| 210 | (False, EngineShutdownState.SHUTTING_DOWN, False, 0, None, ["shutdown"]), |
| 211 | (True, EngineShutdownState.RUNNING, False, 0, None, ["shutdown"]), |
| 212 | (True, EngineShutdownState.SHUTTING_DOWN, True, 0, None, ["shutdown"]), |
| 213 | (True, EngineShutdownState.SHUTTING_DOWN, False, 7, None, ["shutdown"]), |
| 214 | (True, EngineShutdownState.SHUTTING_DOWN, False, 0, 1, ["shutdown"]), |
| 215 | ], |
| 216 | ) |
| 217 | def test_freeze_gc_after_clean_rocm_engine_core_shutdown( |
| 218 | monkeypatch: pytest.MonkeyPatch, |
| 219 | is_rocm: bool, |
| 220 | shutdown_state: EngineShutdownState, |
| 221 | has_work: bool, |
| 222 | shutdown_timeout: int, |
| 223 | exit_code: int | None, |
| 224 | expected_calls: list[str], |
| 225 | ): |
| 226 | calls: list[str] = [] |
| 227 | vllm_config = SimpleNamespace(shutdown_timeout=shutdown_timeout) |
| 228 | proc = SimpleNamespace( |
| 229 | shutdown_state=EngineShutdownState.RUNNING, |
| 230 | has_work=lambda: has_work, |
| 231 | vllm_config=vllm_config, |
| 232 | ) |
| 233 | |
| 234 | def run_busy_loop(): |
| 235 | proc.shutdown_state = shutdown_state |
| 236 | raise SystemExit(exit_code) |
| 237 | |
| 238 | proc.run_busy_loop = run_busy_loop |
| 239 | proc.shutdown = lambda: calls.append("shutdown") |
| 240 | parallel_config = SimpleNamespace( |
| 241 | data_parallel_size=1, |
| 242 | numa_bind=False, |
| 243 | reconfigure_for_independent_dp_rank=lambda: None, |
| 244 | ) |
| 245 | vllm_config.parallel_config = parallel_config |
| 246 | |
| 247 | for name in ( |
| 248 | "maybe_register_config_serialize_by_value", |
| 249 | "set_process_title", |
| 250 | "maybe_init_worker_tracer", |
| 251 | "decorate_logs", |
| 252 | ): |
| 253 | monkeypatch.setattr(core_module, name, lambda *args, **kwargs: None) |
| 254 | monkeypatch.setattr(core_module, "EngineCoreProc", lambda *args, **kwargs: proc) |
| 255 | monkeypatch.setattr( |
| 256 | core_module, |
| 257 | "SignalCallback", |
| 258 | lambda callback: SimpleNamespace(trigger=lambda: None, stop=lambda: None), |
| 259 | ) |
| 260 | monkeypatch.setattr(core_module.signal, "signal", lambda *args: None) |
| 261 | monkeypatch.setattr( |
| 262 | platforms, "current_platform", SimpleNamespace(is_rocm=lambda: is_rocm) |
| 263 | ) |
| 264 | monkeypatch.setattr(core_module.gc, "freeze", lambda: calls.append("freeze")) |
| 265 | |
| 266 | with pytest.raises(SystemExit): |
| 267 | EngineCoreProc.run_engine_core(vllm_config=vllm_config) |
| 268 | |
| 269 | assert calls == expected_calls |
| 270 | |
| 271 | |
| 272 | class _FinishedProcess: |
| 273 | name = "RustFrontend" |
| 274 | |
| 275 | def __init__(self, sentinel): |
| 276 | self.sentinel = sentinel |
| 277 | |
| 278 | @property |
| 279 | def exitcode(self): |
| 280 | return 1 |
| 281 | |
| 282 | |
| 283 | def test_wait_for_engine_startup_reports_watched_process_exit(): |
| 284 | ctx = zmq.Context() |
| 285 | handshake_socket = ctx.socket(zmq.ROUTER) |
| 286 | recv, send = connection.Pipe(duplex=False) |
| 287 | send.close() |
| 288 | |
| 289 | parallel_config = SimpleNamespace( |
| 290 | data_parallel_size_local=1, |
| 291 | data_parallel_hybrid_lb=False, |
| 292 | data_parallel_external_lb=False, |
| 293 | ) |
| 294 | |
| 295 | try: |
| 296 | launch = CoreEngineLaunch( |
| 297 | engine_manager=None, |
| 298 | coordinator=None, |
| 299 | addresses=EngineZmqAddresses(inputs=[], outputs=[]), |
| 300 | tensor_queue=None, |
| 301 | ) |
| 302 | launch.watched_frontend_processes = [_FinishedProcess(recv)] |
| 303 | with pytest.raises(RuntimeError) as exc_info: |
| 304 | wait_for_engine_startup( |
| 305 | handshake_socket, |
| 306 | [CoreEngine()], |
| 307 | parallel_config, |
| 308 | coordinated_dp=False, |
| 309 | cache_config=None, |
| 310 | launch=launch, |
| 311 | ) |
| 312 | finally: |
| 313 | recv.close() |
| 314 | handshake_socket.close(linger=0) |
| 315 | ctx.term() |
| 316 | |
| 317 | assert "Frontend process failed during engine core initialization" in str( |
| 318 | exc_info.value |
| 319 | ) |
| 320 | assert "Failed frontend proc(s): {'RustFrontend': 1}" in str(exc_info.value) |