如何在Slurm集群中让Python程序跨多节点使用CPU核心

现有一个单机多核心运行的Python程序，未使用MPI或OpenMP，通过CreateProcess类的机制（类似multiprocessing）调用CPU核心/线程。希望在Slurm集群中跨多节点使用核心，比如用5个节点共100核运行程序，但提交的sbatch脚本仅让程序使用了第一个节点的20核，其余节点资源闲置。

提交的sbatch脚本如下：

#!/bin/bash
#SBATCH -n 100                      # Number of CPU cores     
#SBATCH -t 0-24:00                  # wall time (D-HH:MM)      set to 24h to test
#SBATCH -p slgrid                    # partition
#SBATCH -o myoutput.%j.out          # STDOUT (%j = JobId)
#SBATCH -e myoutput.%j.err          # STDERR (%j = JobId)
#SBATCH --mail-type=ALL           # Send when job starts, stops, or fails

source /scratch/xxx/.venv/bin/activate     # 激活虚拟环境
python -m darwin.run_search template.txt tokens.json options.json    # Python运行命令

核心原因 #

你的Python程序依赖单机多进程机制，只会在启动它的节点（Slurm分配的第一个节点）上创建子进程，没有跨节点通信和任务调度能力，因此无法利用其他节点的资源。要实现跨节点使用核心，需要借助分布式计算框架或修改任务分发方式。

可选方案 #

1. 修改代码适配分布式计算框架

选择支持跨节点的Python并行库，修改程序逻辑以实现任务的分布式调度：

• mpi4py：基于MPI标准，适合编写分布式并行程序。需要将任务拆分为多个子任务，通过MPI的通信接口在不同节点的进程间传递数据和任务，每个节点的进程处理分配到的任务。
• Dask：可以将原有的单机多进程代码平滑迁移到分布式环境。用Slurm启动一个Dask集群后，让程序连接到该集群，将计算任务提交到集群的所有节点执行，无需大幅修改原有逻辑。
• Ray：轻量级分布式计算框架，支持任务并行和数据并行。只需在代码中添加少量装饰器（如@ray.remote），即可将函数转为分布式可执行任务，Ray会自动处理跨节点的调度和通信。

2. 用Slurm的srun分发独立任务

如果你的程序可以将整体任务拆分为多个独立的子任务（无依赖关系），可以通过srun直接在所有分配的核心上启动独立的Python进程，每个进程处理一个子任务：
修改sbatch脚本中的运行命令，例如：

srun --ntasks=100 python -m darwin.run_search_single_task template.txt tokens.json options.json <task_id>

其中darwin.run_search_single_task是修改后的单任务版本程序，<task_id>用于区分不同子任务，需要提前将任务拆分好并传递给每个进程。

注意事项 #

• 无论选择哪种方案，都需要确保虚拟环境在所有Slurm分配的节点上都能正常访问（比如将虚拟环境放在共享存储中），或者在每个节点上安装相同的依赖包。
• 如果使用MPI类框架，需要确保集群已安装MPI环境（如OpenMPI），且mpi4py是基于该MPI编译的。

内容的提问来源于stack exchange，提问作者Quantum Monte Carlo