[Core] Bound UniProc EngineCore startup threads to available CPUs (#58946)
Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com>
Andreas Karatzas committed
Sep 30, 2026 at 01:02 UTC
866fa130fac1e1252330a680fbed45ad05eba64f
2 files changed
+189
-28
tests/v1/engine/test_startup_watch_processes.py
+123
@@ -1,14 +1,18 @@
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,10 +23,129 @@ from vllm.v1.engine.utils import (
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
[
vllm/v1/engine/utils.py
+66
-28
@@ -28,9 +28,9 @@ from vllm.utils.network_utils import (
28
get_tcp_uri,
29
zmq_socket_ctx,
30
)
31
-from vllm.utils.system_utils import get_mp_context
31
+from vllm.utils.system_utils import get_mp_context, set_env_var
32
from vllm.v1.engine.coordinator import DPCoordinator
33
-from vllm.v1.executor import Executor
33
+from vllm.v1.executor import Executor, UniProcExecutor
34
from vllm.v1.executor.ray_utils import WORKER_SPECIFIC_ENV_VARS
35
from vllm.v1.utils import _SubprocessWrapper, get_engine_client_zmq_addr, shutdown
36
@@ -166,6 +166,42 @@ def _node_ip_from_resources(node_resources: dict) -> str | None:
166
return None
167
168
169
+@contextlib.contextmanager
170
+def _configure_uniproc_startup_threads(
171
+ executor_class: type[Executor], local_engine_count: int
172
+) -> Iterator[None]:
173
+ if (
174
+ not issubclass(executor_class, UniProcExecutor)
175
+ or current_platform.is_cpu()
176
+ or "OMP_NUM_THREADS" in os.environ
177
+ ):
178
+ yield
179
+ return
180
+
181
+ import torch
182
+
183
+ from vllm.utils.torch_utils import (
184
+ OMP_NUM_THREADS_SET_BY_VLLM,
185
+ set_default_torch_num_threads,
186
+ startup_omp_num_threads,
187
+ )
188
+
189
+ original_threads = torch.get_num_threads()
190
+ num_threads = min(original_threads, startup_omp_num_threads(local_engine_count))
191
+ torch_threads = (
192
+ set_default_torch_num_threads(num_threads)
193
+ if num_threads < original_threads
194
+ else contextlib.nullcontext()
195
+ )
196
+ # Fork inherits Torch's setting; spawn reads OMP_NUM_THREADS on import.
197
+ with (
198
+ set_env_var("OMP_NUM_THREADS", str(num_threads)),
199
+ set_env_var(OMP_NUM_THREADS_SET_BY_VLLM, "1"),
200
+ torch_threads,
201
+ ):
202
+ yield
203
+
204
+
205
class CoreEngineProcManager:
206
"""Utility class to handle creation, readiness, and shutdown
207
of background processes used by the AsyncLLM and LLMEngine.
@@ -229,34 +265,36 @@ class CoreEngineProcManager:
265
# pickles process args at start() time, sequentially per rank.
266
user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids
267
try:
232
- for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
233
- # Populate the logical-to-physical GPU mapping in DP for
234
- # platforms that cannot rely on
235
- # torch.accelerator.set_device_index(), and for Ray.
236
- needs_device_env_isolation = not (
237
- current_platform.is_cuda_alike() or current_platform.is_xpu()
238
- )
239
- if is_dp and (
240
- needs_device_env_isolation or vllm_config.parallel_config.use_ray
241
- ):
242
- set_assigned_physical_gpu_ids_for_dp_rank(
243
- vllm_config, local_dp_rank, user_assigned_gpu_ids
268
+ with _configure_uniproc_startup_threads(executor_class, local_engine_count):
269
+ for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
270
+ # Populate the logical-to-physical GPU mapping in DP for
271
+ # platforms that cannot rely on
272
+ # torch.accelerator.set_device_index(), and for Ray.
273
+ needs_device_env_isolation = not (
274
+ current_platform.is_cuda_alike() or current_platform.is_xpu()
275
)
276
+ if is_dp and (
277
+ needs_device_env_isolation
278
+ or vllm_config.parallel_config.use_ray
279
+ ):
280
+ set_assigned_physical_gpu_ids_for_dp_rank(
281
+ vllm_config, local_dp_rank, user_assigned_gpu_ids
282
+ )
283
246
- with numa_utils.configure_subprocess(
247
- # EngineCore itself does not have a TP/PP-local rank.
248
- # When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank()
249
- # populates the logical-to-physical mapping for this DP
250
- # shard, so local_rank=0 means "the first local GPU in
251
- # this shard". The actual TP/PP worker processes spawned
252
- # by the executor are bound separately with their own
253
- # local_rank values.
254
- vllm_config,
255
- local_rank=0,
256
- dp_local_rank=local_dp_rank,
257
- process_kind="EngineCore",
258
- ):
259
- proc.start()
284
+ with numa_utils.configure_subprocess(
285
+ # EngineCore itself does not have a TP/PP-local rank.
286
+ # set_assigned_physical_gpu_ids_for_dp_rank() populates
287
+ # the logical-to-physical mapping for this DP
288
+ # shard, so local_rank=0 means "the first local GPU in
289
+ # this shard". The actual TP/PP worker processes spawned
290
+ # by the executor are bound separately with their own
291
+ # local_rank values.
292
+ vllm_config,
293
+ local_rank=0,
294
+ dp_local_rank=local_dp_rank,
295
+ process_kind="EngineCore",
296
+ ):
297
+ proc.start()
298
finally:
299
# Kill other procs if not all are running.
300
if self.finished_procs():