未设置 --gpus 时防止在 SLURM 中使用 GPU
Prevent GPU usage in SLURM when --gpus is not set
我们正在使用 SLURM 来管理小型本地集群。我们正在管理的一个关键资源是 GPU。当用户通过 --gpus=2
请求 GPU 时,CUDA_VISIBLE_DEVICES
环境变量设置为 SLURM 分配给用户的 GPU。
$ 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
当未指定 --gpus
时,如何配置 SLURM 以将 CUDA_VISIBLE_DEVICES
设置为 -1
?
如果 Slurm 未设置 TaskProlog
变量,您可以使用 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
.
我们正在使用 SLURM 来管理小型本地集群。我们正在管理的一个关键资源是 GPU。当用户通过 --gpus=2
请求 GPU 时,CUDA_VISIBLE_DEVICES
环境变量设置为 SLURM 分配给用户的 GPU。
$ 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
当未指定 --gpus
时,如何配置 SLURM 以将 CUDA_VISIBLE_DEVICES
设置为 -1
?
如果 Slurm 未设置 TaskProlog
变量,您可以使用 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
.