已合并
新增 Host 诊断调优工具 #109
rozwel_dx创建于 3月5日
新增 Host 诊断调优工具 #109
已合并
共 6 个文件变更+995-0
| @@ -0,0 +1,160 @@ | |||
| 1 | +# Host 诊断调优工具 | ||
| 2 | + | ||
| 3 | +## 简介 | ||
| 4 | + | ||
| 5 | +当前 AI 训练、推理业务场景中,Host侧(CPU)的任务下发(如算子调度、内存分配)与Device侧(NPU)的任务执行是异步进行的。当Host侧任务下发耗时超过Device侧任务执行耗时时,Device会因等待新任务而处于空闲状态,形成性能瓶颈,即HostBound问题。 | ||
| 6 | + | ||
| 7 | +针对以上问题,我们设计了Host侧的诊断调优工具,提供简单、易用的绑核能力,通过将进程和线程分别绑定在不同CPU上执行以减少互相之间的干扰与资源竞争 | ||
| 8 | + | ||
| 9 | +## 使用前准备 | ||
| 10 | + | ||
| 11 | +1. 获取仓库中提供的参数解析、绑核脚本:[entrance.py](./entrance.py),[cpu_binder.py](./cpu_binder/cpu_binder.py) | ||
| 12 | + | ||
| 13 | +## 功能介绍 | ||
| 14 | + | ||
| 15 | +### 自定义绑核功能 | ||
| 16 | + | ||
| 17 | +#### 功能说明 | ||
| 18 | + | ||
| 19 | +提供自定义绑核能力,根据用户输入json中的配置方案完成进线程级绑核;缺省输入时以经验最优方案进行绑核: | ||
| 20 | + | ||
| 21 | + - 为每张卡的关键线程 acl_thread/release_thread 单独分配一个 CPU 核,dev[i]_sq_task 单独分配一个 CPU 核(宿主机可绑,未查询到时跳过),算子相关中断 sq_send_trigger_irq 和 cq_update_irq 各单独分配一个 CPU 核(需要拥有/proc目录写权限,权限不足时跳过),其余推理进线程共同分配到 6 个 CPU 核上,即每张 NPU 卡绑定到 11 个 CPU 核(分配 CPU 时考虑 NPU 亲和及跨 NUMA 内存访问时延) | ||
| 22 | + | ||
| 23 | +#### 命令格式 | ||
| 24 | + | ||
| 25 | +```bash | ||
| 26 | +python3 entrance.py bind [-l] [-c <config_path>] | ||
| 27 | +``` | ||
| 28 | + | ||
| 29 | +#### 参数说明 | ||
| 30 | + | ||
| 31 | +| 参数 | 可选/必选 | 说明 | | ||
| 32 | +|----------------|-------|-------------------------------------------------| | ||
| 33 | +| -l/--log-level | 可选 | 指定打印时的日志等级,类型为int,可选值为[0, 1, 2, 3],默认值为1 | | ||
| 34 | +| -c/--config | 可选 | 指定自定义绑核所需的json配置文件,缺省时按默认策略进行绑核,类型为str,默认值为None | | ||
| 35 | + | ||
| 36 | +#### json配置文件参数说明 | ||
| 37 | + | ||
| 38 | +| 参数 | 可选/必选 | 说明 | | ||
| 39 | +|------------------|-------|---------------------------------------------------------------------------------------------------------------------------------------------| | ||
| 40 | +| custom_bind | 必选 | json键值,类型为str,对应value需要为List[Dict],将每一组绑核对象存放在List中 | | ||
| 41 | +| process_name | 可选 | 需要绑定的进程或线程名称,类型为str,由is_thread参数决定是进程或线程,指定非NPU卡对应进程或线程时进程数或线程数应等于cpu_list参数长度 <br/>默认值为None | | ||
| 42 | +| pid | 可选 | 需要绑定的PID号,可以是一个或多个PID,该参数长度应等于cpu_list参数长度,类型为List[int] <br/>默认值为[] | | ||
| 43 | +| cpu_list | 必选 | 绑定CPU列表,可以是一个或多个CPU区间,类型为List[str] <br/>默认值为[] | | ||
| 44 | +| mem_bind | 可选 | 选择是否需要在绑定CPU后将内存迁移至对应NUMA节点,类型为bool,取值为:<br/>• true:表示绑核后将使用内存迁移至对应NUMA节点<br/>• false:表示绑核后不需要将使用内存迁移至对应NUMA节点<br/>默认值为false | | ||
| 45 | +| is_thread | 可选 | 选择绑定的ID是进程ID还是线程ID,类型为bool,取值为:<br/>• true:表示绑定线程ID<br/>• false:表示绑定进程ID<br/>默认值为false | | ||
| 46 | +| is_irq | 可选 | 选择绑定的是否为硬件中断,类型为bool,取值为:<br/>• true:表示绑定的是硬件中断<br/>• false:表示绑定的不是硬件中断<br/>默认值为false | | ||
| 47 | +| irq_id | 可选 | 需要绑定的硬件中断的中断号,类型为List[int] | | ||
| 48 | +| bind_sub_process | 可选 | 选择绑核时是否需要绑定进程ID或线程ID的所有子线程,类型为bool,取值为:<br/>• true:表示绑核时需要同时绑定进程ID或线程ID的所有子线程<br/>• false:表示绑核时不需要绑定进程ID或线程ID的所有子线程<br/>默认值为false | | ||
| 49 | + | ||
| 50 | +#### json配置文件示例 | ||
C | |||
| 51 | + | ||
| 52 | +```json | ||
| 53 | +{ | ||
| 54 | + "custom_bind": [ | ||
| 55 | + { | ||
| 56 | + "process_name": "VLLM::Worker_TP", | ||
| 57 | + "cpu_list": ["4-10","16-22","52-58","64-70"], | ||
| 58 | + "bind_sub_process": true | ||
| 59 | + }, | ||
| 60 | + { | ||
| 61 | + "process_name": "acl_thread", | ||
| 62 | + "cpu_list": ["11","23","59","71"], | ||
| 63 | + "mem_bind": true, | ||
| 64 | + "is_thread": true | ||
| 65 | + }, | ||
| 66 | + { | ||
| 67 | + "process_name": "release_thread", | ||
| 68 | + "cpu_list": ["12,13","24,25","60,61","72,73"], | ||
| 69 | + "is_thread": true | ||
| 70 | + }, | ||
| 71 | + { | ||
| 72 | + "process_name": "VLLM::EngineCore", | ||
| 73 | + "cpu_list": ["44"], | ||
| 74 | + "bind_sub_process": false, | ||
| 75 | + "mem_bind": true | ||
| 76 | + }, | ||
| 77 | + { | ||
| 78 | + "pid": [110351, 110352, 110353, 110354], | ||
| 79 | + "cpu_list": ["32-35"] | ||
| 80 | + }, | ||
| 81 | + { | ||
| 82 | + "irq_id": [3008, 3009, 3264, 3265, 2496, 2497, 2752, 2753], | ||
| 83 | + "cpu_list": ["15", "16", "27", "28", "64", "65", "76", "77"], | ||
| 84 | + "is_irq": true | ||
| 85 | + } | ||
| 86 | + ] | ||
| 87 | +} | ||
| 88 | +``` | ||
| 89 | + | ||
| 90 | +#### 使用示例 | ||
| 91 | + | ||
| 92 | +- 使用示例1 | ||
| 93 | + | ||
| 94 | + ```python | ||
| 95 | + python3 entrance.py bind -c ./bind_design.json | ||
| 96 | + ``` | ||
| 97 | + | ||
| 98 | + 按照[json配置文件](#json配置文件示例)中的方案进行绑核,假设当前仅NPU4,5,6,7在运行,则: | ||
| 99 | + 1. 以进程名"VLLM::Worker_TP"匹配NPU4,5,6,7中的该进程,并分别绑定到CPU "4-10","16-22","52-58","64-70",绑定时会同时绑定该进程下的所有子线程 | ||
| 100 | + 2. 以线程名"acl_thread"匹配NPU4,5,6,7中的该线程,并分别绑定到CPU "11","23","59","71" | ||
| 101 | + 3. 以线程名"release_thread"匹配NPU4,5,6,7中的该线程,并分别绑定到CPU "12,13","24,25","60,61","72,73" | ||
| 102 | + 4. 以进程名"VLLM::EngineCore"匹配环境中的该进程,并绑定到CPU "44" | ||
| 103 | + 5. 分别绑定pid "110351","110352","110353","110354"到CPU "14","26","62,63","74-75" | ||
| 104 | + 6. 分别绑定中断 "3008","3009","3264","3265","2496","2497","2752","2753"到CPU "15", "16", "27", "28", "64", "65", "76", "77" | ||
| 105 | +- 使用示例2 | ||
| 106 | + | ||
| 107 | + ```python | ||
| 108 | + python3 entrance.py bind | ||
| 109 | + ``` | ||
| 110 | + | ||
| 111 | + 1. 结合NPU亲和性以及可访问CPU数量,为每张NPU卡分配对应的CPU区间(每张卡分配的CPU区间都会在一个NUMA节点内) | ||
| 112 | + 2. 绑定每张NPU卡对应的sq_send_trigger_irq和cq_update_irq(算子下发硬件中断)到其CPU区间的前两个CPU核,需要有/proc/irq/\<irq_id\>/smp_affinity文件的写权限,否则会执行失败 | ||
| 113 | + 3. 绑定每张NPU卡对应的dev[i]_sq_task(NPU驱动进程)到其CPU区间的第三个CPU核,需要在宿主机才能查询到此进程,否则会执行失败 | ||
| 114 | + 4. 绑定每张NPU卡对应的主进程和所有子线程到其CPU区间的第四到第九个CPU核 | ||
| 115 | + 5. 绑定每张NPU卡对应的acl_thread(算子下发线程)到其CPU区间的第十个CPU核 | ||
| 116 | + 6. 绑定每张NPU卡对应的release_thread(资源释放线程)到其CPU区间的第十一个CPU核 | ||
| 117 | + | ||
| 118 | +#### 输出示例 | ||
| 119 | + | ||
| 120 | +以使用示例1为例,会输出以下信息展示将绑定结果: | ||
| 121 | + | ||
| 122 | + ```text | ||
| 123 | + [2026-03-13 10:34:27,406] [INFO]:Start binding core round 1: {"process_name": "VLLM::Worker_TP", "cpu_list": ["4-10","16-22","52-58","64-70"], "bind_sub_process": true} | ||
| 124 | + [2026-03-13 10:34:28,490] [INFO]:Bind the target (pid=1603713) to CPU4,5,6,7,8,9,10 | ||
| 125 | + [2026-03-13 10:34:28,498] [INFO]:Bind the target (pid=1603714) to CPU16,17,18,19,20,21,22 | ||
| 126 | + [2026-03-13 10:34:28,506] [INFO]:Bind the target (pid=1603715) to CPU52,53,54,55,56,57,58 | ||
| 127 | + [2026-03-13 10:34:28,514] [INFO]:Bind the target (pid=1603716) to CPU64,65,66,67,68,69,70 | ||
| 128 | + [2026-03-13 10:34:28,522] [INFO]:===== Round 1 of core binding has ended ===== | ||
| 129 | + [2026-03-13 10:34:28,533] [INFO]:Start binding core round 2: {"process_name": "acl_thread", "cpu_list": ["11","23","59","71"], "mem_bind": true, "is_thread": true} | ||
| 130 | + [2026-03-13 10:34:29,592] [INFO]:Bind the target (pid=1648512) to CPU11 | ||
| 131 | + [2026-03-13 10:34:29,603] [INFO]:Bind the target (pid=1648576) to CPU23 | ||
| 132 | + [2026-03-13 10:34:29,613] [INFO]:Bind the target (pid=1648615) to CPU59 | ||
| 133 | + [2026-03-13 10:34:29,623] [INFO]:Bind the target (pid=1648667) to CPU71 | ||
| 134 | + [2026-03-13 10:34:29,633] [INFO]:===== Round 2 of core binding has ended ===== | ||
| 135 | + [2026-03-13 10:34:29,633] [INFO]:Start binding core round 3: {"process_name": "release_thread", "cpu_list": ["12","24","60","72"], "is_thread": true} | ||
| 136 | + [2026-03-13 10:34:30,644] [INFO]:Bind the target (pid=1648513) to CPU12 | ||
| 137 | + [2026-03-13 10:34:30,729] [INFO]:Bind the target (pid=1648577) to CPU24 | ||
| 138 | + [2026-03-13 10:34:30,736] [INFO]:Bind the target (pid=1648616) to CPU60 | ||
| 139 | + [2026-03-13 10:34:30,743] [INFO]:Bind the target (pid=1648668) to CPU72 | ||
| 140 | + [2026-03-13 10:34:30,750] [INFO]:===== Round 3 of core binding has ended ===== | ||
| 141 | + [2026-03-13 10:34:30,757] [INFO]:Start binding core round 4: {"process_name": "VLLM::EngineCore", "cpu_list": ["44"], "bind_sub_process": false, "mem_bind": true} | ||
| 142 | + [2026-03-13 10:34:31,490] [INFO]:Bind the target (pid=1603113) to CPU44 | ||
| 143 | + [2026-03-13 10:34:31,522] [INFO]:===== Round 4 of core binding has ended ===== | ||
| 144 | + [2026-03-13 10:34:31,348] [INFO]:Start binding core round 5: {"pid": [110351, 110352, 110353, 110354], "cpu_list": ["32-35"]} | ||
| 145 | + [2026-03-13 10:34:32,356] [INFO]:Bind the target (pid=110351) to CPU32,33,34,35 | ||
| 146 | + [2026-03-13 10:34:32,363] [INFO]:Bind the target (pid=110352) to CPU32,33,34,35 | ||
| 147 | + [2026-03-13 10:34:32,370] [INFO]:Bind the target (pid=110353) to CPU32,33,34,35 | ||
| 148 | + [2026-03-13 10:34:32,378] [INFO]:Bind the target (pid=110354) to CPU32,33,34,35 | ||
| 149 | + [2026-03-13 10:34:32,385] [INFO]:===== Round 5 of core binding has ended ===== | ||
| 150 | + [2026-03-13 10:34:32,406] [INFO]:Start binding core round 6: {"irq_id": [3008, 3009, 3264, 3265, 2496, 2497, 2752, 2753], "cpu_list": ["15", "16", "27", "28", "64", "65", "76", "77"], "is_irq": true} | ||
| 151 | + [2026-03-13 10:34:33,540] [INFO]:Bind the interrupt of IRQ-3008 to CPU15 | ||
| 152 | + [2026-03-13 10:34:33,540] [INFO]:Bind the interrupt of IRQ-3009 to CPU16 | ||
| 153 | + [2026-03-13 10:34:33,540] [INFO]:Bind the interrupt of IRQ-3264 to CPU27 | ||
| 154 | + [2026-03-13 10:34:33,540] [INFO]:Bind the interrupt of IRQ-3265 to CPU28 | ||
| 155 | + [2026-03-13 10:34:33,540] [INFO]:Bind the interrupt of IRQ-2496 to CPU64 | ||
| 156 | + [2026-03-13 10:34:33,541] [INFO]:Bind the interrupt of IRQ-2497 to CPU65 | ||
| 157 | + [2026-03-13 10:34:33,541] [INFO]:Bind the interrupt of IRQ-2752 to CPU76 | ||
| 158 | + [2026-03-13 10:34:33,541] [INFO]:Bind the interrupt of IRQ-2753 to CPU77 | ||
| 159 | + [2026-03-13 10:34:33,542] [INFO]:===== Round 6 of core binding has ended ===== | ||
| 160 | + ``` | ||
| @@ -0,0 +1,587 @@ | |||
| 1 | +import argparse | ||
| 2 | +import logging | ||
| 3 | +import os | ||
| 4 | +import re | ||
| 5 | +import shutil | ||
| 6 | +import subprocess | ||
| 7 | +from collections import defaultdict | ||
| 8 | +from typing import List, Dict, Tuple | ||
| 9 | + | ||
| 10 | +from misc.gil_tracer.file_manager import FileManager | ||
| 11 | + | ||
| 12 | +CPU_MASK_BIT = 32 | ||
| 13 | +MAIN_PROCESS_RANGE = 6 | ||
| 14 | +ACL_THREAD_RANGE = 1 | ||
| 15 | +RELEASE_THREAD_RANGE = 1 | ||
| 16 | +ALLOWED_CPUS_PATH = "/proc/self/status" | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +def execute_command(cmd: List[str]) -> Tuple[str, int]: | ||
| 20 | + with subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as p: | ||
| 21 | + out, err = p.communicate(timeout=1000) | ||
| 22 | + if err: | ||
| 23 | + logging.debug(f"Command stderr while running {cmd}: {err.decode()}") | ||
| 24 | + return out.decode(), p.returncode | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +def expand_cpu_list(cpu_str: str) -> List[int]: | ||
| 28 | + cpus = [] | ||
| 29 | + try: | ||
| 30 | + for part in cpu_str.split(","): | ||
| 31 | + if "-" in part: | ||
| 32 | + start, end = map(int, part.split("-")) | ||
| 33 | + cpus.extend(range(start, end + 1)) | ||
| 34 | + else: | ||
| 35 | + cpus.append(int(part)) | ||
| 36 | + except ValueError: | ||
[codereview][一般] execute_command未检查communicate(timeout=1000)超时异常;建议捕获subprocess.TimeoutExpired并主动终止进程,避免子进程泄漏。 ![]() ![]() | |||
| 37 | + raise RuntimeError(f"The cpu_list parameter must consist of digits, ',' and '-', which is '{cpu_str}'.") | ||
| 38 | + return cpus | ||
| 39 | + | ||
| 40 | + | ||
| 41 | +class DeviceInfo: | ||
| 42 | + def __init__(self): | ||
| 43 | + self.main_pid_list: List[List[int]] = [] | ||
| 44 | + self.npu_map_info: Dict[str, Dict[str, str]] = self.get_npu_map_info() | ||
| 45 | + self.allowed_cpus: List[int] = self.parse_allowed_cpus() | ||
| 46 | + self.running_npu_list: List[int] = self.get_running_npus() | ||
| 47 | + self.npu_affinity: Dict[int, List[int]] = self.parse_topo_affinity() | ||
| 48 | + | ||
| 49 | + | ||
| 50 | + def get_npu_map_info() -> Dict[str, Dict[str, str]]: | ||
| 51 | + npu_map_info: Dict[str, Dict[str, str]] = {} | ||
| 52 | + npu_info, _ = execute_command(["npu-smi", "info", "-m"]) | ||
| 53 | + npu_map = npu_info.strip().split("\n")[1:] | ||
| 54 | + for line in npu_map: | ||
| 55 | + parts = line.strip().split() | ||
| 56 | + if len(parts) < 3: | ||
| 57 | + continue | ||
| 58 | + npu_id, chip_id, chip_logic_id = parts[:3] | ||
| 59 | + if chip_logic_id.isdigit(): | ||
| 60 | + npu_map_info.setdefault(npu_id, {})[chip_id] = chip_logic_id | ||
| 61 | + logging.debug(f"build npu_map_info: {npu_map_info}") | ||
| 62 | + return npu_map_info | ||
| 63 | + | ||
| 64 | + | ||
| 65 | + def parse_allowed_cpus() -> List[int]: | ||
| 66 | + if not os.path.exists(ALLOWED_CPUS_PATH): | ||
| 67 | + return [] | ||
| 68 | + with open(ALLOWED_CPUS_PATH) as f: | ||
| 69 | + for line in f: | ||
| 70 | + if line.startswith("Cpus_allowed_list"): | ||
| 71 | + allowed_cpu_list = expand_cpu_list(line.split()[1]) | ||
| 72 | + logging.debug(f"Cpus_allowed_list: {allowed_cpu_list}") | ||
| 73 | + return allowed_cpu_list | ||
| 74 | + return [] | ||
| 75 | + | ||
| 76 | + | ||
| 77 | + def parse_topo_affinity() -> Dict[int, List[int]]: | ||
| 78 | + chip_logic_id = 0 | ||
| 79 | + affinity: Dict[int, List[int]] = {} | ||
| 80 | + affinity_message, _ = execute_command(["npu-smi", "info", "-t", "topo"]) | ||
| 81 | + for line in affinity_message.splitlines(): | ||
[codereview][一般] parse_topo_affinity使用自增chip_logic_id作为NPU标识,强依赖npu-smi输出顺序;建议解析真实逻辑ID字段,避免设备顺序变化导致CPU亲和性映射错误。 ![]() ![]() | |||
| 82 | + if line.startswith("NPU"): | ||
| 83 | + last_part = line.split()[-1] | ||
| 84 | + if last_part != "Affinity": | ||
| 85 | + affinity[chip_logic_id] = expand_cpu_list(last_part) | ||
| 86 | + chip_logic_id += 1 | ||
| 87 | + logging.debug(f"build affinity map: {affinity}") | ||
| 88 | + return affinity | ||
| 89 | + | ||
| 90 | + def get_running_npus(self) -> List[int]: | ||
| 91 | + npu_message, _ = execute_command(["npu-smi", "info"]) | ||
| 92 | + in_proc_section = False | ||
| 93 | + running_npu_set = set() | ||
| 94 | + chip_pid_map: Dict[int, List[Tuple[int, int]]] = {} | ||
| 95 | + for line in npu_message.splitlines(): | ||
| 96 | + line = line.strip() | ||
| 97 | + if line.startswith("| NPU") and "Process id" in line: | ||
| 98 | + in_proc_section = True | ||
| 99 | + continue | ||
| 100 | + if not in_proc_section or not line.startswith("| "): | ||
| 101 | + continue | ||
| 102 | + parts = [p.strip() for p in line.strip("|").split("|")] | ||
| 103 | + if len(parts) < 4 or not parts[1].isdigit(): | ||
| 104 | + continue | ||
| 105 | + pid = int(parts[1]) | ||
| 106 | + try: | ||
| 107 | + mem = int(parts[3]) | ||
| 108 | + except ValueError: | ||
| 109 | + mem = 0 | ||
| 110 | + npu_id, chip_id = parts[0].split()[:2] | ||
[codereview][严重] parts[0].split()[:2]默认认为字段一定包含NPU ID和Chip ID;建议增加长度校验,避免输出格式变化触发ValueError ![]() ![]() | |||
| 111 | + chip_logic_id = self.npu_map_info.get(npu_id, {}).get(chip_id) | ||
| 112 | + if chip_logic_id and chip_logic_id.isdigit(): | ||
| 113 | + chip_logic_id = int(chip_logic_id) | ||
| 114 | + chip_pid_map.setdefault(chip_logic_id, []).append((pid, mem)) | ||
| 115 | + running_npu_set.add(chip_logic_id) | ||
| 116 | + | ||
| 117 | + self.main_pid_list = [] | ||
| 118 | + running_npu_list = sorted(running_npu_set) | ||
| 119 | + for npu in running_npu_list: | ||
| 120 | + pid_mem_list = chip_pid_map.get(npu, []) | ||
| 121 | + if pid_mem_list: | ||
| 122 | + max_pid = max(pid_mem_list, key=lambda x: x[1])[0] | ||
| 123 | + self.main_pid_list.append([max_pid]) | ||
| 124 | + logging.debug(f"identifying the running NPU card: {running_npu_set}") | ||
| 125 | + return running_npu_list | ||
| 126 | + | ||
| 127 | + | ||
| 128 | +class CpuAlloc: | ||
| 129 | + def __init__(self): | ||
| 130 | + self.device_info: DeviceInfo = DeviceInfo() | ||
| 131 | + self.cpu_node: Dict[int, int] = {} | ||
| 132 | + self.numa_to_cpu_map: Dict[int, List[int]] = defaultdict(list) | ||
| 133 | + self.npu_cpu_pool: Dict[int, List[int]] = {} | ||
| 134 | + self.npu_cpu_pool_all: Dict[int, List[int]] = {} | ||
| 135 | + self.assign_main: Dict[int, List[int]] = {} | ||
| 136 | + self.assign_acl: Dict[int, List[int]] = {} | ||
| 137 | + self.assign_rel: Dict[int, List[int]] = {} | ||
| 138 | + | ||
| 139 | + | ||
| 140 | + def average_distribute(groups: Dict[str, List[int]], pool: Dict[int, List[int]]) -> Dict[int, List[int]]: | ||
| 141 | + result: Dict[int, List[int]] = {} | ||
| 142 | + for key, npu_list in groups.items(): | ||
| 143 | + cpu_list = sorted(pool[npu_list[0]]) | ||
| 144 | + cpu_num_per_npu = len(cpu_list) // len(npu_list) | ||
| 145 | + for i, npu in enumerate(npu_list): | ||
| 146 | + start_index = i * cpu_num_per_npu | ||
| 147 | + end_index = (i + 1) * cpu_num_per_npu if i < len(npu_list) - 1 else len(cpu_list) | ||
| 148 | + result[npu] = cpu_list[start_index:end_index] | ||
| 149 | + return result | ||
| 150 | + | ||
| 151 | + | ||
| 152 | + def get_acl_main_threads() -> List[int]: | ||
| 153 | + thread_message, _ = execute_command(["ps", "-Te"]) | ||
| 154 | + pids: List[int] = [] | ||
| 155 | + acl_threads_set = set() | ||
| 156 | + for line in thread_message.splitlines(): | ||
| 157 | + if "acl_thread" in line: | ||
| 158 | + pid = line.split()[0] | ||
| 159 | + if pid not in acl_threads_set: | ||
| 160 | + acl_threads_set.add(pid) | ||
| 161 | + pids.append(int(pid)) | ||
| 162 | + return pids | ||
| 163 | + | ||
| 164 | + def dev_alloc(self) -> tuple[List[int], List[str]]: | ||
| 165 | + dev_pid_list: List[int] = [] | ||
| 166 | + dev_cpu_list: List[str] = [] | ||
| 167 | + out, _ = execute_command(["ps", "aux"]) | ||
| 168 | + for line in out.splitlines(): | ||
| 169 | + m = re.search(r"dev(\d+)_sq_task", line) | ||
| 170 | + if not m: | ||
| 171 | + continue | ||
| 172 | + dev_id = int(m.group(1)) | ||
| 173 | + pid = int(line.split()[1]) | ||
| 174 | + cpus = self.npu_cpu_pool_all.get(dev_id, []) | ||
| 175 | + if cpus: | ||
| 176 | + core = cpus[2] if len(cpus) >= 3 else cpus[0] | ||
| 177 | + dev_pid_list.append(pid) | ||
| 178 | + dev_cpu_list.append(str(core)) | ||
| 179 | + return dev_pid_list, dev_cpu_list | ||
| 180 | + | ||
| 181 | + def irq_alloc(self) -> tuple[List[int], List[str]]: | ||
| 182 | + sq_irqs = [] | ||
| 183 | + irq_id_list: List[int] = [] | ||
| 184 | + irq_cpu_list: List[str] = [] | ||
| 185 | + try: | ||
| 186 | + with open("/proc/interrupts") as f: | ||
| 187 | + for line in f: | ||
| 188 | + if "sq_send_trigger_irq" in line: | ||
| 189 | + irq = line.split(":")[0].strip() | ||
| 190 | + sq_irqs.append(irq) | ||
| 191 | + except (IOError, PermissionError) as e: | ||
| 192 | + raise RuntimeError(f"Can't open /proc/interrupts for: {e}") | ||
| 193 | + | ||
| 194 | + for npu in sorted(self.npu_cpu_pool_all.keys()): | ||
| 195 | + cpus = self.npu_cpu_pool_all[npu] | ||
| 196 | + if len(cpus) < 2: | ||
| 197 | + continue | ||
| 198 | + | ||
| 199 | + info, _ = execute_command(["npu-smi", "info", "-t", "board", "-i", str(npu)]) | ||
| 200 | + pci_addr = "" | ||
| 201 | + for line in info.splitlines(): | ||
| 202 | + if "PCIe Bus Info" in line: | ||
| 203 | + pci_addr = line.split()[-1].lower() | ||
| 204 | + break | ||
| 205 | + if not pci_addr: | ||
| 206 | + raise RuntimeError(f"Can't find PCI address of NPU{npu} .") | ||
| 207 | + | ||
| 208 | + msi_irq_dir = f"/sys/bus/pci/devices/{pci_addr}/msi_irqs/" | ||
| 209 | + if not os.path.exists(msi_irq_dir): | ||
| 210 | + raise RuntimeError(f"Can't find MSI interrupt directory of NPU{npu} .") | ||
| 211 | + | ||
| 212 | + npu_irq_list = sorted( | ||
| 213 | + os.listdir(f"/sys/bus/pci/devices/{pci_addr}/msi_irqs/"), | ||
| 214 | + key=lambda x: int(x) | ||
| 215 | + ) | ||
| 216 | + for irq in sq_irqs: | ||
| 217 | + if irq in npu_irq_list: | ||
| 218 | + irq_id_list.extend([int(irq), int(irq) + 1]) | ||
| 219 | + irq_cpu_list.extend([str(cpus[0]), str(cpus[1])]) | ||
| 220 | + break | ||
| 221 | + return irq_id_list, irq_cpu_list | ||
| 222 | + | ||
| 223 | + def build_cpu_node_map(self) -> None: | ||
| 224 | + cpu_numa_map, _ = execute_command(["lscpu", "-e=CPU,NODE"]) | ||
| 225 | + for line in cpu_numa_map.splitlines(): | ||
| 226 | + line = line.strip() | ||
| 227 | + if not line or not line[0].isdigit(): | ||
| 228 | + continue | ||
| 229 | + cpu_str, node_str = line.split() | ||
| 230 | + cpu = int(cpu_str) | ||
| 231 | + node = int(node_str) | ||
| 232 | + self.cpu_node[cpu] = node | ||
| 233 | + self.numa_to_cpu_map[node].append(cpu) | ||
| 234 | + if not self.numa_to_cpu_map: | ||
| 235 | + raise RuntimeError("The output of 'lscpu' is incorrect and no NUMA node is detected.") | ||
| 236 | + | ||
| 237 | + def extend_numa(self, cpu_list: List[int]) -> List[int]: | ||
| 238 | + if not cpu_list: | ||
| 239 | + return [] | ||
| 240 | + nodes = {self.cpu_node[c] for c in cpu_list} | ||
| 241 | + if len(nodes) != 1: | ||
| 242 | + return cpu_list | ||
| 243 | + node = list(nodes)[0] | ||
| 244 | + next_node = (node + 1) % len(self.numa_to_cpu_map) | ||
| 245 | + extended = cpu_list[:] | ||
| 246 | + for cpu in self.numa_to_cpu_map[next_node]: | ||
| 247 | + if cpu in self.device_info.allowed_cpus: | ||
| 248 | + extended.append(cpu) | ||
| 249 | + return sorted(set(extended)) | ||
| 250 | + | ||
| 251 | + def handle_no_affinity(self) -> None: | ||
| 252 | + num_running_npu = len(self.device_info.running_npu_list) | ||
| 253 | + num_numa_node = len(self.numa_to_cpu_map) | ||
| 254 | + if num_numa_node == 0 or num_running_npu == 0: | ||
| 255 | + return | ||
| 256 | + npu_num_per_node = (num_running_npu // num_numa_node) + (1 if num_running_npu % num_numa_node else 0) | ||
| 257 | + index = 0 | ||
| 258 | + for node in sorted(self.numa_to_cpu_map): | ||
| 259 | + cpus = [c for c in self.numa_to_cpu_map[node] if c in self.device_info.allowed_cpus] | ||
| 260 | + if not cpus: | ||
| 261 | + continue | ||
| 262 | + npu_num_this_node = min(npu_num_per_node, num_running_npu - index) | ||
| 263 | + total_cpu_num = len(cpus) | ||
| 264 | + base_cpu_num = total_cpu_num // npu_num_this_node | ||
| 265 | + extra_cpu_num = total_cpu_num % npu_num_this_node | ||
| 266 | + start_index = 0 | ||
| 267 | + for i in range(npu_num_this_node): | ||
| 268 | + take_cpu_num = base_cpu_num + (1 if i < extra_cpu_num else 0) | ||
| 269 | + end_index = start_index + take_cpu_num | ||
| 270 | + select_cpus_list = cpus[start_index:end_index] | ||
| 271 | + if index < num_running_npu: | ||
| 272 | + npu = self.device_info.running_npu_list[index] | ||
| 273 | + self.npu_cpu_pool[npu] = select_cpus_list | ||
| 274 | + index += 1 | ||
| 275 | + start_index = end_index | ||
| 276 | + | ||
| 277 | + def build_cpu_pools_all(self) -> None: | ||
| 278 | + raw_pool: Dict[int, List[int]] = {} | ||
| 279 | + | ||
| 280 | + if self.device_info.npu_affinity: | ||
| 281 | + for npu, cpus in self.device_info.npu_affinity.items(): | ||
| 282 | + filtered = [c for c in cpus if c in self.device_info.allowed_cpus] | ||
| 283 | + raw_pool[npu] = filtered | ||
| 284 | + else: | ||
| 285 | + self.handle_no_affinity() | ||
| 286 | + raw_pool = self.npu_cpu_pool.copy() | ||
| 287 | + | ||
| 288 | + groups: Dict[str, List[int]] = defaultdict(list) | ||
| 289 | + for npu, cpus in raw_pool.items(): | ||
| 290 | + groups[str(cpus)].append(npu) | ||
| 291 | + | ||
| 292 | + final_pool: Dict[int, List[int]] = {} | ||
| 293 | + for key, npu_list in groups.items(): | ||
| 294 | + if len(npu_list) == 1: | ||
| 295 | + final_pool[npu_list[0]] = raw_pool[npu_list[0]] | ||
| 296 | + else: | ||
| 297 | + final_pool.update(self.average_distribute({key: npu_list}, raw_pool)) | ||
| 298 | + logging.debug(f"npu_cpu_pool_all: {final_pool}") | ||
| 299 | + self.npu_cpu_pool_all = final_pool | ||
| 300 | + | ||
| 301 | + def build_cpu_pools_running(self) -> None: | ||
| 302 | + self.build_cpu_node_map() | ||
| 303 | + raw_pool: Dict[int, List[int]] = {} | ||
| 304 | + | ||
| 305 | + if self.device_info.npu_affinity: | ||
| 306 | + for npu in self.device_info.running_npu_list: | ||
| 307 | + cpus = self.device_info.npu_affinity.get(npu, []) | ||
| 308 | + filtered = [c for c in cpus if c in self.device_info.allowed_cpus] | ||
| 309 | + raw_pool[npu] = filtered | ||
| 310 | + else: | ||
| 311 | + self.handle_no_affinity() | ||
| 312 | + for npu in self.device_info.running_npu_list: | ||
| 313 | + if npu in self.npu_cpu_pool: | ||
| 314 | + raw_pool[npu] = self.npu_cpu_pool[npu] | ||
| 315 | + | ||
| 316 | + groups: Dict[str, List[int]] = defaultdict(list) | ||
| 317 | + for npu, cpus in raw_pool.items(): | ||
| 318 | + groups[str(cpus)].append(npu) | ||
| 319 | + | ||
| 320 | + final_pool: Dict[int, List[int]] = {} | ||
| 321 | + for key, npu_list in groups.items(): | ||
| 322 | + if len(npu_list) == 1: | ||
| 323 | + final_pool[npu_list[0]] = raw_pool[npu_list[0]] | ||
| 324 | + else: | ||
| 325 | + final_pool.update(self.average_distribute({key: npu_list}, raw_pool)) | ||
| 326 | + logging.debug(f"npu_cpu_pool: {final_pool}") | ||
| 327 | + self.npu_cpu_pool = final_pool | ||
| 328 | + | ||
| 329 | + def allocate(self, main_range: int, acl_range: int, rel_range: int) -> None: | ||
| 330 | + for npu, pool in self.npu_cpu_pool.items(): | ||
| 331 | + usable_pool = pool[3:] | ||
| 332 | + need = main_range + acl_range + rel_range | ||
| 333 | + if len(usable_pool) < need: | ||
| 334 | + raise RuntimeError(f"The numaber of CPUs on NPU{npu} is insufficient. " | ||
| 335 | + f"The default solution requires at least {need} CPUs.") | ||
| 336 | + self.assign_main[npu] = usable_pool[:main_range] | ||
| 337 | + self.assign_acl[npu] = usable_pool[main_range:main_range + acl_range] | ||
| 338 | + self.assign_rel[npu] = usable_pool[main_range + acl_range:main_range + acl_range + rel_range] | ||
| 339 | + | ||
| 340 | + | ||
| 341 | +class CustomBind: | ||
| 342 | + def __init__(self, process_name: str = "", cpu_list: List[str] = None, | ||
| 343 | + bind_sub_process: bool = False, is_thread: bool = False, is_irq: bool = False, | ||
| 344 | + mem_bind: bool = False, pid: List[int] = None, irq_id: List[int] = None): | ||
| 345 | + self.process_name = process_name | ||
| 346 | + self.bind_sub_process = bind_sub_process | ||
| 347 | + self.is_thread = is_thread | ||
| 348 | + self.is_irq = is_irq | ||
| 349 | + self.mem_bind = mem_bind | ||
| 350 | + self.pid = pid or [] | ||
| 351 | + self.irq_id = irq_id or [] | ||
| 352 | + self.cpu_list = [expand_cpu_list(seg) for seg in (cpu_list or [])] | ||
| 353 | + | ||
| 354 | + | ||
| 355 | + def cpu_to_mask(cpus: List[int]) -> str: | ||
[codereview][严重] cpu_to_mask在cpus为空时执行max(groups.keys())会抛出异常;建议增加空列表保护逻辑。 ![]() ![]() | |||
| 356 | + groups = defaultdict(int) | ||
| 357 | + for cpu in cpus: | ||
| 358 | + group = cpu // CPU_MASK_BIT | ||
| 359 | + bit = cpu % CPU_MASK_BIT | ||
| 360 | + groups[group] |= (1 << bit) | ||
| 361 | + | ||
| 362 | + max_group = max(groups.keys()) | ||
| 363 | + mask_parts = [] | ||
| 364 | + for group in reversed(range(max_group + 1)): | ||
| 365 | + mask_parts.append(f"{groups.get(group, 0):08x}") | ||
| 366 | + return ",".join(mask_parts) | ||
| 367 | + | ||
| 368 | + | ||
| 369 | + def get_main_pid_from_docker(pid: int) -> int: | ||
[codereview][一般] get_main_pid_from_docker通过grep解析文件内容额外创建进程开销较大;建议直接读取/proc/ ![]() ![]() | |||
| 370 | + pid_file = f"/proc/{pid}/status" | ||
| 371 | + if not os.path.exists(pid_file): | ||
| 372 | + return 0 | ||
| 373 | + out, return_code = execute_command(["grep", "Ngid", pid_file]) | ||
| 374 | + if return_code != 0: | ||
| 375 | + return 0 | ||
| 376 | + parts = out.strip().split() | ||
| 377 | + if parts[-1] != "0": | ||
| 378 | + return int(parts[-1]) | ||
| 379 | + else: | ||
| 380 | + return 0 | ||
| 381 | + | ||
| 382 | + def get_real_main_pid_list(self, pid_list: List[Tuple[int, int]], | ||
| 383 | + main_pid_list: List[List[int]]) -> List[List[int]]: | ||
| 384 | + real_main_pid_list: List[List[int]] = [] | ||
| 385 | + for pid, ppid in pid_list: | ||
| 386 | + per_real_pid_list: List[int] = [] | ||
| 387 | + for pids in main_pid_list: | ||
| 388 | + if pid in pids: | ||
| 389 | + per_real_pid_list.append(pid) | ||
| 390 | + continue | ||
| 391 | + elif ppid in pids: | ||
| 392 | + per_real_pid_list.append(ppid) | ||
| 393 | + continue | ||
| 394 | + real_pid = self.get_main_pid_from_docker(pid) | ||
| 395 | + if real_pid in pids: | ||
| 396 | + per_real_pid_list.append(pid) | ||
| 397 | + continue | ||
| 398 | + real_ppid = self.get_main_pid_from_docker(ppid) | ||
| 399 | + if real_ppid in pids: | ||
| 400 | + per_real_pid_list.append(ppid) | ||
| 401 | + if per_real_pid_list: | ||
| 402 | + real_main_pid_list.append(per_real_pid_list) | ||
| 403 | + unique_list = list(dict.fromkeys(tuple(lst) for lst in real_main_pid_list)) | ||
| 404 | + return [list(x) for x in unique_list] | ||
| 405 | + | ||
| 406 | + def find_threads(self) -> List[Tuple[int, int]]: | ||
| 407 | + if self.pid: | ||
| 408 | + pid_list = [] | ||
| 409 | + for p in self.pid: | ||
| 410 | + try: | ||
| 411 | + ppid = int(subprocess.check_output(["ps", "-o", "ppid=", "-p", str(p)], text=True).strip()) | ||
| 412 | + except subprocess.CalledProcessError: | ||
| 413 | + ppid = -1 | ||
| 414 | + except ValueError: | ||
| 415 | + ppid = -1 | ||
| 416 | + pid_list.append((p, ppid)) | ||
| 417 | + return pid_list | ||
| 418 | + | ||
| 419 | + select_idx = 1 if self.is_thread else 0 | ||
| 420 | + out, _ = execute_command(["ps", "-Te"]) if self.is_thread else execute_command(["ps", "-eo", "pid,ppid,cmd"]) | ||
| 421 | + pid_list = [] | ||
| 422 | + for line in out.splitlines(): | ||
| 423 | + if self.process_name in line: | ||
| 424 | + parts = line.split() | ||
| 425 | + if len(parts) >= 2 and parts[0].isdigit() and parts[1].isdigit(): | ||
| 426 | + pid = int(parts[select_idx]) | ||
| 427 | + ppid = int(parts[1 - select_idx]) | ||
| 428 | + pid_list.append((pid, ppid)) | ||
| 429 | + if not pid_list: | ||
| 430 | + raise RuntimeError(f"No process whose name contains {self.process_name} is found.") | ||
| 431 | + return pid_list | ||
| 432 | + | ||
| 433 | + def irq_bind(self) -> None: | ||
[codereview][严重] real_main_pid_list.index(pids)位于循环内部,每次都会线性查找;建议提前构建PID到CPU列表索引映射,避免O(n²)性能退化。 ![]() ![]() | |||
| 434 | + if not shutil.which("systemctl"): | ||
| 435 | + logging.warning("The systemctl command cannot be used in the current environment.If the irqbalance " | ||
| 436 | + "service is enabled, manually disable the irqbalance service.Otherwise, the " | ||
| 437 | + "interrupt-core binding cannot take effect.") | ||
| 438 | + else: | ||
| 439 | + out, return_code = execute_command(["systemctl", "list-unit-files"]) | ||
| 440 | + if return_code == 0 and "irqbalance.service" in out: | ||
| 441 | + _, return_code = execute_command(["systemctl", "is-active", "--quiet", "irqbalance"]) | ||
| 442 | + if return_code == 0: | ||
| 443 | + logging.info("The irqbalance service is running and has been stopped.") | ||
| 444 | + _, return_code = execute_command(["systemctl", "stop", "irqbalance"]) | ||
| 445 | + if return_code != 0: | ||
| 446 | + logging.warning("The irqbalance service cannot be stopped.You need to manually stop it." | ||
| 447 | + "Otherwise, the interrupt-core binding cannot take effect.") | ||
| 448 | + | ||
| 449 | + for irq_id, target_cpu_list in zip(self.irq_id, self.cpu_list): | ||
| 450 | + affinity_file_path = f"/proc/irq/{irq_id}/smp_affinity" | ||
| 451 | + FileManager.check_directory_path_writeable(affinity_file_path) | ||
| 452 | + with open(affinity_file_path, "w") as f: | ||
| 453 | + f.write(self.cpu_to_mask(target_cpu_list)) | ||
| 454 | + logging.info(f"Bind the interrupt of IRQ-{irq_id} to CPU{target_cpu_list}") | ||
| 455 | + | ||
| 456 | + def execute_bind(self, pid: int, cpu_list_str: str, process_type:str, | ||
| 457 | + source_numa: str, cpu_node: Dict[int, int]) -> None: | ||
| 458 | + cmd = ["taskset", "-acp" if self.bind_sub_process else "-cp", cpu_list_str, str(pid)] | ||
| 459 | + logging.info(f"Bind the {self.process_name or 'target'} ({process_type}={pid}) to CPU{cpu_list_str}") | ||
| 460 | + _, return_code = execute_command(cmd) | ||
| 461 | + if return_code != 0: | ||
| 462 | + raise RuntimeError(f"Failed to execute the command: {' '.join(cmd)}") | ||
| 463 | + if self.mem_bind and shutil.which("numactl"): | ||
| 464 | + target_numa = cpu_node.get(int(cpu_list_str.split(",")[0])) | ||
| 465 | + cmd = ["migratepages", str(pid), source_numa, str(target_numa)] | ||
| 466 | + _, return_code = execute_command(cmd) | ||
| 467 | + if return_code != 0: | ||
| 468 | + logging.warning(f"Failed to execute the command: {' '.join(cmd)}") | ||
| 469 | + | ||
| 470 | + def bind(self, source_numa: str, cpu_allocer: CpuAlloc) -> None: | ||
| 471 | + process_type = "pid" if not self.is_thread else "tid" | ||
| 472 | + pid_list = self.find_threads() | ||
| 473 | + if not pid_list: | ||
| 474 | + return | ||
| 475 | + real_main_pid_list = self.get_real_main_pid_list(pid_list, cpu_allocer.device_info.main_pid_list) | ||
| 476 | + cpu_index = -1 | ||
| 477 | + for pid, ppid in pid_list: | ||
| 478 | + cpu_list_str = "" | ||
| 479 | + if len(self.cpu_list) == 1: | ||
| 480 | + cpu_list_str = ",".join(map(str, self.cpu_list[0])) | ||
| 481 | + self.execute_bind(pid, cpu_list_str, process_type, source_numa, cpu_allocer.cpu_node) | ||
| 482 | + continue | ||
| 483 | + if len(pid_list) == len(self.cpu_list): | ||
| 484 | + cpu_index += 1 | ||
| 485 | + cpu_list_str = ",".join(map(str, self.cpu_list[cpu_index])) | ||
| 486 | + for pids in real_main_pid_list: | ||
| 487 | + if pid in pids or ppid in pids: | ||
| 488 | + cpu_list_str = ",".join(map(str, self.cpu_list[real_main_pid_list.index(pids)])) | ||
| 489 | + if not cpu_list_str: | ||
| 490 | + logging.warning(f"Failed to bind process (pid: {pid}, ppid: {ppid}) to CPU {self.cpu_list}. Please " | ||
| 491 | + f"ensure that the number of processes to be bound is the same as the number of " | ||
| 492 | + f"cpu_list you have entered, or ensure that the Ngid field in the /proc/<pid>/status " | ||
| 493 | + f"file is not 0. It is recommended that the script be executed on the host machine.") | ||
| 494 | + continue | ||
| 495 | + self.execute_bind(pid, cpu_list_str, process_type, source_numa, cpu_allocer.cpu_node) | ||
| 496 | + | ||
| 497 | + | ||
| 498 | +def export_bind_config(cpu_alloc: CpuAlloc) -> Dict: | ||
| 499 | + main_cpu_list: List[str] = [] | ||
| 500 | + acl_cpu_list: List[str] = [] | ||
| 501 | + rel_cpu_list: List[str] = [] | ||
| 502 | + config = {"custom_bind": []} | ||
| 503 | + cpu_alloc.build_cpu_pools_all() | ||
| 504 | + cpu_alloc.allocate(MAIN_PROCESS_RANGE, ACL_THREAD_RANGE, RELEASE_THREAD_RANGE) | ||
| 505 | + main_pid_list = cpu_alloc.get_acl_main_threads() | ||
| 506 | + dev_pid_list, dev_cpu_list = cpu_alloc.dev_alloc() | ||
| 507 | + irq_id_list, irq_cpu_list = cpu_alloc.irq_alloc() | ||
| 508 | + | ||
| 509 | + for npu in sorted(cpu_alloc.device_info.running_npu_list): | ||
| 510 | + main_cpu_list.append(",".join(map(str, cpu_alloc.assign_main[npu]))) | ||
| 511 | + acl_cpu_list.append(",".join(map(str, cpu_alloc.assign_acl[npu]))) | ||
| 512 | + rel_cpu_list.append(",".join(map(str, cpu_alloc.assign_rel[npu]))) | ||
| 513 | + | ||
| 514 | + config["custom_bind"].append({ | ||
| 515 | + "pid": main_pid_list, | ||
| 516 | + "cpu_list": main_cpu_list, | ||
| 517 | + "bind_sub_process": True | ||
| 518 | + }) | ||
| 519 | + config["custom_bind"].append({ | ||
| 520 | + "process_name": "acl_thread", | ||
| 521 | + "cpu_list": acl_cpu_list, | ||
| 522 | + "is_thread": True | ||
| 523 | + }) | ||
| 524 | + config["custom_bind"].append({ | ||
| 525 | + "process_name": "release_thread", | ||
| 526 | + "cpu_list": rel_cpu_list, | ||
[codereview][提示] bind.__dict__直接打印对象全部属性可能泄露PID及系统信息;建议按需输出关键字段,降低日志噪声并避免敏感信息泄露。 ![]() ![]() | |||
| 527 | + "is_thread": True | ||
| 528 | + }) | ||
| 529 | + config["custom_bind"].append({ | ||
| 530 | + "pid": dev_pid_list, | ||
| 531 | + "cpu_list": dev_cpu_list | ||
| 532 | + }) | ||
| 533 | + config["custom_bind"].append({ | ||
| 534 | + "irq_id": irq_id_list, | ||
| 535 | + "cpu_list": irq_cpu_list, | ||
| 536 | + "is_irq": True | ||
| 537 | + }) | ||
| 538 | + | ||
| 539 | + return config | ||
| 540 | + | ||
| 541 | + | ||
| 542 | +def load_custom_bind(data: Dict) -> List[CustomBind]: | ||
| 543 | + binders = [] | ||
| 544 | + for item in data.get("custom_bind", []): | ||
| 545 | + binders.append(CustomBind( | ||
| 546 | + process_name=item.get("process_name", ""), | ||
| 547 | + cpu_list=item["cpu_list"], | ||
| 548 | + bind_sub_process=item.get("bind_sub_process", False), | ||
| 549 | + is_thread=item.get("is_thread", False), | ||
| 550 | + is_irq=item.get("is_irq", False), | ||
| 551 | + mem_bind=item.get("mem_bind", False), | ||
| 552 | + pid=item.get("pid", []), | ||
| 553 | + irq_id=item.get("irq_id", []) | ||
| 554 | + )) | ||
| 555 | + return binders | ||
| 556 | + | ||
| 557 | + | ||
| 558 | +def run(args: argparse.Namespace) -> None: | ||
| 559 | + loop_count = 0 | ||
| 560 | + cpu_allocer = CpuAlloc() | ||
| 561 | + cpu_allocer.build_cpu_pools_running() | ||
| 562 | + all_numa_nodes = ",".join(map(str, cpu_allocer.numa_to_cpu_map.keys())) | ||
| 563 | + if not args.config: | ||
| 564 | + default_json = export_bind_config(cpu_allocer) | ||
| 565 | + logging.info(f"No configuration file is detected." | ||
| 566 | + f"The default configuration is used for core binding: {default_json}") | ||
| 567 | + binder_list = load_custom_bind(default_json) | ||
| 568 | + else: | ||
| 569 | + if not os.path.exists(args.config): | ||
| 570 | + logging.error(f"The {args.config} file does not exist.Please check and try again.") | ||
| 571 | + return | ||
| 572 | + input_data = FileManager.read_json_file(args.config) | ||
| 573 | + binder_list = load_custom_bind(input_data) | ||
| 574 | + for bind in binder_list: | ||
| 575 | + loop_count += 1 | ||
| 576 | + logging.info(f"Start binding core round {loop_count}: {bind.__dict__}") | ||
| 577 | + try: | ||
| 578 | + if bind.is_irq: | ||
| 579 | + bind.irq_bind() | ||
| 580 | + else: | ||
| 581 | + if not bind.pid and not bind.process_name: | ||
| 582 | + logging.error(f"No input bound object. One of 'pid, process_name, irq_id' are required.") | ||
| 583 | + continue | ||
| 584 | + bind.bind(all_numa_nodes, cpu_allocer) | ||
| 585 | + except RuntimeError as e: | ||
| 586 | + logging.error(f"Error occurred while binding: {e}") | ||
| 587 | + logging.info(f"===== Round {loop_count} of core binding has ended =====") | ||
| @@ -0,0 +1,35 @@ | |||
| 1 | +import argparse | ||
| 2 | +import logging | ||
| 3 | + | ||
| 4 | +from misc.host_analyzer.cpu_binder import cpu_binder | ||
| 5 | + | ||
| 6 | +LOG_MAP = { | ||
| 7 | + 0: logging.DEBUG, | ||
| 8 | + 1: logging.INFO, | ||
| 9 | + 2: logging.WARNING, | ||
| 10 | + 3: logging.ERROR | ||
| 11 | +} | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +def main(): | ||
| 15 | + parser = argparse.ArgumentParser(description="总入口:支持多个子命令") | ||
| 16 | + subparsers = parser.add_subparsers(dest="command", required=True) | ||
| 17 | + | ||
| 18 | + bind_parser = subparsers.add_parser("bind", help="执行绑核逻辑") | ||
| 19 | + bind_parser.add_argument("-c", "--config", type=str, help="绑定进线程配置文件路径,自定义绑定推理进线程时使用") | ||
| 20 | + bind_parser.add_argument("-l", "--log-level", type=int, | ||
| 21 | + default=1, choices=[0, 1, 2, 3], help="日志级别") | ||
| 22 | + | ||
| 23 | + args = parser.parse_args() | ||
| 24 | + | ||
| 25 | + log_level = LOG_MAP.get(args.log_level, logging.INFO) | ||
| 26 | + logging.basicConfig(level=log_level, format='[%(asctime)s] [%(levelname)s]:%(message)s') | ||
| 27 | + | ||
| 28 | + if args.command == "bind": | ||
| 29 | + cpu_binder.run(args) | ||
| 30 | + else: | ||
| 31 | + parser.print_help() | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +if __name__ == "__main__": | ||
| 35 | + main() | ||
| @@ -0,0 +1,9 @@ | |||
| 1 | +set(MISC_CODE ${TOP_DIR}/abl/msprof/misc) | ||
| 2 | +set(MISC_DB ${TOP_DIR}/llt/abl/msprof/misc_python/ut/testcase) | ||
| 3 | + | ||
| 4 | +run_python_llt_test( | ||
| 5 | + TARGET msprof_python_utest | ||
| 6 | + SRC_FILES_DIR ${TOP_DIR}/abl/msprof/misc/ | ||
| 7 | + TEST_FILES_DIR ${TOP_DIR}/llt/abl/msprof/misc_python/ut/testcase | ||
| 8 | + EXPORT_PYTHONPATH "${MISC_CODE}:${MISC_DB}" | ||
| 9 | +) | ||
The file is empty
| @@ -0,0 +1,204 @@ | |||
| 1 | +import unittest | ||
| 2 | +from unittest.mock import patch | ||
| 3 | + | ||
| 4 | +from host_analyzer.cpu_binder.cpu_binder import CpuAlloc, DeviceInfo, CustomBind, expand_cpu_list | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +class TestDeviceInfo(unittest.TestCase): | ||
| 8 | + | ||
| 9 | + | ||
| 10 | + def setUp(self, mock_execute_command): | ||
| 11 | + mock_execute_command.side_effect = [ | ||
| 12 | + ("NPU ID Chip ID Chip Logic ID Chip Name\n0 0 0 Ascend\n0 1 - Mcu\n1 0 1 Ascend", 0), | ||
| 13 | + ("| NPU Chip | Process id |\n| 0 0 | 1234 | vllm | 56000 |\n| 1 0 | 1235 | vllm | 56000 |", 0), | ||
| 14 | + ("", 0) | ||
| 15 | + ] | ||
| 16 | + self.device_info = DeviceInfo() | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + def test_get_npu_map_info(self, mock_execute_command): | ||
| 20 | + execute_result_list = [ | ||
| 21 | + ("NPU ID Chip ID Chip Logic ID Chip Phy-ID Chip Name\n0 0 0 0 Ascend\n0 1 1 1 Ascend\n0 2 - - Mcu", 0), | ||
| 22 | + ("NPU ID Chip ID Chip Logic ID Chip Name\n8 0 0 Ascend\n8 1 - Mcu\n9 0 1 Ascend", 0) | ||
| 23 | + ] | ||
| 24 | + result_list = [{'0': {'0': '0', '1': '1'}}, {'8': {'0': '0'}, '9': {'0': '1'}}] | ||
| 25 | + for result in execute_result_list: | ||
| 26 | + mock_execute_command.return_value = result | ||
| 27 | + npu_map_info = self.device_info.get_npu_map_info() | ||
| 28 | + expected = result_list.pop(0) | ||
| 29 | + self.assertEqual(npu_map_info, expected) | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + def test_get_running_npus(self, mock_execute_command): | ||
| 33 | + mock_execute_command.side_effect = [ | ||
| 34 | + ("| NPU Chip | Process id |\n| 0 1 | 1236 | vllm | 56000 |", 0), | ||
| 35 | + ("", 0), | ||
| 36 | + ("| NPU Chip | Process id |\n| 1 0 | 1236 | vllm | 56000 |", 0) | ||
| 37 | + ] | ||
| 38 | + expected_result_list = [[], [], [1]] | ||
| 39 | + for expected_result in expected_result_list: | ||
| 40 | + running_npu_list = self.device_info.get_running_npus() | ||
| 41 | + self.assertEqual(running_npu_list, expected_result) | ||
| 42 | + | ||
| 43 | + | ||
| 44 | + def test_parse_topo_affinity(self, mock_execute_command): | ||
| 45 | + mock_execute_command.side_effect = [ | ||
| 46 | + ("NPU0 X HCCS HCCS HCCS HCCS HCCS HCCS HCCS 0-3", 0), | ||
| 47 | + ("GPU0 X HCCS HCCS HCCS HCCS HCCS HCCS HCCS 0-3", 0) | ||
| 48 | + ] | ||
| 49 | + expected_result_list = [{0: [0, 1, 2, 3]}, {}] | ||
| 50 | + for expected_result in expected_result_list: | ||
| 51 | + affinity = self.device_info.parse_topo_affinity() | ||
| 52 | + self.assertEqual(affinity, expected_result) | ||
| 53 | + | ||
| 54 | + def test_expand_cpu_list(self): | ||
| 55 | + result = expand_cpu_list("0-2, 4, 6-8") | ||
| 56 | + self.assertEqual(result, [0, 1, 2, 4, 6, 7, 8]) | ||
| 57 | + with self.assertRaises(RuntimeError): | ||
| 58 | + expand_cpu_list("0/1") | ||
| 59 | + | ||
| 60 | + | ||
| 61 | +class TestCpuAlloc(unittest.TestCase): | ||
| 62 | + | ||
| 63 | + | ||
| 64 | + def setUp(self, mock_execute_command): | ||
| 65 | + mock_execute_command.side_effect = [ | ||
| 66 | + ("NPU ID Chip ID Chip Logic ID Chip Name\n0 0 0 Ascend\n0 1 - Mcu\n1 0 1 Ascend", 0), | ||
| 67 | + ("| NPU Chip | Process id |\n| 0 0 | 1234 | vllm | 56000 |\n| 1 0 | 1235 | vllm | 56000 |", 0), | ||
| 68 | + ("", 0) | ||
| 69 | + ] | ||
| 70 | + self.cpu_alloc = CpuAlloc() | ||
| 71 | + | ||
| 72 | + def test_average_distribute(self): | ||
| 73 | + npu_cpu_pool = { | ||
| 74 | + 0: [10, 11, 12, 13], | ||
| 75 | + 1: [10, 11, 12, 13] | ||
| 76 | + } | ||
| 77 | + groups = {"[10, 11, 12, 13]": [0, 1]} | ||
| 78 | + result = self.cpu_alloc.average_distribute(groups, npu_cpu_pool) | ||
| 79 | + self.assertEqual(result, {0: [10, 11], 1: [12, 13]}) | ||
| 80 | + npu_cpu_pool = { | ||
| 81 | + 0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], | ||
| 82 | + 1: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], | ||
| 83 | + 2: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] | ||
| 84 | + } | ||
| 85 | + groups = {"[0, 1, 2, 3, 4, 5]": [0, 1, 2]} | ||
| 86 | + result = self.cpu_alloc.average_distribute(groups, npu_cpu_pool) | ||
| 87 | + self.assertEqual(result, { | ||
| 88 | + 0: [0, 1, 2, 3], | ||
| 89 | + 1: [4, 5, 6, 7], | ||
| 90 | + 2: [8, 9, 10, 11, 12, 13] | ||
| 91 | + }) | ||
| 92 | + | ||
| 93 | + def test_extend_numa(self): | ||
| 94 | + result = self.cpu_alloc.extend_numa([]) | ||
| 95 | + self.assertEqual(result, []) | ||
| 96 | + self.cpu_alloc.cpu_node = {0: 0, 1: 0, 2: 1, 3: 1} | ||
| 97 | + self.cpu_alloc.numa_to_cpu_map = {0: [0, 1], 1: [2, 3]} | ||
| 98 | + self.cpu_alloc.device_info.allowed_cpus = [0, 1, 2, 3] | ||
| 99 | + result = self.cpu_alloc.extend_numa([0, 1]) | ||
| 100 | + self.assertEqual(result, [0, 1, 2, 3]) | ||
| 101 | + self.cpu_alloc.device_info.allowed_cpus = [0, 1, 3] | ||
| 102 | + result = self.cpu_alloc.extend_numa([0, 1]) | ||
| 103 | + self.assertEqual(result, [0, 1, 3]) | ||
| 104 | + | ||
| 105 | + | ||
| 106 | + def test_build_cpu_node_map(self, mock_execute_command): | ||
| 107 | + mock_execute_command.return_value = ("", 0) | ||
| 108 | + with self.assertRaises(RuntimeError): | ||
| 109 | + self.cpu_alloc.build_cpu_node_map() | ||
| 110 | + mock_execute_command.return_value = ("0 0\n1 1\n2 0\n3 1", 0) | ||
| 111 | + self.cpu_alloc.build_cpu_node_map() | ||
| 112 | + expected_cpu_node = {0: 0, 1: 1, 2: 0, 3: 1} | ||
| 113 | + expected_numa_to_cpu_map = {0: [0, 2], 1: [1, 3]} | ||
| 114 | + self.assertEqual(self.cpu_alloc.cpu_node, expected_cpu_node) | ||
| 115 | + self.assertEqual(self.cpu_alloc.numa_to_cpu_map, | ||
| 116 | + expected_numa_to_cpu_map) | ||
| 117 | + | ||
| 118 | + | ||
| 119 | + def test_handle_no_affinity(self, mock_execute_command): | ||
| 120 | + mock_execute_command.side_effect = [("0 0\n1 1", 0), ("0 0\n1 1", 0)] | ||
| 121 | + self.cpu_alloc.device_info.running_npu_list = [0, 1] | ||
| 122 | + self.cpu_alloc.device_info.allowed_cpus = [0, 1, 2, 3] | ||
| 123 | + self.cpu_alloc.numa_to_cpu_map = {0: [4, 5], 1: [6, 7]} | ||
| 124 | + self.cpu_alloc.handle_no_affinity() | ||
| 125 | + self.assertEqual(self.cpu_alloc.npu_cpu_pool, {}) | ||
| 126 | + self.cpu_alloc.numa_to_cpu_map = {0: [0, 1, 2], 1: [3, 4, 5]} | ||
| 127 | + self.cpu_alloc.handle_no_affinity() | ||
| 128 | + self.assertEqual(self.cpu_alloc.npu_cpu_pool, {0: [0, 1, 2], 1: [3]}) | ||
| 129 | + | ||
| 130 | + def test_allocate(self): | ||
| 131 | + self.cpu_alloc.device_info.running_npu_list = [0] | ||
| 132 | + self.cpu_alloc.npu_cpu_pool = {0: [0, 1, 2, 3, 4, 5, 6]} | ||
| 133 | + self.cpu_alloc.allocate(2, 1, 1) | ||
| 134 | + self.assertEqual(self.cpu_alloc.assign_main[0], [3, 4]) | ||
| 135 | + self.assertEqual(self.cpu_alloc.assign_acl[0], [5]) | ||
| 136 | + self.assertEqual(self.cpu_alloc.assign_rel[0], [6]) | ||
| 137 | + self.cpu_alloc.npu_cpu_pool = {0: [0, 1]} | ||
| 138 | + with self.assertRaises(RuntimeError): | ||
| 139 | + self.cpu_alloc.allocate(6, 1, 1) | ||
| 140 | + | ||
| 141 | + | ||
| 142 | +class TestCustomBind(unittest.TestCase): | ||
| 143 | + | ||
| 144 | + def setUp(self): | ||
| 145 | + self.binder = CustomBind(process_name="test_process", cpu_list=["0-3"], bind_sub_process=True) | ||
| 146 | + | ||
| 147 | + def test_cpu_to_mask(self): | ||
| 148 | + cpu_list = [0, 1, 30, 39] | ||
| 149 | + mask_str = self.binder.cpu_to_mask(cpu_list) | ||
| 150 | + self.assertEqual(mask_str, "00000080,40000003") | ||
| 151 | + cpu_list = [0] | ||
| 152 | + mask_str = self.binder.cpu_to_mask(cpu_list) | ||
| 153 | + self.assertEqual(mask_str, "00000001") | ||
| 154 | + | ||
| 155 | + | ||
| 156 | + | ||
| 157 | + def test_get_main_pid_from_docker(self, mock_execute_command, mock_exists): | ||
| 158 | + mock_execute_command.return_value = ("Ngid:\t123", 0) | ||
| 159 | + pid = self.binder.get_main_pid_from_docker(1000) | ||
| 160 | + self.assertEqual(pid, 123) | ||
| 161 | + mock_execute_command.return_value = ("", 1) | ||
| 162 | + pid = self.binder.get_main_pid_from_docker(1000) | ||
| 163 | + self.assertEqual(pid, 0) | ||
| 164 | + | ||
| 165 | + | ||
| 166 | + def test_get_real_main_pid_list(self, mock_get_main_pid_from_docker): | ||
| 167 | + mock_get_main_pid_from_docker.return_value = 0 | ||
| 168 | + pid_list = [(1111, 2111), (2112, 3112), (8888, 9999)] | ||
| 169 | + main_pid_list = [[2111], [2112]] | ||
| 170 | + real_main_pid_list = self.binder.get_real_main_pid_list(pid_list, main_pid_list) | ||
| 171 | + self.assertEqual(real_main_pid_list, main_pid_list) | ||
| 172 | + pid_list = [(8888, 9999)] | ||
| 173 | + real_main_pid_list = self.binder.get_real_main_pid_list(pid_list, main_pid_list) | ||
| 174 | + self.assertEqual(real_main_pid_list, []) | ||
| 175 | + mock_get_main_pid_from_docker.return_value = 2111 | ||
| 176 | + pid_list = [(8888, 9999)] | ||
| 177 | + real_main_pid_list = self.binder.get_real_main_pid_list(pid_list, main_pid_list) | ||
| 178 | + self.assertEqual(real_main_pid_list, [[8888]]) | ||
| 179 | + | ||
| 180 | + | ||
| 181 | + | ||
| 182 | + def test_find_threads(self, mock_execute_command, mock_check_output): | ||
| 183 | + mock_execute_command.return_value = ("", 0) | ||
| 184 | + with self.assertRaises(RuntimeError): | ||
| 185 | + self.binder.find_threads() | ||
| 186 | + mock_execute_command.return_value = ("1113 1114 test_process\n", 0) | ||
| 187 | + pid_list = self.binder.find_threads() | ||
| 188 | + self.assertEqual(pid_list, [(1113, 1114)]) | ||
| 189 | + self.binder.pid = [1111] | ||
| 190 | + pid_list = self.binder.find_threads() | ||
| 191 | + self.assertEqual(pid_list, [(1111, 1112)]) | ||
| 192 | + | ||
| 193 | + | ||
| 194 | + | ||
| 195 | + def test_execute_bind(self, mock_execute_command, mock_which): | ||
| 196 | + mock_execute_command.return_value = ("", 1) | ||
| 197 | + with self.assertRaises(RuntimeError): | ||
| 198 | + self.binder.execute_bind(123, "0, 1", "pid", "0", {0: 0}) | ||
| 199 | + mock_execute_command.return_value = ("", 0) | ||
| 200 | + self.binder.execute_bind(123, "0, 1", "pid", "0", {0: 0}) | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +if __name__ == '__main__': | ||
| 204 | + unittest.main() | ||


【review】【资料】此处的配置和说明,难以形象理解为何如此配置,建议增加cpu、numa的分布介绍图,更好的理解配置的合理性,建议后续完善。