SLURM集群中Python subprocess调用mpirun出现内存不足问题求助

在SLURM集群中编写Python包装器用于同步多个应用程序执行时，遇到内存分配异常问题：

• 直接通过SLURM运行单个应用程序时，作业可正常执行，sacct命令显示其最大内存占用（MaxRSS）为8.2GB，远低于节点可用的96GB，对应的sbatch命令如下：

mpirun -np 4 --map-by ppr:1:node:pe=20 executable

• 使用Python包装器执行python3 simple_wrapper.py时，可执行程序出现内存分配错误，sacct查询到的内存占用仅为5GB，包装器代码如下：

import subprocess
import sys
with open('log', 'w') as out:
    cmd = 'mpirun -np 4 --map-by ppr:1:node:pe=20 executable'
    subprocess.check_call(cmd.split(), bufsize=-1, stdout=out, stderr=subprocess.STDOUT)

请问这是集群中用Python3运行MPI应用的常见问题吗？该如何解决？

这是Python调用MPI应用时的常见问题，核心原因是SLURM的资源环境变量未被Python子进程正确继承，导致MPI进程只能获取默认的小内存配额，而非SLURM分配给作业的全部资源。以下是具体解决步骤：

• 显式传递环境变量
  默认情况下subprocess.check_call会继承父进程环境，但部分SLURM相关变量可能被Python启动时过滤。修改代码显式传递当前环境变量：

  import subprocess
  import sys
  import os
  
  with open('log', 'w') as out:
      cmd = 'mpirun -np 4 --map-by ppr:1:node:pe=20 executable'
      subprocess.check_call(cmd.split(), bufsize=-1, stdout=out, stderr=subprocess.STDOUT, env=os.environ)
  

• 使用列表形式定义命令
  cmd.split()在复杂参数场景下易出错（比如含空格的路径），建议直接用列表定义命令，同时确保MPI能正确识别SLURM资源分配：

  import subprocess
  import sys
  import os
  
  with open('log', 'w') as out:
      cmd = ['mpirun', '-np', '4', '--map-by', 'ppr:1:node:pe=20', 'executable']
      subprocess.check_call(cmd, bufsize=-1, stdout=out, stderr=subprocess.STDOUT, env=os.environ)
  

• 强制MPI绑定SLURM资源
  在mpirun命令中添加--use-hwthread-cpus参数，确保MPI进程能获取到SLURM分配的全部CPU和内存资源：

  cmd = ['mpirun', '-np', '4', '--map-by', 'ppr:1:node:pe=20', '--use-hwthread-cpus', 'executable']
  

• 确认SLURM作业脚本的资源配置
  确保sbatch脚本中正确指定内存配额，避免Python进程自身被限制内存：

  #SBATCH --mem=96G
  #SBATCH --nodes=4
  #SBATCH --ntasks-per-node=1
  #SBATCH --cpus-per-task=20
  python3 simple_wrapper.py
  

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