wehub-resource-sync
290 lines
12 KiB
Python
290 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""OMP Aware Multiprocessing manager for running multiprocessing.Process()
|
|
Copyright (c) 2026 Red Hat Inc
|
|
Copyright (c) 2026 Cambridge Greys Ltd
|
|
"""
|
|
|
|
import os
|
|
from collections.abc import Callable
|
|
from contextlib import contextmanager
|
|
from typing import TYPE_CHECKING
|
|
|
|
import vllm.utils.cpu_resource_utils as cr_utils
|
|
from vllm import envs
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms import CpuArchEnum, current_platform
|
|
from vllm.utils.cpu_resource_utils import LogicalCPUInfo
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config import VllmConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OMPProcessManager:
|
|
def __init__(self, config: "VllmConfig"):
|
|
if not current_platform.is_cpu():
|
|
return
|
|
|
|
self.local_world_size = config.parallel_config.local_world_size
|
|
self.local_dp_rank = config.parallel_config.data_parallel_rank_local
|
|
# This is a bit tricky because the internal DP size
|
|
# is always 1 for non-MoE models
|
|
self.internal_dp_size = config.parallel_config._api_process_count
|
|
|
|
self.simulate_multi_node = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"
|
|
ld_preload_str = os.getenv("LD_PRELOAD", "")
|
|
self.use_iomp = "libiomp" in ld_preload_str or "libomp" in ld_preload_str
|
|
self.use_gomp = "libgomp" in ld_preload_str
|
|
|
|
assert not (self.use_iomp and self.use_gomp)
|
|
|
|
# at least reserve 1/local_world_size(for ARM/RISC-V) core for scheduler
|
|
# proc as always use MP executor
|
|
# TODO: make scheduler proc sleep when idle
|
|
self.reserve_cpu_num = (
|
|
self.local_world_size
|
|
if current_platform.get_cpu_architecture()
|
|
in (CpuArchEnum.ARM, CpuArchEnum.RISCV)
|
|
else 1
|
|
)
|
|
# reserve at one more core for nixl_connector under p/d case
|
|
if config.kv_transfer_config:
|
|
self.reserve_cpu_num += 1
|
|
|
|
if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None:
|
|
if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU:
|
|
msg = (
|
|
f"VLLM_CPU_NUM_OF_RESERVED_CPU is less than "
|
|
"the minimum requirement"
|
|
f": {self.reserve_cpu_num} cores"
|
|
)
|
|
logger.warning(msg=msg)
|
|
self.reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU
|
|
|
|
self._parse_omp_threads_bind_env()
|
|
|
|
assert not self.simulate_multi_node or self.auto_setup
|
|
|
|
@contextmanager
|
|
def configure_omp_envs(self, rank: int, local_rank: int):
|
|
if not current_platform.is_cpu() or self.skip_setup:
|
|
yield
|
|
return
|
|
|
|
envs_dict = {}
|
|
cpu_list = [str(i) for i in self.cpu_lists[local_rank]]
|
|
envs_dict["OMP_NUM_THREADS"] = str(len(cpu_list))
|
|
if self.use_iomp:
|
|
# set IOMP envs
|
|
cpu_list_str = ",".join(cpu_list)
|
|
envs_dict["KMP_AFFINITY"] = (
|
|
f"granularity=fine,explicit,proclist=[{cpu_list_str}]"
|
|
)
|
|
# The time(milliseconds) that a thread should wait after
|
|
# completing the execution of a parallel region, before sleeping.
|
|
# A value of 5 masks thread underutilization.
|
|
# Set to 1 when debugging thread utilization issues.
|
|
envs_dict["KMP_BLOCKTIME"] = "5"
|
|
# Prevents the CPU to run into low performance state
|
|
envs_dict["KMP_TPAUSE"] = "0"
|
|
elif self.use_gomp:
|
|
# set GOMP envs
|
|
# likes '0 1 2 ...'
|
|
cpu_list_str = " ".join(cpu_list)
|
|
envs_dict["GOMP_CPU_AFFINITY"] = cpu_list_str
|
|
else:
|
|
# set OMP envs
|
|
# likes '{0,1,2,...}'
|
|
cpu_list_str = ",".join(cpu_list)
|
|
envs_dict["OMP_PLACES"] = f"{{{cpu_list_str}}}"
|
|
envs_dict["OMP_PROC_BIND"] = "true"
|
|
|
|
# backup envs
|
|
old_envs_dict = {}
|
|
for k in envs_dict:
|
|
old_envs_dict[k] = os.environ.get(k)
|
|
|
|
try:
|
|
# set envs
|
|
for k, v in envs_dict.items():
|
|
os.environ[k] = v
|
|
yield
|
|
finally:
|
|
# restore old envs
|
|
for k, v in old_envs_dict.items(): # type: ignore
|
|
if v is None:
|
|
os.environ.pop(k, None)
|
|
else:
|
|
os.environ[k] = v
|
|
|
|
def _parse_omp_threads_bind_env(self):
|
|
vllm_mask = envs.VLLM_CPU_OMP_THREADS_BIND
|
|
self.skip_setup = vllm_mask == "nobind"
|
|
self.auto_setup = vllm_mask == "auto"
|
|
self.reserved_cpu_list = []
|
|
self.cpu_lists = []
|
|
|
|
if self.auto_setup:
|
|
# auto generate CPU lists
|
|
cpu_arch = current_platform.get_cpu_architecture()
|
|
if cpu_arch == CpuArchEnum.POWERPC:
|
|
# For POWERPC SMT-8/4/2
|
|
cpu_list, reserve_list = self._get_autobind_cpu_ids(
|
|
lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4]
|
|
)
|
|
elif cpu_arch in (CpuArchEnum.X86, CpuArchEnum.S390X):
|
|
# For x86/S390X SMT-2, use 1 logical CPU per physical core
|
|
cpu_list, reserve_list = self._get_autobind_cpu_ids(
|
|
lambda cpus: cpus[-1:]
|
|
)
|
|
elif cpu_arch in (CpuArchEnum.ARM, CpuArchEnum.RISCV):
|
|
# For AArch64 / RISC-V, no SMT, use all logical CPUs
|
|
cpu_list, reserve_list = self._get_autobind_cpu_ids(lambda cpus: cpus)
|
|
else:
|
|
cpu_list, reserve_list = [], []
|
|
raise RuntimeError(f"{cpu_arch} doesn't support auto CPU binding.")
|
|
|
|
for item in cpu_list:
|
|
self.cpu_lists.append([x.id for x in item])
|
|
self.reserved_cpu_list = [x.id for x in reserve_list]
|
|
elif not self.skip_setup:
|
|
# user defined CPU lists
|
|
omp_cpuids_list = vllm_mask.split("|")
|
|
if self.local_dp_rank is not None:
|
|
local_dp_rank = self.local_dp_rank
|
|
world_size = self.local_world_size
|
|
# Rank mapping [DP, PP, TP]
|
|
omp_cpuids_list = omp_cpuids_list[
|
|
local_dp_rank * world_size : (local_dp_rank + 1) * world_size
|
|
]
|
|
|
|
assert len(omp_cpuids_list) == self.local_world_size, (
|
|
"Given "
|
|
f"number of CPU id list {omp_cpuids_list} doesn't match "
|
|
f"local world size {self.local_world_size}."
|
|
)
|
|
|
|
# parse CPU list strings like "5,2-4" to [5, 2, 3, 4]
|
|
self.cpu_lists = [cr_utils.parse_id_list(s) for s in omp_cpuids_list]
|
|
else:
|
|
# skip
|
|
self.cpu_lists = []
|
|
|
|
msg = (
|
|
"OpenMP thread binding info: \n"
|
|
f"\tVLLM_CPU_OMP_THREADS_BIND={vllm_mask!r}, "
|
|
f"auto_setup={self.auto_setup}, skip_setup={self.skip_setup}\n"
|
|
f"\tlocal_world_size={self.local_world_size}, "
|
|
f"reserve_cpu_num={self.reserve_cpu_num}\n"
|
|
)
|
|
for i, cpus in enumerate(self.cpu_lists):
|
|
msg += f"\tlocal_rank={i}, core ids={cpus}\n"
|
|
msg += f"\treserved_cpus={self.reserved_cpu_list}"
|
|
logger.info(msg)
|
|
|
|
def _get_autobind_cpu_ids(
|
|
self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
|
|
) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]:
|
|
"""
|
|
Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for
|
|
other processes.
|
|
Currently for rank N, only CPU ids on the N-th node in available NUMA
|
|
node list will be selected.
|
|
Args:
|
|
cpu_selector: a callable object to select CPUs from a CPU list
|
|
of a physical core. The input is a LogicalCPUInfo list contains
|
|
logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id.
|
|
A selected LogicalCPUInfo list should be returned.
|
|
"""
|
|
|
|
# this memory node list has been sliced for DP offset
|
|
allowed_numa_nodes = cr_utils.get_visible_memory_node()
|
|
logical_cpu_list = cr_utils.get_allowed_cpu_list()
|
|
|
|
local_world_size = self.local_world_size
|
|
assert (
|
|
len(allowed_numa_nodes) >= local_world_size or self.simulate_multi_node
|
|
), (
|
|
f"Not enough allowed NUMA nodes to bind threads of "
|
|
f"{local_world_size} local CPUWorkers. "
|
|
f"Allowed NUMA nodes are {allowed_numa_nodes}. "
|
|
"Please try to bind threads manually or decrease DP/TP/PP."
|
|
)
|
|
|
|
# Generate OMP CPU list for each rank
|
|
cpu_lists_of_ranks = []
|
|
reserved_cpu_list = []
|
|
total_cpu_num = 0
|
|
for local_rank in range(self.local_world_size):
|
|
if not self.simulate_multi_node:
|
|
selected_numa_node = allowed_numa_nodes[local_rank]
|
|
selected_logical_cpu_list = [
|
|
x for x in logical_cpu_list if x.numa_node == selected_numa_node
|
|
]
|
|
else:
|
|
world_size_across_dp = self.local_world_size * self.internal_dp_size
|
|
assert len(logical_cpu_list) >= world_size_across_dp
|
|
selected_logical_cpu_list = sorted(
|
|
logical_cpu_list, key=lambda x: x.numa_node
|
|
)
|
|
sim_cpu_num_per_node = (
|
|
len(selected_logical_cpu_list) // world_size_across_dp
|
|
)
|
|
assert self.local_dp_rank is not None
|
|
start_idx = (
|
|
local_rank + self.local_world_size * self.local_dp_rank
|
|
) * sim_cpu_num_per_node
|
|
selected_logical_cpu_list = selected_logical_cpu_list[
|
|
start_idx : (start_idx + sim_cpu_num_per_node)
|
|
]
|
|
|
|
# Select logical CPUs on same physical cores via cpu_selector
|
|
core_to_cpus: dict[int, list[LogicalCPUInfo]] = {}
|
|
for cpu_info in selected_logical_cpu_list:
|
|
if cpu_info.physical_core not in core_to_cpus:
|
|
core_to_cpus[cpu_info.physical_core] = []
|
|
core_to_cpus[cpu_info.physical_core].append(cpu_info)
|
|
selected_logical_cpu_list = []
|
|
for cpu_list in core_to_cpus.values():
|
|
cpu_list = sorted(cpu_list, key=lambda x: x.id)
|
|
selected_logical_cpu_list.extend(cpu_selector(cpu_list))
|
|
|
|
# sort selected cores based on core id
|
|
selected_logical_cpu_list = sorted(
|
|
selected_logical_cpu_list, key=lambda x: x.id
|
|
)
|
|
|
|
cpu_lists_of_ranks.append(selected_logical_cpu_list)
|
|
total_cpu_num += len(selected_logical_cpu_list)
|
|
|
|
# Reserve CPUs for other processes
|
|
if total_cpu_num <= self.reserve_cpu_num:
|
|
logger.warning(
|
|
"Selected CPU core number (%s) "
|
|
"should be greater than reserved CPU core "
|
|
"number (%s).",
|
|
total_cpu_num,
|
|
self.reserve_cpu_num,
|
|
)
|
|
return cpu_lists_of_ranks, []
|
|
|
|
reserve_num_per_rank = [
|
|
self.reserve_cpu_num // self.local_world_size
|
|
] * self.local_world_size
|
|
# last rank first
|
|
for i in range(
|
|
self.local_world_size - 1,
|
|
self.local_world_size - 1 - self.reserve_cpu_num % self.local_world_size,
|
|
-1,
|
|
):
|
|
reserve_num_per_rank[i] += 1
|
|
for i in range(self.local_world_size):
|
|
num = reserve_num_per_rank[i]
|
|
if num > 0:
|
|
reserved_cpu_list.extend(cpu_lists_of_ranks[i][-num:])
|
|
cpu_lists_of_ranks[i] = cpu_lists_of_ranks[i][:-num]
|
|
|
|
return cpu_lists_of_ranks, reserved_cpu_list
|