经常有小伙伴和我抱怨说拿不到计算资源，于是这不就来了吗。这是一个后台GPU排队脚本，主要是为了解决实验室中的显卡使用/占用问题。
✨阅前须知Talk is cheap, show me the code. 废话少说，直接给我代码！
脚本代码跳转链接：脚本在这里
说给小白白的一些话：
- 本文只针对Nvidia显卡，依赖nvidia-smi查看显卡状态的命令。
- 本文提出的方法使用的是python作为终端脚本的启动器，默认python即可。
- 本文的方法是多显卡的贪心脚本，一个空窗期内有几张显卡就会用几张，小心使用。
- 本文方法无法实现显卡累加的操作，如果你想要那种闲下一张卡就抢过来的脚本，可以试试在本文的基础上进行修改。
对理解GPU多卡排队脚本有帮助的内容：
- 在服务器上训练人工智能模型的时候往往是启动一个后台任务，启动后台任务的方法如下：
#!/bin/bash nohup 【命令】 &
- 对于使用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
- 对于多卡训练任务，往往是通过修改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):
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}Wn "
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（活跃于知识区生活区和动画区）
