如何让含Numba并行函数的Python脚本利用集群多节点全部核心

1. 核心限制说明 #

Numba的parallel=True默认依赖OpenMP/TBB这类线程库，仅能在单个节点内的CPU核心上并行，无法直接跨节点调度任务。要利用多节点的全部150核，必须结合分布式并行框架（如MPI），将任务拆分到各节点的进程中，每个进程再通过Numba的线程并行利用本地节点核心。

2. 调整Slurm作业配置 #

根据集群节点20核的规格，修改SBATCH脚本，通过MPI启动Python进程，匹配申请的150核资源：

方案A：每个核心对应1个MPI进程（线程数设为1） #

#SBATCH -n 150
#SBATCH --ntasks-per-node=20  # 匹配单节点20核的硬件规格
mpirun -np 150 NUMBA_NUM_THREADS=1 python main.py

方案B：每个节点对应1个MPI进程（线程数设为20） #

更高效，减少进程间通信开销：

#SBATCH -n 8  # 8个节点（7个满节点140核+1个10核节点）
#SBATCH --ntasks-per-node=1
mpirun -np 8 NUMBA_NUM_THREADS=20 python main.py

3. 修改Python脚本（结合mpi4py） #

通过mpi4py实现跨节点任务拆分与结果聚合：

from mpi4py import MPI
import numba

# 初始化MPI环境
comm = MPI.COMM_WORLD
rank = comm.Get_rank()  # 当前进程编号
total_processes = comm.Get_size()  # 总进程数

# 示例：拆分全局任务数据（需根据你的实际任务逻辑实现）
def split_global_data(global_data, rank, total_processes):
    chunk_size = len(global_data) // total_processes
    start = rank * chunk_size
    end = start + chunk_size if rank != total_processes -1 else len(global_data)
    return global_data[start:end]

# 全局任务数据（替换为你的实际数据）
global_data = list(range(150000))
local_data = split_global_data(global_data, rank, total_processes)

# Numba并行处理函数
@numba.njit(parallel=True)
def process_local_data(data):
    result = []
    for i in numba.prange(len(data)):
        result.append(data[i] ** 2)
    return result

# 每个进程处理本地数据
local_result = process_local_data(local_data)

# 收集所有进程结果到主进程（rank=0）
all_results = comm.gather(local_result, root=0)

# 主进程合并结果
if rank == 0:
    final_result = []
    for res in all_results:
        final_result.extend(res)
    print(f"处理完成，结果长度：{len(final_result)}")

4. 关键配置注意事项 #

• 避免资源冲突：MPI进程数 × 每个进程的Numba线程数，需等于你申请的总核心数150，防止资源浪费或超配。
• NUMBA_NUM_THREADS的合理设置：单节点内的进程线程数总和不能超过节点核心数（20），否则会引发Slurm资源报错。
• 任务拆分逻辑：需根据你的计算任务类型（如数组运算、批量任务）实现对应的全局数据拆分与结果合并逻辑。

内容的提问来源于stack exchange，提问作者velenos14