SLURM环境下Python multiprocessing多进程性能异常低下问题

问题背景 #

• 本地4×2核机器上，Python multiprocessing并行任务运行正常，每个进程耗时约5分钟（与串行单任务耗时一致）
• 集群为2台AMD节点（每节点64核），目标是在单个节点使用总计128核资源（64物理核×2逻辑核）
• 串行SLURM脚本（不调用multiprocessing）运行耗时约5分钟，配置如下：

  #!/bin/bash
  
  #SBATCH --mem-per-cpu=2G
  #SBATCH --cpus_per_task=64
  #SBATCH --nodes=1
  #SBATCH --ntasks=1
  #SBATCH --nstaks-per-node=1  # 存在拼写错误
  
  module load python
  python 0 myscript.py 
  

• 并行SLURM脚本（调用multiprocessing.Pool(64)）运行时，每个进程耗时约30分钟（是本地的6倍），配置如下：

  #!/bin/bash
  
  #SBATCH --mem-per-cpu=2G
  #SBATCH --cpus_per_task=64
  #SBATCH --nodes=1
  #SBATCH --ntasks=1
  #SBATCH --nstaks-per-node=1  # 拼写错误
  
  module load python
  
  export MKL_NUM_THREADS=1
  export NUMEXPR_NUM_THREADS=1
  export OMP_NUM_THREADS=1
  
  python 1 64 myscript.py 
  

• 额外现象：sacct显示SLURM实际分配128核，而非脚本指定的64核；补充测试中，设置MKL_NUM_THREADS=2用16进程处理32任务，比32进程处理32任务耗时更短。

核心原因分析 #

1. SLURM脚本拼写错误导致资源分配异常 #

脚本中--nstaks-per-node是拼写错误，正确参数为--ntasks-per-node。该错误会导致SLURM资源调度逻辑混乱，出现实际分配核数与指定值不符的情况（如分配128逻辑核）。

2. CPU未绑定引发上下文切换开销 #

集群核心数远多于本地机器，若进程未绑定到固定核心，会频繁在不同核心间切换上下文，大幅增加运行耗时。本地核心少，切换成本可忽略，但集群环境下该开销会被放大数倍。

3. 超线程与进程数不匹配 #

AMD节点默认开启超线程，64物理核对应128逻辑核。若仅配置64个进程，SLURM可能将多个进程分配到同一物理核的逻辑核上，导致缓存竞争；或进程数未充分利用逻辑核，同时进程切换成本高于线程切换。

4. Fork模式的资源继承问题 #

multiprocessing默认的fork模式会继承父进程的所有资源，包括未正确配置的SLURM CPU亲和性设置，导致子进程无法获得独立的核心资源。

解决方案 #

1. 修正SLURM脚本错误 #

将--nstaks-per-node=1改为--ntasks-per-node=1，确保资源调度逻辑正常。

2. 配置CPU绑定 #

在SLURM脚本中添加CPU绑定参数，强制进程固定到核心：

#SBATCH --cpu-bind=cores  # 绑定到物理核
# 或绑定到逻辑核：#SBATCH --cpu-bind=threads

添加--cpu-bind=verbose可打印绑定详情，验证每个进程是否分配到独立核心。

3. 匹配进程数与核心数 #

若目标是使用单个节点的全部128逻辑核：

• 修改SLURM脚本：#SBATCH --cpus_per_task=128
• 修改Python调用参数：python 1 128 myscript.py（创建128个进程）

若仅使用64物理核：

• 保持#SBATCH --cpus_per_task=64，并配合--cpu-bind=cores，确保每个进程绑定到一个物理核。

4. 优化multiprocessing启动方式 #

改用forkserver启动方式，避免继承父进程的不必要资源：

import multiprocessing

if __name__ == '__main__':
    multiprocessing.set_start_method('forkserver')
    # 后续创建Pool的代码

或在fork后手动设置子进程的CPU亲和性：

import os
import multiprocessing

def worker_func(core_id):
    # 绑定当前进程到指定核心
    os.sched_setaffinity(0, [core_id])
    # 实际任务代码

if __name__ == '__main__':
    num_cores = int(os.environ.get('SLURM_CPUS_PER_TASK', 64))
    with multiprocessing.Pool(num_cores) as pool:
        pool.map(worker_func, range(num_cores))

5. 验证线程环境变量生效 #

在Python代码中打印环境变量，确保MKL/OMP线程数限制生效：

import os
print(f"MKL_NUM_THREADS: {os.environ.get('MKL_NUM_THREADS')}")
print(f"OMP_NUM_THREADS: {os.environ.get('OMP_NUM_THREADS')}")

补充测试疑问解答 #

设置MKL_NUM_THREADS=2用16进程处理32任务更快，核心原因是进程上下文切换成本远高于线程切换：

• fork创建的进程拥有独立内存空间、进程表项，切换时需保存/恢复更多系统状态；线程共享进程内存空间，切换成本极低。
• AMD超线程的逻辑核共享同一物理核的L1/L2缓存，线程间数据交互更高效，避免了进程间的缓存失效问题。
  即使没有进程间共享数据，多进程的系统开销仍远大于多线程，因此在超线程环境下，结合线程池（如MKL/OMP的线程）比纯多进程更高效。

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