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

分布式訓練與 Kubeflow

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

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

1. 在 Kubernetes 集群上創建 Pod 以拉起各個訓練行程
2. 配置用作服務發現的資訊（如 TF_CONFIG）以及創建相關 Kubernetes 資源（如 Service）
3. 監控并更新整個任務的狀態

事實上，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 的專案維護依然存在許多挑戰，

主要包括：

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

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

1. 用戶需要安裝多個 operator 組件才能支持多種訓練 APIs
2. 各種 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 的未來的發展，我們認為存在以下幾個領域值得持續投入：

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

關于我們

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

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

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

標籤：其他

上一篇：歷時90天，從外包公司到位元組軟體測開工程師，就靠這份測開檔案

下一篇：第三章 類檔案結構