[筆記] Python multi-process task in Docker Container

Python multi-process task in Docker Container

container setting

啟動 container 時要做一些設定

假設我的機器有個 8 核心 CPU，以及多張 GPU
要讓 container 使用全部運算資源的話可以這樣設定

docker run --gpus all --cpus 8 --cpuset-cpus 0-7  (後略...)

參數設定

• --cpuset-cpus 0-7
  • 8 核心的 cpu 編號從 0 開始到 7
  • 如果想指定 cpu 特定某幾個核心的話可以這樣改:
    --cpuset-cpus 0,7 使用第一個跟最後一個核心
    --cpuset-cpus 0,2,4,6 使用第 1,3,5,7 個核心

ref: https://docs.docker.com/config/containers/resource_constraints/

python 使用 multiprocessing

查看 cpu 核心數

import multiprocessing

cpus = multiprocessing.cpu_count()
print(cpus) # 8

假設我們有個任務名稱叫做 sleep_task，執行每個運算時都會睡一秒
如果數據量超大，使用單個核心會等到天荒地老

import time

# task 內容: 睡一秒後再做放大兩倍的運算
def sleep_task(input_list):
    res = []
    for x in input_list:
        time.sleep(1)
        res.append(x*2)
    return res

if __name__ == '__main__':
    # Single-process start
    start_time = time.time()    
    input_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
    sp_result = sleep_task(input_list)  
    # Single-process result
    print('sp_result', sp_result)
    # 因為 input_list 有 16 個元素，每個元素運算需要 1 秒，所以大概會花 16 秒左右
    print('sp_result time cost:', time.time()-start_time) 
    
    # Multi-process start
    start_time = time.time()
    p = multiprocessing.Pool()
    
    results = []
    for i in range(8):  # 使用 8 核心 cpu
        # 欲使每個核心分別處理 input_list 當中的 2 個資料
        # 可以手動將資料切成小塊透過 args 送進去 sleep_task
        result = p.apply_async(sleep_task, args=(input_list[i*2:(i+1)*2], )) 
        results.append(result)
        
    p.close()
    
    mp_result = [] 
    for r in results:
        mp_result += r.get()
        
    # Multi-process result
    print('mp_result', mp_result)
    # 因為 input_list 有 16 個元素，每個元素運算需要 1 秒，但分給 8 個核心去願算，所以大概會花 2 秒左右
    print('mp_result time cost', time.time()-start_time)

執行結果

使用同樣的步驟，在 12 核心機器上執行任務
可以用 top 指令監看各個 cpu 使用率
 輸入 top 指令後按下 1 就可以看到各個 cpu 目前使用情形

(按下 q 退出)