【星光03】GPU多卡排队/抢占/贪心脚本，实验室必备

经常有小伙伴和我抱怨说拿不到计算资源，于是这不就来了吗。这是一个后台GPU排队脚本，主要是为了解决实验室中的显卡使用/占用问题。

✨阅前须知

Talk is cheap, show me the code. 废话少说，直接给我代码！
脚本代码跳转链接：脚本在这里

说给小白白的一些话：

1. 本文只针对Nvidia显卡，依赖nvidia-smi查看显卡状态的命令。
2. 本文提出的方法使用的是python作为终端脚本的启动器，默认python即可。
3. 本文的方法是多显卡的贪心脚本，一个空窗期内有几张显卡就会用几张，小心使用。
4. 本文方法无法实现显卡累加的操作，如果你想要那种闲下一张卡就抢过来的脚本，可以试试在本文的基础上进行修改。

✨预备知识

对理解GPU多卡排队脚本有帮助的内容：

1. 在服务器上训练人工智能模型的时候往往是启动一个后台任务，启动后台任务的方法如下：

   #!/bin/bash
   nohup 【命令】 &
   

2. 对于使用GPU的后台任务，如果没有好好的退出，会一直占用GPU资源。所以需要执行如下命令退出：

   # 执行以下命令在全部线程中寻找你的后台线程：
   ps -ef | grep 【刚才启动的命令，支持正则表达式】
   # 确定上述搜索命令中的全部结果都是期望结束的线程后，执行：
   ps -ef | grep 【同上】 | grep -v grep |cut -c 9-15 | xargs kill
   # 补充：kill 命令的默认参数是15，意为正常退出。如果出现无法退出的情况，可以给参数9，强制退出，如下：
   # ps -ef | grep 【同上】 | grep -v grep |cut -c 9-15 | xargs kill -9
   

3. 对于多卡训练任务，往往是通过修改CUDA_VISIBLE_DEVICES变量和NUM_GPUS变量。所以思路很简单，寻找全部可以使用的显卡，然后修改这两个变量即可。

✨脚本

我们这里使用常用到显卡的人工智能训练任务作为例子，下面是一个多卡训练的命令：

GPUS="0,1,2,3,4" # 前两行为Parameter
NUM_GPUS=5
CUDA_VISIBLE_DEVICES=${GPUS} tools/dist_train.sh ${NUM_GPUS}

⭐后台启动脚本

创建train.bash终端脚本文件，实现后台启动python脚本的功能，内容如下:

#!/bin/bash
nohup \
    train.py \
&

注：\ 是折行的意思，方便理解超级长的命令。

⭐GPU排队脚本

创建train.pypython脚本文件，内容如下：

#! /usr/bin/python3
import os
import time
import sys


class GPUGet:
    def __init__(self,
                 min_gpu_number,
                 time_interval):
        self.min_gpu_number = min_gpu_number
        self.time_interval = time_interval

    def get_gpu_info(self):
        gpu_status = os.popen('nvidia-smi | grep %').read().split('|')[1:]
        gpu_dict = dict()
        for i in range(len(gpu_status) // 4):
            index = i * 4
            gpu_state = str(gpu_status[index].split('   ')[2].strip())
            gpu_power = int(gpu_status[index].split('   ')[-1].split('/')[0].split('W')[0].strip())
            gpu_memory = int(gpu_status[index + 1].split('/')[0].split('M')[0].strip())
            gpu_dict[i] = (gpu_state, gpu_power, gpu_memory)
        return gpu_dict

    def loop_monitor(self):
        available_gpus = []
        while True:
            gpu_dict = self.get_gpu_info()
            for i, (gpu_state, gpu_power, gpu_memory) in gpu_dict.items():
                if gpu_state == "P8" and gpu_power <= 40 and gpu_memory <= 1000:  # 设置GPU选用条件，当前适配的是Nvidia-RTX3090
                    gpu_str = f"GPU/id: {i}, GPU/state: {gpu_state}, GPU/memory: {gpu_memory}MiB, GPU/power: {gpu_power}W\n "
                    sys.stdout.write(gpu_str)
                    sys.stdout.flush()
                    available_gpus.append(i)
            if len(available_gpus) >= self.min_gpu_number:
                return available_gpus
            else:
                available_gpus = []
                time.sleep(self.time_interval)

    def run(self, cmd_parameter, cmd_command):
        available_gpus = self.loop_monitor()
        gpu_list_str = ",".join(map(str, available_gpus))
        # 构建终端命令
        cmd_parameter = fr"""{cmd_parameter}
                          NUM_GPUS={len(available_gpus)} ; \ """  # 一定要有 `; \ `
        cmd_command = fr"""CUDA_VISIBLE_DEVICES={gpu_list_str} \ 
                         {cmd_command}"""
        command = fr"""{cmd_parameter} {cmd_command}"""
        print(command)
        os.system(command)


if __name__ == '__main__':
    min_gpu_number = 3  # 最小GPU数量，多于这个数值才会开始执行训练任务。
    time_interval = 5  # 监控GPU状态的频率，单位秒。
    gpu_get = GPUGet(min_gpu_number, time_interval)

    cmd_parameter = r""""""  # 命令会使用到的参数，使用 `;` 连接。
    cmd_command = r"""tools/dist_train.sh ${NUM_GPUS} \ """
    gpu_get.run(cmd_parameter, cmd_command)

✨脚本执行

直接在终端中启动train.bash即可。

注：不要忘记修改.bash和.py两个脚本的运行权限，可以执行命令：chmod 764 【文件名】

✨参考

• 封面画师：溪涧殇
• GPU排队脚本——一旦空闲就触发python脚本

---

🍀碎碎念🍀
Hello米娜桑，这里是英国留学中的杨丝儿。我的博客的关键词集中在编程、算法、机器人、人工智能、数学等等，点个关注吧，持续高质量输出中。
🌸唠嗑QQ群：兔叽的魔术工房 (942848525)
⭐️B站账号：YangSierCode000（活跃于知识区生活区和动画区）

---