未设置 --gpus 时防止 SLURM 中的 GPU 使用

Question

我们正在使用 SLURM 来管理一个小型内部部署集群。我们正在管理的一个关键资源是 GPU。当用户经由请求的GPU--gpus=2的CUDA_VISIBLE_DEVICES环境变量被设置与所述GPU SLURM分配给用户。

$ srun --gpus=2 bash -c 'echo $CUDA_VISIBLE_DEVICES'
0,1

我们有一个小团队，可以相信我们的用户不会滥用系统（他们很容易覆盖环境变量），所以这很好用。但是，意外绕过它有点太容易了，因为何时--gpus未指定$CUDA_VISIBLE_DEVICES未设置，因此用户可以使用任何 GPU（我们通常使用 PyTorch）。

换句话说，以下命令可以正常工作（只要它位于 GPU 节点上），但我希望它失败（因为没有请求 GPU）。

srun sudo docker run -e CUDA_VISIBLE_DEVICES --runtime=nvidia pytorch/pytorch:1.1.0-cuda10.0-cudnn7.5-runtime python -c 'import torch; print(torch.tensor([1., 2.], device=torch.device("cuda:0")))'

如果$CUDA_VISIBLE_DEVICES设置为，它将失败-1。

$ CUDA_VISIBLE_DEVICES=-1 sudo docker run -e CUDA_VISIBLE_DEVICES --runtime=nvidia pytorch/pytorch:1.1.0-cuda10.0-cudnn7.5-runtime python -c 'import torch; print(torch.tensor([1., 2.], device=torch.device("cuda:0")))'
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/THCGeneral.cpp line=51 error=38 : no CUDA-capable device is detected
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/opt/conda/lib/python3.6/site-packages/torch/cuda/__init__.py", line 163, in _lazy_init
    torch._C._cuda_init()
RuntimeError: cuda runtime error (38) : no CUDA-capable device is detected at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THC/THCGeneral.cpp:51

如何配置SLURM到组CUDA_VISIBLE_DEVICES到-1时--gpus未指定？

Answer 1

如果 Slurm 未设置变量，您可以使用TaskProlog脚本将其设置为。$CUDA_VISIBLE_DEVICES-1

在slurm.conf，配置TaskProlog=/path/to/prolog.sh和设置以下内容prolog.sh。

#! /bin/bash

if [[ -z $CUDA_VISIBLE_DEVICES]]; then
echo export CUDA_VISIBLE_DEVICES=-1
fi

该echo export ...部件将注入CUDA_VISIBLE_DEVICES=-1作业环境。

确保/path/to从所有计算节点都可见。但这不会阻止用户播放系统并从 Python 脚本中重新定义变量。真正阻止访问需要配置cgroups.