*作者：陳裘凱（ 求索）*

## 前言

KubeDL 是阿里開源的基於 Kubernetes 的 AI 工作負載管理框架，取自"Kubernetes-Deep-Learning"的縮寫，希望能夠依託阿里巴巴的場景，將大規模機器學習作業排程與管理的經驗反哺社群。目前 KubeDL 已經進入 CNCF Sandbox 專案孵化，我們會不斷探索雲原生 AI 場景中的最佳實踐，助力演算法科學家們簡單高效地實現創新落地。

KubeDL 為分散式訓練作業帶來了 HostNetwork 網路模式，支援計算節點之間透過宿主機網路相互通訊以提升網路效能，同時適應 RDMA/SCC 等新型高效能資料中心架構的網路環境，此外，KubeDL 針對 HostNetwork 模式帶來的 FailOver 後新埠互相感知等問題也帶來了新的解決思路。

Github 地址：
[https://github.com/kubedl-io/kubedl](https://github.com/kubedl-io/kubedl)

網站：
[https://kubedl.io/model/intro/](https://kubedl.io/model/intro/)

## Overlay 不是銀彈

Kubernetes 原生的容器網路模型定義了一系列不依賴 NAT 的"Pod-Pod"間通訊規約，基於 VxLAN 組建的 Overlay 網路很好地實現了這一模型（如經典的 Flannel）並解決了諸多大規模容器編排系統中的網路管理的痛點：

- Pod 的無感遷移：Overlay 網路是基於物理網路構建的虛擬二層網路，Pod IP 並不與任何節點繫結，當節點當機或發生其他硬體異常時，對應的服務 Pod 可以透過相同的 IP 在其他節點上重新啟動，只要底層的物理網路連通不中斷就不影響服務的可用性。在大規模的分散式機器學習訓練中。KubeDL 也是基於“Pod 可能漂移，但 Service 是固定的”這一前提實現的計算節點故障轉移（FailOver）；

- 網路節點的規模：經典的 NAT 地址解析通常透過 ARP 廣播協議來自動學習鄰接節點 IP 與 MAC 地址的對映，但當節點規模龐大時，一次廣播很容易造成 ARP 風暴並引起網路擁塞，而基於隧道穿越的 Overlay 網路只需知道少數的 VTEP 節點的 MAC 地址即能實現資料包的轉發，極大的降低了網路的壓力；

- 租戶網路隔離：Kubernetes 強大的網路外掛擴充套件性配合 VxLAN 的協議設計，很容易實現虛擬網路的再劃分從而實現租戶之間的網路隔離；

這些都是虛擬容器網路帶來的好處，但虛擬化的代價是網路效能的損耗：Pod 與主機網路透過一對 Veth 虛擬網橋裝置連線來實現網路 namespace 的互相隔離，每一次"Pod-Pod"間通訊的資料包都需要經過”封包-路由-乙太網-路由-拆包“等流程才能到達對端的 Pod，拖慢網路效能的同時還會增加宿主機核心網路棧的處理壓力從而提升負載。

![在這裡插入圖片描述](https://img-blog.csdnimg.cn/c7c1b775c9eb4de9ab7b8160ab302a2d.png?x-oss-process=image/watermark,type_d3F5LXplbmhlaQ,shadow_50,text_Q1NETiBA6Zi_6YeM5be05be05LqR5Y6f55Sf,size_20,color_FFFFFF,t_70,g_se,x_16#pic_center)

隨著多模態模型訓練、大規模稠密引數模型訓練等分散式訓練模式的興起，以及資料集規模、特徵引數的爆炸，網路通訊已然成為分散式訓練效率的一個“水桶短板”。最直接的最佳化網路效能的方法即使用主機網路（HostNetwork）通訊，免去容器網路虛擬化的開銷。同時隨著 RDMA（RoCE），Nvidia GPU Direct 等技術的成熟，這些新型的高效能網路技術逐漸被應用於大規模的商業生產環境來大幅提升模型訓練的效率，透過旁路核心網路棧的開銷和零複製直讀資料等技術充分利用網路頻寬，Efficiency Is Money！這些原生的高效能網路通訊庫原語（如 RDMA_CM）也同樣依賴主機網路實現，無法直接基於 Pod 虛擬網路通訊。

KubeDL 在支援分散式訓練基於標準容器網路通訊的基礎上擴充套件了主機網路的通訊模型，同時解決了埠衝突和 FailOver 後新埠互相感知等分散式訓練中的常見問題，實現高效能網路的輕鬆使能。

## 使能 Host 高效能網路

### 標準容器網路拓撲

在標準的容器網路通訊模型中，Master/Worker/PS 等不同 Workload 角色的之間透過 Headless Service 實現服務發現，Pod 之間透過恆定的域名相互通訊，由 CoreDNS 實現域名到 Pod IP 的解析，由於 Pod 是可以漂移的但 Service 及其附屬的域名是恆定的，即使部分 Pod 執行時異常了也能很好地實現 FailOver，在異常 Pod 重新拉起之後與其他 Pod 重連線。

```
apiVersion: training.kubedl.io/v1alpha1
kind: "TFJob"
metadata:
  name: "mnist"
  namespace: kubedl
spec:
  cleanPodPolicy: None
  tfReplicaSpecs:
    PS:
      replicas: 2
      restartPolicy: Never
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubedl/tf-mnist-with-summaries:1.0
              command:
                - "python"
                - "/var/tf_mnist/mnist_with_summaries.py"
                - "--log_dir=/train/logs"
                - "--learning_rate=0.01"
                - "--batch_size=150"
              volumeMounts:
                - mountPath: "/train"
                  name: "training"
              resources:
                limits:
                  cpu: 2048m
                  memory: 2Gi
                requests:
                  cpu: 1024m
                  memory: 1Gi
          volumes:
            - name: "training"
              hostPath:
                path: /tmp/data
                type: DirectoryOrCreate
    Worker:
      replicas: 3
      restartPolicy: ExitCode
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubedl/tf-mnist-with-summaries:1.0
              command:
                - "python"
                - "/var/tf_mnist/mnist_with_summaries.py"
                - "--log_dir=/train/logs"
                - "--learning_rate=0.01"
                - "--batch_size=150"
              volumeMounts:
                - mountPath: "/train"
                  name: "training"
              resources:
                limits:
                  cpu: 2048m
                  memory: 2Gi
                requests:
                  cpu: 1024m
                  memory: 1Gi
          volumes:
            - name: "training"
              hostPath:
                path: /tmp/data
                type: DirectoryOrCreate
```

以一個經典 PS-Worker 架構的 Tensorflow 分散式訓練作業為例，Worker 負責計算引數的梯度，由 PS 負責聚合、更新並廣播引數，因此每個 PS 都可能和所有 Worker 建立連線並通訊，反之亦是。

在 Tensorflow 框架的實現中，這樣一個作業間拓撲結構由一個 TF Cluster Spec 結構來描述，每個 Role（PS or Worker）例項都包含一個 Index 標識自身索引號，可以透過Role+Index 獲取自身或其他Role例項的服務地址，即可建立連線開始通訊。在標準容器網路模式中，使用者提交以下 TFJob，KubeDL 會生成 TF Cluster Spec 並以環境變數的形式傳入並被框架接收，同時為每個 Role 例項都準備好 Headless Service，它的 Endpoint 域名地址即對應 TF Cluster Spec 中的服務地址，每個 Pod 都擁有一份獨立的 Linux Network Namespace，Pod 的埠地址空間也相互隔離，因此排程到相同的 Node 上也可以使用相同的容器埠。

至此不同 Role 的例項間就能透過 Tensorflow 原生的方式開始分散式訓練及通訊。

![在這裡插入圖片描述](https://img-blog.csdnimg.cn/515f0e5396054027836709fd2efe86ab.png?x-oss-process=image/watermark,type_d3F5LXplbmhlaQ,shadow_50,text_Q1NETiBA6Zi_6YeM5be05be05LqR5Y6f55Sf,size_20,color_FFFFFF,t_70,g_se,x_16#pic_center)

標準容器網路的好處顯而易見，簡單直觀的網路設定，FailOver 友好的網路容錯，都使得這一方案能夠滿足大多數場景下的需求。但對高效能網路有訴求的場景下又該如何運轉呢？KubeDL 給出了主機網路的解決方案。

### Host 容器網路拓撲

沿用以上的例子，啟用主機網路的方式很簡單，只要給 TFJob 追加一個 annotation 即可，其餘的作業配置都無需特殊改造，如下所示：

```
apiVersion: training.kubedl.io/v1alpha1
kind: "TFJob"
metadata:
  name: "mnist"
  namespace: kubedl
  annotations:
    kubedl.io/network-mode: host
spec:
  cleanPodPolicy: None
  tfReplicaSpecs:
    PS:
    ...
    Worker:
    ...
```

當 KubeDL 發現該作業宣告瞭使用主機網路後，會透過以下步驟完成網路的連線設定：

- 建立 Pod 時不再使用固定埠，而是在一定埠範圍內隨機出一個主機埠，並設定對應暴露的容器埠號，透過上下文的方式傳遞到後續的控制流中；

- 對 Pod 啟用 HostNetwork 並設定 DNS 解析策略為 Host 優先；

- 不再建立 Headless Service，取而代之的是一個正常的流量轉發 Service，暴露埠為原先的恆定值，目標埠為 Pod 的真實值；

- 生成的 TF Cluster Spec 中，自身對應的 Role+Index 可見 Local 地址埠為真實的主機埠，其他 Role 例項的地址埠都是恆定的，無論對方的 Pod 如何漂移都能透過 Service 正確轉發；

- 當發生 FailOver 時，KubeDL 會為重建後的 Pod 重新選擇埠，新啟動的 Pod 會透過 TF_CONFIG 得到新的 Local 地址埠，同時 KubeDL 保證對應 Service 的目標埠得到正確更新，其他與之相連的 Role 也能在 Service 目標埠更新後繼續通訊；

這樣一個根據訓練作業拓撲結構搭建的主機網路就準備換好了，與之前的不同之處在於，所有的 Pod 都與主機共用了一個 Network Namespace，因此也共享了主機的埠號，而 Pod 之間的通訊也從原先透過解析域名為 Pod IP 並建立連線，變成了透過 Service 實現流量的轉發，另一方面 TF Cluster Spec 發生了變化但沒有改變原生 Tensorflow 的模式，當前 Pod 直接獲得 Local Port 監聽，而其他的 Pod 地址看起來都是恆定的 Service 對應的域名和暴露的埠永遠恆定，只有目標埠可能隨著 FailOver 不斷改變，這一切都透過 KubeDL 處理變得無感。

![在這裡插入圖片描述](https://img-blog.csdnimg.cn/8d327dba7bfc47b6895a1c8b2ae9e1ab.png?x-oss-process=image/watermark,type_d3F5LXplbmhlaQ,shadow_50,text_Q1NETiBA6Zi_6YeM5be05be05LqR5Y6f55Sf,size_20,color_FFFFFF,t_70,g_se,x_16#pic_center)

我們以 Tensorflow 作為主機網路的例子，因為它的 Cluster Spec 複雜性更具代表性，但 KubeDL 的內建工作負載（如 PyTorch，XGBoost 等）我們也都針對其框架的行為實現了對應主機網路模式的網路拓撲設定。

## 總結

KubeDL 透過擴充套件現有的分散式訓練作業標準容器網路通訊模式，實現了基於原生主機網路的通訊模式，在常見訓練場景下獲得網路效能增益的同時，也完美適應了 RDMA/SCC 等高效能網路架構的環境，助力分散式訓練作業執行效率的大幅提升，這一通訊模式已經在阿里巴巴內部的生產叢集中廣泛使用，比如達摩院在雲棲大會最新發布的 AliceMind 超大模型就是透過 KubeDL 主機網路+RDMA 在高效能運算叢集中訓練的產物。我們期待更多開發者參與 KubeDL 社群的建設，一起最佳化深度學習工作負載的排程及執行時效率！

戳[此處](https://kubedl.io/model/intro/)，立即瞭解 KubeDL 專案！

KubeDL HostNetwork：加速分散式訓練通訊效率

相關文章