![頭圖.png](https://ucc.alicdn.com/pic/developer-ecology/a6e2e19e5348456eafd6336c22c8fe78.png)

作者 | 徐曉舟（蕭元）
來源|阿[裡巴巴雲原生公眾號](https://mp.weixin.qq.com/s/LQfa0R9JruXxqDEKDWBBZQ)

# 背景

由於雲端計算在資源成本和彈性擴容方面的天然優勢，越來越多客戶願意在雲上構建 AI 系統，而以容器、Kubernetes 為代表的雲原生技術，已經成為釋放雲價值的最短路徑， 在雲上基於 Kubernetes 構建 AI 平臺已經成為趨勢。

當面臨較複雜的模型訓練或者資料量大時，單機的計算能力往往無法滿足算力要求。透過使用阿里的 AiACC 或者社群的 [horovod](https://github.com/uber/horovod) 等分散式訓練框架，僅需修改幾行程式碼，就能將一個單機的訓練任務擴充套件為支援分散式的訓練任務。在 Kubernetes 上常見的是 kubeflow 社群的 tf-operator 支援 Tensorflow PS 模式，或者 mpi-operator 支援 horovod 的 mpi allreduce 模式。

# 現狀

Kubernetes 和雲端計算提供敏捷性和伸縮性，我們可以透過 cluster-AutoScaler 等元件為訓練任務設定彈性策略，利用 Kubernetes 的彈效能力，按需建立，減少 GPU 裝置空轉。

但這種伸縮模式面對訓練這種離線任務還是略有不足：

- 不支援容錯，當部分 Worker 由於裝置原因失敗，整個任務需要停止重來。
- 訓練任務一般時間較長，佔用算力大，任務缺少彈效能力。當資源不足時，除非任務終止，無法按需為其他業務騰出資源。
- 訓練任務時間較長，不支援 worker 動態配置， 無法安全地使用搶佔例項，發揮雲上最大價效比

如何給訓練任務賦予彈效能力，是提高價效比的關鍵路徑。近期 horovod 等分散式框架逐漸支援了 Elastic Training，即彈性訓練能力。也就是允許一個訓練任務在執行的過程中動態的擴容或者縮容訓練 worker， 從不會引起訓練任務的中斷。需要在程式碼中做少量修改適配，可參考：[https://horovod.readthedocs.io/en/stable/elastic_include.html](https://horovod.readthedocs.io/en/stable/elastic_include.html)。

對 Elastic training 的實現原理感興趣可以看這篇 [Elastic Horovod 設計文件](https://docs.google.com/document/d/15ZoHA5AeSI_boeyIBapg9WPXKrYXMRvPytPzQWTCTn4/edit#)， 本文不詳細介紹。

![1.png](https://ucc.alicdn.com/pic/developer-ecology/8a760c717a594407b33fda33d70dd9f0.png)

在 mpi-operator 中，參與訓練的 Worker 都是作為靜態資源設計和維護，支援彈性訓練模式後，給任務增加了靈活性，同時也給運維層帶來了挑戰，例如：

- 必須透過 horovod 提供的 horovordrun 作為入口，horovod 中 launcher 透過 ssh 登陸 worker，需要打通 launcher 和 worker 之間的登陸隧道。

- 負責計算彈性的 Elastic Driver 模組透過指定 discover_host 指令碼獲取最新 worker 拓撲資訊，從而拉起或停止 worker 例項。當 worker 變化時，首先要更新 discover_host 指令碼的返回值。

- 在搶佔或價格計算等場景中，有時需要指定 worker 縮容，K8s 原生的編排元語 deployment，statefulset 無法滿足指定縮容的場景。

# 解決方法

針對以上問題，我們設計開發了 et-operator，提供 TrainingJob CRD 描述訓練任務, ScaleOut 和 ScaleIn  CRD 描述擴容和縮容操作， 透過它們的組合，使我們的訓練任務更具有彈性。將這個方案開源，歡迎大家提需求、交流、吐槽。

開源方案地址：[https://github.com/AliyunContainerService/et-operator](https://github.com/AliyunContainerService/et-operator)

# 設計

TrainingJob Controller 主要有以下功能：

- 維護 TrainingJob 的建立/刪除生命週期，以及子資源管理。
- 執行擴縮容操作。
- 容錯，當 worker 被驅逐，建立新的 worker 加入到訓練中。

## 1. 資源建立

TrainingJob 子資源建立順序如下：

- 建立打通 ssh 所需的金鑰對， 建立 secret。
- 建立 workers，包含 service 和 pod，掛載 secret 公鑰。
- 建立 configmap， 包含 discover_host 指令碼 , hostfile 檔案。
- 建立 launcher，掛載 configmap。由於 hostfile 後續會隨著拓撲關係修改，所以 hostfile 單獨透過 initcontainer 從 configmap 複製到單獨目錄。

TrainingJob 相關資源：

![2.png](https://ucc.alicdn.com/pic/developer-ecology/fd9a6d7b61254e26b2994445f4370b90.png)

TrainingJob CR 的配置分為 Lanucher 和 Worker。在 Launcher 中指定任務的映象和啟動執行， 預設 et-operator 會根據 worker 分配情況，生成一個 hostfile 檔案和 discover_host 指令碼，discover_host 指令碼掛載到 Launcher 的 /etc/edl/discover_hosts.sh 檔案， 在入口指令碼的 horovodrun 執行中透過 --host-discovery-script 引數指定。在 Worker 設定中指定 worker 的映象和 GPU 佔用 ，並可以透過 maxReplicas / minReplicas 指定 workers 的副本數允許範圍。

```
apiVersion: kai.alibabacloud.com/v1alpha1
kind: TrainingJob
metadata:
  name: elastic-training
  namespace: default
spec:
  cleanPodPolicy: Running
  etReplicaSpecs:
    launcher:
      replicas: 1
      template:
        spec:
          containers:
          - command:
            - sh
            - -c
            - horovodrun -np 2 --min-np 1 --max-np 9 --host-discovery-script
              /etc/edl/discover_hosts.sh python /examples/elastic/tensorflow2_mnist_elastic.py
            image: registry.cn-huhehaote.aliyuncs.com/lumo/horovod:master-tf2.1.0-torch1.4.0-mxnet-py3.6-gpu
            imagePullPolicy: Always
            name: mnist-elastic
    worker:
      maxReplicas: 9
      minReplicas: 1
      replicas: 2
      template:
        spec:
          containers:
          - image: registry.cn-huhehaote.aliyuncs.com/lumo/horovod:master-tf2.1.0-torch1.4.0-mxnet-py3.6-gpu
            imagePullPolicy: Always
            name: mnist-elastic
            resources:
              limits:
                nvidia.com/gpu: "1"
              requests:
                nvidia.com/gpu: "1"
status:
  currentWorkers:
  - elastic-training-worker-0
  - elastic-training-worker-1
  - elastic-training-worker-2
  - elastic-training-worker-3
  phase: Succeeded
  replicaStatuses:
    Launcher:
      active: 1
      succeeded: 1
    Worker:
      active: 4
```

![3.png](https://ucc.alicdn.com/pic/developer-ecology/76abbe8bd0d64318af93e52b411ca078.png)

## 2. Worker 擴容 / 縮容

除了 TrainingJob 外，et-operator 同時支援 ScaleOut 和 ScaleIn 兩種 CRD，下發訓練任務擴容和縮容操作。

當下發一個 ScaleOut CR，ScaleOutController 觸發 Reconcile， 這裡工作很簡單，根據 ScaleOut CR 中的 Selector 欄位，找到 Scaler 對應的 TrainingJob，設定到 CR 的 OwnerReferences 上。

以一個 ScaleOut 操作舉例：

```
- apiVersion: kai.alibabacloud.com/v1alpha1
  kind: ScaleOut
  metadata:
    creationTimestamp: "2020-11-04T13:54:26Z
    name: scaleout-ptfnk
    namespace: default
    ownerReferences:
    - apiVersion: kai.alibabacloud.com/v1alpha1
      blockOwnerDeletion: true
      controller: true
      kind: TrainingJob
      name: elastic-training // 指向擴容物件TrainingJob
      uid: 075b9c4a-22f9-40ce-83c7-656b329a2b9e
  spec:
  selector:
    name: elastic-training
  toAdd:
    count: 2
```

TrainingJobController 中監聽到屬於 TrainingJob  的 ScaleOut CR 有更新， 觸發 TrainingJob 的 Reconcile，遍歷過濾 TrainingJob 下 OwnerReference 指向的 ScaleIn 和 ScaleOut， 根據建立時間和狀態時間決定執行的擴容或者縮容。

```
apiVersion: kai.alibabacloud.com/v1alpha1
kind: TrainingJob
metadata:
  name: elastic-training
  namespace: default
spec: 
  // ...... Launcher and Worker spec
status:
  currentScaler: ScaleIn:default/scaleout-ptfnk
  phase: Scaling
  currentWorkers:
  - elastic-training-worker-0
  - elastic-training-worker-1
```

ScaleOut 任務 CR：

![4.png](https://ucc.alicdn.com/pic/developer-ecology/302954134f5349c3892e2d3d9173ce62.png)

ScaleIn 任務 CR：

![5.png](https://ucc.alicdn.com/pic/developer-ecology/c7e5b10ab8e34a24acf951cbd3195eae.png)

詳細工作過程：

![6.png](https://ucc.alicdn.com/pic/developer-ecology/615ef2a5cc4041368f90b1ec3ae8237d.png)

# 執行

## 1. 安裝 ET-Operator

```
mkdir -p $(go env GOPATH)/src/github.com/aliyunContainerService
cd $(go env GOPATH)/src/github.com/aliyunContainerService
git clone https://http://github.com/aliyunContainerService/et-operator
cd et-operator
kubectl create -f deploy/all_in_one.yaml 
```

檢測 crd 的安裝：

```
# kubectl get crd
NAME                                    CREATED AT
scaleins.kai.alibabacloud.com           2020-11-11T11:16:13Z
scaleouts.kai.alibabacloud.com          2020-11-11T11:16:13Z
trainingjobs.kai.alibabacloud.com       2020-11-11T11:16:13Z
```

檢測 controller 的執行狀態，預設安裝在 kube-ai 中：

```
# kubectl -n kube-ai get po
NAME                                         READY   STATUS              RESTARTS   AGE
et-operator-controller-manager-7877968489-c5kv4   0/2     ContainerCreating   0          5s
```

## 2. 執行 TrainingJob

執行事先已準備好的示例：

```
kubectl apply -f examples/training_job.yaml
```

檢測執行狀態：

```
# kubectl get trainingjob
NAME                          PHASE     AGE
elastic-training              Running   77s

# kubectl get po
NAME                                      READY   STATUS             RESTARTS   AGE
elastic-training-launcher                 1/1     Running            0          7s
elastic-training-worker-0                 1/1     Running            0          10s
elastic-training-worker-1                 1/1     Running            0          9s
```

## 3. 縮容訓練任務 Worker

執行縮容時，可以透過 ScaleIn CR 中的 spec.toDelete.count  或 spec.toDelete.podNames  欄位指定縮容的 worker。

透過 count 配置縮容的數量，則透過 index 計算由高到低縮容 Worker。

```
apiVersion: kai.alibabacloud.com/v1alpha1
kind: ScaleIn
metadata:
  name: scalein-workers
spec:
  selector:
    name: elastic-training
  toDelete:
    count: 1
```

如果想要縮容特定的 Worker，可以配置 podNames：

```
apiVersion: kai.alibabacloud.com/v1alpha1
kind: ScaleIn
metadata:
  name: scalein-workers
spec:
  selector:
    name: elastic-training
  toDelete:
    podNames:
    - elastic-training-worker-1
```

執行一個縮容示例，指定數量縮容 1 個 worker：

```
kubectl create -f examples/scale_in_count.yaml
```

檢測縮容執行狀態和訓練任務：

```
# kubectl get scalein
NAME                                     PHASE            AGE
scalein-sample-t8jxd                     ScaleSucceeded   11s

# kubectl get po
NAME                                      READY   STATUS             RESTARTS   AGE
elastic-training-launcher                 1/1     Running            0          47s
elastic-training-worker-0                 1/1     Running            0          50s
```

## 4. 擴容訓練任務

在 ScaleOut CR 中，透過 spec.toAdd.count 欄位指定擴容的 worker 數：

```
apiVersion: kai.alibabacloud.com/v1alpha1
  kind: ScaleOut
  metadata:
    name: elastic-training-scaleout-9dtmw
    namespace: default
  spec:
    selector:
      name: elastic-training
    timeout: 300
    toAdd:
      count: 2
```

執行示例：

```
kubectl create -f examples/scale_out.yaml
```

檢測縮容執行狀態和訓練任務：

```
kubectl get scaleout
NAME                                     PHASE            AGE
elastic-training-scaleout-9dtmw          ScaleSucceeded   30s
kubectl get po
NAME                                      READY   STATUS             RESTARTS   AGE
elastic-training-launcher                 1/1     Running            0          2m5s
elastic-training-worker-0                 1/1     Running            0          2m8s
elastic-training-worker-1                 1/1     Running            0          40s
elastic-training-worker-2                 1/1     Running            0          40s
```

# 總結

ET-Operator 提供一組訓練和擴縮容 CRD 和 Controller， 讓我們在 Kubernetes 上方便地執行彈性分散式訓練，支援下發分散式訓練任務，並透過和分散式框架的整合聯動，在訓練任務執行過程中動態地擴容和縮容參與運算的 Workers。使我們的訓練任務具有彈效能力，結合搶佔例項，能夠更好的利用雲上的資源彈性和價效比優勢。

在 Kubernetes 上彈性深度學習訓練利器 -- Elastic Training Operator

相關文章