diff --git a/tests/v1/engine/test_startup_watch_processes.py b/tests/v1/engine/test_startup_watch_processes.py index 0b346efb65..feb1e04268 100644 --- a/tests/v1/engine/test_startup_watch_processes.py +++ b/tests/v1/engine/test_startup_watch_processes.py @@ -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"), [ diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index e786cf2c5b..9df749f365 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -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():