在SBATCH/SLURM中运行含ptemcee的Python脚本卡住如何解决？

问题核心原因 #

ptemcee/emcee默认采用的多进程启动模式、底层数学库的多线程抢占、srun命令的进程管理规则，和SLURM的资源分配机制存在冲突，导致进程死锁无响应。

具体修改步骤 #

• 1. 修改Python脚本的多进程启动模式 #

在teste.py的最开头（所有其他导入语句之前）加入如下代码，将多进程启动方式从默认的fork改为兼容性更好的spawn：

import multiprocessing
multiprocessing.set_start_method('spawn', force=True)

• 2. 调整SBATCH提交脚本 #

删除srun包装，添加环境变量禁用numpy、MKL等底层库的自动多线程，避免和ptemcee的多进程抢占CPU资源：

#!/usr/bin/env bash
#SBATCH -J Exemplo            # 任务名
#SBATCH --cpus-per-task=15
#SBATCH --ntasks 1           # 总进程数
#SBATCH --partition=batch 

# 禁用底层库自动多线程，避免资源冲突
export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1

module load python
# 去掉srun直接运行Python脚本
python teste.py
date

• 3. 显式指定ptemcee的并行进程数 #

在创建ptemcee采样器时，显式传入processes参数，数值和你申请的--cpus-per-task值保持一致，避免自动检测CPU数出错：

# 示例代码，替换你原有创建sampler的对应部分
sampler = ptemcee.EnsembleSampler(
    n_walkers, dim, your_log_prob_function,
    processes=15 # 和--cpus-per-task=15对应
)

排查建议 #

如果修改后仍然出现卡顿，可以在Python脚本的采样逻辑前后添加日志打印，确认卡住的具体节点，也可以在提交脚本中加入squeue -j $SLURM_JOB_ID 查看任务的实时资源占用状态。

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