807 lines
32 KiB
Python
807 lines
32 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
from collections import defaultdict
|
|
|
|
import psutil
|
|
import regex as re
|
|
from vllm.logger import logger
|
|
|
|
from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type
|
|
|
|
MASK_BIT = 32 # Number of bits in a CPU affinity mask group
|
|
MIN_CPUS_PER_NPU = 5 # 2(IRQ) + 1(main, at least 1 CPU) + 1(acl) + 1(release) = 5 CPUs per NPU
|
|
MIN_CPUS_PER_NPU_WITHOUT_IRQ = 3 # 1(main, at least 1 CPU) + 1(acl) + 1(release)
|
|
ASCEND_950_PHYSICAL_CPUS_PER_CLUSTER = 8
|
|
ALLOWED_CPUS_PATH = "/proc/self/status"
|
|
ASCEND_RT_VISIBLE_DEVICES = os.getenv("ASCEND_RT_VISIBLE_DEVICES")
|
|
|
|
TOPO_AFFINITY_MODE = "topo_affinity"
|
|
GLOBAL_SLICE_MODE = "global_slice"
|
|
|
|
DEVICE_BINDING_MODE: dict["AscendDeviceType", str] = {
|
|
AscendDeviceType.A2: TOPO_AFFINITY_MODE,
|
|
AscendDeviceType.A3: GLOBAL_SLICE_MODE,
|
|
AscendDeviceType._310P: TOPO_AFFINITY_MODE,
|
|
AscendDeviceType.A5: TOPO_AFFINITY_MODE,
|
|
}
|
|
NO_IRQ_BINDING_DEVICE_TYPES = {AscendDeviceType.A5}
|
|
|
|
|
|
def is_arm_cpu() -> bool:
|
|
arch = platform.machine().lower()
|
|
if arch in {"x86_64", "amd64", "i386", "i686"}:
|
|
return False
|
|
if arch in {"aarch64", "arm64"} or arch.startswith("arm"):
|
|
return True
|
|
logger.warning(
|
|
"Unknown CPU architecture. arch=%s, action: disabling CPU binding.",
|
|
arch,
|
|
)
|
|
return False
|
|
|
|
|
|
def execute_command(cmd: list[str]) -> tuple[str, int]:
|
|
# Force a C locale so subprocess output remains parseable on localized OSes.
|
|
env = os.environ.copy()
|
|
env["LC_ALL"] = "C"
|
|
env["LANG"] = "C"
|
|
env["LC_MESSAGES"] = "C"
|
|
with subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env) as p:
|
|
try:
|
|
out, _ = p.communicate(timeout=1000)
|
|
except subprocess.TimeoutExpired:
|
|
p.kill()
|
|
out, _ = p.communicate()
|
|
return out.decode(), p.returncode
|
|
|
|
|
|
class DeviceInfo:
|
|
def __init__(self):
|
|
self.npu_map_info: dict[str, dict[str, str]] = self.get_npu_map_info()
|
|
self.allowed_cpus: list[int] = self.parse_allowed_cpus()
|
|
self.running_npu_list: list[int] = self.get_running_npus()
|
|
self.npu_affinity: dict[int, list[int]] = self.parse_topo_affinity()
|
|
self.all_logic_npus: list[int] = self.get_all_logic_npus()
|
|
self.total_logic_npus: int = len(self.all_logic_npus)
|
|
|
|
@staticmethod
|
|
def split_npu_smi_header(line: str) -> list[str]:
|
|
return [item.strip() for item in re.split(r"\s{2,}", line.strip()) if item.strip()]
|
|
|
|
@staticmethod
|
|
def is_cpu_list(cpu_list_str: str) -> bool:
|
|
return bool(re.fullmatch(r"\d+(?:-\d+)?(?:,\s*\d+(?:-\d+)?)*", cpu_list_str))
|
|
|
|
@staticmethod
|
|
def expand_cpu_list(allowed_list_str: str) -> list[int]:
|
|
allowed_cpus_list: list[int] = []
|
|
for per_range in allowed_list_str.split(","):
|
|
if "-" in per_range:
|
|
start_cpu, end_cpu = map(int, per_range.split("-"))
|
|
allowed_cpus_list.extend(range(start_cpu, end_cpu + 1))
|
|
else:
|
|
allowed_cpus_list.append(int(per_range))
|
|
return allowed_cpus_list
|
|
|
|
def get_all_logic_npus(self) -> list[int]:
|
|
"""Collect all logical NPU IDs from the NPU mapping.
|
|
|
|
self.npu_map_info maps a board_id (A3) or npu_id (A2) to a per-chip map.
|
|
The per-chip map uses chip_id as the key and the logical NPU ID string
|
|
as the value.
|
|
"""
|
|
logic_ids: set[int] = set()
|
|
for _, chip_map in self.npu_map_info.items():
|
|
for _, logic_str in chip_map.items():
|
|
if logic_str and logic_str.isdigit():
|
|
logic_ids.add(int(logic_str))
|
|
return sorted(logic_ids)
|
|
|
|
@staticmethod
|
|
def get_npu_map_info() -> dict[str, dict[str, str]]:
|
|
npu_map_info: dict[str, dict[str, str]] = {}
|
|
npu_info, _ = execute_command(["npu-smi", "info", "-m"])
|
|
npu_map = [line.strip() for line in npu_info.splitlines() if line.strip()]
|
|
|
|
header = DeviceInfo.split_npu_smi_header(npu_map[0])
|
|
npu_id_idx = header.index("NPU ID")
|
|
chip_id_idx = header.index("Chip ID") if "Chip ID" in header else None
|
|
chip_logic_id_idx = header.index("Chip Logic ID") if "Chip Logic ID" in header else None
|
|
|
|
for line in npu_map[1:]:
|
|
parts = line.split()
|
|
npu_id = parts[npu_id_idx]
|
|
|
|
chip_id = "0"
|
|
if chip_id_idx is not None:
|
|
chip_id = parts[chip_id_idx]
|
|
|
|
if chip_logic_id_idx is not None:
|
|
chip_logic_id = parts[chip_logic_id_idx]
|
|
else:
|
|
chip_logic_id = npu_id
|
|
if not chip_logic_id.isdigit():
|
|
continue
|
|
if npu_id not in npu_map_info:
|
|
npu_map_info[npu_id] = {}
|
|
npu_map_info[npu_id][chip_id] = chip_logic_id
|
|
return npu_map_info
|
|
|
|
def resolve_logic_id(self, npu_id: str, chip_id: str | None) -> int:
|
|
chip_map = self.npu_map_info.get(npu_id, {})
|
|
if chip_id is not None:
|
|
chip_logic_id = chip_map.get(chip_id)
|
|
elif len(chip_map) == 1:
|
|
chip_logic_id = next(iter(chip_map.values()))
|
|
else:
|
|
raise RuntimeError(
|
|
"Failed to resolve chip_logic_id because the process table does not contain chip_id "
|
|
f"and NPU {npu_id} has {len(chip_map)} mapped chips."
|
|
)
|
|
if not chip_logic_id or not chip_logic_id.isdigit():
|
|
raise RuntimeError("Failed to get correct chip_logic_id from command 'npu-smi info -m'.")
|
|
return int(chip_logic_id)
|
|
|
|
def get_running_npus(self) -> list[int]:
|
|
npu_message, _ = execute_command(["npu-smi", "info"])
|
|
in_proc_section = False
|
|
proc_npu_field = ""
|
|
running_npu_set = set()
|
|
for line in npu_message.splitlines():
|
|
line = line.strip()
|
|
if line.startswith("| NPU") and "Process id" in line:
|
|
parts = [p.strip() for p in line.strip("|").split("|")]
|
|
proc_npu_field = " ".join(parts[0].split())
|
|
in_proc_section = True
|
|
continue
|
|
if not in_proc_section:
|
|
continue
|
|
if line.startswith("| "):
|
|
parts = [p.strip() for p in line.strip("|").split("|")]
|
|
if len(parts) < 2:
|
|
continue
|
|
npu_chip_parts = parts[0].split()
|
|
npu_id = npu_chip_parts[0]
|
|
chip_id = None
|
|
if proc_npu_field == "NPU Chip":
|
|
chip_id = npu_chip_parts[1]
|
|
if not npu_id.isdigit() or (chip_id is not None and not chip_id.isdigit()):
|
|
continue
|
|
running_npu_set.add(self.resolve_logic_id(npu_id, chip_id))
|
|
if ASCEND_RT_VISIBLE_DEVICES:
|
|
devices_str = ASCEND_RT_VISIBLE_DEVICES
|
|
devices_list = [int(x) for x in devices_str.split(",")]
|
|
running_npu_set = set(devices_list) & running_npu_set
|
|
if not running_npu_set:
|
|
raise RuntimeError("Can not get running npu info.")
|
|
return sorted(running_npu_set)
|
|
|
|
def parse_allowed_cpus(self) -> list[int]:
|
|
if not os.path.exists(ALLOWED_CPUS_PATH):
|
|
return []
|
|
with open(ALLOWED_CPUS_PATH) as f:
|
|
for line in f:
|
|
if line.startswith("Cpus_allowed_list"):
|
|
return self.expand_cpu_list(line.split()[1])
|
|
raise RuntimeError("Can not found specific 'Cpus_allowed_list' in the '/proc/self/status' file.")
|
|
|
|
def parse_topo_affinity(self) -> dict[int, list[int]]:
|
|
affinity: dict[int, list[int]] = {}
|
|
affinity_message, _ = execute_command(["npu-smi", "info", "-t", "topo"])
|
|
for line in affinity_message.splitlines():
|
|
if line.startswith("NPU"):
|
|
parts = line.split()
|
|
npu_match = re.fullmatch(r"NPU(\d+)", parts[0])
|
|
if not npu_match:
|
|
continue
|
|
last_part = parts[-1]
|
|
if self.is_cpu_list(last_part):
|
|
affinity[int(npu_match.group(1))] = self.expand_cpu_list(last_part)
|
|
return affinity
|
|
|
|
|
|
class CpuAlloc:
|
|
def __init__(self, rank_id: int):
|
|
self.rank_id = rank_id
|
|
self.device_info: DeviceInfo = DeviceInfo()
|
|
self.cpu_node: dict[int, int] = {}
|
|
self.numa_to_cpu_map: dict[int, list[int]] = defaultdict(list)
|
|
self.npu_cpu_pool: dict[int, list[int]] = {}
|
|
self.assign_main: dict[int, list[int]] = {}
|
|
self.assign_acl: dict[int, list[int]] = {}
|
|
self.assign_rel: dict[int, list[int]] = {}
|
|
self.uvb_cpu_pool: list[int] = []
|
|
|
|
@staticmethod
|
|
def cpu_to_mask(cpu: int) -> str:
|
|
group = cpu // MASK_BIT
|
|
bit = cpu % MASK_BIT
|
|
value = 1 << bit
|
|
mask = f"{value:08x}"
|
|
for _ in range(1, group + 1):
|
|
mask = f"{mask},{'0' * (MASK_BIT // 4)}"
|
|
return mask
|
|
|
|
@staticmethod
|
|
def get_threads_map(thread_message: str) -> dict[str, dict[str, list[str]]]:
|
|
threads_map: dict[str, dict[str, list[str]]] = {}
|
|
for line in thread_message.splitlines():
|
|
parts = line.split()
|
|
if len(parts) < 2:
|
|
continue
|
|
main_pid, sub_pid = parts[0], parts[1]
|
|
if "acl_thread" in line:
|
|
key = "acl_thread"
|
|
elif "release_thread" in line:
|
|
key = "release_thread"
|
|
else:
|
|
continue
|
|
if main_pid not in threads_map:
|
|
threads_map[main_pid] = {"acl_thread": [], "release_thread": []}
|
|
threads_map[main_pid][key].append(sub_pid)
|
|
return threads_map
|
|
|
|
@staticmethod
|
|
def bind(pid: str, cpus: list[int], bind_sub_thread: bool) -> None:
|
|
if cpus:
|
|
cpu_list = ",".join(map(str, cpus))
|
|
if bind_sub_thread:
|
|
bind_result, return_code = execute_command(["taskset", "-acp", cpu_list, pid])
|
|
else:
|
|
bind_result, return_code = execute_command(["taskset", "-cp", cpu_list, pid])
|
|
if return_code != 0:
|
|
raise RuntimeError(f"Failed to bind {pid} to CPU {cpu_list}.")
|
|
|
|
def average_distribute(self, groups: dict[str, list[int]]) -> dict[int, list[int]]:
|
|
result: dict[int, list[int]] = {}
|
|
for key, npu_list in groups.items():
|
|
cpu_list = sorted(self.npu_cpu_pool[npu_list[0]])
|
|
cpu_num_per_npu = len(cpu_list) // len(npu_list)
|
|
for i, npu in enumerate(npu_list):
|
|
start_index = i * cpu_num_per_npu
|
|
end_index = (i + 1) * cpu_num_per_npu if i < len(npu_list) - 1 else len(cpu_list)
|
|
result[npu] = cpu_list[start_index:end_index]
|
|
return result
|
|
|
|
def extend_numa(self, cpu_list: list[int]) -> list[int]:
|
|
if not cpu_list:
|
|
return []
|
|
nodes = {self.cpu_node[c] for c in cpu_list}
|
|
if len(nodes) != 1:
|
|
return cpu_list
|
|
node = list(nodes)[0]
|
|
next_node = (node + 1) % len(self.numa_to_cpu_map)
|
|
extended = cpu_list[:]
|
|
for cpu in self.numa_to_cpu_map[next_node]:
|
|
if cpu in self.device_info.allowed_cpus:
|
|
extended.append(cpu)
|
|
return sorted(set(extended))
|
|
|
|
def build_cpu_node_map(self) -> None:
|
|
cpu_numa_map, _ = execute_command(["lscpu", "-e=CPU,NODE"])
|
|
for line in cpu_numa_map.splitlines():
|
|
line = line.strip()
|
|
if not line or not line[0].isdigit():
|
|
continue
|
|
cpu_str, node_str = line.split()
|
|
cpu = int(cpu_str)
|
|
node = int(node_str)
|
|
self.cpu_node[cpu] = node
|
|
self.numa_to_cpu_map[node].append(cpu)
|
|
if len(self.numa_to_cpu_map) == 0:
|
|
raise RuntimeError("lscpu command output error, no NUMA node available. Please check!")
|
|
|
|
@staticmethod
|
|
def parse_threads_per_core(lscpu_message: str) -> int | None:
|
|
for line in lscpu_message.splitlines():
|
|
if "Thread(s) per core:" not in line:
|
|
continue
|
|
value = line.split(":", 1)[1].strip()
|
|
if value.isdigit():
|
|
return int(value)
|
|
return None
|
|
|
|
@staticmethod
|
|
def get_uvb_poll_window_threads(thread_message: str) -> list[str]:
|
|
thread_ids = []
|
|
for line in thread_message.splitlines():
|
|
if "uvb_poll_window_thread" not in line:
|
|
continue
|
|
parts = line.split()
|
|
if len(parts) >= 2 and parts[1].isdigit():
|
|
thread_ids.append(parts[1])
|
|
return thread_ids
|
|
|
|
def get_ascend_950_cluster_size(self) -> int | None:
|
|
lscpu_message, _ = execute_command(["lscpu"])
|
|
threads_per_core = self.parse_threads_per_core(lscpu_message)
|
|
if threads_per_core not in {1, 2}:
|
|
logger.warning(
|
|
"[cpu_bind_ascend_950] unsupported Thread(s) per core: %s. action: skip worker CPU binding.",
|
|
threads_per_core,
|
|
)
|
|
return None
|
|
return ASCEND_950_PHYSICAL_CPUS_PER_CLUSTER * threads_per_core
|
|
|
|
def get_single_numa_node(self, cpu_list: list[int]) -> int | None:
|
|
if not cpu_list:
|
|
return None
|
|
nodes = {self.cpu_node[cpu] for cpu in cpu_list}
|
|
if len(nodes) != 1:
|
|
return None
|
|
return next(iter(nodes))
|
|
|
|
def bind_uvb_poll_window_threads(self) -> None:
|
|
self.uvb_cpu_pool = [
|
|
cpu for cpu in self.numa_to_cpu_map.get(0, []) if cpu in self.device_info.allowed_cpus and cpu != 0
|
|
]
|
|
if not self.uvb_cpu_pool:
|
|
logger.info("[cpu_bind_ascend_950] no available NUMA0 CPUs for uvb_poll_window_thread.")
|
|
return
|
|
|
|
thread_message, _ = execute_command(["ps", "-Te"])
|
|
uvb_threads = self.get_uvb_poll_window_threads(thread_message)
|
|
if not uvb_threads:
|
|
logger.warning(
|
|
"[cpu_bind_ascend_950] uvb_poll_window_thread not found. "
|
|
"If running in Docker, start the container with '--pid=host'."
|
|
)
|
|
return
|
|
|
|
bound_threads = []
|
|
for thread_id in uvb_threads:
|
|
try:
|
|
self.bind(thread_id, self.uvb_cpu_pool, False)
|
|
bound_threads.append(thread_id)
|
|
except RuntimeError as e:
|
|
logger.warning(
|
|
"[cpu_bind_ascend_950] failed to bind uvb_poll_window_thread %s: %s. "
|
|
"If running in Docker, start the container with '--pid=host'.",
|
|
thread_id,
|
|
e,
|
|
)
|
|
if bound_threads:
|
|
logger.info(
|
|
"[cpu_bind_ascend_950] uvb_poll_window_thread tids=[%s] cpus=[%s]",
|
|
" ".join(bound_threads),
|
|
" ".join(map(str, self.uvb_cpu_pool)),
|
|
)
|
|
|
|
# TODO: Drop the vLLM Ascend implementation for Ascend 950 CPU binding and
|
|
# switch to the CANN affinity tool once that tool is visible in external
|
|
# commercial releases.
|
|
def build_ascend_950_cpu_pools(self) -> bool:
|
|
if not self.device_info.npu_affinity:
|
|
logger.warning("[cpu_bind_ascend_950] NPU topo affinity not found. action: skip worker CPU binding.")
|
|
return False
|
|
|
|
cluster_size = self.get_ascend_950_cluster_size()
|
|
if cluster_size is None:
|
|
return False
|
|
|
|
allowed_cpu_set = set(self.device_info.allowed_cpus)
|
|
uvb_cpu_set = set(self.uvb_cpu_pool)
|
|
candidate_npus = [
|
|
npu
|
|
for npu in self.device_info.all_logic_npus
|
|
if npu in self.device_info.running_npu_list
|
|
or any(cpu in allowed_cpu_set for cpu in self.device_info.npu_affinity.get(npu, []))
|
|
]
|
|
|
|
numa_to_npus: dict[int, list[int]] = defaultdict(list)
|
|
for npu in candidate_npus:
|
|
base_cpu_list = [
|
|
cpu for cpu in self.device_info.npu_affinity.get(npu, []) if cpu in self.device_info.allowed_cpus
|
|
]
|
|
numa_node = self.get_single_numa_node(base_cpu_list)
|
|
if numa_node is None:
|
|
logger.warning(
|
|
"[cpu_bind_ascend_950] invalid topo affinity. npu=%s, cpus=%s. action: skip worker CPU binding.",
|
|
npu,
|
|
base_cpu_list,
|
|
)
|
|
return False
|
|
numa_to_npus[numa_node].append(npu)
|
|
|
|
final: dict[int, list[int]] = {}
|
|
for numa_node, npu_list in numa_to_npus.items():
|
|
cpu_list = [
|
|
cpu
|
|
for cpu in sorted(self.numa_to_cpu_map.get(numa_node, []))
|
|
if cpu in self.device_info.allowed_cpus and cpu not in uvb_cpu_set
|
|
]
|
|
clusters = [
|
|
cpu_list[i : i + cluster_size]
|
|
for i in range(0, len(cpu_list) - len(cpu_list) % cluster_size, cluster_size)
|
|
]
|
|
npu_list = sorted(npu_list)
|
|
if len(clusters) < len(npu_list):
|
|
logger.warning(
|
|
"[cpu_bind_ascend_950] insufficient CPU clusters. numa=%s, clusters=%s, npus=%s. "
|
|
"action: skip worker CPU binding.",
|
|
numa_node,
|
|
len(clusters),
|
|
npu_list,
|
|
)
|
|
return False
|
|
for index, npu in enumerate(npu_list):
|
|
final[npu] = clusters[index]
|
|
|
|
try:
|
|
self.npu_cpu_pool = {npu: final[npu] for npu in self.device_info.running_npu_list}
|
|
except KeyError as e:
|
|
logger.warning(
|
|
"[cpu_bind_ascend_950] failed to build CPU pool for visible NPU %s. action: skip worker CPU binding.",
|
|
e,
|
|
)
|
|
return False
|
|
return True
|
|
|
|
def build_global_slice_cpu_pool(self) -> None:
|
|
"""
|
|
Build per-NPU CPU pools by slicing allowed_cpus using GLOBAL logical NPU ids.
|
|
|
|
Why:
|
|
- Multiple processes/DP groups may share the SAME cpuset (same allowed_cpus).
|
|
- If each process slices only its visible NPUs, CPU ranges overlap across processes.
|
|
- Global slicing ensures deterministic, non-overlapping CPU partitions per logical NPU id.
|
|
|
|
Notes:
|
|
- This strategy does NOT rely on npu-smi topo affinity.
|
|
- NUMA locality is achieved only if CPU numbering aligns with NUMA layout.
|
|
- Requires enough CPUs for each device's role split.
|
|
"""
|
|
running_npus = list(self.device_info.running_npu_list)
|
|
if not running_npus:
|
|
return
|
|
|
|
allowed_cpus = sorted(set(self.device_info.allowed_cpus))
|
|
total_cpu_count = len(allowed_cpus)
|
|
if total_cpu_count == 0:
|
|
return
|
|
|
|
# Prefer mapping info (npu-smi info -m), fallback to topo keys, then visible list
|
|
if self.device_info.total_logic_npus > 0:
|
|
total_npus = self.device_info.total_logic_npus
|
|
elif self.device_info.npu_affinity:
|
|
total_npus = len(self.device_info.npu_affinity)
|
|
else:
|
|
total_npus = len(running_npus)
|
|
|
|
# Compute global per-NPU slicing
|
|
base_cpu_count = total_cpu_count // total_npus
|
|
extra_cpu_count = total_cpu_count % total_npus
|
|
|
|
logger.debug(
|
|
"[cpu_global_slice] rank:%s ASCEND_RT_VISIBLE_DEVICES=%s running_npu_list:%s total_npus:%s "
|
|
"allowed_cpus:%s base:%s extra:%s allowed_cpus_head:%s allowed_cpus_tail:%s",
|
|
self.rank_id,
|
|
ASCEND_RT_VISIBLE_DEVICES,
|
|
running_npus,
|
|
total_npus,
|
|
total_cpu_count,
|
|
base_cpu_count,
|
|
extra_cpu_count,
|
|
allowed_cpus[:16],
|
|
allowed_cpus[-16:],
|
|
)
|
|
|
|
min_cpus_per_npu = self._min_cpus_per_npu()
|
|
# Enforce the minimum per-NPU slice length.
|
|
# Because with remainder distribution, some NPUs may get 'base' cores and some get 'base+1'.
|
|
# The minimum slice size is 'base'.
|
|
if base_cpu_count < min_cpus_per_npu:
|
|
raise RuntimeError(
|
|
"Insufficient CPUs for binding with CPU role reservations: "
|
|
f"total_allowed={total_cpu_count}, total_npus={total_npus}, "
|
|
f"min_per_npu={base_cpu_count} (<{min_cpus_per_npu}). "
|
|
f"Need at least {total_npus * min_cpus_per_npu} CPUs in cpuset."
|
|
)
|
|
|
|
def _slice_for_npu(global_npu_id: int) -> list[int]:
|
|
# start = global_npu_id*base + min(global_npu_id, extra)
|
|
start = global_npu_id * base_cpu_count + (
|
|
global_npu_id if global_npu_id < extra_cpu_count else extra_cpu_count
|
|
)
|
|
take = base_cpu_count + (1 if global_npu_id < extra_cpu_count else 0)
|
|
end = start + take
|
|
return allowed_cpus[start:end]
|
|
|
|
for npu in running_npus:
|
|
if npu < 0 or npu >= total_npus:
|
|
raise RuntimeError(f"Invalid NPU id {npu}, total_npus={total_npus}.")
|
|
self.npu_cpu_pool[npu] = _slice_for_npu(npu)
|
|
|
|
@staticmethod
|
|
def _binding_mode() -> str:
|
|
device_type = get_ascend_device_type()
|
|
return DEVICE_BINDING_MODE.get(device_type, TOPO_AFFINITY_MODE)
|
|
|
|
@staticmethod
|
|
def _is_ascend_950() -> bool:
|
|
return get_ascend_device_type() == AscendDeviceType.A5
|
|
|
|
@staticmethod
|
|
def _reserve_irq_cpus() -> bool:
|
|
return get_ascend_device_type() not in NO_IRQ_BINDING_DEVICE_TYPES
|
|
|
|
@staticmethod
|
|
def _min_cpus_per_npu() -> int:
|
|
if CpuAlloc._reserve_irq_cpus():
|
|
return MIN_CPUS_PER_NPU
|
|
return MIN_CPUS_PER_NPU_WITHOUT_IRQ
|
|
|
|
def build_cpu_pools(self) -> bool:
|
|
self.build_cpu_node_map()
|
|
|
|
mode = self._binding_mode()
|
|
logger.info(
|
|
"[cpu_bind_mode] mode=%s rank=%s visible_npus=%s", mode, self.rank_id, self.device_info.running_npu_list
|
|
)
|
|
if self._is_ascend_950():
|
|
self.bind_uvb_poll_window_threads()
|
|
return self.build_ascend_950_cpu_pools()
|
|
|
|
if mode == GLOBAL_SLICE_MODE:
|
|
self.build_global_slice_cpu_pool()
|
|
return True
|
|
|
|
self._build_topo_affinity_cpu_pool()
|
|
return True
|
|
|
|
def _build_topo_affinity_cpu_pool(self) -> None:
|
|
# topo_affinity mode
|
|
if not self.device_info.npu_affinity:
|
|
logger.warning("NPU topo affinity not found. action: fallback to global-slice CPU binding.")
|
|
self.build_global_slice_cpu_pool()
|
|
return
|
|
|
|
allowed_cpu_set = set(self.device_info.allowed_cpus)
|
|
# Consider all logical NPUs to avoid CPU distribution overlap across worker processes.
|
|
candidate_npus = [
|
|
npu
|
|
for npu in self.device_info.all_logic_npus
|
|
# Always keep visible NPUs for strict conflict checks; include
|
|
# non-visible NPUs only when they overlap this process's cpuset.
|
|
if npu in self.device_info.running_npu_list
|
|
or any(cpu in allowed_cpu_set for cpu in self.device_info.npu_affinity.get(npu, []))
|
|
]
|
|
for npu in candidate_npus:
|
|
base_cpu_list = [
|
|
cpu for cpu in self.device_info.npu_affinity.get(npu, []) if cpu in self.device_info.allowed_cpus
|
|
]
|
|
if not base_cpu_list:
|
|
raise RuntimeError("CPUs available in 'Cpus_allowed_list' conflict with NUMA affinity.")
|
|
extra_cpu_list = self.extend_numa(base_cpu_list)
|
|
self.npu_cpu_pool[npu] = extra_cpu_list
|
|
|
|
groups = defaultdict(list)
|
|
for npu, cpus in self.npu_cpu_pool.items():
|
|
groups[str(cpus)].append(npu)
|
|
|
|
final: dict[int, list[int]] = {}
|
|
for key, npu_list in groups.items():
|
|
if len(npu_list) == 1:
|
|
final[npu_list[0]] = self.npu_cpu_pool[npu_list[0]]
|
|
else:
|
|
final.update(self.average_distribute({key: npu_list}))
|
|
# Keep only visible NPUs in the final binding pool. Non-visible NPUs are used only to avoid overlap.
|
|
self.npu_cpu_pool = {npu: final[npu] for npu in self.device_info.running_npu_list}
|
|
|
|
def allocate(self) -> None:
|
|
if self._is_ascend_950():
|
|
self._allocate_ascend_950_roles()
|
|
return
|
|
self._allocate_default_roles()
|
|
|
|
def _allocate_ascend_950_roles(self) -> None:
|
|
self.assign_main = {npu: cpu_pool for npu, cpu_pool in self.npu_cpu_pool.items()}
|
|
self.assign_acl = {npu: [] for npu in self.npu_cpu_pool}
|
|
self.assign_rel = {npu: [] for npu in self.npu_cpu_pool}
|
|
|
|
def _allocate_default_roles(self) -> None:
|
|
reserve_irq_cpus = self._reserve_irq_cpus()
|
|
min_cpus_per_npu = self._min_cpus_per_npu()
|
|
for npu, cpu_pool in self.npu_cpu_pool.items():
|
|
if len(cpu_pool) >= min_cpus_per_npu:
|
|
if reserve_irq_cpus:
|
|
main = cpu_pool[2:-2]
|
|
else:
|
|
main = cpu_pool[:-2]
|
|
acl = [cpu_pool[-2]]
|
|
rel = [cpu_pool[-1]]
|
|
else:
|
|
raise RuntimeError(
|
|
f"The number of CPUs is insufficient. Each NPU requires at least {min_cpus_per_npu} CPUs."
|
|
)
|
|
self.assign_main[npu] = main
|
|
self.assign_acl[npu] = acl
|
|
self.assign_rel[npu] = rel
|
|
|
|
def print_plan(self) -> None:
|
|
logger.info("The CPU allocation plan is as follows:")
|
|
current_npu = self.device_info.running_npu_list[self.rank_id]
|
|
if self._is_ascend_950():
|
|
self._print_ascend_950_plan(current_npu)
|
|
return
|
|
self._print_default_plan(current_npu)
|
|
|
|
def _print_ascend_950_plan(self, current_npu: int) -> None:
|
|
main = " ".join(map(str, self.assign_main[current_npu]))
|
|
logger.info("Ascend 950 NPU%s: worker=[%s]", current_npu, main)
|
|
|
|
def _print_default_plan(self, current_npu: int) -> None:
|
|
main = " ".join(map(str, self.assign_main[current_npu]))
|
|
acl = " ".join(map(str, self.assign_acl[current_npu]))
|
|
rel = str(self.assign_rel[current_npu]) if self.assign_rel[current_npu] else ""
|
|
logger.info("NPU%s: main=[%s] acl=[%s] release=[%s]", current_npu, main, acl, rel)
|
|
|
|
def bind_memory(self, pid: str, npu: int) -> None:
|
|
def _get_npu_numa_node(npu_id: int) -> int | None:
|
|
cpu_pool = self.npu_cpu_pool.get(npu_id, [])
|
|
if not cpu_pool:
|
|
return None
|
|
anchor_cpu = cpu_pool[0]
|
|
return self.cpu_node.get(anchor_cpu)
|
|
|
|
if not shutil.which("migratepages"):
|
|
logger.info("The 'migratepages' command is not available, skipping memory binding.")
|
|
return
|
|
target_numa = _get_npu_numa_node(npu)
|
|
if target_numa is None:
|
|
logger.warning(
|
|
"[migrate] NPU has no CPU pool. rank=%s, npu=%s.",
|
|
self.rank_id,
|
|
npu,
|
|
)
|
|
return
|
|
all_numa_nodes = sorted(self.numa_to_cpu_map.keys())
|
|
if target_numa not in all_numa_nodes:
|
|
logger.warning("[migrate] NUMA node not found. npu=%s, numa=%s. ", npu, target_numa)
|
|
return
|
|
# Bind memory to the NPU's NUMA node only to minimize cross-NUMA traffic.
|
|
logger.info("[migrate] NPU:%s -> NUMA [%s]", npu, target_numa)
|
|
execute_command(
|
|
[
|
|
"migratepages",
|
|
pid,
|
|
",".join(map(str, all_numa_nodes)),
|
|
str(target_numa),
|
|
]
|
|
)
|
|
|
|
def bind_threads(self) -> None:
|
|
if self._is_ascend_950():
|
|
self.bind_ascend_950_threads()
|
|
return
|
|
self._bind_default_threads()
|
|
|
|
def _bind_default_threads(self) -> None:
|
|
thread_message, _ = execute_command(["ps", "-Te"])
|
|
threads_map = self.get_threads_map(thread_message)
|
|
main_pid = str(psutil.Process().pid)
|
|
current_npu = self.device_info.running_npu_list[self.rank_id]
|
|
self.bind(main_pid, self.assign_main[current_npu], True)
|
|
for thread_id in threads_map.get(main_pid, {}).get("acl_thread", []):
|
|
self.bind(thread_id, self.assign_acl[current_npu], False)
|
|
for thread_id in threads_map.get(main_pid, {}).get("release_thread", []):
|
|
self.bind(thread_id, self.assign_rel[current_npu], False)
|
|
# Migrate memory once for the whole process, after all threads are pinned.
|
|
self.bind_memory(main_pid, current_npu)
|
|
|
|
def bind_ascend_950_threads(self) -> None:
|
|
main_pid = str(psutil.Process().pid)
|
|
current_npu = self.device_info.running_npu_list[self.rank_id]
|
|
self.bind(main_pid, self.assign_main[current_npu], True)
|
|
self.bind_memory(main_pid, current_npu)
|
|
|
|
def bind_npu_irq(self) -> None:
|
|
if not self._reserve_irq_cpus():
|
|
logger.info("[irq] IRQ binding skipped on Ascend 950.")
|
|
return
|
|
|
|
if not os.access("/proc/irq", os.W_OK):
|
|
return
|
|
|
|
# Only bind IRQ for current rank's NPU to avoid multi-process overwrite.
|
|
current_npu = self.device_info.running_npu_list[self.rank_id]
|
|
if current_npu not in self.npu_cpu_pool:
|
|
logger.warning("[irq] NPU has no CPU pool. rank=%s, npu=%s. ", self.rank_id, current_npu)
|
|
return
|
|
|
|
if shutil.which("systemctl"):
|
|
output, _ = execute_command(["systemctl", "list-unit-files"])
|
|
if "irqbalance.service" in output:
|
|
_, return_code = execute_command(["systemctl", "is-active", "--quiet", "irqbalance"])
|
|
if return_code == 0:
|
|
logger.warning(
|
|
"The irqbalance service is running and has been stopped. "
|
|
"action: service stopped to enable IRQ binding. "
|
|
)
|
|
execute_command(["systemctl", "stop", "irqbalance"])
|
|
|
|
sq_irqs = []
|
|
with open("/proc/interrupts") as f:
|
|
for line in f:
|
|
if "sq_send_trigger_irq" in line:
|
|
irq = line.split(":")[0].strip()
|
|
sq_irqs.append(irq)
|
|
|
|
npu = current_npu
|
|
cpus = self.npu_cpu_pool[npu]
|
|
if len(cpus) < 2:
|
|
logger.warning("[irq] CPU pool too small. npu=%s, cpu_count=%d, min_required=2. ", npu, len(cpus))
|
|
return
|
|
|
|
sq_cpu, cq_cpu = cpus[0], cpus[1] # Reserved for IRQ binding
|
|
pci_addr = ""
|
|
|
|
device_type = get_ascend_device_type()
|
|
if device_type == AscendDeviceType.A3:
|
|
# A3: logical npu_id = card_id*2 + chip_id
|
|
card_id = npu // 2
|
|
chip_id = npu % 2
|
|
info, _ = execute_command(["npu-smi", "info", "-t", "board", "-i", str(card_id), "-c", str(chip_id)])
|
|
else:
|
|
# A2 / others: logical npu_id is card id
|
|
info, _ = execute_command(["npu-smi", "info", "-t", "board", "-i", str(npu)])
|
|
|
|
for line in info.splitlines():
|
|
if "PCIe Bus Info" in line:
|
|
pci_addr = line.split()[-1].lower()
|
|
break
|
|
|
|
if not pci_addr:
|
|
logger.warning("Can't find PCI address. npu=%s. ", npu)
|
|
return
|
|
|
|
try:
|
|
npu_irq_list = sorted(os.listdir(f"/sys/bus/pci/devices/{pci_addr}/msi_irqs/"), key=lambda x: int(x))
|
|
except FileNotFoundError:
|
|
logger.warning("The msi_irqs folder cannot be found. pci_addr=%s. ", pci_addr)
|
|
return
|
|
|
|
sq_irq, cq_irq = "", ""
|
|
for irq in sq_irqs:
|
|
if irq in npu_irq_list:
|
|
sq_irq = irq
|
|
cq_irq = str(int(irq) + 1)
|
|
break
|
|
if not sq_irq:
|
|
logger.warning("The sq_send_trigger_irq is not found. npu=%s. ", npu)
|
|
return
|
|
|
|
logger.info(
|
|
"NPU%s(PCI %s): sq_send_trigger_irq IRQ_ID=%s -> CPU%s, cq_update_irq IRQ_ID=%s -> CPU%s",
|
|
npu,
|
|
pci_addr,
|
|
sq_irq,
|
|
sq_cpu,
|
|
cq_irq,
|
|
cq_cpu,
|
|
)
|
|
with open(f"/proc/irq/{sq_irq}/smp_affinity", "w") as f:
|
|
f.write(self.cpu_to_mask(sq_cpu))
|
|
with open(f"/proc/irq/{cq_irq}/smp_affinity", "w") as f:
|
|
f.write(self.cpu_to_mask(cq_cpu))
|
|
|
|
def run_all(self) -> None:
|
|
if not self.build_cpu_pools():
|
|
return
|
|
self.allocate()
|
|
self.print_plan()
|
|
self.bind_threads()
|
|
self.bind_npu_irq()
|
|
|
|
|
|
def bind_cpus(rank_id: int) -> None:
|
|
if not is_arm_cpu():
|
|
logger.info("CPU binding skipped: non-ARM CPU detected.")
|
|
return
|
|
binder = CpuAlloc(rank_id)
|
|
binder.run_all()
|