已合并
新增 Host 诊断调优工具 #109
rozwel_dx创建于 3月5日
新增 Host 诊断调优工具 #109
已合并
rozwel_dx创建于 3月5日
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/>&bull; true:表示绑核后将使用内存迁移至对应NUMA节点<br/>&bull; false:表示绑核后不需要将使用内存迁移至对应NUMA节点<br/>默认值为false |
45+| is_thread | 可选 | 选择绑定的ID是进程ID还是线程ID,类型为bool,取值为:<br/>&bull; true:表示绑定线程ID<br/>&bull; false:表示绑定进程ID<br/>默认值为false |
46+| is_irq | 可选 | 选择绑定的是否为硬件中断,类型为bool,取值为:<br/>&bull; true:表示绑定的是硬件中断<br/>&bull; false:表示绑定的不是硬件中断<br/>默认值为false |
47+| irq_id | 可选 | 需要绑定的硬件中断的中断号,类型为List[int] |
48+| bind_sub_process | 可选 | 选择绑核时是否需要绑定进程ID或线程ID的所有子线程,类型为bool,取值为:<br/>&bull; true:表示绑核时需要同时绑定进程ID或线程ID的所有子线程<br/>&bull; false:表示绑核时不需要绑定进程ID或线程ID的所有子线程<br/>默认值为false |
49+ 
50+#### json配置文件示例
C
Cchenhao_12093月25日

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

likedislike
rozwel_dx
rozwel_dx
3月25日 评论:
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:
xieanran
xieanranxieanran6月1日

[codereview][一般] execute_command未检查communicate(timeout=1000)超时异常;建议捕获subprocess.TimeoutExpired并主动终止进程,避免子进程泄漏。

likedislike
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+ @staticmethod
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+ @staticmethod
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+ @staticmethod
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():
xieanran
xieanranxieanran6月1日

[codereview][一般] parse_topo_affinity使用自增chip_logic_id作为NPU标识,强依赖npu-smi输出顺序;建议解析真实逻辑ID字段,避免设备顺序变化导致CPU亲和性映射错误。

likedislike
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]
xieanran
xieanranxieanran6月1日

[codereview][严重] parts[0].split()[:2]默认认为字段一定包含NPU ID和Chip ID;建议增加长度校验,避免输出格式变化触发ValueError

likedislike
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+ @staticmethod
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+ @staticmethod
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+ @staticmethod
355+ def cpu_to_mask(cpus: List[int]) -> str:
xieanran
xieanranxieanran6月1日

[codereview][严重] cpu_to_mask在cpus为空时执行max(groups.keys())会抛出异常;建议增加空列表保护逻辑。

likedislike
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+ @staticmethod
369+ def get_main_pid_from_docker(pid: int) -> int:
xieanran
xieanranxieanran6月1日

[codereview][一般] get_main_pid_from_docker通过grep解析文件内容额外创建进程开销较大;建议直接读取/proc//status文件解析,提高执行效率。

likedislike
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:
xieanran
xieanranxieanran6月1日

[codereview][严重] real_main_pid_list.index(pids)位于循环内部,每次都会线性查找;建议提前构建PID到CPU列表索引映射,避免O(n²)性能退化。

likedislike
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,
xieanran
xieanranxieanran6月1日

[codereview][提示] bind.__dict__直接打印对象全部属性可能泄露PID及系统信息;建议按需输出关键字段,降低日志噪声并避免敏感信息泄露。

likedislike
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+)
@@ -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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('os.path.exists', return_value=True)
156+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.CustomBind.get_main_pid_from_docker')
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+ @patch('subprocess.check_output', return_value="1112")
181+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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+ @patch('shutil.which', return_value=True)
194+ @patch('misc.host_analyzer.cpu_binder.cpu_binder.execute_command')
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()