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

✨阅前须知

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

说给小白白的一些话:

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

✨预备知识

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

  1. 在服务器上训练人工智能模型的时候往往是启动一个后台任务,启动后台任务的方法如下:
    1
    2
    #!/bin/bash
    nohup 【命令】 &
  2. 对于使用GPU的后台任务,如果没有好好的退出,会一直占用GPU资源。所以需要执行如下命令退出:
    1
    2
    3
    4
    5
    6
    # 执行以下命令在全部线程中寻找你的后台线程:
    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变量。所以思路很简单,寻找全部可以使用的显卡,然后修改这两个变量即可。

✨脚本

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

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

⭐✨后台启动脚本

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

1
2
3
4
#!/bin/bash
nohup \
train.py \
&

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

⭐✨GPU排队脚本

创建train.pypython脚本文件,内容如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#! /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):
gpu_dict = self.get_gpu_info()
available_gpus = []
while True:
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 【文件名】

✨参考


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