【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 = '
')
        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
RUN ON GPU : %d
Started @ %s
'%(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 = '
')

    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 ! ')

随时监测是否有GPU空闲,若有,则将任务添加上去,直至所有任务计算完毕。

实验结果:

 

实验结果表明可以达到我们的目的。

参考和引用:

1、https://github.com/QuantumLiu/tf_gpu_manager

2、https://github.com/calico/basenji/blob/master/basenji/util.py

tz@croplab, HZAU

2020-9-16

原文地址:https://www.cnblogs.com/acm-icpcer/p/13679508.html