背景：

有大量的GPU任務需要在多GPU伺服器上執行，每個任務理論上僅使用單張GPU卡。在不依賴叢集排程程度的基礎上，並考慮伺服器其他使用者爭搶GPU資源的可能性，此程式碼庫提供可以序列或並行地部署多GPU任務到多GPU卡，並動態的將隊列當中的等待任務前赴後繼地新增到隨時空閒出來的GPU上的解決方案。

PS：目前僅能做到通過空餘視訊記憶體數量來判斷GPU是否空閒。

原始碼：

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