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。