如何正确使用Slurm sbatch与Python Multiprocessing限制CPU数量

在Slurm架构服务器上运行多进程Python代码，期望限制可用CPU数量，让代码为每个CPU创建子进程，但无论怎么调整参数，代码始终使用服务器最大CPU数（272）。相关代码、脚本及执行结果如下：

Python代码 #

def Func(ins) : 
  ###
  things...
  ###
return var

if __name__ == '__main__' :
  from multiprocessing import Pool
  from multiprocessing import active_children
  from multiprocessing import cpu_count

  p = Pool()
  print("active cpus = ", cpu_count())
  print("open process = ", p._processes)
  print("active_children = ", len(active_children()))
  results = p.map(Func, range(2000))
  p.close()

  exit()

SBATCH调度脚本 #

#!/bin/bash

#SBATCH --time=1:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=48
#SBATCH --mem=40000 # Memory per node (in MB).

module load python 
conda activate myenv
python3 test.py

echo 'done!'

执行结果 #

active cpus =  272
open process =  272
active_children =  272
done!

作业通过sbatch job.sh提交。

1. cpu_count()不识别Slurm分配：Python的multiprocessing.cpu_count()返回的是服务器物理CPU总核数，不会读取Slurm通过--cpus-per-task分配的CPU配额。
2. Pool()默认行为忽略配额：无参数的Pool()会直接使用cpu_count()的结果创建进程池，完全无视Slurm的CPU限制配置。
3. 未读取Slurm环境变量：Slurm分配CPU后会设置SLURM_CPUS_PER_TASK环境变量，但你的代码没有利用这个变量来控制进程池大小。

可以通过两种方式修正：

方式一：读取Slurm环境变量设置进程池大小 #

修改Python代码，读取Slurm提供的SLURM_CPUS_PER_TASK变量，以此指定进程池的进程数：

def Func(ins) : 
  ###
  things...
  ###
return var

if __name__ == '__main__' :
  import os
  from multiprocessing import Pool
  from multiprocessing import active_children
  from multiprocessing import cpu_count

  # 优先读取Slurm分配的CPU数，未获取到则回退到物理核数
  cpu_num = int(os.getenv('SLURM_CPUS_PER_TASK', cpu_count()))
  p = Pool(processes=cpu_num)
  print("allocated cpus = ", cpu_num)
  print("open process = ", p._processes)
  print("active_children = ", len(active_children()))
  results = p.map(Func, range(2000))
  p.close()
  p.join()  # 建议添加join()等待所有子进程执行完毕

  exit()

方式二：用taskset绑定CPU（辅助限制） #

在SBATCH脚本中，通过taskset强制Python进程仅使用分配的CPU核心，即使进程池创建了多余进程，操作系统也会限制其CPU使用：

#!/bin/bash

#SBATCH --time=1:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=48
#SBATCH --mem=40000 # Memory per node (in MB).

module load python 
conda activate myenv
# 绑定到分配的CPU核心区间
taskset -c 0-$((SLURM_CPUS_PER_TASK-1)) python3 test.py

echo 'done!'

建议：将方式一和方式二结合使用，既从代码层面控制进程数，又通过系统层面限制CPU使用，避免资源浪费。

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