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

隨時監測是否有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