防止在单节点 slurm 设置中直接使用 GPU

Question

场景：我在一台 Ubuntu 机器上安装了 slurm。用户可以登录到该机器来执行一些GPU和非GPU任务。

目标：我想阻止这些用户直接使用 GPU。只有从 slurm 开始的作业才应该能够使用 GPU。

我所做的：为了防止直接使用 GPU，我更改了 /dev/nvidiaN ( chown root.gpu /dev/nvidia*) 的所有者，现在只有组内的用户gpu可以使用它们。这样可行。问题是：slurm 作业当然以启动时的同一用户身份运行，因此这些提交的作业也无法使用任何 GPU。

有什么办法可以实现这个目标吗？我考虑过使用镜像用户（user->user-slurm在组内gpu但没有登录），sudo 和一些sbatch --uid=UID --guid=GID ...白名单命令，但这看起来非常复杂。

Answer 1

我也有同样的问题。在我的实验室中，我们想要在单节点 GPU 中设置本地 SLURM 。我们想要一个简单的解决方案，以便只有 SLURM 可以使用 GPU。这似乎是一个比乍看起来更困难的提议。

我采用的解决方案是（有时）终止任何使用尚未通过 SLURM 启动的 GPU 的进程。这可以通过一个相对简单的 Python 脚本来完成。

#!/usr/bin/env python3
# Kills all processes using /dev/nvidia* that have not
# been initiated by SLURM.

import os, subprocess, psutil, syslog

# get all nvidia devices
devs = []
i = 0
while os.path.exists(f'/dev/nvidia{i}'):
    devs.append(f'/dev/nvidia{i}')
    i += 1

# get all processes using any of those devices
out = subprocess.check_output(['lsof', '-t'] + devs)

# kill processes whose parent is not SLURM (slurmstepd)
for pid in out.split():
    p = process = psutil.Process(int(pid))
    has_slurm_parent = False
    while p := p.parent():
        if p.name() == 'slurmstepd':
            has_slurm_parent = True
            break
    if not has_slurm_parent:
        syslog.syslog(f'Killing process {pid} {process.name()} ({process.username()}) - cannot use GPU outside of SLURM')
        process.kill()

然后，只需将其添加到 cron ( sudo crontab -e) 中，以每 5 分钟左右运行一次：

*/5 * * * * python3 /path/to/script.py

此外，由于此脚本只是最后的防御手段，因此最好通过禁用 cuda 可见设备来防止这种情况发生/etc/environment：

CUDA_VISIBLE_DEVICES=""

我们还可以在 bash 用户 ( ) 的情况下锁定此变量/etc/profile.d/xxxx.sh以及一条不错的欢迎消息：

export CUDA_VISIBLE_DEVICES=""
readonly CUDA_VISIBLE_DEVICES
echo "Use srun or sbatch to access the GPU."