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

[Core] Bound UniProc EngineCore startup threads to available CPUs (#58946)

Browse Source 
Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com>
This commit is contained in:
Andreas Karatzas
2026-09-29 23:02:55 -07:00
committed by GitHub
parent 5faf81a429
commit 866fa130fa
2 changed files with 189 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/engine/test_startup_watch_processes.py
+123
View File
@@ -1,14 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from contextlib import nullcontext
from multiprocessing import connection
from threading import Event
from types import SimpleNamespace
import pytest
import torch
import zmq
import vllm.platforms as platforms
from vllm.utils import torch_utils
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
@@ -19,10 +23,129 @@ from vllm.v1.engine.utils import (
EngineZmqAddresses,
wait_for_engine_startup,
)
from vllm.v1.executor import UniProcExecutor, multiproc_executor
pytestmark = pytest.mark.skip_global_cleanup
@pytest.mark.parametrize(
(
"executor_class",
"local_engine_count",
"initial_threads",
"user_threads",
"fail_start",
"expected",
),
[
(UniProcExecutor, 1, 192, None, False, (12, "12", "1")),
(UniProcExecutor, 2, 192, None, False, (6, "6", "1")),
(UniProcExecutor, 2, 192, None, True, (6, "6", "1")),
(UniProcExecutor, 2, 5, None, False, (5, "5", "1")),
(UniProcExecutor, 2, 1, None, False, (1, "1", "1")),
(UniProcExecutor, 2, 5, "5", False, (5, "5", None)),
(
multiproc_executor.MultiprocExecutor,
2,
192,
None,
False,
(192, None, None),
),
],
)
def test_engine_core_startup_threads_are_scoped_to_launch(
monkeypatch: pytest.MonkeyPatch,
executor_class,
local_engine_count: int,
initial_threads: int,
user_threads: str | None,
fail_start: bool,
expected: tuple[int, str | None, str | None],
):
"""UniProc workers inherit their CPU share and restore the parent."""
marker = torch_utils.OMP_NUM_THREADS_SET_BY_VLLM
monkeypatch.delenv(marker, raising=False)
monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
if user_threads is not None:
monkeypatch.setenv("OMP_NUM_THREADS", user_threads)
monkeypatch.setattr(torch_utils, "available_cpu_count", lambda: 12)
platform = SimpleNamespace(
is_cpu=lambda: False, is_cuda_alike=lambda: True, is_xpu=lambda: False
)
monkeypatch.setattr(engine_utils, "current_platform", platform)
threads = SimpleNamespace(count=initial_threads)
setter_calls = []
def set_threads(count):
setter_calls.append(count)
threads.count = count
monkeypatch.setattr(torch, "get_num_threads", lambda: threads.count)
monkeypatch.setattr(torch, "set_num_threads", set_threads)
start_states = []
def thread_state():
return (
threads.count,
os.environ.get("OMP_NUM_THREADS"),
os.environ.get(marker),
)
def start_process():
start_states.append(thread_state())
if fail_start:
raise RuntimeError("start failed")
def get_context():
assert thread_state() == (initial_threads, user_threads, None)
return SimpleNamespace(
Process=lambda **kwargs: SimpleNamespace(
name=kwargs["name"],
exitcode=1 if fail_start else None,
start=start_process,
)
)
monkeypatch.setattr(engine_utils, "get_mp_context", get_context)
monkeypatch.setattr(engine_utils, "shutdown", lambda *_args, **_kwargs: None)
monkeypatch.setattr(
engine_utils.numa_utils,
"configure_subprocess",
lambda *args, **kwargs: nullcontext(),
)
with (
pytest.raises(RuntimeError, match="start failed")
if fail_start
else nullcontext()
):
manager = CoreEngineProcManager(
local_engine_count=local_engine_count,
start_index=0,
local_start_index=0,
vllm_config=SimpleNamespace(
shutdown_timeout=0,
parallel_config=SimpleNamespace(
data_parallel_size=local_engine_count,
assigned_physical_gpu_ids=None,
use_ray=False,
),
),
local_client=True,
handshake_address="unused",
executor_class=executor_class,
log_stats=False,
)
manager._finalizer.detach()
assert start_states == [expected] * (1 if fail_start else local_engine_count)
assert thread_state() == (initial_threads, user_threads, None)
assert setter_calls == (
[expected[0], initial_threads]
if executor_class is UniProcExecutor and expected[0] < initial_threads
else []
)
@pytest.mark.parametrize(
("is_rocm", "request_timeout", "manager_timeout", "process_timeout"),
[
vllm/v1/engine/utils.py
+66 -28
View File
@@ -28,9 +28,9 @@ from vllm.utils.network_utils import (
get_tcp_uri,
zmq_socket_ctx,
)
from vllm.utils.system_utils import get_mp_context
from vllm.utils.system_utils import get_mp_context, set_env_var
from vllm.v1.engine.coordinator import DPCoordinator
from vllm.v1.executor import Executor
from vllm.v1.executor import Executor, UniProcExecutor
from vllm.v1.executor.ray_utils import WORKER_SPECIFIC_ENV_VARS
from vllm.v1.utils import _SubprocessWrapper, get_engine_client_zmq_addr, shutdown
@@ -166,6 +166,42 @@ def _node_ip_from_resources(node_resources: dict) -> str | None:
return None
@contextlib.contextmanager
def _configure_uniproc_startup_threads(
executor_class: type[Executor], local_engine_count: int
) -> Iterator[None]:
if (
not issubclass(executor_class, UniProcExecutor)
or current_platform.is_cpu()
or "OMP_NUM_THREADS" in os.environ
):
yield
return
import torch
from vllm.utils.torch_utils import (
OMP_NUM_THREADS_SET_BY_VLLM,
set_default_torch_num_threads,
startup_omp_num_threads,
)
original_threads = torch.get_num_threads()
num_threads = min(original_threads, startup_omp_num_threads(local_engine_count))
torch_threads = (
set_default_torch_num_threads(num_threads)
if num_threads < original_threads
else contextlib.nullcontext()
)
# Fork inherits Torch's setting; spawn reads OMP_NUM_THREADS on import.
with (
set_env_var("OMP_NUM_THREADS", str(num_threads)),
set_env_var(OMP_NUM_THREADS_SET_BY_VLLM, "1"),
torch_threads,
):
yield
class CoreEngineProcManager:
"""Utility class to handle creation, readiness, and shutdown
of background processes used by the AsyncLLM and LLMEngine.
@@ -229,34 +265,36 @@ class CoreEngineProcManager:
# pickles process args at start() time, sequentially per rank.
user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids
try:
for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
# Populate the logical-to-physical GPU mapping in DP for
# platforms that cannot rely on
# torch.accelerator.set_device_index(), and for Ray.
needs_device_env_isolation = not (
current_platform.is_cuda_alike() or current_platform.is_xpu()
)
if is_dp and (
needs_device_env_isolation or vllm_config.parallel_config.use_ray
):
set_assigned_physical_gpu_ids_for_dp_rank(
vllm_config, local_dp_rank, user_assigned_gpu_ids
with _configure_uniproc_startup_threads(executor_class, local_engine_count):
for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
# Populate the logical-to-physical GPU mapping in DP for
# platforms that cannot rely on
# torch.accelerator.set_device_index(), and for Ray.
needs_device_env_isolation = not (
current_platform.is_cuda_alike() or current_platform.is_xpu()
)
if is_dp and (
needs_device_env_isolation
or vllm_config.parallel_config.use_ray
):
set_assigned_physical_gpu_ids_for_dp_rank(
vllm_config, local_dp_rank, user_assigned_gpu_ids
)
with numa_utils.configure_subprocess(
# EngineCore itself does not have a TP/PP-local rank.
# When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank()
# populates the logical-to-physical mapping for this DP
# shard, so local_rank=0 means "the first local GPU in
# this shard". The actual TP/PP worker processes spawned
# by the executor are bound separately with their own
# local_rank values.
vllm_config,
local_rank=0,
dp_local_rank=local_dp_rank,
process_kind="EngineCore",
):
proc.start()
with numa_utils.configure_subprocess(
# EngineCore itself does not have a TP/PP-local rank.
# set_assigned_physical_gpu_ids_for_dp_rank() populates
# the logical-to-physical mapping for this DP
# shard, so local_rank=0 means "the first local GPU in
# this shard". The actual TP/PP worker processes spawned
# by the executor are bound separately with their own
# local_rank values.
vllm_config,
local_rank=0,
dp_local_rank=local_dp_rank,
process_kind="EngineCore",
):
proc.start()
finally:
# Kill other procs if not all are running.
if self.finished_procs():