经常有小伙伴和我抱怨说拿不到计算资源,于是这不就来了吗。这是一个后台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.py
python脚本文件,内容如下:
#! /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 【文件名】
✨参考
🍀碎碎念🍀
Hello米娜桑,这里是英国留学中的杨丝儿。我的博客的关键词集中在编程、算法、机器人、人工智能、数学等等,点个关注吧,持续高质量输出中。
🌸唠嗑QQ群:兔叽的魔术工房 (942848525)
⭐️B站账号:YangSierCode000(活跃于知识区生活区和动画区)