GPU资源排队脚本，支持多个GPU资源(可指定GPU或指定排队运行次数)

由于GPU资源在大部分时间会被较多人使用，而部分时间可能较少人或者没有人使用，这样就会在一定程度上造成GPU资源的浪费，因此出于这个目的写了GPU资源的排队脚本，这里主要是针对多个GPU资源来实现的：

参考：https://blog.csdn.net/leviopku/article/details/102958166
代码如下：
下面展示一些。

import os
import sys
import time

cmd = 'python xxx.py（py文件名）'

def get_info():
    gpu_memory_list = []
    gpu_status = os.popen('nvidia-smi | grep %').readlines()
    for list_ in gpu_status:
        list_ = list_.split('|')
        gpu_memory = int(list_[2].split('/')[0].split('M')[0].strip())
        gpu_memory_list.append(gpu_memory)
        # 主要是针对memory操作，如果对power需要操作的话同理
        # gpu_power = int(list_[1].split('  ')[-1].split('/')[0].split('W')[0].strip())
    return  gpu_memory_list
    
def setup_python(interval=1):
    gpu_memory_list = get_info()
    i = 0
    while(i!=7):
        if gpu_memory_list[i] <= 2000:  # 设置条件，满足则运行python程序
            if i==0 or i==5 or i==6:  # 指定特定的GPU
                print('\n'+cmd)
                os.system(cmd)  # 运行python脚本
                i += 1
                if i%7 == 0 and i!=0:  # GPU的数量
                    i = 0
                    gpu_memory_list = get_info()  # 遍历一次GPU资源后重新查询
            else:
                i += 1
            
        else:
            gpu_memory_str = 'gpu memory:%d MiB'% gpu_memory_list[i]
            sys.stdout.write('\r'+str(i)+' '+gpu_memory_str)  # 写入缓冲
            sys.stdout.flush()  # 输出并清空缓冲
            time.sleep(interval)
            i += 1
            if i%7 == 0 and i!=0: # 遍历一次GPU资源后重新查询
                i = 0
                gpu_memory_list = get_info()

            
if __name__ =='__main__':
    setup_python()

上述的代码是无限次数排队以及根据占用GPU程序的大小来实现

接下来的是根据 特定次数排队（或无限） 以及根据当前排队的GPU程序的大小来实现（以下的是完整版）：

### 需求： ###
# 1、排队运行
# 2、限制运行次数（无限次或者自定义次数都可以）control_time决定：None为无限次否则为自定义次数


import os
import sys
import time

cmd = 'python xxx.py' # 文件名

def get_info(): # 得到当前被占用的空间
    gpu_memory_list = []
    gpu_status = os.popen('nvidia-smi | grep %').readlines()
    for list_ in gpu_status:
        list_ = list_.split('|')
        gpu_memory = int(list_[2].split('/')[0].split('M')[0].strip())
        gpu_memory_list.append(gpu_memory)
        # 主要是针对memory操作，如果对power需要操作的话同理
        # gpu_power = int(list_[1].split('  ')[-1].split('/')[0].split('W')[0].strip())
    return  gpu_memory_list

def get_info_v1():  # 得到的是差值，即当前占用剩余空间
    gpu_memory_list = []
    gpu_status = os.popen('nvidia-smi | grep %').readlines()
    for list_ in gpu_status:
        list_ = list_.split('|')
        gpu_memory = int(list_[2].split('/')[0].split('M')[0].strip())
        total_memory = int(list_[2].split('/')[1].split('M')[0].strip())
        gpu_memory_list.append(total_memory-gpu_memory)   # memory 操作
    return  gpu_memory_list
    
def setup_python(interval=1,gpu_num=2,gpu_store=5000,control_time=2):
    # gpu_memory_list = get_info()
    gpu_memory_list = get_info_v1()
    i = 0
    control_mark = True
    
    while(i!=gpu_num):
        # if gpu_memory_list[i] <= gpu_store:  # 设置条件，满足则运行python程序  
        if gpu_memory_list[i] >= gpu_store:  # 此时是排队程序的大小
            if i==0 or i==1 :  # 指定特定的GPU
                print('\n'+cmd)
                os.system(cmd)  # 运行python脚本
                i += 1
                if control_time is not None:
                    control_time -= 1
                if control_time == 0:
                    control_mark = False
                if not control_mark:
                    break
                else:
                    if i%gpu_num == 0 and i!=0 :  # GPU的数量
                        i = 0
                        # gpu_memory_list = get_info()  # 遍历一次GPU资源后重新查询
                        gpu_memory_list = get_info_v1()
            else:
                i += 1
            
        else:
            gpu_memory_str = 'gpu memory:%d MiB'% gpu_memory_list[i]
            sys.stdout.write('\r'+str(i)+' '+gpu_memory_str)  # 写入缓冲
            sys.stdout.flush()  # 输出并清空缓冲
            time.sleep(interval)
            i += 1
            if i%gpu_num == 0 and i!=0: # 遍历一次GPU资源后重新查询
                i = 0
                # gpu_memory_list = get_info()
                gpu_memory_list = get_info_v1()

            
if __name__ =='__main__':
    gpu_num = 2  # GPU的数量
    gpu_store = 2000  # 运行程序的占用大小
    setup_python(gpu_num=gpu_num,gpu_store=gpu_store,control_time=None)  # None 表示无限次  否则是自定义系数