create a-sched-MindStudio-Boost-AtomGit

已合并

create a-sched #1

baij创建于 6月3日从baij77/Ascend-msboost： master 合入到 Ascend/msboost： master

create a-sched #1

已合并

baij创建于 6月3日

从baij77/Ascend-msboost master合入到Ascend/msboost master

共 15 个文件变更+3352-0

A.gitignore+14-0

@@ -0,0 +1,14 @@
	1	+# ------------------------------------------------------------
	2	+# Build artifacts
	3	+# ------------------------------------------------------------
	4	+build/
	5	+dist/
	6	+
	7	+# ------------------------------------------------------------
	8	+# Python
	9	+# ------------------------------------------------------------
	10	+__pycache__/
	11	+*.egg
	12	+*.egg-info/
	13	+.mypy_cache/
	14	+.pytest_cache/

ALICENSE+127-0

@@ -0,0 +1,127 @@
	1	+ 木兰宽松许可证, 第2版
	2	+
	3	+ 木兰宽松许可证，第2版
	4	+ 2020年1月 http://license.coscl.org.cn/MulanPSL2
	5	+
	6	+
	7	+ 您对“软件”的复制、使用、修改及分发受木兰宽松许可证，第2版（“本许可证”）的如下条款的约束：
	8	+
	9	+ 0. 定义
	10	+
	11	+ “软件”是指由“贡献”构成的许可在“本许可证”下的程序和相关文档的集合。
	12	+
	13	+ “贡献”是指由任一“贡献者”许可在“本许可证”下的受版权法保护的作品。
	14	+
	15	+ “贡献者”是指将受版权法保护的作品许可在“本许可证”下的自然人或“法人实体”。
	16	+
	17	+ “法人实体”是指提交贡献的机构及其“关联实体”。
	18	+
	19	+ “关联实体”是指，对“本许可证”下的行为方而言，控制、受控制或与其共同受控制的机构，此处的控制是指有受控方或共同受控方至少50%直接或间接的投票权、资金或其他有价证券。
	20	+
	21	+ 1. 授予版权许可
	22	+
	23	+ 每个“贡献者”根据“本许可证”授予您永久性的、全球性的、免费的、非独占的、不可撤销的版权许可，您可以复制、使用、修改、分发其“贡献”，不论修改与否。
	24	+
	25	+ 2. 授予专利许可
	26	+
	27	+ 每个“贡献者”根据“本许可证”授予您永久性的、全球性的、免费的、非独占的、不可撤销的（根据本条规定撤销除外）专利许可，供您制造、委托制造、使用、许诺销售、销售、进口其“贡献”或以其他方式转移其“贡献”。前述专利许可仅限于“贡献者”现在或将来拥有或控制的其“贡献”本身或其“贡献”与许可“贡献”时的“软件”结合而将必然会侵犯的专利权利要求，不包括对“贡献”的修改或包含“贡献”的其他结合。如果您或您的“关联实体”直接或间接地，就“软件”或其中的“贡献”对任何人发起专利侵权诉讼（包括反诉或交叉诉讼）或其他专利维权行动，指控其侵犯专利权，则“本许可证”授予您对“软件”的专利许可自您提起诉讼或发起维权行动之日终止。
	28	+
	29	+ 3. 无商标许可
	30	+
	31	+ “本许可证”不提供对“贡献者”的商品名称、商标、服务标志或产品名称的商标许可，但您为满足第4条规定的声明义务而必须使用除外。
	32	+
	33	+ 4. 分发限制
	34	+
	35	+ 您可以在任何媒介中将“软件”以源程序形式或可执行形式重新分发，不论修改与否，但您必须向接收者提供“本许可证”的副本，并保留“软件”中的版权、商标、专利及免责声明。
	36	+
	37	+ 5. 免责声明与责任限制
	38	+
	39	+ “软件”及其中的“贡献”在提供时不带任何明示或默示的担保。在任何情况下，“贡献者”或版权所有者不对任何人因使用“软件”或其中的“贡献”而引发的任何直接或间接损失承担责任，不论因何种原因导致或者基于何种法律理论，即使其曾被建议有此种损失的可能性。
	40	+
	41	+ 6. 语言
	42	+ “本许可证”以中英文双语表述，中英文版本具有同等法律效力。如果中英文版本存在任何冲突不一致，以中文版为准。
	43	+
	44	+ 条款结束
	45	+
	46	+ 如何将木兰宽松许可证，第2版，应用到您的软件
	47	+
	48	+ 如果您希望将木兰宽松许可证，第2版，应用到您的新软件，为了方便接收者查阅，建议您完成如下三步：
	49	+
	50	+ 1，请您补充如下声明中的空白，包括软件名、软件的首次发表年份以及您作为版权人的名字；
	51	+
	52	+ 2，请您在软件包的一级目录下创建以“LICENSE”为名的文件，将整个许可证文本放入该文件中；
	53	+
	54	+ 3，请将如下声明文本放入每个源文件的头部注释中。
	55	+
	56	+ Copyright (c) [Year] [name of copyright holder]
	57	+ [Software Name] is licensed under Mulan PSL v2.
	58	+ You can use this software according to the terms and conditions of the Mulan PSL v2.
	59	+ You may obtain a copy of Mulan PSL v2 at:
	60	+ http://license.coscl.org.cn/MulanPSL2
	61	+ THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	62	+ See the Mulan PSL v2 for more details.
	63	+
	64	+
	65	+ Mulan Permissive Software License，Version 2
	66	+
	67	+ Mulan Permissive Software License，Version 2 (Mulan PSL v2)
	68	+ January 2020 http://license.coscl.org.cn/MulanPSL2
	69	+
	70	+ Your reproduction, use, modification and distribution of the Software shall be subject to Mulan PSL v2 (this License) with the following terms and conditions:
	71	+
	72	+ 0. Definition
	73	+
	74	+ Software means the program and related documents which are licensed under this License and comprise all Contribution(s).
	75	+
	76	+ Contribution means the copyrightable work licensed by a particular Contributor under this License.
	77	+
	78	+ Contributor means the Individual or Legal Entity who licenses its copyrightable work under this License.
	79	+
	80	+ Legal Entity means the entity making a Contribution and all its Affiliates.
	81	+
	82	+ Affiliates means entities that control, are controlled by, or are under common control with the acting entity under this License, ‘control’ means direct or indirect ownership of at least fifty percent (50%) of the voting power, capital or other securities of controlled or commonly controlled entity.
	83	+
	84	+ 1. Grant of Copyright License
	85	+
	86	+ Subject to the terms and conditions of this License, each Contributor hereby grants to you a perpetual, worldwide, royalty-free, non-exclusive, irrevocable copyright license to reproduce, use, modify, or distribute its Contribution, with modification or not.
	87	+
	88	+ 2. Grant of Patent License
	89	+
	90	+ Subject to the terms and conditions of this License, each Contributor hereby grants to you a perpetual, worldwide, royalty-free, non-exclusive, irrevocable (except for revocation under this Section) patent license to make, have made, use, offer for sale, sell, import or otherwise transfer its Contribution, where such patent license is only limited to the patent claims owned or controlled by such Contributor now or in future which will be necessarily infringed by its Contribution alone, or by combination of the Contribution with the Software to which the Contribution was contributed. The patent license shall not apply to any modification of the Contribution, and any other combination which includes the Contribution. If you or your Affiliates directly or indirectly institute patent litigation (including a cross claim or counterclaim in a litigation) or other patent enforcement activities against any individual or entity by alleging that the Software or any Contribution in it infringes patents, then any patent license granted to you under this License for the Software shall terminate as of the date such litigation or activity is filed or taken.
	91	+
	92	+ 3. No Trademark License
	93	+
	94	+ No trademark license is granted to use the trade names, trademarks, service marks, or product names of Contributor, except as required to fulfill notice requirements in Section 4.
	95	+
	96	+ 4. Distribution Restriction
	97	+
	98	+ You may distribute the Software in any medium with or without modification, whether in source or executable forms, provided that you provide recipients with a copy of this License and retain copyright, patent, trademark and disclaimer statements in the Software.
	99	+
	100	+ 5. Disclaimer of Warranty and Limitation of Liability
	101	+
	102	+ THE SOFTWARE AND CONTRIBUTION IN IT ARE PROVIDED WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED. IN NO EVENT SHALL ANY CONTRIBUTOR OR COPYRIGHT HOLDER BE LIABLE TO YOU FOR ANY DAMAGES, INCLUDING, BUT NOT LIMITED TO ANY DIRECT, OR INDIRECT, SPECIAL OR CONSEQUENTIAL DAMAGES ARISING FROM YOUR USE OR INABILITY TO USE THE SOFTWARE OR THE CONTRIBUTION IN IT, NO MATTER HOW IT’S CAUSED OR BASED ON WHICH LEGAL THEORY, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
	103	+
	104	+ 6. Language
	105	+
	106	+ THIS LICENSE IS WRITTEN IN BOTH CHINESE AND ENGLISH, AND THE CHINESE VERSION AND ENGLISH VERSION SHALL HAVE THE SAME LEGAL EFFECT. IN THE CASE OF DIVERGENCE BETWEEN THE CHINESE AND ENGLISH VERSIONS, THE CHINESE VERSION SHALL PREVAIL.
	107	+
	108	+ END OF THE TERMS AND CONDITIONS
	109	+
	110	+ How to Apply the Mulan Permissive Software License，Version 2 (Mulan PSL v2) to Your Software
	111	+
	112	+ To apply the Mulan PSL v2 to your work, for easy identification by recipients, you are suggested to complete following three steps:
	113	+
	114	+ i Fill in the blanks in following statement, including insert your software name, the year of the first publication of your software, and your name identified as the copyright owner;
	115	+
	116	+ ii Create a file named “LICENSE” which contains the whole context of this License in the first directory of your software package;
	117	+
	118	+ iii Attach the statement to the appropriate annotated syntax at the beginning of each source file.
	119	+
	120	+
	121	+ Copyright (c) [Year] [name of copyright holder]
	122	+ [Software Name] is licensed under Mulan PSL v2.
	123	+ You can use this software according to the terms and conditions of the Mulan PSL v2.
	124	+ You may obtain a copy of Mulan PSL v2 at:
	125	+ http://license.coscl.org.cn/MulanPSL2
	126	+ THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	127	+ See the Mulan PSL v2 for more details.

Aa-sched/a_sched/__init__.py+30-0

@@ -0,0 +1,30 @@
	1	+from a_sched.api import (
	2	+ group_create,
	3	+ group_destroy,
	4	+ group_add_process,
	5	+ group_remove_process,
	6	+ group_add_thread,
	7	+ group_remove_thread,
	8	+ thread_set_high_priority,
	9	+ process_bind_npu,
	10	+ run_affinity,
	11	+ print_affinity,
	12	+ restore_affinity,
	13	+ set_exclude_cpu,
	14	+)
	15	+
	16	+__version__ = "0.1.0"
	17	+__all__ = [
	18	+ "group_create",
	19	+ "group_destroy",
	20	+ "group_add_process",
	21	+ "group_remove_process",
	22	+ "group_add_thread",
	23	+ "group_remove_thread",
	24	+ "thread_set_high_priority",
	25	+ "process_bind_npu",
	26	+ "run_affinity",
	27	+ "print_affinity",
	28	+ "restore_affinity",
	29	+ "set_exclude_cpu",
	30	+]

Aa-sched/a_sched/affinity_domain.py+352-0

@@ -0,0 +1,352 @@
	1	+from __future__ import annotations
	2	+from abc import ABC, abstractmethod
	3	+from enum import Enum
	4	+from collections import defaultdict
	5	+from dataclasses import dataclass, field
	6	+import os
	7	+import re
	8	+
	9	+from a_sched.utils import CPUMask
	10	+from a_sched.config import AffinityConfig
	11	+import a_sched.utils as utils
	12	+
	13	+
	14	+class DomainLevel(Enum):
	15	+ """亲和域层级（模仿调度器层级）"""
	16	+
	17	+ SOCKET = 0 # 最高层：物理CPU插槽
	18	+ NUMA = 1 # 中间层：NUMA内存节点
	19	+ CLUSTER = 2 # 最底层：CPU性能集群
	20	+ THREAD = 3 # 额外：线程级（可选）
	21	+
	22	+
	23	+class AffinityDomain(ABC):
	24	+ """
	25	+ 亲和域抽象基类
	26	+ 模仿Linux调度器sched_domain的设计
	27	+ """
	28	+
	29	+ def __init__(self, level: DomainLevel, domain_id: int):
	30	+ self.level = level
	31	+ self.domain_id = domain_id
	32	+
	33	+ # 核心拓扑属性（模仿sched_domain）
	34	+ self.cpus = CPUMask() # 该域包含的CPU位图
	35	+ self.parent: AffinityDomain \| None = None # 父域
	36	+ self.children: list[AffinityDomain] = [] # 子域
	37	+
	38	+ # 层级名称
	39	+ self._level_names = {
	40	+ DomainLevel.SOCKET: "SOCKET",
	41	+ DomainLevel.NUMA: "NUMA",
	42	+ DomainLevel.CLUSTER: "CLUSTER",
	43	+ DomainLevel.THREAD: "THREAD",
	44	+ }
	45	+
	46	+ @abstractmethod
	47	+ def detect_from_hardware(self) -> bool:
	48	+ """从硬件检测域信息"""
	49	+ pass
	50	+
	51	+ def add_child(self, child: AffinityDomain):
	52	+ child.parent = self
	53	+ self.children.append(child)
	54	+
	55	+ def get_all_children(self) -> list[AffinityDomain]:
	56	+ return self.children
	57	+
	58	+ def get_all_children_id(self) -> list[int]:
	59	+ children_id = []
	60	+ for child in self.children:
	61	+ children_id.append(child.domain_id)
	62	+ return children_id
	63	+
	64	+ def get_children_num(self) -> int:
	65	+ return len(self.children)
	66	+
	67	+ def __str__(self) -> str:
	68	+ level_name = self._level_names.get(self.level, "UNKNOWN")
	69	+ return f"{level_name}[{self.domain_id}]: CPUs={self.cpus}"
	70	+
	71	+
	72	+class SocketDomain(AffinityDomain):
	73	+ """Socket域 - 物理CPU插槽（最高层）"""
	74	+
	75	+ def __init__(self, socket_id: int):
	76	+ super().__init__(level=DomainLevel.SOCKET, domain_id=socket_id)
	77	+
	78	+ def detect_from_hardware(self) -> bool:
	79	+ return True
	80	+
	81	+
	82	+class NumaDomain(AffinityDomain):
	83	+ """NUMA域 - 内存节点（中间层）"""
	84	+
	85	+ def __init__(self, node_id: int):
	86	+ super().__init__(level=DomainLevel.NUMA, domain_id=node_id)
	87	+
	88	+ def detect_from_hardware(self) -> bool:
	89	+ return True
	90	+
	91	+
	92	+class ClusterDomain(AffinityDomain):
	93	+ """Cluster域 - CPU性能集群（最底层）"""
	94	+
	95	+ def __init__(self, cluster_id: int):
	96	+ super().__init__(level=DomainLevel.CLUSTER, domain_id=cluster_id)
	97	+
	98	+ def detect_from_hardware(self) -> bool:
	99	+ return True
	100	+
	101	+
	102	+@dataclass
	103	+class CpuCore:
	104	+ physical_package_id: int = -1
	105	+ cluster_id: int = -1
	106	+
	107	+
	108	+@dataclass
	109	+class NumaNode:
	110	+ cpulist: list[int] = field(default_factory=list)
	111	+
	112	+
	113	+class AffinityDomainBuilder:
	114	+ """
	115	+ 亲和域构建
	116	+ """
	117	+
	118	+ def __init__(self, config: AffinityConfig) -> None:
	119	+ self.socket_domains: list[SocketDomain] = []
	120	+ self.numa_domains: list[NumaDomain] = []
	121	+ self.cluster_domains: list[ClusterDomain] = []
	122	+
	123	+ self._exclude_cpus = config.exclude_cpus
	124	+ self._container_cpus = utils.get_allowed_cpu_list()
	125	+
	126	+ self._cpu_dict: dict[int, CpuCore] = {}
	127	+ self._numa_dict: dict[int, NumaNode] = {}
	128	+
	129	+ self._socket_to_cpus: dict[int, list[int]] = defaultdict(list)
	130	+ self._numa_to_cpus: dict[int, list[int]] = defaultdict(list)
	131	+ self._cluster_to_cpus: dict[int, list[int]] = defaultdict(list)
	132	+ self._socket_to_numas: dict[int, set[int]] = defaultdict(set)
	133	+ self._numa_to_clusters: dict[int, set[int]] = defaultdict(set)
	134	+
	135	+ def build_affinity_domain(self) -> bool:
	136	+ try:
	137	+ self._read_cpu()
	138	+ self._read_numa()
	139	+ self._build_topo()
	140	+ self._build_socket_domains()
	141	+ for socket in self.socket_domains:
	142	+ self._build_numa_domains_for_socket(socket)
	143	+ for numa in self.numa_domains:
	144	+ self._build_cluster_domains_for_numa(numa)
	145	+ return True
	146	+
	147	+ except Exception as e:
	148	+ print(f"Build affinity domain fail: {e}.")
	149	+ return False
	150	+
	151	+ def _read_cpu(self) -> None:
	152	+ cpu_root_dir = "/sys/devices/system/cpu/"
	153	+ cpu_pattern = re.compile(r"^cpu(\d+)$")
	154	+
	155	+ for entry in utils.safe_listdir(cpu_root_dir):
	156	+ match = cpu_pattern.match(entry)
	157	+ if not match:
	158	+ continue
	159	+ cpu_id = int(match.group(1))
	160	+
	161	+ # 过滤掉不可用的cpu
	162	+ if not self._is_cpu_available(cpu_id):
	163	+ continue
	164	+
	165	+ # 过滤掉不在线的cpu
	166	+ if not utils.is_cpu_online(cpu_id):
	167	+ continue
	168	+
	169	+ cpu_dir = os.path.join(cpu_root_dir, f"cpu{cpu_id}")
	170	+ topology_dir = os.path.join(cpu_dir, "topology/")
	171	+
	172	+ cpu_core = CpuCore(
	173	+ physical_package_id=utils.read_int_param(topology_dir, "physical_package_id"),
	174	+ cluster_id=utils.read_int_param(topology_dir, "cluster_id"),
	175	+ )
	176	+ self._cpu_dict[cpu_id] = cpu_core
	177	+
	178	+ def _read_numa(self) -> None:
	179	+ node_root_dir = "/sys/devices/system/node/"
	180	+ node_pattern = re.compile(r"^node(\d+)$")
	181	+
	182	+ for entry in utils.safe_listdir(node_root_dir):
	183	+ match = node_pattern.match(entry)
	184	+ if not match:
	185	+ continue
	186	+ node_id = int(match.group(1))
	187	+ node_dir = os.path.join(node_root_dir, f"node{node_id}")
	188	+ numa = NumaNode(cpulist=utils.read_list_param(node_dir, "cpulist")[0])
	189	+ self._numa_dict[node_id] = numa
	190	+
	191	+ def _is_cpu_available(self, cpu: int) -> bool:
	192	+ """
	193	+ cpu是否可用
	194	+ 1) 容器场景如果不在容器绑定的cpu范围需要剔除
	195	+ 2) 在通过set_exclude_cpu设置的排除范围需要剔除
	196	+ """
	197	+
	198	+ if self._container_cpus:
	199	+ if cpu not in self._container_cpus:
	200	+ return False
	201	+ if self._exclude_cpus:
	202	+ if cpu in self._exclude_cpus:
	203	+ return False
	204	+ return True
	205	+
	206	+ def _build_topo(self) -> None:
	207	+ for cpu_id, cpu_core in self._cpu_dict.items():
	208	+ if cpu_core.physical_package_id != -1:
	209	+ self._socket_to_cpus[cpu_core.physical_package_id].append(cpu_id)
	210	+ # 部分硬件不存在cluster，按照cluster等于cpu进行处理，保持cluster级抽象
	211	+ if cpu_core.cluster_id == -1:
	212	+ cpu_core.cluster_id = cpu_id
	213	+ self._cluster_to_cpus[cpu_core.cluster_id].append(cpu_id)
	214	+
	215	+ for numa_id, numa in self._numa_dict.items():
	216	+ for cpu_id in numa.cpulist:
	217	+ cpu = self._cpu_dict.get(cpu_id)
	218	+ if cpu is None:
	219	+ continue
	220	+ self._numa_to_cpus[numa_id].append(cpu_id)
	221	+ if cpu.physical_package_id != -1:
	222	+ self._socket_to_numas[cpu.physical_package_id].add(numa_id)
	223	+ if cpu.cluster_id != -1:
	224	+ self._numa_to_clusters[numa_id].add(cpu.cluster_id)
	225	+
	226	+ def _build_socket_domains(self) -> None:
	227	+ for socket_id, cpu_list in sorted(self._socket_to_cpus.items()):
	228	+ socket = SocketDomain(socket_id)
	229	+ socket.detect_from_hardware()
	230	+ socket.cpus.from_list(cpu_list)
	231	+ self.socket_domains.append(socket)
	232	+
	233	+ def _build_numa_domains_for_socket(self, socket: SocketDomain) -> None:
	234	+ socket_numas = self._socket_to_numas.get(socket.domain_id)
	235	+ if socket_numas is None:
	236	+ return
	237	+ for numa_id in sorted(socket_numas):
	238	+ numa = NumaDomain(numa_id)
	239	+ numa.detect_from_hardware()
	240	+ cpu_list = self._numa_to_cpus.get(numa_id, [])
	241	+ numa.cpus.from_list(cpu_list)
	242	+ self.numa_domains.append(numa)
	243	+ socket.add_child(numa)
	244	+
	245	+ def _build_cluster_domains_for_numa(self, numa: NumaDomain) -> None:
	246	+ numa_clusters = self._numa_to_clusters.get(numa.domain_id)
	247	+ if numa_clusters is None:
	248	+ return
	249	+ for cluster_id in sorted(numa_clusters):
	250	+ cluster = ClusterDomain(cluster_id)
	251	+ cluster.detect_from_hardware()
	252	+ cpu_list = self._cluster_to_cpus.get(cluster_id, [])
	253	+ cluster.cpus.from_list(cpu_list)
	254	+ self.cluster_domains.append(cluster)
	255	+ numa.add_child(cluster)
	256	+
	257	+
	258	+class AffinityDomainManager:
	259	+ """
	260	+ 亲和域管理
	261	+ """
	262	+
	263	+ def __init__(self, config: AffinityConfig) -> None:
	264	+ self.socket_domains: list[SocketDomain] = []
	265	+ self.numa_domains: list[NumaDomain] = []
	266	+ self.cluster_domains: list[ClusterDomain] = []
	267	+ self._config = config
	268	+
	269	+ def build_affinity_domain(self) -> None:
	270	+ builder = AffinityDomainBuilder(self._config)
	271	+ if builder.build_affinity_domain():
	272	+ self.socket_domains = builder.socket_domains
	273	+ self.numa_domains = builder.numa_domains
	274	+ self.cluster_domains = builder.cluster_domains
	275	+
	276	+ def get_socket_num(self) -> int:
	277	+ return len(self.socket_domains)
	278	+
	279	+ def get_socket_domain(self, socket_id: int) -> SocketDomain \| None:
	280	+ for socket in self.socket_domains:
	281	+ if socket.domain_id == socket_id:
	282	+ return socket
	283	+ return None
	284	+
	285	+ def get_numa_domain(self, numa_id: int) -> NumaDomain \| None:
	286	+ for numa in self.numa_domains:
	287	+ if numa.domain_id == numa_id:
	288	+ return numa
	289	+ return None
	290	+
	291	+ def get_all_numas_id(self) -> list:
	292	+ numa_list: list = []
	293	+ for numa in self.numa_domains:
	294	+ numa_list.append(numa.domain_id)
	295	+ return numa_list
	296	+
	297	+ def get_core_num(self) -> int:
	298	+ core_num: int = 0
	299	+ for socket in self.socket_domains:
	300	+ core_num += socket.cpus.count()
	301	+ return core_num
	302	+
	303	+ def get_cluster_domain(self, cluster_id: int) -> ClusterDomain \| None:
	304	+ for cluster in self.cluster_domains:
	305	+ if cluster.domain_id == cluster_id:
	306	+ return cluster
	307	+ return None
	308	+
	309	+ def print_all(self) -> None:
	310	+ print(
	311	+ f"TOTAL: sockets={self.get_socket_num()}, "
	312	+ f"numas={len(self.numa_domains)}, "
	313	+ f"clusters={len(self.cluster_domains)}, "
	314	+ f"cores={self.get_core_num()}"
	315	+ )
	316	+ for socket in self.socket_domains:
	317	+ self.print_domain(socket)
	318	+
	319	+ def print_domain(self, domain: AffinityDomain, indent: int = 0) -> None:
	320	+ prefix = " " * indent
	321	+ print(f"{prefix}{domain}")
	322	+
	323	+ # 打印子域
	324	+ for child in domain.get_all_children():
	325	+ self.print_domain(child, indent + 1)
	326	+
	327	+ def get_sockets_of_cpus(self, cpus: list[int]) -> list:
	328	+ socket_list: list[int] = []
	329	+ for socket in self.socket_domains:
	330	+ for cpu in cpus:
	331	+ if socket.cpus.isset(cpu=cpu):
	332	+ socket_list.append(socket.domain_id)
	333	+ break
	334	+ return sorted(set(socket_list))
	335	+
	336	+ def get_numas_of_cpus(self, cpus: list[int]) -> list:
	337	+ numa_list: list[int] = []
	338	+ for numa in self.numa_domains:
	339	+ for cpu in cpus:
	340	+ if numa.cpus.isset(cpu=cpu):
	341	+ numa_list.append(numa.domain_id)
	342	+ break
	343	+ return sorted(set(numa_list))
	344	+
	345	+ def get_clusters_of_cpus(self, cpus: list[int]) -> list:
	346	+ cluster_list: list[int] = []
	347	+ for cluster in self.cluster_domains:
	348	+ for cpu in cpus:
	349	+ if cluster.cpus.isset(cpu=cpu):
	350	+ cluster_list.append(cluster.domain_id)
	351	+ break
	352	+ return sorted(set(cluster_list))

Aa-sched/a_sched/api.py+173-0

@@ -0,0 +1,173 @@
	1	+from __future__ import annotations
	2	+from a_sched.engine import affinity_engine as affinity
	3	+
	4	+
	5	+def group_create(name: str = "") -> int:
	6	+ """
	7	+ 创建新组并返回组ID
	8	+
	9	+ Returns:
	10	+ int: 新创建的组ID，如果创建失败返回-1
	11	+
	12	+ Raises:
	13	+ RuntimeError: 当组创建失败时可能抛出异常
	14	+ """
	15	+ return affinity.task.group_create(name=name)
	16	+
	17	+
	18	+def group_destroy(group_id: int) -> None:
	19	+ """
	20	+ 销毁指定组
	21	+
	22	+ Args:
	23	+ group_id: 要销毁的组ID
	24	+
	25	+ Returns:
	26	+ None: 不返回值
	27	+
	28	+ Note:
	29	+ 如果组不存在，静默返回（不抛异常）
	30	+ """
	31	+ affinity.task.destory_group(group_id=group_id)
	32	+
	33	+
	34	+def group_add_thread(
	35	+ group_id: int,
	36	+ tid: int \| None = None,
	37	+ thread_name: str \| None = None,
	38	+ pid: int \| None = None,
	39	+ process_name: str \| None = None,
	40	+) -> None:
	41	+ """
	42	+ 将线程添加到指定任务组中。
	43	+
	44	+ Args:
	45	+ group_id: 目标任务组的ID，必须是已存在的有效组ID。
	46	+ tid: 要添加的线程ID，必须是正在运行线程的有效TID。
	47	+ thread_name: 线程的唯一名称，用于标识和监控。
	48	+
	49	+ Raises:
	50	+ 当添加失败时抛出异常
	51	+ """
	52	+ affinity.task.group_add_thread(
	53	+ group_id=group_id, tid=tid, thread_name=thread_name, pid=pid, process_name=process_name
	54	+ )
	55	+
	56	+
	57	+def group_remove_thread(group_id: int, tid: int \| None = None, thread_name: str \| None = None) -> None:
	58	+ """
	59	+ 从组中移除线程
	60	+
	61	+ Args:
	62	+ group_id: 组ID
	63	+ tid: 线程ID
	64	+ thread_name: 线程名称
	65	+
	66	+ Raises:
	67	+ 当移除失败时抛出异常
	68	+ """
	69	+ affinity.task.group_remove_thread(group_id=group_id, tid=tid, thread_name=thread_name)
	70	+
	71	+
	72	+def group_add_process(
	73	+ group_id: int, pid: int \| None = None, process_name: str \| None = None, parent_name: str \| None = None
	74	+) -> None:
	75	+ """
	76	+ 将进程添加到指定任务组中。
	77	+
	78	+ Args:
	79	+ group_id: 目标任务组的ID。必须是已存在的有效组ID。
	80	+ pid: 要添加的进程ID。必须是正在运行进程的有效PID。
	81	+ process_name: 进程的唯一名称。用于标识和监控。
	82	+ parent_name: 父进程的名称，用于在多个同名进程中过滤指定父进程的子进程。
	83	+
	84	+ Raises:
	85	+ 当添加失败时抛出异常
	86	+ """
	87	+ affinity.task.group_add_process(group_id=group_id, pid=pid, process_name=process_name, parent_name=parent_name)
	88	+
	89	+
	90	+def group_remove_process(
	91	+ group_id: int, pid: int \| None = None, process_name: str \| None = None, parent_name: str \| None = None
	92	+) -> None:
	93	+ """
	94	+ 从组中移除进程
	95	+
	96	+ Args:
	97	+ group_id: 组ID
	98	+ pid: 进程ID
	99	+ process_name: 进程名称
	100	+ parent_name: 父进程的名称，用于在多个同名进程中过滤指定父进程的子进程。
	101	+
	102	+ Raises:
	103	+ 当移除失败时抛出异常
	104	+ """
	105	+ affinity.task.group_remove_process(group_id=group_id, pid=pid, process_name=process_name, parent_name=parent_name)
	106	+
	107	+
	108	+def thread_set_high_priority(
	109	+ tid: int \| None = None, thread_name: str \| None = None, pid: int \| None = None, process_name: str \| None = None
	110	+) -> None:
	111	+ """
	112	+ 指定线程设置高优先级
	113	+
	114	+ Args:
	115	+ tid: 线程ID
	116	+ thread_name: 线程名称
	117	+
	118	+ Raises:
	119	+ 当设置失败时抛出异常
	120	+ """
	121	+ affinity.task.thread_set_high_priority(tid=tid, thread_name=thread_name, pid=pid, process_name=process_name)
	122	+
	123	+
	124	+def process_bind_npu(
	125	+ npu_id: int, pid: int \| None = None, process_name: str \| None = None, parent_name: str \| None = None
	126	+) -> None:
	127	+ """
	128	+ 进程绑定NPU
	129	+
	130	+ Args:
	131	+ npu_id: npu的id
	132	+ pid: 进程ID
	133	+ process_name: 进程名
	134	+ parent_name: 父进程的名称，用于在多个同名进程中过滤指定父进程的子进程。
	135	+
	136	+ Raises:
	137	+ 失败时抛出异常
	138	+ """
	139	+ affinity.task.process_bind_npu(npu_id=npu_id, pid=pid, process_name=process_name, parent_name=parent_name)
	140	+
	141	+
	142	+def run_affinity(dry_run: bool = False) -> None:
	143	+ """
	144	+ 运行亲和调度
	145	+
	146	+ Args:
	147	+ dry_run: True表示试运行，仅输出亲和方案，不做亲和方案执行；False表示运行亲和调度全流程
	148	+ """
	149	+ affinity.run(dry_run=dry_run)
	150	+
	151	+
	152	+def print_affinity() -> None:
	153	+ """
	154	+ 打印亲和组中进程/线程当前实际的亲和信息
	155	+ """
	156	+ affinity.print_affinity()
	157	+
	158	+
	159	+def restore_affinity() -> None:
	160	+ """
	161	+ 恢复亲合组中进程/线程原始的亲和信息
	162	+ """
	163	+ affinity.restore_affinity()
	164	+
	165	+
	166	+def set_exclude_cpu(cpu_str: str) -> None:
	167	+ """
	168	+ 设置不参与亲和调度的cpu
	169	+
	170	+ Args:
	171	+ cpu_str: 字符串格式的cpu列表，示例：0-3, 10, 20
	172	+ """
	173	+ affinity.config.set_exclude_cpu(cpu_str=cpu_str)

Aa-sched/a_sched/backup.py+344-0

@@ -0,0 +1,344 @@
	1	+import glob
	2	+import psutil
	3	+import json
	4	+import os
	5	+from datetime import datetime
	6	+
	7	+from a_sched.task import TaskManager
	8	+from a_sched.affinity_domain import AffinityDomainManager
	9	+import a_sched.utils as utils
	10	+
	11	+# 亲和信息备份文件名前缀，完整文件名形如 cpu_affinity_backup_20260525_143022.json
	12	+AFFINITY_BACKUP_FILE_PREFIX = "cpu_affinity_backup_"
	13	+AFFINITY_BACKUP_FILE_SUFFIX = ".json"
	14	+
	15	+
	16	+class AffinityBackup:
	17	+ """亲和信息备份恢复"""
	18	+
	19	+ def __init__(self, task: TaskManager, domain: AffinityDomainManager):
	20	+ self._task = task
	21	+ self._domain = domain
	22	+
	23	+ def backup_affinity(self) -> None:
	24	+ affinity_data = {
	25	+ "cpu_bind_data": self.build_cpu_bind_data(),
	26	+ "background_bind_data": self.build_background_bind_data(),
	27	+ "irq_bind_data": self.build_irq_bind_data(),
	28	+ "irq_service_status": self.build_irq_service_status(),
	29	+ "dev_sq_bind_data": self.build_dev_sq_bind_data(),
	30	+ }
	31	+
	32	+ # 写入文件
	33	+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	34	+ backup_file = f"{AFFINITY_BACKUP_FILE_PREFIX}{timestamp}{AFFINITY_BACKUP_FILE_SUFFIX}"
	35	+ try:
	36	+ with open(backup_file, "w", encoding="utf-8") as f:
	37	+ json.dump(affinity_data, f, indent=4, ensure_ascii=False)
	38	+ os.chmod(backup_file, 0o444)
	39	+ print(f"save affinity to {backup_file} SUCCESS!")
	40	+ except Exception as e:
	41	+ print(f"save affinity to {backup_file} FAIL! {str(e)}")
	42	+
	43	+ def build_cpu_bind_data(self) -> dict:
	44	+ cpu_bind_data = {}
	45	+ for group in self._task.groups.values():
	46	+ for pid, process_task in group.process_tasks.items():
	47	+ try:
	48	+ process = psutil.Process(pid)
	49	+ threads = []
	50	+ for thread in process.threads():
	51	+ try:
	52	+ threads.append(
	53	+ {
	54	+ "tid": thread.id,
	55	+ "thread_name": utils.get_thread_name_by_tid(tid=thread.id, pid=pid),
	56	+ "cpu_affinity": utils.get_thread_cpus_by_tid(tid=thread.id, pid=pid),
	57	+ }
	58	+ )
	59	+ except Exception as e:
	60	+ print(f"thread [{thread.id}] get cpu affinity failed, {str(e)}")
	61	+
	62	+ try:
	63	+ cpu_bind_data[pid] = {
	64	+ "pid": pid,
	65	+ "process_name": process.name(),
	66	+ "cpu_affinity": process.cpu_affinity(),
	67	+ "threads": threads,
	68	+ }
	69	+ except Exception as e:
	70	+ print(f"process [{pid}-{process_task.name}] get cpu affinity failed, {str(e)}")
	71	+
	72	+ except Exception as e:
	73	+ print(f"process [{pid}-{process_task.name}] not found, {str(e)}")
	74	+ return cpu_bind_data
	75	+
	76	+ def build_irq_service_status(self) -> bool:
	77	+ _, _, ret = utils.execute_command(["systemctl", "is-active", "--quiet", "irqbalance"])
	78	+ return True if ret == 0 else False
	79	+
	80	+ def build_irq_bind_data(self) -> dict:
	81	+ irq_bind_data = {}
	82	+ for group in self._task.groups.values():
	83	+ for npu in group.npu_tasks.values():
	84	+ irqs = []
	85	+ if npu._sq_irq is not None:
	86	+ irqs.append((npu._sq_irq, npu.SQ_IRQ))
	87	+ for cq in npu._cq_irqs:
	88	+ irqs.append((cq, npu.CQ_IRQ))
	89	+ if npu._trs_mbox_irq is not None:
	90	+ irqs.append((npu._trs_mbox_irq, npu.trs_mbox_name))
	91	+ for irq_id, irq_name in irqs:
	92	+ try:
	93	+ irq_bind_data[irq_id] = {
	94	+ "irq_id": irq_id,
	95	+ "irq_name": irq_name,
	96	+ "cpu_affinity": utils.get_irq_cpus_by_irq_id(irq_id=irq_id),
	97	+ }
	98	+ except Exception as e:
	99	+ print(f"irq [{irq_id}-{irq_name}] get cpu affinity failed, {str(e)}")
	100	+ return irq_bind_data
	101	+
	102	+ def build_background_bind_data(self) -> dict:
	103	+ backgroud_bind_data = {}
	104	+ for pid, name in self._task.background_processes.items():
	105	+ try:
	106	+ proc = psutil.Process(pid)
	107	+ backgroud_bind_data[pid] = {
	108	+ "pid": pid,
	109	+ "process_name": proc.name(),
	110	+ "cpu_affinity": proc.cpu_affinity(),
	111	+ }
	112	+ except (psutil.NoSuchProcess, psutil.AccessDenied):
	113	+ continue
	114	+ return backgroud_bind_data
	115	+
	116	+ def build_dev_sq_bind_data(self) -> dict:
	117	+ dev_sq_bind_data = {}
	118	+ for group in self._task.groups.values():
	119	+ for npu_id, npu in group.npu_tasks.items():
	120	+ entry = {"npu_id": npu_id}
	121	+ if npu._dev_sq_task is not None:
	122	+ try:
	123	+ entry["dev_sq_task"] = {
	124	+ "pid": npu._dev_sq_task,
	125	+ "name": npu.dev_sq_task_name,
	126	+ "cpu_affinity": utils.get_process_cpus_by_pid(pid=npu._dev_sq_task),
	127	+ }
	128	+ except Exception as e:
	129	+ print(f"npu[{npu_id}] {npu.dev_sq_task_name} get cpu affinity failed, {str(e)}")
	130	+ if npu._dev_sq_send_wq is not None:
	131	+ try:
	132	+ entry["dev_sq_send_wq"] = {
	133	+ "wq_name": npu.dev_sq_send_wq_name,
	134	+ "cpu_affinity": utils.get_npu_work_queue_cpus_by_name(wq_name=npu.dev_sq_send_wq_name),
	135	+ }
	136	+ except Exception as e:
	137	+ print(f"npu[{npu_id}] {npu.dev_sq_send_wq_name} get cpu affinity failed, {str(e)}")
	138	+ if len(entry) > 1:
	139	+ dev_sq_bind_data[npu_id] = entry
	140	+ return dev_sq_bind_data
	141	+
	142	+ def restore_affinity(self) -> None:
	143	+ # 1. 查找最早的备份文件
	144	+ backup_files = sorted(glob.glob(f"{AFFINITY_BACKUP_FILE_PREFIX}*{AFFINITY_BACKUP_FILE_SUFFIX}"))
	145	+ if not backup_files:
	146	+ print(f"no backup file matching {AFFINITY_BACKUP_FILE_PREFIX}*{AFFINITY_BACKUP_FILE_SUFFIX} found")
	147	+ return
	148	+ backup_file = backup_files[0]
	149	+ print(f"restore from earliest backup file: {backup_file}")
	150	+
	151	+ # 2. 读取保存数据
	152	+ try:
	153	+ with open(backup_file, "r", encoding="utf-8") as f:
	154	+ saved_data = json.load(f)
	155	+ if not saved_data:
	156	+ print("data is empty")
	157	+ return
	158	+ except json.JSONDecodeError:
	159	+ print(f"json decode error: {backup_file}")
	160	+ return
	161	+
	162	+ # 3. 批量恢复
	163	+ self.restore_cpu_bind_data(saved_data.get("cpu_bind_data", {}))
	164	+ self.restore_irq_bind_data(saved_data.get("irq_bind_data", {}))
	165	+ self.restore_irq_service_status(saved_data.get("irq_service_status", False))
	166	+ self.restore_dev_sq_bind_data(saved_data.get("dev_sq_bind_data", {}))
	167	+ self.restore_background_bind_data(saved_data.get("background_bind_data", {}))
	168	+
	169	+ def restore_background_bind_data(self, background_bind_data: dict) -> None:
	170	+ restored = 0
	171	+ for pid_str, data in background_bind_data.items():
	172	+ try:
	173	+ pid = int(pid_str)
	174	+ proc = psutil.Process(pid)
	175	+ orig_affinity = data["cpu_affinity"]
	176	+ proc.cpu_affinity(orig_affinity)
	177	+ for thread in proc.threads():
	178	+ try:
	179	+ psutil.Process(thread.id).cpu_affinity(orig_affinity)
	180	+ except (psutil.NoSuchProcess, psutil.AccessDenied):
	181	+ pass
	182	+ restored += 1
	183	+ except psutil.NoSuchProcess:
	184	+ pass
	185	+ except Exception as e:
	186	+ print(f"background process [{pid_str}] restore failed, {str(e)}")
	187	+
	188	+ if restored:
	189	+ print(f"restore {restored} background processes affinity SUCCESS!")
	190	+
	191	+ def restore_cpu_bind_data(self, cpu_bind_data: dict) -> None:
	192	+ for pid, data in cpu_bind_data.items():
	193	+ try:
	194	+ pid = int(pid)
	195	+ proc = psutil.Process(pid)
	196	+ proc_name = data["process_name"]
	197	+ proc.cpu_affinity(data["cpu_affinity"])
	198	+ cpus = utils.CPUMask().from_list(proc.cpu_affinity())
	199	+ print(f"restore process-{pid}-{data['process_name']}, cpu_affinity:{cpus}")
	200	+ for thread_data in data["threads"]:
	201	+ try:
	202	+ thread = psutil.Process(thread_data["tid"])
	203	+ thread.cpu_affinity(thread_data["cpu_affinity"])
	204	+ cpus = utils.CPUMask().from_list(thread.cpu_affinity())
	205	+ print(
	206	+ f"restore thread-{thread_data['tid']}-{thread_data['thread_name']}(process-{proc_name}), cpu_affinity:{cpus}"
	207	+ )
	208	+ except Exception as e:
	209	+ print(f"thread-{thread_data['thread_name']}(process-{proc_name}) restore failed, {str(e)}")
	210	+
	211	+ except Exception as e:
	212	+ print(f"process [{pid}] restore failed, {str(e)}")
	213	+
	214	+ def restore_irq_service_status(self, should_start: bool) -> None:
	215	+ if should_start:
	216	+ utils.execute_command(["systemctl", "start", "irqbalance"])
	217	+ print("restore irqbalance service to active.")
	218	+
	219	+ def restore_irq_bind_data(self, irq_bind_data: dict) -> None:
	220	+ for irq_id, data in irq_bind_data.items():
	221	+ try:
	222	+ irq_id = int(irq_id)
	223	+ irq_name = data["irq_name"]
	224	+ cpus = data["cpu_affinity"]
	225	+ utils.bind_irq_to_cpus(irq_id=irq_id, cpus=cpus, irq_name=irq_name)
	226	+ print(f"restore irq-{irq_id}-{irq_name}, cpu_affinity: {utils.CPUMask().from_list(cpus)}")
	227	+ except Exception as e:
	228	+ print(f"irq [{irq_id}] restore failed, {str(e)}")
	229	+
	230	+ def restore_dev_sq_bind_data(self, dev_sq_bind_data: dict) -> None:
	231	+ for npu_id, data in dev_sq_bind_data.items():
	232	+ task_info = data.get("dev_sq_task")
	233	+ if task_info:
	234	+ try:
	235	+ pid = task_info["pid"]
	236	+ cpus = task_info["cpu_affinity"]
	237	+ utils.bind_process_to_cpus(pid=pid, cpus=cpus)
	238	+ print(
	239	+ f"restore npu[{npu_id}] {task_info['name']} - {pid}, "
	240	+ f"cpu_affinity: {utils.CPUMask().from_list(cpus)}"
	241	+ )
	242	+ except Exception as e:
	243	+ print(f"npu[{npu_id}] {task_info.get('name')} restore failed, {str(e)}")
	244	+
	245	+ wq_info = data.get("dev_sq_send_wq")
	246	+ if wq_info:
	247	+ try:
	248	+ cpus = wq_info["cpu_affinity"]
	249	+ utils.bind_npu_sq_send_wq_to_cpus(npu_id=int(npu_id), cpus=cpus)
	250	+ print(f"restore npu[{npu_id}] {wq_info['wq_name']}, cpu_affinity: {utils.CPUMask().from_list(cpus)}")
	251	+ except Exception as e:
	252	+ print(f"npu[{npu_id}] {wq_info.get('wq_name')} restore failed, {str(e)}")
	253	+
	254	+ def print_affinity(self) -> None:
	255	+ """打印亲和任务中进程/线程当前实际的亲和信息"""
	256	+
	257	+ for group in self._task.groups.values():
	258	+ print(f"TaskGroup[{group.group_id}]: {f'name={group.name}' if group.name else ''}")
	259	+ for pid, process in group.process_tasks.items():
	260	+ priority = process._priority_names.get(process.priority, "UNKNOWN")
	261	+ cpus = utils.get_process_cpus_by_pid(pid=pid)
	262	+ print(
	263	+ f" - PROCESS[{pid}]: name={process.name}, priority={priority}, "
	264	+ f"socket={self._domain.get_sockets_of_cpus(cpus)}, "
	265	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	266	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	267	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	268	+ )
	269	+ for tid, thread in group.thread_tasks.items():
	270	+ priority = thread._priority_names.get(thread.priority, "UNKNOWN")
	271	+ cpus = utils.get_thread_cpus_by_tid(tid=tid)
	272	+ print(
	273	+ f" - THREAD[{tid}]: name={thread.name}, priority={priority}, "
	274	+ f"socket={self._domain.get_sockets_of_cpus(cpus)}, "
	275	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	276	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	277	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	278	+ )
	279	+ for npu_id, npu in group.npu_tasks.items():
	280	+ print(f" - NPU[{npu_id}]: ")
	281	+ if npu._dev_sq_task is not None:
	282	+ cpus = utils.get_process_cpus_by_pid(pid=npu._dev_sq_task)
	283	+ print(
	284	+ f" - Process[{npu._dev_sq_task}]({npu.dev_sq_task_name}): "
	285	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	286	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	287	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	288	+ )
	289	+ if npu._dev_sq_send_wq is not None:
	290	+ cpus = utils.get_npu_work_queue_cpus_by_name(wq_name=npu.dev_sq_send_wq_name)
	291	+ print(
	292	+ f" - Process[{npu._dev_sq_send_wq}]({npu.dev_sq_send_wq_name}): "
	293	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	294	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	295	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	296	+ )
	297	+ if npu._acl_thread is not None:
	298	+ cpus = utils.get_thread_cpus_by_tid(tid=npu._acl_thread)
	299	+ print(
	300	+ f" - Thread[{npu._acl_thread}]({npu.ACL_THREAD}): "
	301	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	302	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	303	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	304	+ )
	305	+ if npu._release_thread is not None:
	306	+ cpus = utils.get_thread_cpus_by_tid(tid=npu._release_thread)
	307	+ print(
	308	+ f" - Thread[{npu._release_thread}]({npu.RELEASE_THREAD}): "
	309	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	310	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	311	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	312	+ )
	313	+ if npu._rt_recycle_thread is not None:
	314	+ cpus = utils.get_thread_cpus_by_tid(tid=npu._rt_recycle_thread)
	315	+ print(
	316	+ f" - Thread[{npu._rt_recycle_thread}]({npu.RT_RECYCLE_THREAD}): "
	317	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	318	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	319	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	320	+ )
	321	+ if npu._sq_irq is not None:
	322	+ cpus = utils.get_irq_cpus_by_irq_id(irq_id=npu._sq_irq)
	323	+ print(
	324	+ f" - Irq[{npu._sq_irq}]({npu.SQ_IRQ}): "
	325	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	326	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	327	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	328	+ )
	329	+ for cq in npu._cq_irqs:
	330	+ cpus = utils.get_irq_cpus_by_irq_id(irq_id=cq)
	331	+ print(
	332	+ f" - Irq[{cq}]({npu.CQ_IRQ}): "
	333	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	334	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	335	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	336	+ )
	337	+ if npu._trs_mbox_irq is not None:
	338	+ cpus = utils.get_irq_cpus_by_irq_id(irq_id=npu._trs_mbox_irq)
	339	+ print(
	340	+ f" - Irq[{npu._trs_mbox_irq}]({npu.trs_mbox_name}): "
	341	+ f"numa=[{utils.compress_continuous(self._domain.get_numas_of_cpus(cpus))}], "
	342	+ f"cluster=[{utils.compress_continuous(self._domain.get_clusters_of_cpus(cpus))}], "
	343	+ f"cpu=[{utils.CPUMask().from_list(cpus)}]"
	344	+ )

Aa-sched/a_sched/config.py+10-0

@@ -0,0 +1,10 @@
	1	+import a_sched.utils as utils
	2	+
	3	+
	4	+class AffinityConfig:
	5	+
	6	+ def __init__(self) -> None:
	7	+ self.exclude_cpus: list[int] = []
	8	+
	9	+ def set_exclude_cpu(self, cpu_str: str) -> None:
	10	+ self.exclude_cpus = utils.parse_cpu_affinity_string(cpu_str)

Aa-sched/a_sched/engine.py+188-0

@@ -0,0 +1,188 @@
	1	+from a_sched.affinity_domain import AffinityDomainManager
	2	+from a_sched.task import TaskManager
	3	+from a_sched.backup import AffinityBackup
	4	+from a_sched.config import AffinityConfig
	5	+from a_sched.scheduler import Scheduler
	6	+from a_sched.strategy.hierarchical_balance import HierarchicalBalanceScheduler
	7	+import a_sched.utils as utils
	8	+
	9	+
	10	+class AffinityEngine:
	11	+ """亲和调度引擎"""
	12	+
	13	+ def __init__(self) -> None:
	14	+ self._scheduler: Scheduler \| None = None
	15	+ self._init_engine()
	16	+
	17	+ def _init_engine(self) -> None:
	18	+ self.config = AffinityConfig()
	19	+ self.task = TaskManager()
	20	+ self.domain = AffinityDomainManager(config=self.config)
	21	+ self.backup = AffinityBackup(task=self.task, domain=self.domain)
	22	+
	23	+ def run(self, dry_run: bool = False) -> None:
	24	+ """
	25	+ 运行亲和调度
	26	+
	27	+ Args:
	28	+ dry_run: True表示试运行，仅输出亲和方案，不做亲和方案执行；False表示运行亲和调度全流程
	29	+ """
	30	+
	31	+ print("\n-------------------------- Affinity Schedule Begin --------------------------")
	32	+
	33	+ try:
	34	+ print("\nStarting build affinity domain...")
	35	+ self.domain.build_affinity_domain()
	36	+
	37	+ print("The affinity domain is as follows:")
	38	+ self._print_affinity_domain()
	39	+
	40	+ print("\nThe affinity info before schedule is as follows:")
	41	+ self.print_affinity()
	42	+
	43	+ print("\nStarting plan affinity...")
	44	+ self._plan()
	45	+
	46	+ print("The affinity plan is as follows:")
	47	+ self._print_affinity_plan()
	48	+
	49	+ if not dry_run:
	50	+ print("\nStarting execute affinity...")
	51	+ self._execute()
	52	+
	53	+ print("\nThe affinity info after schedule is as follows:")
	54	+ self.print_affinity()
	55	+
	56	+ print("\nAffinity schedule SUCCESS!")
	57	+
	58	+ except Exception as e:
	59	+ print(f"\nAffinity schedule FAILED! Error: {e}.")
	60	+
	61	+ print("---------------------------- Affinity Schedule End ----------------------------")
	62	+
	63	+ def _plan(self) -> None:
	64	+ # 刷新高优先级线程和npu互相绑定关系
	65	+ self.task.update_high_prio_thread_bind_npu()
	66	+
	67	+ # 扫描背景任务
	68	+ self.task.scan_background_tasks()
	69	+
	70	+ # 使用分层均衡亲和调度策略决策亲和方案
	71	+ self._scheduler = HierarchicalBalanceScheduler(self.config, self.domain, self.task)
	72	+ if not self._scheduler.schedule():
	73	+ raise RuntimeError("Plan affinity failed!")
	74	+
	75	+ def _execute(self) -> None:
	76	+ # 执行前先备份当前亲和信息
	77	+ self.backup_affinity()
	78	+ # 停止CPU硬件中断自动均衡
	79	+ self._stop_irq_balance()
	80	+ # 根据亲和方案绑定cpu
	81	+ self._bind_cpus()
	82	+ # 迁移进程内存到新numa节点
	83	+ self._bind_memory()
	84	+
	85	+ def _print_affinity_domain(self) -> None:
	86	+ print("------------------------------- Affinity Domain -------------------------------")
	87	+ self.domain.print_all()
	88	+
	89	+ def _print_affinity_plan(self) -> None:
	90	+ print("-------------------------------- Affinity Plan --------------------------------")
	91	+ self.task.print_all()
	92	+
	93	+ def _stop_irq_balance(self) -> None:
	94	+ print("\nStopping irqbalance service...")
	95	+ _, _, return_code = utils.execute_command(["systemctl", "is-active", "--quiet", "irqbalance"])
	96	+ if return_code == 0:
	97	+ utils.execute_command(["systemctl", "stop", "irqbalance"])
	98	+ print("the irqbalance service has been stopped.")
	99	+
	100	+ def _bind_cpus(self) -> None:
	101	+ print("\nStarting bind cpus...")
	102	+
	103	+ # 绑定亲和组任务
	104	+ for _, group in self.task.groups.items():
	105	+ for task in group.get_all_tasks():
	106	+ task.bind_cpu()
	107	+
	108	+ # 绑定背景任务
	109	+ self._bind_background_tasks()
	110	+
	111	+ def _bind_background_tasks(self) -> None:
	112	+ normal_cpus = self.task.background_tasks_cpus
	113	+ if not normal_cpus:
	114	+ print("[BackgrondTask] No target CPUs found, skipping background tasks binding.")
	115	+ return
	116	+
	117	+ failed = []
	118	+ bound = 0
	119	+
	120	+ for pid, name in self.task.background_processes.items():
	121	+ try:
	122	+ utils.bind_process_to_cpus(pid, normal_cpus)
	123	+ bound += 1
	124	+ except Exception as e:
	125	+ failed.append((pid, name, str(e)))
	126	+
	127	+ if failed:
	128	+ print(f"[BackgroundTask] Failed to bind {len(failed)} processes.")
	129	+ for pid, name, err in failed:
	130	+ print(f" - process[{pid}]({name}): {err}")
	131	+
	132	+ total = len(self.task.background_processes)
	133	+ print(
	134	+ f"[BackgroundTask] Bind {bound}/{total} processes to CPUs {utils.compress_continuous(normal_cpus)}"
	135	+ f" (failed: {len(failed)})"
	136	+ )
	137	+
	138	+ def _bind_memory(self) -> None:
	139	+ print("\nStarting bind memory...")
	140	+ for _, group in self.task.groups.items():
	141	+ for _, process in group.process_tasks.items():
	142	+ if not process.numa:
	143	+ print(f"can not get target numa of process [{process.task_id}]")
	144	+ continue
	145	+ tgt_numa = process.numa[0]
	146	+ src_numa = self.domain.get_all_numas_id()
	147	+ if not src_numa:
	148	+ print(f"can not get source numa of process [{process.task_id}]")
	149	+ continue
	150	+ print(
	151	+ f"migrating pages of process [{process.task_id}] from source numa {src_numa} to target numa [{tgt_numa}]"
	152	+ )
	153	+ try:
	154	+ utils.migrate_process_pages(pid=process.task_id, src_numa=src_numa, tgt_numa=tgt_numa)
	155	+ except Exception as e:
	156	+ print(f"Failed to migrate pages for process [{process.task_id}]: {str(e)}")
	157	+
	158	+ def print_affinity(self) -> None:
	159	+ """打印亲和任务中进程/线程当前实际的亲和信息"""
	160	+
	161	+ print("--------------------------- Current Affinity Status ---------------------------")
	162	+ self.backup.print_affinity()
	163	+
	164	+ def backup_affinity(self) -> None:
	165	+ """备份亲和任务中进程/线程当前亲和信息"""
	166	+
	167	+ print("\nStarting backup current affinity...")
	168	+ self.backup.backup_affinity()
	169	+
	170	+ def restore_affinity(self) -> None:
	171	+ """恢复亲和任务中进程/线程原始亲和信息"""
	172	+
	173	+ print("\n-------------------------- Affinity Resotre Begin ---------------------------")
	174	+ print("Starting restore affinity...")
	175	+
	176	+ self.backup.restore_affinity()
	177	+
	178	+ print("\nThe affinity after restore is as follows:")
	179	+ self.print_affinity()
	180	+
	181	+ print(f"\nAffinity restore SUCCESS!")
	182	+ print("---------------------------- Affinity Resotre End -----------------------------")
	183	+
	184	+ def reset(self) -> None:
	185	+ self._init_engine()
	186	+
	187	+
	188	+affinity_engine = AffinityEngine()

Aa-sched/a_sched/scheduler.py+18-0

@@ -0,0 +1,18 @@
	1	+from abc import ABC, abstractmethod
	2	+
	3	+from a_sched.affinity_domain import AffinityDomainManager
	4	+from a_sched.config import AffinityConfig
	5	+from a_sched.task import TaskManager
	6	+
	7	+
	8	+class Scheduler(ABC):
	9	+ """调度器抽象基类"""
	10	+
	11	+ def __init__(self, config: AffinityConfig, domain: AffinityDomainManager, task: TaskManager):
	12	+ self.config = config
	13	+ self.domain = domain
	14	+ self.task = task
	15	+
	16	+ @abstractmethod
	17	+ def schedule(self) -> bool:
	18	+ return True

Aa-sched/a_sched/strategy/__init__.py+0-0

The file is empty

Aa-sched/a_sched/strategy/hierarchical_balance.py+591-0

@@ -0,0 +1,591 @@
	1	+from __future__ import annotations
	2	+from collections import defaultdict
	3	+
	4	+from a_sched.affinity_domain import AffinityDomainManager, SocketDomain, NumaDomain
	5	+from a_sched.task import TaskManager, TaskGroup, Task
	6	+from a_sched.config import AffinityConfig
	7	+from a_sched.scheduler import Scheduler
	8	+import a_sched.utils as utils
	9	+
	10	+
	11	+class HierarchicalBalanceScheduler(Scheduler):
	12	+ """分层均衡亲和调度器"""
	13	+
	14	+ def __init__(self, config: AffinityConfig, domain: AffinityDomainManager, task: TaskManager):
	15	+ super().__init__(config=config, domain=domain, task=task)
	16	+
	17	+ # 待调度的亲和组，非空，至少有一个task
	18	+ self._schedule_task_groups: list[int] = []
	19	+ # 可被调度的socket，需要剔除部分资源不足的socket
	20	+ self._schedule_sockets: list[int] = []
	21	+
	22	+ # 调度到各socket的亲和组
	23	+ self._socket_to_task_groups: dict[int, list[int]] = defaultdict(list)
	24	+ # 每个socket中待调度的高优先级task
	25	+ self._socket_to_high_prio_tasks: dict[int, list[Task]] = defaultdict(list)
	26	+ # socket除隔离域numa外的其他正常numa
	27	+ self._socket_to_normal_numas: dict[int, list[int]] = defaultdict(list)
	28	+ # socket用于隔离域的numa
	29	+ self._socket_to_isolate_numa: dict[int, int] = {}
	30	+ # 单NUMA场景下，高优先级任务占用的CPU列表
	31	+ self._socket_to_isolate_cpus: dict[int, list[int]] = defaultdict(list)
	32	+
	33	+ # 调度到各numa的亲和组
	34	+ self._numa_to_task_groups: dict[int, list[int]] = defaultdict(list)
	35	+ # 每个numa预留给背景任务的cluster
	36	+ self._numa_to_background_clusters: dict[int, list[int]] = defaultdict(list)
	37	+ # 分配给每个亲和组的cluster
	38	+ self._task_group_to_clusters: dict[int, list[int]] = defaultdict(list)
	39	+
	40	+ def schedule(self) -> bool:
	41	+ # 获取待调度的非空亲和组
	42	+ self._get_task_groups_to_schedule()
	43	+ if len(self._schedule_task_groups) == 0:
	44	+ print("[Error] no valid task group found to schedule.")
	45	+ return False
	46	+
	47	+ # 获取可参与调度的socket
	48	+ self._get_schedule_sockets()
	49	+ if len(self._schedule_sockets) == 0:
	50	+ print("[Error] no valid socket found to schedule.")
	51	+ return False
	52	+
	53	+ # socket层对等均分
	54	+ if not self._schedule_task_groups_to_sockets():
	55	+ return False
	56	+
	57	+ # 每个socket内，对task group进行numa层的对等均分
	58	+ for socket_id in self._socket_to_task_groups:
	59	+ if not self._schedule_task_groups_of_socket(socket_id):
	60	+ return False
	61	+
	62	+ # 每个numa内，对task group进行cluster的对等均分
	63	+ for numa_id in self._numa_to_task_groups:
	64	+ if not self._schedule_numa_task_groups_to_clusters(numa_id):
	65	+ return False
	66	+
	67	+ # 为每个task group刷新affinity
	68	+ for _, group in self.task.groups.items():
	69	+ if not self._update_task_group_affinity(group):
	70	+ return False
	71	+
	72	+ # 每个task group内，对每个task（进程/线程）进行分配
	73	+ for _, group in self.task.groups.items():
	74	+ if not self._schedule_tasks_affinity_of_group(group):
	75	+ return False
	76	+
	77	+ # 刷新高优先级task的亲和信息给所在的group
	78	+ for _, group in self.task.groups.items():
	79	+ self._update_task_group_with_high_prio_tasks(group)
	80	+
	81	+ # 背景任务分配cpu
	82	+ self._alloc_cpus_for_background_tasks()
	83	+
	84	+ return True
	85	+
	86	+ def _get_task_groups_to_schedule(self) -> None:
	87	+ groups_to_schedule: list[int] = []
	88	+ for group_id, group in self.task.groups.items():
	89	+ if group.get_all_tasks_num() != 0:
	90	+ groups_to_schedule.append(group_id)
	91	+ self._schedule_task_groups = groups_to_schedule
	92	+
	93	+ def _get_schedule_sockets(self) -> None:
	94	+ schedule_sockets: list[int] = []
	95	+ for socket in self.domain.socket_domains:
	96	+ schedule_sockets.append(socket.domain_id)
	97	+ self._schedule_sockets = schedule_sockets
	98	+
	99	+ def _schedule_task_groups_to_sockets(self) -> bool:
	100	+ total_task_group_num = len(self._schedule_task_groups)
	101	+ total_socket_num = len(self._schedule_sockets)
	102	+ if total_socket_num == 0:
	103	+ print("no valid socket found.")
	104	+ return False
	105	+
	106	+ base = total_task_group_num // total_socket_num
	107	+ extra = total_task_group_num % total_socket_num
	108	+
	109	+ start = 0
	110	+ for i, socket_id in enumerate(self._schedule_sockets):
	111	+ take = base + (1 if i < extra else 0)
	112	+ end = start + take
	113	+ take_task_groups = self._schedule_task_groups[start:end]
	114	+ if not take_task_groups:
	115	+ break
	116	+ self._socket_to_task_groups[socket_id].extend(take_task_groups)
	117	+ start = end
	118	+
	119	+ return True
	120	+
	121	+ def _schedule_task_groups_of_socket(self, socket_id: int) -> bool:
	122	+ socket = self.domain.get_socket_domain(socket_id)
	123	+ if socket is None:
	124	+ print(f"socket domain [{socket_id}] not found.")
	125	+ return False
	126	+
	127	+ # numa隔离预处理
	128	+ if not self._schedule_isolate_on_numa_of_socket(socket):
	129	+ return False
	130	+
	131	+ # 高优先级任务分配, 基于是否存在 isolate_numa 自动分流
	132	+ if not self._dispatch_high_prio_tasks_by_topology(socket):
	133	+ return False
	134	+
	135	+ # numa非隔离域内对task group进行调度
	136	+ socket_normal_numas = self._socket_to_normal_numas[socket_id]
	137	+ if not self._schedule_socket_task_groups_to_numas(socket=socket, numas=socket_normal_numas):
	138	+ return False
	139	+
	140	+ return True
	141	+
	142	+ def _schedule_isolate_on_numa_of_socket(self, socket: SocketDomain) -> bool:
	143	+ socket_numas = sorted(socket.get_all_children_id())
	144	+ socket_numa_num = len(socket_numas)
	145	+ if socket_numa_num == 0:
	146	+ print(f"no numa with online CPU found in socket [{socket.domain_id}]")
	147	+ return False
	148	+
	149	+ socket_high_prio_tasks: list[Task] = []
	150	+ socket_task_groups = self._socket_to_task_groups[socket.domain_id]
	151	+ for group_id in socket_task_groups:
	152	+ group_high_prio_tasks = self.task.get_high_prio_tasks_of_group(group_id)
	153	+ socket_high_prio_tasks.extend(group_high_prio_tasks)
	154	+
	155	+ # 没有高优先级任务，不需要进行numa隔离
	156	+ if len(socket_high_prio_tasks) == 0:
	157	+ self._socket_to_normal_numas[socket.domain_id] = socket_numas
	158	+ return True
	159	+
	160	+ # 存在高优先级任务，且numa数量小于2时，没有多余的numa可以隔离，不进行numa隔离，改为在cpu层面进行隔离
	161	+ if socket_numa_num < 2:
	162	+ print(f"[Warning] no enough numa for isolate in socket [{socket.domain_id}].")
	163	+ self._socket_to_high_prio_tasks[socket.domain_id] = socket_high_prio_tasks
	164	+ self._socket_to_normal_numas[socket.domain_id] = socket_numas
	165	+ return True
	166	+
	167	+ self._socket_to_high_prio_tasks[socket.domain_id] = socket_high_prio_tasks
	168	+
	169	+ # 使用socket中最后一个numa作为隔离域numa，其他的numa作为非隔离域numa
	170	+ self._socket_to_normal_numas[socket.domain_id].extend(socket_numas[:-1])
	171	+ self._socket_to_isolate_numa[socket.domain_id] = socket_numas[-1]
	172	+
	173	+ return True
	174	+
	175	+ def _dispatch_high_prio_tasks_by_topology(self, socket: SocketDomain) -> bool:
	176	+ socket_high_prio_tasks = self._socket_to_high_prio_tasks.get(socket.domain_id, [])
	177	+ if not socket_high_prio_tasks:
	178	+ return True
	179	+
	180	+ isolate_numa_id = self._socket_to_isolate_numa.get(socket.domain_id)
	181	+
	182	+ if isolate_numa_id is not None:
	183	+ # 有预留的隔离numa → numa 级别隔离
	184	+ return self._schedule_socket_high_prio_tasks_to_isolate_numa(socket)
	185	+ else:
	186	+ # 1 socket 1 numa → CPU 级别隔离
	187	+ return self._schedule_high_prio_tasks_to_isolate_cpu(socket, socket_high_prio_tasks)
	188	+
	189	+ def _schedule_socket_high_prio_tasks_to_isolate_numa(self, socket: SocketDomain) -> bool:
	190	+ socket_high_prio_tasks = self._socket_to_high_prio_tasks[socket.domain_id]
	191	+ task_num = len(socket_high_prio_tasks)
	192	+ if task_num == 0:
	193	+ return True
	194	+
	195	+ isolate_numa_id = self._socket_to_isolate_numa.get(socket.domain_id)
	196	+ if isolate_numa_id is None:
	197	+ return True
	198	+
	199	+ isolate_numa = self.domain.get_numa_domain(isolate_numa_id)
	200	+ if isolate_numa is None:
	201	+ print(f"isolate numa domain [{isolate_numa_id}] not found.")
	202	+ return False
	203	+
	204	+ clusters = isolate_numa.get_all_children_id()
	205	+ cluster_num = len(clusters)
	206	+
	207	+ task_needed_clusters = 0
	208	+ for task in socket_high_prio_tasks:
	209	+ task_needed_clusters += task.min_cluster
	210	+
	211	+ if len(clusters) == 0:
	212	+ print(f"isolate numa [{isolate_numa_id}] has no available cluster")
	213	+ return False
	214	+
	215	+ # cluster数量大于等于task需要的cluster数量时，按照task的实际cluster需求分配
	216	+ # 需验证每个task分配到的cluster是否有足够的CPU
	217	+ if cluster_num >= task_needed_clusters:
	218	+ self._alloc_high_prio_tasks_by_cluster(
	219	+ socket=socket,
	220	+ isolate_numa=isolate_numa,
	221	+ clusters=clusters,
	222	+ socket_high_prio_tasks=socket_high_prio_tasks,
	223	+ )
	224	+ return True
	225	+
	226	+ print(
	227	+ f"[Warning] no enough clusters in isolate numa [{isolate_numa_id}], "
	228	+ f"need {task_needed_clusters}, actual {cluster_num}."
	229	+ )
	230	+
	231	+ self._alloc_high_prio_tasks_by_cpu(
	232	+ socket=socket, isolate_numa=isolate_numa, socket_high_prio_tasks=socket_high_prio_tasks
	233	+ )
	234	+ return True
	235	+
	236	+ def _alloc_high_prio_tasks_by_cluster(
	237	+ self, socket: SocketDomain, isolate_numa: NumaDomain, clusters: list[int], socket_high_prio_tasks: list[Task]
	238	+ ) -> None:
	239	+ start = 0
	240	+ cluster_count = len(clusters)
	241	+
	242	+ for task in socket_high_prio_tasks:
	243	+ if start >= cluster_count:
	244	+ print(
	245	+ f"[Warning] no more clusters for task(id={task.task_id}, name={task.name}), "
	246	+ f"fallback to CPU-level allocation."
	247	+ )
	248	+ self._alloc_high_prio_tasks_by_cpu(
	249	+ socket=socket, isolate_numa=isolate_numa, socket_high_prio_tasks=socket_high_prio_tasks
	250	+ )
	251	+ return
	252	+
	253	+ end = min(start + task.min_cluster, cluster_count)
	254	+ take_cluster = clusters[start:end]
	255	+ cpu_list = self._get_cpu_list_of_clusters(take_cluster)
	256	+
	257	+ while len(cpu_list) < task.min_cpu and end < cluster_count:
	258	+ end += 1
	259	+ take_cluster = clusters[start:end]
	260	+ cpu_list = self._get_cpu_list_of_clusters(take_cluster)
	261	+
	262	+ if len(cpu_list) < task.min_cpu:
	263	+ print(
	264	+ f"[Warning] clusters {take_cluster} only have {len(cpu_list)} CPUs, "
	265	+ f"task(id={task.task_id}, name={task.name}) needs {task.min_cpu} CPUs, "
	266	+ f"fallback to CPU-level allocation."
	267	+ )
	268	+ self._alloc_high_prio_tasks_by_cpu(
	269	+ socket=socket, isolate_numa=isolate_numa, socket_high_prio_tasks=socket_high_prio_tasks
	270	+ )
	271	+ return
	272	+
	273	+ self._update_task_affinity(task, [socket.domain_id], [isolate_numa.domain_id], take_cluster, cpu_list)
	274	+ task.do_isolate = True
	275	+ start = end
	276	+
	277	+ def _alloc_high_prio_tasks_by_cpu(
	278	+ self, socket: SocketDomain, isolate_numa: NumaDomain, socket_high_prio_tasks: list[Task]
	279	+ ) -> None:
	280	+ task_needed_cpus = sum(task.min_cpu for task in socket_high_prio_tasks)
	281	+ cpu_list = isolate_numa.cpus.to_list()
	282	+ cpu_num = len(cpu_list)
	283	+
	284	+ if cpu_num >= task_needed_cpus:
	285	+ start = 0
	286	+ for task in socket_high_prio_tasks:
	287	+ end = start + task.min_cpu
	288	+ take_core = cpu_list[start:end]
	289	+ cluster_ids = self.domain.get_clusters_of_cpus(take_core)
	290	+ self._update_task_affinity(task, [socket.domain_id], [isolate_numa.domain_id], cluster_ids, take_core)
	291	+ task.do_isolate = True
	292	+ start = end
	293	+ return
	294	+
	295	+ print(
	296	+ f"[Warning] no enough cpus in isolate numa [{isolate_numa.domain_id}], "
	297	+ f"need {task_needed_cpus}, actual {cpu_num}."
	298	+ )
	299	+
	300	+ clusters = isolate_numa.get_all_children_id()
	301	+ for task in socket_high_prio_tasks:
	302	+ self._update_task_affinity(task, [socket.domain_id], [isolate_numa.domain_id], clusters, cpu_list)
	303	+ task.do_isolate = True
	304	+
	305	+ def _schedule_high_prio_tasks_to_isolate_cpu(
	306	+ self, socket: SocketDomain, socket_high_prio_tasks: list[Task]
	307	+ ) -> bool:
	308	+ """单NUMA场景下,为高优先级任务分配独立CPU"""
	309	+ socket_numas = self._socket_to_normal_numas[socket.domain_id]
	310	+ if not socket_numas:
	311	+ print(f"[Error] no numa found for socket [{socket.domain_id}]")
	312	+ return False
	313	+
	314	+ all_clusters = []
	315	+ for numa_id in socket_numas:
	316	+ numa = self.domain.get_numa_domain(numa_id)
	317	+ if numa is not None:
	318	+ all_clusters.extend(numa.get_all_children_id())
	319	+ all_clusters = sorted(all_clusters)
	320	+ cluster_num = len(all_clusters)
	321	+
	322	+ if cluster_num == 0:
	323	+ print(f"[Error] no clusters found in socket [{socket.domain_id}]")
	324	+ return False
	325	+
	326	+ # 计算高优先级任务需要的CPU数量（每个任务1个CPU）
	327	+ task_needed_cpus = 0
	328	+ for task in socket_high_prio_tasks:
	329	+ task_needed_cpus += task.min_cpu
	330	+
	331	+ # 按CPU数量均分：高优先级任务占用前N个CPU，剩余给普通任务
	332	+ cpu_list = []
	333	+ for cluster_id in all_clusters:
	334	+ cluster_domain = self.domain.get_cluster_domain(cluster_id)
	335	+ if cluster_domain:
	336	+ cpu_list.extend(cluster_domain.cpus.to_list())
	337	+ cpu_list = sorted(cpu_list)
	338	+
	339	+ if len(cpu_list) < task_needed_cpus:
	340	+ print(f"[Warning] not enough CPUs for high priority tasks: need {task_needed_cpus}, actual {len(cpu_list)}")
	341	+ # 如果CPU不够，所有高优先级任务共享所有CPU
	342	+ for task in socket_high_prio_tasks:
	343	+ self._update_task_affinity(task, [socket.domain_id], socket_numas, all_clusters, cpu_list)
	344	+ task.do_isolate = True
	345	+ return True
	346	+
	347	+ start = 0
	348	+ for task in socket_high_prio_tasks:
	349	+ end = start + task.min_cpu
	350	+ take_cpu = cpu_list[start:end]
	351	+ cluster_ids = self.domain.get_clusters_of_cpus(take_cpu)
	352	+ self._update_task_affinity(task, [socket.domain_id], socket_numas, cluster_ids, take_cpu)
	353	+ task.do_isolate = True
	354	+ start = end
	355	+
	356	+ self._socket_to_isolate_cpus[socket.domain_id] = cpu_list[:task_needed_cpus]
	357	+ return True
	358	+
	359	+ def _schedule_socket_task_groups_to_numas(self, socket: SocketDomain, numas: list[int]) -> bool:
	360	+ socket_task_groups = self._socket_to_task_groups[socket.domain_id]
	361	+ task_group_num = len(socket_task_groups)
	362	+
	363	+ numa_num = len(numas)
	364	+ if numa_num == 0:
	365	+ print(f"no numa to schedule for socket{socket.domain_id}")
	366	+ return False
	367	+
	368	+ base = task_group_num // numa_num
	369	+ extra = task_group_num % numa_num
	370	+
	371	+ start = 0
	372	+ for i, numa in enumerate(numas):
	373	+ take = base + (1 if i < extra else 0)
	374	+ end = start + take
	375	+ take_task_groups = socket_task_groups[start:end]
	376	+ if not take_task_groups:
	377	+ break
	378	+ self._numa_to_task_groups[numa].extend(take_task_groups)
	379	+ start = end
	380	+
	381	+ return True
	382	+
	383	+ def _schedule_numa_task_groups_to_clusters(self, numa_id: int) -> bool:
	384	+ numa_domain = self.domain.get_numa_domain(numa_id)
	385	+ if numa_domain is None:
	386	+ print(f"numa domain [{numa_id}] not found")
	387	+ return False
	388	+
	389	+ cluster_domains = numa_domain.get_all_children_id()
	390	+
	391	+ current_socket_id = None
	392	+ for socket_id, normal_numas in self._socket_to_normal_numas.items():
	393	+ if numa_id in normal_numas:
	394	+ current_socket_id = socket_id
	395	+ break
	396	+
	397	+ if current_socket_id is not None:
	398	+ isolate_cpus = self._socket_to_isolate_cpus.get(current_socket_id, [])
	399	+ if isolate_cpus:
	400	+ filtered_clusters = []
	401	+ for cluster_id in cluster_domains:
	402	+ cluster_domain = self.domain.get_cluster_domain(cluster_id)
	403	+ if cluster_domain is None:
	404	+ continue
	405	+ cluster_cpus = set(cluster_domain.cpus.to_list())
	406	+ if not cluster_cpus.intersection(isolate_cpus):
	407	+ filtered_clusters.append(cluster_id)
	408	+ if filtered_clusters:
	409	+ cluster_domains = filtered_clusters
	410	+
	411	+ cluster_num = len(cluster_domains)
	412	+
	413	+ numa_task_groups = self._numa_to_task_groups[numa_id]
	414	+ task_group_num = len(numa_task_groups)
	415	+ if task_group_num == 0:
	416	+ return True
	417	+
	418	+ if cluster_num == 0:
	419	+ print(f"numa[{numa_id}] has no available clusters, skip.")
	420	+ return False
	421	+
	422	+ reserve_count = max(1, int(cluster_num * 0.2))
	423	+ self._numa_to_background_clusters[numa_id] = cluster_domains[:reserve_count]
	424	+ cluster_domains = cluster_domains[reserve_count:]
	425	+ cluster_num = len(cluster_domains)
	426	+ print(
	427	+ f"numa[{numa_id}] reserve {reserve_count} cluster(s) for background tasks, "
	428	+ f"{cluster_num} cluster(s) left for normal tasks."
	429	+ )
	430	+
	431	+ if cluster_num < task_group_num:
	432	+ # cluster数量少于group数量时，不再做细粒度拆分，每个亲和组分享该numa下所有cluster
	433	+ print(f"numa [{numa_id}]: cluster number ({cluster_num}) < task group number ({task_group_num})")
	434	+ for _, task_group_id in enumerate(numa_task_groups):
	435	+ self._task_group_to_clusters[task_group_id].extend(cluster_domains)
	436	+ return True
	437	+
	438	+ base = cluster_num // task_group_num
	439	+ start = 0
	440	+ for _, task_group_id in enumerate(numa_task_groups):
	441	+ end = start + base
	442	+ take_clusters = cluster_domains[start:end]
	443	+ if not take_clusters:
	444	+ break
	445	+ self._task_group_to_clusters[task_group_id].extend(take_clusters)
	446	+ start = end
	447	+
	448	+ return True
	449	+
	450	+ def _update_task_group_affinity(self, group: TaskGroup) -> bool:
	451	+ # update task group clusters
	452	+ task_group_clusters = self._task_group_to_clusters.get(group.group_id)
	453	+ if task_group_clusters is None:
	454	+ print(f"no clusters for task group [{group.group_id}]")
	455	+ return False
	456	+ group.cluster = sorted(task_group_clusters)
	457	+
	458	+ # update task group cpus
	459	+ for cluster_id in group.cluster:
	460	+ cluster_domain = self.domain.get_cluster_domain(cluster_id)
	461	+ if cluster_domain is None:
	462	+ print(f"cluster domain [{cluster_id}] not found")
	463	+ return False
	464	+ group.cpus.set_list(cluster_domain.cpus.to_list())
	465	+
	466	+ # update task group numas
	467	+ task_group_numas: list[int] = []
	468	+ for numa_id, numa_task_groups in self._numa_to_task_groups.items():
	469	+ for group_id in numa_task_groups:
	470	+ if group_id == group.group_id:
	471	+ task_group_numas.append(numa_id)
	472	+ break
	473	+ group.numa = sorted(task_group_numas)
	474	+
	475	+ # update socket
	476	+ task_group_sockets: list[int] = []
	477	+ for socked_id, socket_task_groups in self._socket_to_task_groups.items():
	478	+ for group_id in socket_task_groups:
	479	+ if group_id == group.group_id:
	480	+ task_group_sockets.append(socked_id)
	481	+ break
	482	+ group.socket = sorted(task_group_sockets)
	483	+
	484	+ return True
	485	+
	486	+ def _schedule_tasks_affinity_of_group(self, group: TaskGroup) -> bool:
	487	+ schedule_tasks = self._get_tasks_to_schedule_for_group(group)
	488	+ schedule_task_num = len(schedule_tasks)
	489	+ group_cluster_num = len(group.cluster)
	490	+
	491	+ if schedule_task_num == 0:
	492	+ print(f"no task need to schedule in group [{group.group_id}]")
	493	+ return True
	494	+
	495	+ group_cpus = group.cpus.to_list()
	496	+ if len(group_cpus) == 0:
	497	+ print(f"group[{group.group_id}] has no available CPU, cannot schedule.")
	498	+ return False
	499	+
	500	+ for task in schedule_tasks:
	501	+ self._update_task_affinity(task, group.socket, group.numa, group.cluster, group.cpus.to_list())
	502	+
	503	+ return True
	504	+
	505	+ def _get_tasks_to_schedule_for_group(self, group: TaskGroup) -> list[Task]:
	506	+ schedule_tasks: list[Task] = []
	507	+ for task in group.get_all_tasks():
	508	+ if task.cpus.count() == 0: # cpu尚未分配
	509	+ schedule_tasks.append(task)
	510	+ return schedule_tasks
	511	+
	512	+ def _update_task_group_with_high_prio_tasks(self, group: TaskGroup) -> None:
	513	+ for task in group.get_high_prio_tasks():
	514	+ if task.cpus.count() != 0 and task.do_isolate:
	515	+ group.isolate_numa.extend(task.numa)
	516	+ group.isolate_cluster.extend(task.cluster)
	517	+ group.isolate_cpus.set_list(task.cpus.to_list())
	518	+ group.isolate_numa = sorted(set(group.isolate_numa))
	519	+ group.isolate_cluster = sorted(set(group.isolate_cluster))
	520	+
	521	+ def _update_task_affinity(
	522	+ self,
	523	+ task: Task,
	524	+ socket: list[int],
	525	+ numa: list[int],
	526	+ cluster: list[int],
	527	+ cpus: list[int],
	528	+ ):
	529	+ task.socket = list(socket)
	530	+ task.numa = list(numa)
	531	+ task.cluster = list(cluster)
	532	+ task.assign_cpu(cpus)
	533	+
	534	+ def _get_cpu_list_of_clusters(self, clusters: list[int]) -> list:
	535	+ cpu_list: list = []
	536	+ for cluster_id in clusters:
	537	+ cluster_domain = self.domain.get_cluster_domain(cluster_id)
	538	+ if cluster_domain is None:
	539	+ print(f"cluster domain [{cluster_id}] not found")
	540	+ continue
	541	+ cpu_list.extend(cluster_domain.cpus.to_list())
	542	+ return sorted(cpu_list)
	543	+
	544	+ def get_background_task_cpus(self) -> list:
	545	+ """
	546	+ 获取背景任务应绑定的CPU列表。
	547	+ 优先使用空闲NUMA的CPU，若无空闲NUMA则使用各NUMA预留给系统任务的cluster。
	548	+ """
	549	+
	550	+ used_numas = set()
	551	+ for _, group in self.task.groups.items():
	552	+ for task in group.get_all_tasks():
	553	+ if task.numa:
	554	+ used_numas.update(task.numa)
	555	+ all_numas = set(self.domain.get_all_numas_id())
	556	+ unused_numas = all_numas - used_numas
	557	+
	558	+ if unused_numas:
	559	+ cpus = []
	560	+ for numa_id in sorted(unused_numas):
	561	+ numa = self.domain.get_numa_domain(numa_id)
	562	+ if numa:
	563	+ cpus.extend(numa.cpus.to_list())
	564	+ result = sorted(set(cpus))
	565	+ if result:
	566	+ print(
	567	+ f"Using unused NUMA CPUs: {utils.compress_continuous(result)} "
	568	+ f"numa={sorted(unused_numas)}"
	569	+ )
	570	+ return result
	571	+
	572	+ background_cpus = []
	573	+ for numa_id in sorted(self._numa_to_background_clusters):
	574	+ clusters = self._numa_to_background_clusters[numa_id]
	575	+ background_cpus.extend(self._get_cpu_list_of_clusters(clusters))
	576	+
	577	+ result = sorted(set(background_cpus))
	578	+ if result:
	579	+ print(
	580	+ f"Using reserved clusters from each NUMA: {utils.compress_continuous(result)} "
	581	+ f"(numa_reserved={dict(sorted(self._numa_to_background_clusters.items()))})"
	582	+ )
	583	+ return result
	584	+
	585	+ def _alloc_cpus_for_background_tasks(self) -> None:
	586	+ normal_cpus = self.get_background_task_cpus()
	587	+ if normal_cpus:
	588	+ self.task.background_tasks_cpus = normal_cpus
	589	+ print(f"Allocated CPUs: {utils.compress_continuous(normal_cpus)}")
	590	+ else:
	591	+ print("No target CPUs found for background tasks")

Aa-sched/a_sched/task.py+704-0

@@ -0,0 +1,704 @@
	1	+from __future__ import annotations
	2	+from collections import defaultdict
	3	+from enum import Enum, auto
	4	+from datetime import datetime
	5	+
	6	+from a_sched.utils import CPUMask
	7	+import a_sched.utils as utils
	8	+
	9	+
	10	+class PriorityLevel(Enum):
	11	+ NORMAL = auto()
	12	+ HIGH = auto()
	13	+
	14	+
	15	+class TaskType(Enum):
	16	+ THREAD = auto() # 线程
	17	+ PROCESS = auto() # 进程
	18	+ NPU = auto() # NPU
	19	+
	20	+
	21	+class TaskGroupStatus(Enum):
	22	+ UNKNOWN = auto()
	23	+ CREATED = auto()
	24	+
	25	+
	26	+class Task:
	27	+ def __init__(self, task_type: TaskType, group_id: int, task_id: int, name: str \| None = None):
	28	+ # 基本信息
	29	+ self.group_id: int = group_id
	30	+ self.task_id: int = task_id
	31	+ self.name: str \| None = name
	32	+ self.task_type: TaskType = task_type
	33	+ self.priority: PriorityLevel = PriorityLevel.NORMAL # 默认优先级为Normal
	34	+ self.bind_npu: int \| None = None
	35	+
	36	+ # 调度管理
	37	+ self.cpus: CPUMask = CPUMask() # 分配到的cpu
	38	+ self.socket: list[int] = [] # 调度到的socket
	39	+ self.numa: list[int] = [] # 调度到的numa
	40	+ self.cluster: list[int] = [] # 调度到的cluster
	41	+ self.do_isolate: bool = False # 是否做了隔离调度
	42	+
	43	+ self._type_names = {TaskType.THREAD: "THREAD", TaskType.PROCESS: "PROCESS", TaskType.NPU: "NPU"}
	44	+ self._priority_names = {PriorityLevel.NORMAL: "NORMAL", PriorityLevel.HIGH: "HIGH"}
	45	+
	46	+ @property
	47	+ def min_cpu(self) -> int:
	48	+ return 1
	49	+
	50	+ @property
	51	+ def min_cluster(self) -> int:
	52	+ return 1
	53	+
	54	+ def assign_cpu(self, cpus: list[int]) -> None:
	55	+ pass
	56	+
	57	+ def bind_cpu(self) -> None:
	58	+ pass
	59	+
	60	+ def __str__(self) -> str:
	61	+ type_name = self._type_names.get(self.task_type, "UNKNOWN")
	62	+ priority = self._priority_names.get(self.priority, "UNKNOWN")
	63	+ return (
	64	+ f"{type_name}[{self.task_id}]: "
	65	+ f"name={self.name if self.name is not None else ''}, priority={priority}, "
	66	+ f"socket={self.socket}, "
	67	+ f"numa=[{utils.compress_continuous(self.numa)}], "
	68	+ f"cluster=[{utils.compress_continuous(self.cluster)}], "
	69	+ f"cpu=[{self.cpus}]"
	70	+ )
	71	+
	72	+
	73	+class ThreadTask(Task):
	74	+ def __init__(self, group_id: int, tid: int, pid: int, name: str \| None = None):
	75	+ super().__init__(TaskType.THREAD, group_id=group_id, task_id=tid, name=name)
	76	+ self.pid: int = pid
	77	+
	78	+ def bind_cpu(self) -> None:
	79	+ print(f"binding thread[{self.task_id}]({self.name}) to cpu [{self.cpus}]")
	80	+ utils.bind_thread_to_cpus(tid=self.task_id, cpus=self.cpus.to_list())
	81	+
	82	+ def assign_cpu(self, cpus: list[int]) -> None:
	83	+ if not cpus:
	84	+ print(f"[Error] assign cpu for thread[{self.task_id}](name={self.name}) fail, cpus={cpus}")
	85	+ return
	86	+ self.cpus.set_list(cpus[:1])
	87	+
	88	+
	89	+class ProcessTask(Task):
	90	+ def __init__(self, group_id: int, pid: int, name: str \| None = None):
	91	+ super().__init__(TaskType.PROCESS, group_id=group_id, task_id=pid, name=name)
	92	+
	93	+ def bind_cpu(self) -> None:
	94	+ print(f"binding process[{self.task_id}]({self.name}) to cpu [{self.cpus}]")
	95	+ utils.bind_process_to_cpus(pid=self.task_id, cpus=self.cpus.to_list())
	96	+
	97	+ def assign_cpu(self, cpus: list[int]) -> None:
	98	+ self.cpus.set_list(cpus)
	99	+
	100	+
	101	+class NpuTask(Task):
	102	+ ACL_THREAD = "acl_thread"
	103	+ RELEASE_THREAD = "release_thread"
	104	+ RT_RECYCLE_THREAD = "RT_RECYCLE"
	105	+ SQ_IRQ = "sq_send_trigger_irq"
	106	+ CQ_IRQ = "cq_update_irq"
	107	+ TRS_MBOX_IRQ_PREFIX = "trs-mbox"
	108	+
	109	+ def __init__(self, group_id: int, npu_id: int, bind_pid: int, name: str = ""):
	110	+ super().__init__(TaskType.NPU, group_id=group_id, task_id=npu_id, name=name)
	111	+ self.priority = PriorityLevel.HIGH
	112	+
	113	+ self._bind_pid = bind_pid # npu关联的worker进程pid
	114	+ self._acl_thread: int \| None = None # acl线程 tid
	115	+ self._release_thread: int \| None = None # release线程 tid
	116	+ self._rt_recycle_thread: int \| None = None # rt_recycle线程 tid
	117	+ self._sq_irq: int \| None = None # sq_send_trigger_irq中断号
	118	+ self._cq_irqs: list[int] = [] # cq_update_irq中断号，有16个
	119	+ self._trs_mbox_irq: int \| None = None # trs_mbox中断号
	120	+ self._dev_sq_task: int \| None = None # dev_sq_task进程pid
	121	+ self._dev_sq_send_wq: int \| None = None # dev_sq_send_wq进程pid
	122	+
	123	+ self._bind_high_prio_thread: list[ThreadTask] = []
	124	+
	125	+ self.trs_mbox_name = f"{self.TRS_MBOX_IRQ_PREFIX}-{self.task_id}-0"
	126	+ self.dev_sq_task_name = f"dev{self.task_id}_sq_task"
	127	+ self.dev_sq_send_wq_name = f"dev{self.task_id}_sq_send_wq"
	128	+
	129	+ self._init_npu()
	130	+
	131	+ self._acl_thread_cpus: list[int] = []
	132	+ self._release_thread_cpus: list[int] = []
	133	+ self._rt_recycle_thread_cpus: list[int] = []
	134	+ self._sq_irq_cpus: list[int] = []
	135	+ self._cq_irqs_cpus: list[int] = []
	136	+ self._trs_mbox_irq_cpus: list[int] = []
	137	+ self._dev_sq_task_cpus: list[int] = []
	138	+ self._dev_sq_send_wq_cpus: list[int] = []
	139	+
	140	+ def _init_npu(self) -> None:
	141	+ print(f"\nStarting init npu[{self.task_id}]...")
	142	+ self._get_npu_threads()
	143	+ self._get_npu_irqs()
	144	+ self._get_npu_dev_sq()
	145	+
	146	+ @property
	147	+ def min_cpu(self) -> int:
	148	+ need_cpu = 0
	149	+ need_cpu += 1 if self._acl_thread is not None else 0
	150	+ need_cpu += 1 if self._release_thread is not None else 0
	151	+ need_cpu += 1 if self._rt_recycle_thread is not None else 0
	152	+ need_cpu += 1 if self._sq_irq is not None else 0
	153	+ need_cpu += 1 if self._cq_irqs else 0
	154	+ need_cpu += 1 if self._trs_mbox_irq is not None else 0
	155	+ need_cpu += 1 if self._dev_sq_task is not None else 0
	156	+ need_cpu += 1 if self._dev_sq_send_wq is not None else 0
	157	+ need_cpu += len(self._bind_high_prio_thread)
	158	+ return need_cpu
	159	+
	160	+ @property
	161	+ def min_cluster(self) -> int:
	162	+ # todo: 需要根据实际硬件参数来计算cluster数量，暂时按照a3来处理
	163	+ cpu_num_per_cluster = 4
	164	+ base = self.min_cpu // cpu_num_per_cluster
	165	+ extra = self.min_cpu % cpu_num_per_cluster
	166	+ return base + 1 if extra > 0 else base
	167	+
	168	+ def _get_npu_dev_sq(self) -> None:
	169	+ print(f"Starting get npu[{self.task_id}] dev_sq...")
	170	+
	171	+ # dev_sq_task
	172	+ pids = utils.get_pid_by_process_name(process_name=self.dev_sq_task_name)
	173	+ if not pids:
	174	+ print(f"[Warning] cannot found {self.dev_sq_task_name} for npu [{self.task_id}]")
	175	+ else:
	176	+ self._dev_sq_task = pids[0][0]
	177	+ print(f"get npu process (name={self.dev_sq_task_name}, pid={self._dev_sq_task}) success.")
	178	+
	179	+ # dev_sq_send_wq
	180	+ pids = utils.get_pid_by_process_name(process_name=self.dev_sq_send_wq_name[:15])
	181	+ if not pids:
	182	+ print(f"[Warning] cannot found {self.dev_sq_send_wq_name} for npu [{self.task_id}]")
	183	+ else:
	184	+ self._dev_sq_send_wq = pids[0][0]
	185	+ print(f"get npu process (name={self.dev_sq_send_wq_name}, pid={self._dev_sq_send_wq}) success.")
	186	+
	187	+ def _get_npu_threads(self) -> None:
	188	+ print(f"Starting get npu[{self.task_id}] threads...")
	189	+
	190	+ # acl_thread
	191	+ tids = utils.get_tid_by_thread_name(thread_name=self.ACL_THREAD, pid=self._bind_pid)
	192	+ if not tids:
	193	+ print(f"[Warning] cannot found {self.ACL_THREAD} for npu [{self.task_id}]")
	194	+ else:
	195	+ self._acl_thread = tids[0]
	196	+ print(f"get npu thread (name={self.ACL_THREAD}, tid={self._acl_thread}) success.")
	197	+
	198	+ # release_thread
	199	+ tids = utils.get_tid_by_thread_name(thread_name=self.RELEASE_THREAD, pid=self._bind_pid)
	200	+ if not tids:
	201	+ print(f"[Warning] cannot found {self.RELEASE_THREAD} for npu [{self.task_id}]")
	202	+ else:
	203	+ self._release_thread = tids[0]
	204	+ print(f"get npu thread (name={self.RELEASE_THREAD}, tid={self._release_thread}) success.")
	205	+
	206	+ # rt_recycle
	207	+ tids = utils.get_tid_by_thread_name(thread_name=self.RT_RECYCLE_THREAD, pid=self._bind_pid)
	208	+ if not tids:
	209	+ raise ValueError(f"[Warning] cannot found {self.RT_RECYCLE_THREAD} for npu [{self.task_id}]")
	210	+ else:
	211	+ self._rt_recycle_thread = tids[0]
	212	+ print(f"get npu thread (name={self.RT_RECYCLE_THREAD}, tid={self._rt_recycle_thread}) success.")
	213	+
	214	+ def _get_npu_irqs(self) -> None:
	215	+ print(f"Starting get npu[{self.task_id}] irqs...")
	216	+
	217	+ # sq
	218	+ irqs = utils.get_npu_irq_by_name(irq_name=self.SQ_IRQ, npu_id=self.task_id)
	219	+ if not irqs:
	220	+ print(f"[Warning] cannot found {self.SQ_IRQ} of npu [{self.task_id}]")
	221	+ else:
	222	+ self._sq_irq = irqs[0]
	223	+ print(f"get npu irq (name={self.SQ_IRQ}, id={self._sq_irq}) success.")
	224	+
	225	+ # cq
	226	+ irqs = utils.get_npu_irq_by_name(irq_name=self.CQ_IRQ, npu_id=self.task_id)
	227	+ if not irqs:
	228	+ print(f"[Warning] cannot found {self.CQ_IRQ} of npu [{self.task_id}]")
	229	+ else:
	230	+ self._cq_irqs = irqs
	231	+ print(f"get npu irq (name={self.CQ_IRQ}, id={self._cq_irqs}) success.")
	232	+
	233	+ # trs_mbox
	234	+ irqs = utils.get_npu_irq_by_name(irq_name=self.trs_mbox_name, npu_id=self.task_id)
	235	+ if not irqs:
	236	+ print(f"[Warning] cannot found {self.trs_mbox_name} of npu [{self.task_id}]")
	237	+ else:
	238	+ self._trs_mbox_irq = irqs[0]
	239	+ print(f"get npu irq (name={self.trs_mbox_name}, id={self._trs_mbox_irq}) success.")
	240	+
	241	+ def bind_cpu(self) -> None:
	242	+ if self._acl_thread is not None and self._acl_thread_cpus:
	243	+ print(
	244	+ f"binding npu[{self.task_id}] thread[{self._acl_thread}]({self.ACL_THREAD}) "
	245	+ f"to cpu {self._acl_thread_cpus}"
	246	+ )
	247	+ utils.bind_thread_to_cpus(tid=self._acl_thread, cpus=self._acl_thread_cpus)
	248	+
	249	+ if self._release_thread is not None and self._release_thread_cpus:
	250	+ print(
	251	+ f"binding npu[{self.task_id}] thread[{self._release_thread}]({self.RELEASE_THREAD}) "
	252	+ f"to cpu {self._release_thread_cpus}"
	253	+ )
	254	+ utils.bind_thread_to_cpus(tid=self._release_thread, cpus=self._release_thread_cpus)
	255	+
	256	+ if self._rt_recycle_thread is not None and self._rt_recycle_thread_cpus:
	257	+ print(
	258	+ f"binding npu[{self.task_id}] thread[{self._rt_recycle_thread}]({self.RT_RECYCLE_THREAD}) "
	259	+ f"to cpu {self._rt_recycle_thread_cpus}"
	260	+ )
	261	+ utils.bind_thread_to_cpus(tid=self._rt_recycle_thread, cpus=self._rt_recycle_thread_cpus)
	262	+
	263	+ if self._sq_irq is not None and self._sq_irq_cpus:
	264	+ print(f"binding npu[{self.task_id}] irq[{self._sq_irq}]({self.SQ_IRQ}) to cpu {self._sq_irq_cpus}")
	265	+ utils.bind_irq_to_cpus(irq_id=self._sq_irq, cpus=self._sq_irq_cpus, irq_name=self.SQ_IRQ)
	266	+
	267	+ if self._cq_irqs and self._cq_irqs_cpus:
	268	+ print(f"binding npu[{self.task_id}] irq{self._cq_irqs}({self.CQ_IRQ}) to cpu {self._cq_irqs_cpus}")
	269	+ for cq in self._cq_irqs:
	270	+ utils.bind_irq_to_cpus(irq_id=cq, cpus=self._cq_irqs_cpus, irq_name=self.CQ_IRQ)
	271	+
	272	+ if self._trs_mbox_irq is not None and self._trs_mbox_irq_cpus:
	273	+ print(
	274	+ f"binding npu[{self.task_id}] irq[{self._trs_mbox_irq}]({self.trs_mbox_name}) "
	275	+ f"to cpu {self._trs_mbox_irq_cpus}"
	276	+ )
	277	+ utils.bind_irq_to_cpus(irq_id=self._trs_mbox_irq, cpus=self._trs_mbox_irq_cpus, irq_name=self.trs_mbox_name)
	278	+
	279	+ if self._dev_sq_task is not None and self._dev_sq_task_cpus:
	280	+ print(
	281	+ f"binding npu[{self.task_id}] process[{self._dev_sq_task}]({self.dev_sq_task_name}) "
	282	+ f"to cpu {self._dev_sq_task_cpus}"
	283	+ )
	284	+ utils.bind_process_to_cpus(pid=self._dev_sq_task, cpus=self._dev_sq_task_cpus)
	285	+
	286	+ if self._dev_sq_send_wq is not None and self._dev_sq_send_wq_cpus:
	287	+ print(
	288	+ f"binding npu[{self.task_id}] process[{self._dev_sq_send_wq}]({self.dev_sq_send_wq_name}) "
	289	+ f"to cpu {self._dev_sq_send_wq_cpus}"
	290	+ )
	291	+ utils.bind_npu_sq_send_wq_to_cpus(npu_id=self.task_id, cpus=self._dev_sq_send_wq_cpus)
	292	+
	293	+ def assign_cpu(self, cpus: list[int]) -> None:
	294	+ self.cpus.set_list(cpus)
	295	+ cpu_num = len(cpus)
	296	+ if cpu_num < self.min_cpu:
	297	+ print(
	298	+ f"[Error] assign cpu for npu[{self.task_id}] fail, no enough cpus, assigned {cpu_num}, need {self.min_cpu}"
	299	+ )
	300	+ return
	301	+
	302	+ start = 0
	303	+ for thread in self._bind_high_prio_thread:
	304	+ end = start + 1
	305	+ thread.assign_cpu(cpus[start:end])
	306	+ start = end
	307	+ thread.socket = list(self.socket)
	308	+ thread.numa = list(self.numa)
	309	+ thread.cluster = list(self.cluster)
	310	+ if self._acl_thread is not None:
	311	+ end = start + 1
	312	+ self._acl_thread_cpus = cpus[start:end]
	313	+ start = end
	314	+ if self._release_thread is not None:
	315	+ end = start + 1
	316	+ self._release_thread_cpus = cpus[start:end]
	317	+ start = end
	318	+ if self._rt_recycle_thread is not None:
	319	+ end = start + 1
	320	+ self._rt_recycle_thread_cpus = cpus[start:end]
	321	+ start = end
	322	+ if self._dev_sq_task is not None:
	323	+ end = start + 1
	324	+ self._dev_sq_task_cpus = cpus[start:end]
	325	+ start = end
	326	+ if self._dev_sq_send_wq is not None:
	327	+ end = start + 1
	328	+ self._dev_sq_send_wq_cpus = cpus[start:end]
	329	+ start = end
	330	+ if self._sq_irq is not None:
	331	+ end = start + 1
	332	+ self._sq_irq_cpus = cpus[start:end]
	333	+ start = end
	334	+ if self._cq_irqs is not None:
	335	+ end = start + 1
	336	+ self._cq_irqs_cpus = cpus[start:end]
	337	+ start = end
	338	+ if self._trs_mbox_irq is not None:
	339	+ end = start + 1
	340	+ self._trs_mbox_irq_cpus = cpus[start:end]
	341	+ start = end
	342	+
	343	+ def __str__(self) -> str:
	344	+ npu_str = ""
	345	+ type_name = self._type_names.get(self.task_type, "UNKNOWN")
	346	+ priority = self._priority_names.get(self.priority, "UNKNOWN")
	347	+ npu_str = (
	348	+ f"{type_name}[{self.task_id}]: priority={priority}, "
	349	+ f"socket={self.socket}, numa=[{utils.compress_continuous(self.numa)}], cluster=[{utils.compress_continuous(self.cluster)}], "
	350	+ f"cpu=[{self.cpus}]"
	351	+ )
	352	+ npu_str += (
	353	+ (
	354	+ f"\n - Process[{self._dev_sq_task}]({self.dev_sq_task_name}): cpu={utils.compress_continuous(self._dev_sq_task_cpus)}"
	355	+ )
	356	+ if self._dev_sq_task is not None
	357	+ else ""
	358	+ )
	359	+ npu_str += (
	360	+ (
	361	+ f"\n - Process[{self._dev_sq_send_wq}]({self.dev_sq_send_wq_name}): cpu={utils.compress_continuous(self._dev_sq_send_wq_cpus)}"
	362	+ )
	363	+ if self._dev_sq_send_wq is not None
	364	+ else ""
	365	+ )
	366	+ npu_str += (
	367	+ (
	368	+ f"\n - Thread[{self._acl_thread}]({self.ACL_THREAD}): cpu={utils.compress_continuous(self._acl_thread_cpus)}"
	369	+ )
	370	+ if self._acl_thread is not None
	371	+ else ""
	372	+ )
	373	+ npu_str += (
	374	+ (
	375	+ f"\n - Thread[{self._release_thread}]({self.RELEASE_THREAD}): cpu={utils.compress_continuous(self._release_thread_cpus)}"
	376	+ )
	377	+ if self._release_thread is not None
	378	+ else ""
	379	+ )
	380	+ npu_str += (
	381	+ (
	382	+ f"\n - Thread[{self._rt_recycle_thread}]({self.RT_RECYCLE_THREAD}): cpu={utils.compress_continuous(self._rt_recycle_thread_cpus)}"
	383	+ )
	384	+ if self._rt_recycle_thread is not None
	385	+ else ""
	386	+ )
	387	+ npu_str += (
	388	+ (f"\n - Irq[{self._sq_irq}]({self.SQ_IRQ}): cpu={utils.compress_continuous(self._sq_irq_cpus)}")
	389	+ if self._sq_irq is not None
	390	+ else ""
	391	+ )
	392	+ npu_str += (
	393	+ (f"\n - Irq[{self._cq_irqs[0]}]({self.CQ_IRQ}): cpu={utils.compress_continuous(self._cq_irqs_cpus)}")
	394	+ if self._cq_irqs
	395	+ else ""
	396	+ )
	397	+ npu_str += (
	398	+ (
	399	+ f"\n - Irq[{self._trs_mbox_irq}]({self.trs_mbox_name}): cpu={utils.compress_continuous(self._trs_mbox_irq_cpus)}"
	400	+ )
	401	+ if self._trs_mbox_irq is not None
	402	+ else ""
	403	+ )
	404	+ return npu_str
	405	+
	406	+
	407	+class TaskGroup:
	408	+ def __init__(self, group_id: int, name: str = "", desc: str = ""):
	409	+ # 基本信息
	410	+ self.group_id: int = group_id
	411	+ self.name: str = name
	412	+ self.description: str = desc
	413	+
	414	+ # 任务管理
	415	+ self.thread_tasks: dict[int, ThreadTask] = {} # tid -> task
	416	+ self.process_tasks: dict[int, ProcessTask] = {} # pid -> task
	417	+ self.npu_tasks: dict[int, NpuTask] = {} # npu_id -> task
	418	+
	419	+ # 调度管理
	420	+ self.cpus: CPUMask = CPUMask() # 分配到的cpu
	421	+ self.socket: list[int] = [] # 调度到的socket
	422	+ self.numa: list[int] = [] # 调度到的numa
	423	+ self.cluster: list[int] = [] # 调度到的cluster
	424	+ self.isolate_numa: list[int] = []
	425	+ self.isolate_cluster: list[int] = []
	426	+ self.isolate_cpus: CPUMask = CPUMask()
	427	+
	428	+ # 状态和统计
	429	+ self.created_at: datetime = datetime.now()
	430	+ self.status: TaskGroupStatus = TaskGroupStatus.CREATED
	431	+
	432	+ def get_all_tasks_num(self) -> int:
	433	+ return len(self.thread_tasks) + len(self.process_tasks) + len(self.npu_tasks)
	434	+
	435	+ def get_all_tasks(self) -> list[Task]:
	436	+ all_tasks: list[Task] = []
	437	+ for _, process in self.process_tasks.items():
	438	+ all_tasks.append(process)
	439	+ for _, thread in self.thread_tasks.items():
	440	+ all_tasks.append(thread)
	441	+ for _, npu in self.npu_tasks.items():
	442	+ all_tasks.append(npu)
	443	+ return all_tasks
	444	+
	445	+ def get_normal_prio_tasks(self) -> list[Task]:
	446	+ normal_prio_tasks: list[Task] = []
	447	+ for _, process in self.process_tasks.items():
	448	+ if process.priority == PriorityLevel.NORMAL:
	449	+ normal_prio_tasks.append(process)
	450	+ for _, thread in self.thread_tasks.items():
	451	+ if thread.priority == PriorityLevel.NORMAL:
	452	+ normal_prio_tasks.append(thread)
	453	+ return normal_prio_tasks
	454	+
	455	+ def get_high_prio_tasks(self) -> list[Task]:
	456	+ high_prio_tasks: list[Task] = []
	457	+ for _, process in self.process_tasks.items():
	458	+ if process.priority == PriorityLevel.HIGH:
	459	+ high_prio_tasks.append(process)
	460	+ for _, thread in self.thread_tasks.items():
	461	+ # bind_npu的thread被加到npu_task中处理，这里不用再加了
	462	+ if thread.priority == PriorityLevel.HIGH and thread.bind_npu is None:
	463	+ high_prio_tasks.append(thread)
	464	+ for _, npu in self.npu_tasks.items():
	465	+ high_prio_tasks.append(npu)
	466	+ return high_prio_tasks
	467	+
	468	+ def __str__(self) -> str:
	469	+ return (
	470	+ f"TaskGroup[{self.group_id}]: name={self.name}, "
	471	+ f"socket={self.socket}, "
	472	+ f"numa=[{utils.compress_continuous(self.numa)}], "
	473	+ f"cluster=[{utils.compress_continuous(self.cluster)}], "
	474	+ f"cpu=[{self.cpus}]"
	475	+ f"{f', isol_numa=[{utils.compress_continuous(self.isolate_numa)}]' if self.isolate_numa else ''}"
	476	+ f"{f', isol_cluster=[{utils.compress_continuous(self.isolate_cluster)}]' if self.isolate_cluster else ''}"
	477	+ f"{f', isol_cpu=[{self.isolate_cpus}]' if self.isolate_cpus.count() != 0 else ''}"
	478	+ )
	479	+
	480	+
	481	+class TaskManager:
	482	+ """调度任务管理"""
	483	+
	484	+ def __init__(self) -> None:
	485	+ self.groups: dict[int, TaskGroup] = {} # group_id -> task group
	486	+ self.process_to_npu: dict[int, int] = {} # pid -> npu_id
	487	+ self.background_processes: dict[int, str] = {} # 背景进程
	488	+ self.background_tasks_cpus: list[int] = [] # 背景任务分配的CPU
	489	+ self._current_group_index: int = 0
	490	+
	491	+ def group_create(self, name: str = "") -> int:
	492	+ group_id = self._current_group_index
	493	+ self.groups[group_id] = TaskGroup(group_id=group_id, name=name)
	494	+ self._current_group_index += 1
	495	+ return group_id
	496	+
	497	+ def destory_group(self, group_id: int) -> None:
	498	+ self.groups.pop(group_id)
	499	+
	500	+ def group_add_thread(
	501	+ self,
	502	+ group_id: int,
	503	+ tid: int \| None = None,
	504	+ thread_name: str \| None = None,
	505	+ pid: int \| None = None,
	506	+ process_name: str \| None = None,
	507	+ ) -> None:
	508	+ if tid is None and thread_name is None:
	509	+ raise ValueError("add thread failed, either tid or thread_name must be set")
	510	+
	511	+ group = self.groups.get(group_id, None)
	512	+ if group is None:
	513	+ raise ValueError(f"add thread failed, group (id={group_id}) not found")
	514	+
	515	+ if pid is None and process_name is not None:
	516	+ pids = utils.get_pid_by_process_name(process_name)
	517	+ pid = pids[0][0] if pids else None
	518	+
	519	+ if pid is None:
	520	+ raise ValueError(f"add thread failed, process (pid={pid}, name={process_name}) not found")
	521	+
	522	+ if tid is None and thread_name is not None:
	523	+ tids = utils.get_tid_by_thread_name(thread_name=thread_name, pid=pid, process_name=process_name)
	524	+ tid = tids[0] if tids else None
	525	+
	526	+ if tid is None:
	527	+ raise ValueError(f"add thread failed, thread (tid={tid}, name={thread_name}) not found")
	528	+
	529	+ group.thread_tasks[tid] = ThreadTask(group_id=group_id, tid=tid, pid=pid, name=thread_name)
	530	+
	531	+ def group_remove_thread(
	532	+ self,
	533	+ group_id: int,
	534	+ tid: int \| None,
	535	+ thread_name: str \| None = None,
	536	+ pid: int \| None = None,
	537	+ process_name: str \| None = None,
	538	+ ) -> None:
	539	+ if tid is None and thread_name is None:
	540	+ raise ValueError("remove thread failed, either tid or thread_name must be set")
	541	+
	542	+ group = self.groups.get(group_id, None)
	543	+ if group is None:
	544	+ raise ValueError(f"remove thread failed, group({group_id}) not found")
	545	+
	546	+ if tid is None and thread_name is not None:
	547	+ tids = utils.get_tid_by_thread_name(thread_name=thread_name, pid=pid, process_name=process_name)
	548	+ tid = tids[0] if tids else None
	549	+
	550	+ if tid is None:
	551	+ raise ValueError(f"remove thread failed, thread (tid={tid}, name={thread_name}) not found")
	552	+
	553	+ group.thread_tasks.pop(tid, None)
	554	+
	555	+ def group_add_process(
	556	+ self, group_id: int, pid: int \| None = None, process_name: str \| None = None, parent_name: str \| None = None
	557	+ ) -> None:
	558	+ if pid is None and process_name is None:
	559	+ raise ValueError("add process failed, either pid or process_name must be set")
	560	+
	561	+ group = self.groups.get(group_id, None)
	562	+ if group is None:
	563	+ raise ValueError(f"add process failed, group (id={group_id}) not found")
	564	+
	565	+ if pid is None and process_name is not None:
	566	+ pids = utils.get_pid_by_process_name(process_name=process_name, parent_name=parent_name)
	567	+ pid = pids[0][0] if pids else None
	568	+
	569	+ if pid is None:
	570	+ raise ValueError(f"add process failed, process (pid={pid}, name={process_name}) not found")
	571	+
	572	+ group.process_tasks[pid] = ProcessTask(group_id=group_id, pid=pid, name=process_name)
	573	+
	574	+ def group_remove_process(
	575	+ self, group_id: int, pid: int \| None = None, process_name: str \| None = None, parent_name: str \| None = None
	576	+ ) -> None:
	577	+ if pid is None and process_name is None:
	578	+ raise ValueError("remove process failed, either pid or process_name must be set")
	579	+
	580	+ group = self.groups.get(group_id, None)
	581	+ if group is None:
	582	+ raise ValueError(f"remove process failed, group({group_id}) not found")
	583	+
	584	+ if pid is None and process_name is not None:
	585	+ pids = utils.get_pid_by_process_name(process_name=process_name, parent_name=parent_name)
	586	+ pid = pids[0][0] if pids else None
	587	+
	588	+ if pid is None:
	589	+ raise ValueError(f"remove process failed, process (pid={pid}, name={process_name}) not found")
	590	+
	591	+ group.process_tasks.pop(pid)
	592	+
	593	+ def thread_set_high_priority(
	594	+ self,
	595	+ tid: int \| None = None,
	596	+ thread_name: str \| None = None,
	597	+ pid: int \| None = None,
	598	+ process_name: str \| None = None,
	599	+ ) -> None:
	600	+ if tid is None and thread_name is None:
	601	+ raise ValueError("set high priority failed, either tid or thread_name must be set")
	602	+
	603	+ if tid is None and thread_name is not None:
	604	+ tids = utils.get_tid_by_thread_name(thread_name=thread_name, pid=pid, process_name=process_name)
	605	+ tid = tids[0] if tids else None
	606	+
	607	+ if tid is None:
	608	+ raise ValueError(f"set high priority failed, thread (tid={tid}, name={thread_name}) not found")
	609	+
	610	+ for _, group in self.groups.items():
	611	+ task = group.thread_tasks.get(tid)
	612	+ if task is not None:
	613	+ task.priority = PriorityLevel.HIGH
	614	+ break
	615	+
	616	+ def process_bind_npu(
	617	+ self, npu_id: int, pid: int \| None = None, process_name: str \| None = None, parent_name: str \| None = None
	618	+ ) -> None:
	619	+ if pid is None and process_name is None:
	620	+ raise ValueError("process bind npu failed, either pid or process_name must be set")
	621	+
	622	+ if pid is None and process_name is not None:
	623	+ pids = utils.get_pid_by_process_name(process_name=process_name, parent_name=parent_name)
	624	+ pid = pids[0][0] if pids else None
	625	+
	626	+ if pid is None:
	627	+ raise ValueError(f"process bind npu failed, process (pid={pid}, name={process_name}) not found")
	628	+
	629	+ group = self.find_group_by_pid(pid)
	630	+ if group is None:
	631	+ raise ValueError(f"process bind npu failed, not found group of process (pid={pid}, name={process_name})")
	632	+
	633	+ self.process_to_npu[pid] = npu_id
	634	+ group.npu_tasks[npu_id] = NpuTask(group_id=group.group_id, npu_id=npu_id, bind_pid=pid, name=f"npu[{npu_id}]")
	635	+
	636	+ def get_group_num(self) -> int:
	637	+ return len(self.groups)
	638	+
	639	+ def get_all_group_id(self) -> list[int]:
	640	+ task_groups_id: list[int] = []
	641	+ for group_id in self.groups:
	642	+ task_groups_id.append(group_id)
	643	+ return task_groups_id
	644	+
	645	+ def get_group(self, group_id: int) -> TaskGroup \| None:
	646	+ return self.groups.get(group_id, None)
	647	+
	648	+ def find_group_by_pid(self, find_pid: int) -> TaskGroup \| None:
	649	+ for _, group in self.groups.items():
	650	+ for pid in group.process_tasks:
	651	+ if pid == find_pid:
	652	+ return group
	653	+ return None
	654	+
	655	+ def get_high_prio_tasks_of_group(self, group_id: int) -> list[Task]:
	656	+ group = self.get_group(group_id)
	657	+ return group.get_high_prio_tasks() if group is not None else []
	658	+
	659	+ def print_all(self) -> None:
	660	+ for _, group in self.groups.items():
	661	+ self.print_task_group(group)
	662	+
	663	+ def print_task_group(self, group: TaskGroup) -> None:
	664	+ print(f"{group}")
	665	+ for _, process in group.process_tasks.items():
	666	+ print(f" - {process}")
	667	+ for _, thread in group.thread_tasks.items():
	668	+ print(f" - {thread}")
	669	+ for _, npu in group.npu_tasks.items():
	670	+ print(f" - {npu}")
	671	+
	672	+ def update_high_prio_thread_bind_npu(self) -> None:
	673	+ for _, group in self.groups.items():
	674	+ for _, thread in group.thread_tasks.items():
	675	+ if thread.priority != PriorityLevel.HIGH:
	676	+ continue
	677	+ thread.bind_npu = self.process_to_npu.get(thread.pid)
	678	+ if thread.bind_npu is None:
	679	+ continue
	680	+ npu_task = group.npu_tasks.get(thread.bind_npu)
	681	+ if npu_task is not None:
	682	+ npu_task._bind_high_prio_thread.append(thread)
	683	+
	684	+ def scan_background_tasks(self) -> None:
	685	+ """
	686	+ 扫描背景任务
	687	+ """
	688	+
	689	+ user_pids = set()
	690	+ for _, group in self.groups.items():
	691	+ for pid in group.process_tasks.keys():
	692	+ user_pids.add(pid)
	693	+ for _, npu_task in group.npu_tasks.items():
	694	+ if npu_task._dev_sq_task is not None:
	695	+ user_pids.add(npu_task._dev_sq_task)
	696	+ if npu_task._dev_sq_send_wq is not None:
	697	+ user_pids.add(npu_task._dev_sq_send_wq)
	698	+
	699	+ background_tasks = utils.get_all_user_processes()
	700	+ for pid, name in background_tasks:
	701	+ if pid not in user_pids:
	702	+ self.background_processes[pid] = name
	703	+
	704	+ print(f"[BackgroundTask] Scanned {len(self.background_processes)} processes")

Aa-sched/a_sched/utils.py+644-0

Aa-sched/examples/affinity_vllm.py+144-0

Aa-sched/pyproject.toml+13-0