雲原生 AI 前沿：Kubeflow Training Operator 統一雲上 AI 訓練

分散式訓練與 Kubeflow

當開發者想要講深度學習的分散式訓練搬上 Kubernetes 叢集時，首先想到的往往就是 Kubeflow 社群中形形色色的 operators，如 tf-operator、mpi-operator。

這些服務於各種深度學習訓練（TensorFlow、PyTorch、MXNet 等）的 operators 主要的工作包括：

在 Kubernetes 叢集上建立 Pod 以拉起各個訓練程序
配置用作服務發現的資訊（如 TF_CONFIG）以及建立相關 Kubernetes 資源（如 Service）
監控並更新整個任務的狀態

事實上，Kubeflow 的訓練 Operators 已經成為在 Kubernetes 上執行分散式訓練任務的實際標準。

不僅各大公有云廠商都已經基本收錄或集成了 Kubeflow 的訓練 operators，社群上其他與深度學習訓練相關的專案（如用以自動機器學習的 Katib，又如提供自動化編排功能的 Flyte）都對接了 Kubeflow 中的 operators 作為下發建立分散式訓練任務的工具。

Kubeflow Operators 的問題

在 2019 年初，Kubeflow 社群啟動了 kubeflow/common 專案用以維護 operator 之間重複使用的部分程式碼。經過一年多的迭代和重構，在 2020 年中該專案逐漸穩定並開始接入訓練 operator 。當前，tf-operator、mxnet-operator 和 xgboost-operator 即為構建在 kubeflow/common 專案之上的訓練 operators。

然而，整個 Kubeflow 訓練 operators 的專案維護依然存在許多挑戰。

主要包括：

大量開發者的精力耗費在針對不同訓練框架的功能增強和故障修復上
難以將測試和釋出的基礎功能與服務在不同 operators 之間複用
第三方元件需要對接大量不同的 operators
新的訓練框架需要開發完整的對應的 operator 才能使用，開發成本過高
眾多的 operators 對剛剛接觸 Kubeflow 的新人開發者而言學習成本過高

以上問題都是 Kubeflow 的開發者和維護者面對的。除此之外，這些 operator 的使用者同樣面臨一些問題：

使用者需要安裝多個 operator 元件才能支援多種訓練 APIs
各種 Kubeflow Jobs 的 JobSpec 看上去很類似，但是又有些許不同，並沒有提供統一的使用體驗

這問題的原因主要在於每個深度學習框架都對應一個的 operator 獨立在一個 repository 中進行維護。這種分開維護的模式使得諸如構建環境、測試環境、部署方式以及程式碼邏輯都無法做到很好的整合。

儘管深度學習框架的數量處在收斂的過程中，但依然會有源源不斷的新框架希望通過 Kubeflow 可以快速接入 Kubernetes 進行分散式訓練，而這些新的增量使得問題變得更為嚴重。

Proposal：All-in-One

針對上面提到的各項問題，經過社群會議的多次討論，決定嘗試通過融合的方式將多個 Kubeflow 的訓練 operator 程式碼匯聚到一個倉庫。

同時，參照 controller-runtime 中推薦的 One-Manager-Multi-Controller 的模式，讓多個處理不同 API 的 controller 可以共享一個 Manager 及其 cache，在簡化程式碼的同時也減少了在多個 operator 同時部署時冗餘的 APIServer 請求：

    mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{...})

    ...

    for _, s := range enabledSchemes {

        setupFunc, supported := controller_v1.SupportedSchemeReconciler[s]

        if !supported {os.Exit(1)}

        if err = setupFunc(mgr, enableGangScheduling); err != nil {

            setupLog.Error(err, "unable to create controller", "controller", s)

            os.Exit(1)

        }

    }

所有的 Controller（Reconciler）都需要向 SupportedSchemeReconciler 提前完成註冊：

var SupportedSchemeReconciler = map[string]ReconcilerSetupFunc{

    tensorflowv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {

        return tensorflowcontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)

    },

    pytorchv1.Kind: func(mgr manager.Manager, enableGangScheduling bool) error {

        return pytorchcontroller.NewReconciler(mgr, enableGangScheduling).SetupWithManager(mgr)

    },

    ...,

}

使用者可以在啟動 operator 程序時通過 --enable-scheme 來指定需要開啟支援的 API。後續有新的 Controller 接入，按照這種“先註冊後啟動”的方式來選擇性地開啟對應的 controllers。

進展與近期規劃

當前融合已經正式併入 tf-operator 的 master 分支。使用者很快可以在即將釋出的 Kubeflow 1.4 Release 中體驗到融合後的 tf-operator：部署單個 operator 即可支援包括 TFJob、PyTorchJob、MXNetJob 和 XGBoostJob 在內的四種 API 支援。

在程式碼倉庫層面的融合是 Kubeflow Training Operator 邁向下一個階段的第一步。這一步更多地解決了在專案運營層面，包括環境複用、整體程式碼管理上的一致性。而針對開發者的低程式碼開發，包括新功能增強、bug 修復和新 API 接入，將是我們規劃的下一步目標。

根據這樣的設計，開發者只需要修改非常有限的幾個函式即可接入新的 API。

主要包括：

// 根據 ctrl.Request 獲取對應的自定義 Job

GetJob(ctx context.Context, req ctrl.Request) (client.Object, error)

// 從自定義 Job 中以 map[commonv1.ReplicaType]*commonv1.ReplicaSpec 的格式抽取 ReplicasSpecs

ExtractReplicasSpec(job client.Object) (map[commonv1.ReplicaType]*commonv1.ReplicaSpec, error)

// 從自定義 Job 中抽取 RunPolicy

ExtractRunPolicy(job client.Object) (*commonv1.RunPolicy, error)

// 從自定義 Job 中抽取 JobStatus

ExtractJobStatus(job client.Object) (*commonv1.JobStatus, error)

開發者如果需要注入一些用以服務發現的環境變數，可以覆蓋方法 DecoratePod(rtype commonv1.ReplicaType, podTemplate *corev1.PodTemplateSpec, job client.Object) 在 client 向 APIServer 提交建立請求前修改 Pod。

以上低程式碼開發方式的基礎已經以 pkg/reconciler.v1 的形態合入 kubeflow/common 倉庫。很快，我們也將在 tf-operator 上引入基於該 reconciler.v1 包的基礎 API，希望可以在驗證 reconciler.v1 的同時為更多通用的實用案例提供一種更為簡便接入 Kubernetes 的方式。

如果開發者希望以更低層 API 的方式對 controller 進行開發，pkg/controller.v1 包可以滿足這一類開發者的需求。

遠景展望

儘管針對 Kubeflow Training Operator 的優化改造還在進行中，我們並沒有止步於此。對於 Training Operator 的未來的發展，我們認為存在以下幾個領域值得持續投入：

首先是進一步提高 Kubeflow Training Operator 適配定製化需求 Job 時的靈活性。我們計劃提出與深度學習訓練框架解耦的一種 Job API 以支援更廣泛的任務定義，並允許使用者可以藉助 kubeflow/common 中的 controller.v1 和 reconciler.v1 進行定製化開發，但其學習成本和開發成本依然過高。甚至在將來，初級開發者可以不修改 operator 而僅僅新增/修改一些 webhook 或是 decorator server 來實現定製化修改。
第二個方面是進一步增強 Kubeflow Training Operator 和其他第三方元件互動時的便利性。我們希望未來利用 Kubeflow Training Operator 來構建 AI 平臺的開發者可以方便地將其與其他模組對接，實現諸如任務佇列、流水線、超引數搜尋等功能。
最後也是最關鍵的，我們依然希望可以進一步提升 Kubeflow Training Operator 的穩定性。

我們歡迎更多的同學能夠嘗試、體驗 Kubeflow 並且投入到 Kubeflow 專案中來。

參考資料

[1]add reconciler.v1: 【https://github.com/kubeflow/common/pull/141】

[2]

reconciler.v1 implementation: 【https://github.com/kubeflow/common/tree/master/pkg/reconciler.v1/common】

[3] All-in-one Kubeflow Training Operator: 【https://docs.google.com/document/d/1x1JPDQfDMIbnoQRftDH1IzGU0qvHGSU4W6Jl4rJLPhI/edit】

關於我們

更多關於雲原生的案例和知識，可關注同名【騰訊雲原生】公眾號~

福利：公眾號後臺回覆【手冊】，可獲得《騰訊雲原生路線圖手冊》&《騰訊雲原生最佳實踐》~

【騰訊雲原生】雲說新品、雲研新術、雲遊新活、雲賞資訊，掃碼關注同名公眾號，及時獲取更多幹貨！！

雲原生 AI 前沿：Kubeflow Training Operator 統一雲上 AI 訓練