背景

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

當面臨較複雜的模型訓練或者資料量大時，單機的計算能力往往無法滿足算力要求。透過使用阿里的 AiACC 或者社群的 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。

對 Elastic training 的實現原理感興趣可以看這篇 Elastic Horovod 設計文件， 本文不詳細介紹。

在 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

設計

TrainingJob Controller 主要有以下功能：

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

TrainingJob 子資源建立順序如下：

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

TrainingJob 相關資源：

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/v1alpha1kind: TrainingJobmetadata:  name: elastic-training  namespace: defaultspec:  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

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/v1alpha1kind: TrainingJobmetadata:  name: elastic-training  namespace: defaultspec:   // ...... Launcher and Worker specstatus:  currentScaler: ScaleIn:default/scaleout-ptfnk  phase: Scaling  currentWorkers:  - elastic-training-worker-0  - elastic-training-worker-1

ScaleOut 任務 CR：

ScaleIn 任務 CR：

詳細工作過程：

執行1. 安裝 ET-Operator

mkdir -p $(go env GOPATH)/src/github.com/aliyunContainerServicecd $(go env GOPATH)/src/github.com/aliyunContainerServicegit clone https://http://github.com/aliyunContainerService/et-operatorcd et-operatorkubectl create -f deploy/all_in_one.yaml 

檢測 crd 的安裝：

# kubectl get crdNAME                                    CREATED ATscaleins.kai.alibabacloud.com           2020-11-11T11:16:13Zscaleouts.kai.alibabacloud.com          2020-11-11T11:16:13Ztrainingjobs.kai.alibabacloud.com       2020-11-11T11:16:13Z

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

# kubectl -n kube-ai get poNAME                                         READY   STATUS              RESTARTS   AGEet-operator-controller-manager-7877968489-c5kv4   0/2     ContainerCreating   0          5s

2. 執行 TrainingJob

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

kubectl apply -f examples/training_job.yaml

檢測執行狀態：

# kubectl get trainingjobNAME                          PHASE     AGEelastic-training              Running   77s# kubectl get poNAME                                      READY   STATUS             RESTARTS   AGEelastic-training-launcher                 1/1     Running            0          7selastic-training-worker-0                 1/1     Running            0          10selastic-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/v1alpha1kind: ScaleInmetadata:  name: scalein-workersspec:  selector:    name: elastic-training  toDelete:    count: 1

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

apiVersion: kai.alibabacloud.com/v1alpha1kind: ScaleInmetadata:  name: scalein-workersspec:  selector:    name: elastic-training  toDelete:    podNames:    - elastic-training-worker-1

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

kubectl create -f examples/scale_in_count.yaml

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

# kubectl get scaleinNAME                                     PHASE            AGEscalein-sample-t8jxd                     ScaleSucceeded   11s# kubectl get poNAME                                      READY   STATUS             RESTARTS   AGEelastic-training-launcher                 1/1     Running            0          47selastic-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 scaleoutNAME                                     PHASE            AGEelastic-training-scaleout-9dtmw          ScaleSucceeded   30skubectl get poNAME                                      READY   STATUS             RESTARTS   AGEelastic-training-launcher                 1/1     Running            0          2m5selastic-training-worker-0                 1/1     Running            0          2m8selastic-training-worker-1                 1/1     Running            0          40selastic-training-worker-2                 1/1     Running            0          40s

總結

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