如何等待由外部程序后台启动的Slurm作业步骤完成？

原生Slurm脚本的正确等待方式 #

当直接在Slurm批处理脚本中启动后台进程时，wait命令可以阻塞直到所有后台子进程完成，避免作业提前终止：

#!/bin/bash
#SBATCH --ntasks=2
srun --overlap -n1 step_one &
srun --overlap -n2 step_two &
wait

问题：DVC启动后台进程后wait失效 #

如果通过DVC启动相同的后台进程，DVC会在启动进程后提前退出，此时wait无法捕获这些非当前shell子进程的进程，导致作业提前终止：
Slurm脚本：

#!/bin/bash
#SBATCH --ntasks=2
dvc repro
wait  # 无效果

对应的DVC流水线文件：

stages:
  one:
    cmd: srun --overlap -n1 step_one &
  two:
    cmd: srun --overlap -n1 step_two &

现有方案的局限性 #

你尝试用sstat循环检测Slurm步骤的方式存在不稳定风险，比如依赖sstat的输出行数判断，若Slurm步骤命名或输出格式变化，可能导致误判：

#!/bin/bash
#SBATCH --ntasks=2
dvc repro
while [ $(sstat -n -a -j $SLURM_JOB_ID | wc -l) -gt 1 ]
do
    sleep 10
done

可行解决方案 #

方法1：修改DVC阶段，利用DVC的并行管理（推荐） #

去掉DVC阶段命令末尾的&，让DVC负责并行执行和等待，这是最符合DVC设计的方式：
修改后的DVC流水线：

stages:
  one:
    cmd: srun --overlap -n1 step_one
  two:
    cmd: srun --overlap -n1 step_two

对应的Slurm脚本：

#!/bin/bash
#SBATCH --ntasks=2
dvc repro --jobs 2

dvc repro --jobs 2会并行启动两个阶段，并自动等待所有阶段完成，无需额外wait命令。

方法2：捕获后台进程PID并等待 #

如果必须保留DVC阶段的后台运行方式，可在启动进程时将PID写入文件，后续在主脚本中等待这些PID：
修改后的DVC流水线：

stages:
  one:
    cmd: srun --overlap -n1 step_one & echo $! > "${SLURM_JOB_ID}_step_one.pid"
  two:
    cmd: srun --overlap -n1 step_two & echo $! > "${SLURM_JOB_ID}_step_two.pid"

对应的Slurm脚本：

#!/bin/bash
#SBATCH --ntasks=2
# 清理旧PID文件
rm -f "${SLURM_JOB_ID}_step_one.pid" "${SLURM_JOB_ID}_step_two.pid"
# 运行DVC
dvc repro
# 等待指定PID的进程完成
wait $(cat "${SLURM_JOB_ID}_step_one.pid" "${SLURM_JOB_ID}_step_two.pid")
# 清理PID文件
rm -f "${SLURM_JOB_ID}_step_one.pid" "${SLURM_JOB_ID}_step_two.pid"

注：使用$SLURM_JOB_ID作为PID文件名前缀，避免多作业间的文件冲突，确保文件存储在集群节点共享的目录中。

方法3：改进Slurm步骤检测逻辑 #

优化sstat的检测方式，直接排除batch步骤，避免误判：

#!/bin/bash
#SBATCH --ntasks=2
dvc repro
# 等待所有非batch的Slurm步骤完成
while true; do
    # 统计除batch步骤外的活跃步骤数
    active_steps=$(sstat -n -a -j "$SLURM_JOB_ID" | grep -v "${SLURM_JOB_ID}.batch" | wc -l)
    if [ "$active_steps" -eq 0 ]; then
        break
    fi
    sleep 5
done

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