【Python】多GPU服务器自动选择空闲GPU

背景：

有大量的GPU任务需要在多GPU服务器上执行，每个任务理论上仅使用单张GPU卡。在不依赖集群调度程序的基础上，并考虑服务器其他用户争抢GPU资源的可能性，此代码库提供可以串行或并行地部署多GPU任务到多GPU卡、并动态的将队列当中的等待任务前赴后继地添加到随时空闲出来的GPU上的解决方案。

PS：目前仅能做到通过空余显存数量来判断GPU是否空闲。这样做的原因是：长期的实践经验表明，GPU计算任务是否能将提交到GPU上受显存的影响因素最大；即使GPU浮点计算使用率为百分之百，只要有足够的显存，任务还是可以提交上去；但是如果显存不足，即使GPU浮点计算率为0，任务也提交不上去。

源码：

https://github.com/wnm1503303791/Multi_GPU_Runner

测试环境：

代码库可解决的两种情况：

（1）我们仅需要执行一系列串行的GPU任务（一般适用于前后相关联的一系列GPU计算任务）：

from manager import GPUManager
gm=GPUManager()

while(1):
    localtime = time.asctime( time.localtime(time.time()) )
    gpu_index = gm.choose_no_task_gpu()
    if gpu_index >= 0 :
        print('Mission Start Running @ %s'%(localtime));

        # gpu_index = 0

        cmd_1 = 'CUDA_VISIBLE_DEVICES=' + str(gpu_index) + ' ' + 'python ...'
        subprocess.call(cmd_1, shell=True)

        cmd_2 = 'python ...'
        subprocess.call(cmd_2, shell=True)

        break;
    else:
        print('Keep Looking @ %s'%(localtime),end = '\r')
        continue;

print('Mission Complete ! Checking GPU Process Over ! ')

原理很简单，使用while循环持续探测GPU情况，只要有一个GPU被其他用户的进程释放，则立即将我们需要计算的任务部署到空闲的GPU上。串行完成所有计算任务之后打破循环，结束主进程。

（2）有一系列GPU任务，任务之间不相关联，可以动态地并行部署到多GPU卡上，目的是尽早结束所有GPU计算任务：

from manager import GPUManager
gm=GPUManager()

mission_queue = []
#for i in range(3):
if(1):
    #以下的cmd_用于测试目的，真正使用的时候将字符串cmd_的内容换成自己需要执行的GPU任务命令即可
    cmd_ = 'python ./fizzbuzz.py > fizzbuzz_1'
    mission_queue.append(cmd_)
    cmd_ = 'python fizzbuzz.py > fizzbuzz_2'
    mission_queue.append(cmd_)
    cmd_ = 'python ./fizzbuzz.py > fizzbuzz_3'
    mission_queue.append(cmd_)
    cmd_ = 'python fizzbuzz.py > fizzbuzz_4'
    mission_queue.append(cmd_)
    cmd_ = 'python ./fizzbuzz.py > fizzbuzz_5'
    mission_queue.append(cmd_)

p = []
total = len(mission_queue)
finished = 0
running = 0

while(finished + running < total):
    '''
    if len(mission_queue) <= 0 :
        break;
    '''
    localtime = time.asctime( time.localtime(time.time()) )
    gpu_av = gm.choose_no_task_gpu()
    # 在每轮epoch当中仅提交1个GPU计算任务
    if len(gpu_av) > 0 :
        gpu_index = random.sample(gpu_av, 1)[0]#为了保证服务器上所有GPU负载均衡，从所有空闲GPU当中随机选择一个执行本轮次的计算任务
        cmd_ = 'CUDA_VISIBLE_DEVICES=' + str(gpu_index) + ' ' + mission_queue.pop(0)#mission_queue当中的任务采用先进先出优先级策略
        print('Mission : %s\nRUN ON GPU : %d\nStarted @ %s\n'%(cmd_, gpu_index, localtime))
        # subprocess.call(cmd_, shell=True)
        p.append(subprocess.Popen(cmd_, shell=True))
        running += 1
        time.sleep(10)#等待NVIDIA CUDA代码库初始化并启动

    else:#如果服务器上所有GPU都已经满载则不提交GPU计算任务
        print('Keep Looking @ %s'%(localtime), end = '\r')

    new_p = []#用来存储已经提交到GPU但是还没结束计算的进程
    for i in range(len(p)):
        if p[i].poll() != None:
            running -= 1
            finished += 1
        else:
            new_p.append(p[i])

    if len(new_p) == len(p):#此时说明已提交GPU的进程队列当中没有进程被执行完
        time.sleep(1)
    p = new_p

for i in range(len(p)):#mission_queue队列当中的所有GPU计算任务均已提交，等待GPU计算完毕结束主进程
    p[i].wait()

print('Mission Complete ! Checking GPU Process Over ! ')