使用 slurm 运行并行多处理 Python 脚本

Question

我正在尝试使用 slurm 同时运行一个 python 脚本，该脚本使用多处理在内部并行化进程。 python 脚本的目标是多次求解随机微分方程，以便获得足够的统计数据进行分析。为了加快进程，我使用多处理将重复分为 4 个子进程。另一方面，为了使用我的所有资源，我尝试使用不同的输入参数（即方程的不同条件）同时运行此脚本多次。现在，我使用一个具有 16 个 CPU 的节点，因此我想同时运行代码 4 次，为每组条件分配 4 个 CPU，这样 4x4=16。

我的 slurm 实现是一个批处理脚本

run.sh

，如下所示：

#!/bin/bash
#SBATCH -N 1
#SBATCH -n 16
#SBATCH --time=10:00:00 

module purge > /dev/null 2>&1

eval "$(conda shell.bash hook)"
conda activate my_environment


srun --exclusive -n 1 python script.py 0.95 1 1 0 1 0 > nohup1.out 2>&1 & 
srun --exclusive -n 1 python script.py 0.955 1 1 0 1 0 > nohup2.out 2>&1 & 
srun --exclusive -n 1 python script.py 0.96 1 1 0 1 0 > nohup3.out 2>&1 & 
srun --exclusive -n 1 python script.py 0.965 1 1 0 1 0 > nohup4.out 2>&1 & 
srun --exclusive -n 1 python script.py 0.97 1 1 0 1 0 > nohup5.out 2>&1 & 
srun --exclusive -n 1 python script.py 0.975 1 1 0 1 0 > nohup6.out 2>&1 & 

wait

要运行此批处理脚本，我正在使用

sbatch run.sh

。请注意，在示例代码中，我有 6 个参数组合（我计划在实际情况下使用更多参数组合）。我期望 slurm 根据可用资源运行前 4 个，并在完成后继续执行下 2 个。但是，它会同时发送所有脚本。我怀疑这是因为旗帜

--exclusive

。然而，当我删除它时，它只是一一运行脚本。

谢谢！

Answer 1

您正在尝试执行一项称为参数扫描的操作。有一些很棒的工具可以做到这一点，而不是使用大量

srun

自己处理它，据我所知，这可能不起作用。您可以尝试以下方法：

• Facebook 的 Hydra 与 joblib 插件

这给了你很大的控制权，但需要仔细规划。您需要使用单独的记录器，例如 python 记录器或 mlflow 记录器，具体取决于您要保存的内容。您可以在配置文件中定义作业数量，这使您的生活更加轻松。但是，如果您的作业超出了节点时间限制（主要是 2 天），那么您需要手动查找失败的节点。所以要小心。

• 使用wandb扫描管理器

这给你带来了一点点控制权，但它基本上是开箱即用的。有一个在线服务器为每个子作业分配扫描参数。这使得重新启动失败的作业变得非常容易。

您需要探索两者并选择合适的工具。