Tensorflow源碼解析7 -- TensorFlow分佈式運行時

1 概述

TensorFlow架構設計精巧，在後端運行時這一層，除了提供本地運行時外，還提供了分佈式運行時。通過分佈式訓練，在多臺機器上並行執行，大大提高了訓練速度。前端用戶通過session.run()啓動系統執行時，target默認爲空字符串""，對應的是本地運行模式。若target以"grpc://"開頭，則對應的是分佈式運行模式，target指定了要連接的TensorFlow執行引擎。

分佈式運行時同樣分爲client master和worker，只是三者不在同一進程內。分佈式運行時同樣是圍繞計算圖Graph來進行的，流程也與本地運行時幾乎相同。client負責圖的構造，並傳遞給master。master接收後，啓動圖的剪枝和分裂，將分裂後的子圖發送給多個worker進程。worker進程負責執行計算子圖，它會先按照自己所在機器包含的設備，先按照設備進行子圖的二次分裂，然後在每個設備上進行子圖執行。所有設備執行完畢後，從計算圖的終止節點sink中取出數據。

本地運行時通過DirectSession同時管理client master和worker，而分佈式運行時則不同。client對應GrpcSession，master對應MasterSession，worker對應WorkerSession。三者使用同一個句柄session_handle進行協同工作。

2 數據交換

和本地運行時類似，分佈式運行時也存在跨設備的數據依賴。對於跨設備的數據邊，將其分裂，在發送方插入send節點，接收方插入recv節點。如果二者跨進程通信（比如兩臺不同的服務器），則通過GrpcRemoteRendezvous進行數據交換。如果二者是進程內通信（比如同一臺服務器的CPU0和CPU1），則通過IntraProcessRendezvous進行數據交換。上節講過的本地運行時在運行前，就創建了一個IntraProcessRendezvous對象。

3 分佈式集羣結構

TensorFlow爲分佈式運行時，設計了一個精巧的結構。共分爲三級。

1. 集羣cluster，可包含多臺服務器，通過ClusterSpec對象描述。它包含多個job，一個job又包含多個Task。一個Task對應一個server。

2. Job。將目的相同的Task劃歸爲一個job，使用job_id唯一標示。一般存在兩種job

   1. ps：數據存儲，負責存儲和更新模型的參數，比如w和b。比較適合CPU
   2. worker：數據計算，負責train和inference時的數據計算工作。比較適合GPU

一般ps將數據發送給worker，待worker運算完畢後再返回給ps，ps再進行數據更新。

1. Task。Task是提供服務的最小單位，它一般單獨在一個進程內，通過job_id:task_index唯一標示。一個Task對應一個server，提供MasterService和WorkerService兩種服務。

下面是一個集羣配置的例子。

tf.train.ClusterSpec({
    "worker": [
        "worker0:1111", # /job:worker/task:0
        "worker1:2222", # /job:worker/task:1
        "worker2:3333" # /job:worker/task:2
    ],
    "ps": [
        "ps0:1111", # /job:ps/task:0
        "ps1:2222" # /job:ps/task:1
]})

這個集羣cluster內包含2個job，一個ps和一個worker。ps又包含2個task，worker則包含3個task，共計5個task。