SLURM+Docker：如何使用 SLURM scancel 终止 docker 创建的进程

Question

我们目前已经搭建了一个GPU计算集群，并以SLURM作为资源管理器。由于这是一个用于深度学习的集群，我们通过使用 nvidia-docker 镜像来管理依赖关系，以方便不同的框架和 CUDA 版本。

我们的典型用例是使用 srun 分配资源并发出命令来运行 nvidia-docker，该命令按照以下方式运行实验脚本：

srun --gres=gpu:[num gpus required] nvidia-docker run --rm -u $(id -u):$(id -g) /bin/bash -c [python scripts etc..] &

我们发现了一个问题，如果使用 scancel 命令取消 slurm 作业，节点上的 docker 进程将被取消，但在 docker 中启动的任何实验脚本仍然会继续。据我们了解，这不是 SLURM 的错误，而是杀死 docker 进程并不会杀死其生成的进程，只能使用 docker Kill 命令来杀死它们。虽然可能有某种方法可以在 SLURM 序言脚本中执行 docker Kill 命令，但我们想知道其他人是否也遇到过这个问题，以及他们是否已经以某种方式解决了这个问题。为了总结一下，我们想知道：

我们如何确保在 nvidia-docker 容器中启动的进程（该容器又由 SLURM SRUN 启动）被 SCANCEL 终止？

Answer 1

配置 Slurm 使用 cgroups 可能会有所帮助。启用

cgroups

后，属于作业的任何进程都会附加到

cgroup

并在作业结束时销毁。销毁是由内核负责的，因此常规进程无法逃避这一点。

Answer 2

ProctrackType=proctrack/cgroup

插件提供了cgroup的跟踪，它仍然只是向每个进程发送信号，由每个进程负责响应（退出）。

您可以使用包装器脚本通过陷阱捕获终止信号，然后从陷阱处理程序发出

docker kill

命令。

#!/bin/bash
handle_kill() {
    echo >&2 "${0##*/}: Got a kill signal, issuing docker kill ..."
    docker kill $YOUR_DOCKER_KILL_ARGS
    exit $YOUR_RETURN_CODE
}
trap handle_kill CONT TERM
# invoke your docker commands here
docker . . .