Slurm环境下Python多进程运行缓慢问题求助

核心问题 #

你遇到的性能瓶颈源于多层并行嵌套、Slurm资源配置与Python进程模型不匹配，以及第三方库隐式多线程的叠加效应，导致CPU上下文切换过载、资源耗尽。

环境信息 #

SBATCH 提交脚本

#!/bin/bash
##----------------------- Start job description ----------------------- 
#SBATCH --partition=standard
#SBATCH --job-name=multi_Leaspy
#SBATCH --nodes=1
#SBATCH --ntasks=40
#SBATCH --mem-per-cpu=4096
#SBATCH --time=160:00:00
#SBATCH --mail-type=ALL
#SBATCH --output=out-%j.log
#SBATCH --error=err-%j.log
##------------------------ End job description ------------------------ 

Python 并行代码片段

num_workers = ntasks = int(os.environ.get('SLURM_NTASKS', 1))
if __name__ == "__main__":
        with ProcessPoolExecutor(max_workers = num_workers - 1) as executor:       
            arg_list = [(i, row, df_test, id_test, btstrp, classes, classes_auc, 
                      results_table, MAE_table, features_set, vector_range) for i, row in feat_sub.iterrows()]
            for result, MAE_row in executor.map(multiLeaspyIter, arg_list):
                if result is not None and MAE_row is not None:
                    results_table = results_table.append(result, ignore_index=True)
                    MAE_table = MAE_table.append(MAE_row, ignore_index = True)

观测到的异常状态

• htop显示：Tasks=1627, 501 thr, 40 running；Load average=201.3, 180.94, 117.36
• 节点配置：40物理核，启用超线程（80逻辑核）

分步解决措施 #

1. 修正Slurm资源与Python进程的映射逻辑 #

Slurm的--ntasks在单节点场景下默认对应CPU核心数，但Python进程池+内部线程会突破配额。调整配置：

• 修改SBATCH脚本，明确绑定CPU资源（二选一即可）：

  # 方案1：40个独立任务，每个任务占1核
  #SBATCH --ntasks=40
  #SBATCH --cpus-per-task=1
  
  # 方案2：单任务绑定40核（更适合单进程池场景）
  #SBATCH --ntasks=1
  #SBATCH --cpus-per-task=40
  

• Python代码中直接获取可用物理核心数，避免超线程过载：

  import psutil
  # 获取物理核心数，而非逻辑核心（超线程）
  num_workers = psutil.cpu_count(logical=False)
  # 或从Slurm环境变量读取申请的核心数
  num_workers = int(os.environ.get('SLURM_CPUS_ON_NODE', psutil.cpu_count(logical=False)))
  

2. 禁用第三方库的隐式多线程 #

pandas、numpy、Leaspy等库默认会启用多线程优化，与你的进程池形成嵌套并行，直接导致进程/线程爆炸。在SBATCH脚本开头添加环境变量强制单线程：

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export NUMBA_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export LEASPY_NUM_THREADS=1  # 若Leaspy支持此配置

3. 优化数据传递，减少序列化开销 #

arg_list中传递大型DataFrame会触发pickle序列化，产生巨大内存和IO开销。优化方式：

• 改为传递数据文件路径，让子进程自行读取：

  # 先将df_test保存为临时文件
  df_test.to_csv("/tmp/df_test.csv", index=False)
  arg_list = [(i, row, "/tmp/df_test.csv", id_test, btstrp, classes, classes_auc, 
              features_set, vector_range) for i, row in feat_sub.iterrows()]
  # 在multiLeaspyIter中读取：
  def multiLeaspyIter(args):
      i, row, df_test_path, id_test, ... = args
      df_test = pd.read_csv(df_test_path)
      # 后续逻辑
  

• 替换pandas的append为列表收集结果后一次性合并（append多次调用效率极低）：

  results_list = []
  MAE_list = []
  for result, MAE_row in executor.map(multiLeaspyIter, arg_list):
      if result is not None and MAE_row is not None:
          results_list.append(result)
          MAE_list.append(MAE_row)
  # 一次性合并
  results_table = pd.concat(results_list, ignore_index=True)
  MAE_table = pd.concat(MAE_list, ignore_index=True)
  

4. 排查并关闭multiLeaspyIter内部的并行逻辑 #

检查该函数是否内部启动了额外的进程/线程（比如ThreadPoolExecutor、Leaspy的内部并行），如果是，强制关闭：

def multiLeaspyIter(args):
    # 禁用内部多线程
    import os
    os.environ["OMP_NUM_THREADS"] = "1"
    # 若使用了ThreadPoolExecutor，设置max_workers=1
    with ThreadPoolExecutor(max_workers=1) as inner_exec:
        # 内部逻辑
    # 其他逻辑

5. 调整进程池配置，避免内存泄漏 #

改用multiprocessing.Pool并设置maxtasksperchild，定期重启子进程避免内存累积：

from multiprocessing import Pool
if __name__ == "__main__":
    with Pool(processes=num_workers, maxtasksperchild=10) as pool:
        all_results = pool.map(multiLeaspyIter, arg_list)
    # 后续处理结果

验证步骤 #

1. 先添加禁用多线程的环境变量，调整SBATCH的CPU配置。
2. 修改Python代码的进程数计算逻辑，优化数据传递方式。
3. 运行任务后用htop观测：确保总进程/线程数不超过80（逻辑核上限），负载平均值接近40（物理核数）。
4. 对比单进程与多进程的运行时间，确认性能提升。

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