如何为HPC上的Python多进程代码编写正确的SLURM指令

核心概念辨析 #

• ntasks=N：SLURM会创建N个独立的任务实例，每个实例都会完整执行你的Python脚本。这就是你看到重复输出"hello world"、运行时间无变化的原因——8个任务各自跑了一遍全量代码，而非并行执行计算逻辑。
• cpus-per-task=N：给单个任务实例分配N个CPU核心，适合Python多进程场景（因为multiprocessing是由主进程fork子进程，需要在同一个节点内获取多个核心资源）。

正确配置步骤 #

1. SLURM脚本设置 #

使用cpus-per-task指定核心数，同时绑定单个节点（避免跨节点通信开销），示例脚本如下：

#!/bin/bash
#SBATCH --job-name=multiproc_test
#SBATCH --nodes=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=16G  # 按需调整内存
#SBATCH --time=01:00:00

python multiproc1.py

2. Python代码适配 #

必须让代码的进程数与SLURM分配的核心数对齐，推荐通过读取SLURM环境变量动态设置，避免硬编码：

import os
from multiprocessing import Pool

def compute_task():
    # 替换为你的计算逻辑
    print("hello world")
    # 模拟计算耗时
    import time
    time.sleep(10)

if __name__ == "__main__":
    # 从SLURM环境变量获取分配的核心数，默认1个
    core_count = int(os.getenv("SLURM_CPUS_PER_TASK", 1))
    with Pool(core_count) as pool:
        # 根据核心数启动对应数量的进程
        pool.map(compute_task, range(core_count))

问题复盘 #

• 用ntasks=8时，SLURM启动8个独立的Python进程，每个都完整执行脚本，相当于重复跑了8次程序，自然不会缩短计算时间，还会重复输出。
• 之前设置cpus-per-task耗时更长，是因为代码进程数未与分配的核心数匹配，核心资源未被充分利用，甚至因调度冗余增加了开销。

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