如何用mpi4py/multiprocessing在SLURM集群并行多循环任务？

仅用multiprocessing可行吗？ #

完全可行，不需要依赖mpi4py，核心是借助SLURM的**任务数组（Job Array）**来实现跨节点分配外层循环：

• 将外层20次拆分为20个独立任务，通过SLURM任务数组提交，每个任务申请1个节点+12核资源。
• 每个任务内部用multiprocessing并行处理内层循环，逻辑和你现在单节点的写法一致。

示例实现 #

SLURM提交脚本（submit_array.sh）

#!/bin/bash
#SBATCH --job-name=outer_loop_array
#SBATCH --array=0-19  # 对应外层20次循环迭代
#SBATCH --nodes=1
#SBATCH --cpus-per-task=12  # 每个节点分配12核
#SBATCH --output=outer_%A_%a.log

# 加载集群Python环境（根据实际情况调整）
module load python/3.10

# 传递当前迭代索引给Python脚本
python your_task.py $SLURM_ARRAY_TASK_ID

对应Python脚本（your_task.py）

import sys
from multiprocessing import Pool

def inner_task(inner_param):
    # 替换为你的内层循环逻辑
    return inner_param ** 2

def run_outer_iteration(outer_idx):
    # 根据外层索引获取对应参数（替换为你的实际参数逻辑）
    outer_param = outer_idx * 10
    # 准备内层循环参数列表
    inner_params = list(range(10000))
    
    # 用12个进程并行处理内层循环
    with Pool(processes=12) as pool:
        inner_results = pool.map(inner_task, inner_params)
    
    # 保存当前外层迭代的结果
    with open(f"result_outer_{outer_idx}.txt", "w") as f:
        f.write(str(sum(inner_results)))

if __name__ == "__main__":
    # 从SLURM获取当前外层迭代的索引
    outer_index = int(sys.argv[1])
    run_outer_iteration(outer_index)

需要用mpi4py吗？ #

不是必须的，仅在以下场景建议使用：

• 你的任务需要节点间通信（比如实时汇总所有节点的计算结果）；
• 希望用统一的进程管理逻辑，而非拆分多个独立任务。

mpi4py配置示例 #

SLURM提交脚本（submit_mpi.sh）

#!/bin/bash
#SBATCH --job-name=mpi_parallel
#SBATCH --nodes=20
#SBATCH --ntasks-per-node=1  # 每个节点启动1个MPI进程
#SBATCH --cpus-per-task=12
#SBATCH --output=mpi_%A.log

# 加载MPI和Python环境
module load openmpi/4.1.5 python/3.10

# 启动20个MPI进程（对应20个节点）
mpirun -np 20 python mpi_task.py

对应mpi4py脚本（mpi_task.py）

from mpi4py import MPI
from multiprocessing import Pool

def inner_task(inner_param):
    return inner_param ** 2

def main():
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()  # 当前MPI进程的ID（0-19）
    total_ranks = comm.Get_size()  # 总进程数=20

    # 每个MPI进程处理一个外层迭代
    outer_param = rank * 10
    inner_params = list(range(10000))
    
    # 节点内用12核并行处理内层循环
    with Pool(processes=12) as pool:
        inner_sum = sum(pool.map(inner_task, inner_params))
    
    # 汇总所有节点结果到rank=0的进程
    all_results = comm.gather(inner_sum, root=0)
    if rank == 0:
        with open("all_outer_results.txt", "w") as f:
            f.write(str(sum(all_results)))

if __name__ == "__main__":
    main()

方案对比 #

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