KubeDL 加入 CNCF Sandbox，加速 AI 産業雲原生化

作者 | KubeDL 團隊

2021 年 6 月 23 日，雲原生計算基金會（CNCF）宣布通過全球 TOC 投票接納 KubeDL 成為 CNCF Sandbox 項目。KubeDL 是阿裡開源的基于 Kubernetes 的 AI 工作負載管理架構，取自"Kubernetes-Deep-Learning"的縮寫，希望能夠依托阿裡巴巴的場景，将大規模機器學習作業排程與管理的經驗反哺社群。

項目位址：

http://kubedl.io

項目介紹

随着 TensorFlow, PyTorch，XGBoost 等主流 AI 架構的不斷成熟，和以 GPU/TPU 為代表的多種AI異構計算晶片的井噴式湧現，人工智能正快速進入“大規模工業化”落地的階段。從算法工程師着手設計第一層神經網絡結構，到最終上線服務于真實的應用場景，除 AI 算法的研發外還需要大量基礎架構層面的系統支援，包括資料收集和清理、分布式訓練引擎、資源排程與編排、模型管理，推理服務調優，可觀測等。如以下經典圖例所展示，衆多系統元件的協同組成了完整的機器學習流水線。

與此同時，以 Kubernetes 為代表的雲原生技術蓬勃發展，通過優秀的抽象和強大的可擴充性，将應用層與 IaaS（Infrastructure as a Service）層的基礎設施完美解耦：應用能夠以“雲”的範式按需使用資源，無需關注底層基礎設施的複雜性，進而解放生産力并專注于自身領域的創新。

Kubernetes 的出現解決了雲資源如何高效傳遞的問題，但對于 AI 這類本身具備高度複雜性的工作負載還無法做到很好地原生支援，如何整合各類架構的差異并保留其通用性，同時圍繞 AI 工作負載的運作時去建設一系列完善的周邊生态及工具，業界還在不斷探索與嘗試。在實踐中，我們發現了 AI 負載運作在 Kubernetes 生态中面臨着如下挑戰：

• 機器學習架構百花齊放，各自有不同的優化方向和适用場景，但在分布式訓練作業的生命周期管理上又存在着諸多共性，同時針對一些進階特性也有相同的訴求（如網絡模式，鏡像代碼分離，中繼資料持久化，緩存加速等）。為每類架構的負載單獨實作 operater，各自獨立程序無法共享 state，缺乏全局視角，使得全局 Job 層面的排程以及隊列機制難以實作。此外，不利于功能的抽象和複用，在代碼層面存在重複勞動。

• 原生 Kubernetes 無法滿足離線任務多樣的排程需求。Kubernetes 面向 Pod 排程的模型天然适用于微服務等 Long Running 的工作負載，但針對離線任務的高吞吐，Gang Scheduling 排程（All-Or-Nothing），Elastic Capacity 等多種排程訴求，社群演進出了多種排程方案。以機器學習分布式訓練作業排程場景中極為常見的Gang Scheduling為例，社群目前就有YuniKorn，Volcano，Coscheduling 等排程器實作，提供不同的互動協定，我們需要有插件化的手段來啟用對應的排程協定。同時，像 PS/worker 這類根據業務特有屬性，不同 role 之間有啟動依賴的 DAG 編排訴求，需要在控制器中實作；

• 分布式訓練的結果往往以模型作為 output，并存儲在分布式檔案系統中如(阿裡雲 OSS/NAS)，但如何從訓練作業的視角去管理模型，像容器鏡像那樣成為AI服務的“不可變基礎設施”并實作簡單且清晰的版本管理與追溯，業界還缺乏最佳實踐。同時，“訓練”與“推理”兩個階段相對獨立，算法科學家視角中的“訓練->模型->推理”機器學習流水線缺乏斷層，而“模型”作為兩者的中間産物正好能夠充當那個“承前啟後”的角色；

• 分布式訓練尚能大力出奇迹，但推理服務的規格配置卻是一個精細活。顯存量、 CPU 核數、BatchSize、線程數等變量都可能影響推理服務的品質。純粹基于資源水位的容量預估無法反映業務的真實資源需求，因為某些引擎如 TensorFlow 會對顯存進行預占。理論上存在一個服務品質與資源效能的最優平衡點，但它就像黑暗中的幽靈，明知道它的存在卻難以琢磨。随着 GPU 虛拟化技術的成熟，這個平衡點的價值越來越凸顯，更優的規格能顯著提供單 GPU 卡的部署密度，節約大量的成本。

• 推理服務本身是一種特殊的 long running 微服務形态，除了基礎的 deployment 外，針對不同的推理場景還欠缺一些執行個體與流量的管理政策，如：

  1) 算法科學家通常會同時部署兩個甚至多個不同版本的模型執行個體進行 A/B Test 以驗證最佳的服務效果，需要基于權重的精細化流量控制；

  2) 能夠根據流量請求水準和目前推理服務的 metrics 來自動觸發執行個體的擴縮，在充分保障服務可用性的前提下最小化資源成本等等。

KubeDL

針對上述難題，阿裡巴巴雲原生，叢集管理和 PAI 團隊将管理大規模機器學習工作負載的經驗沉澱為通用的運作時管理架構——KubeDL，涵蓋分布式訓練，模型管理，推理服務等機器學習流水線的各階段，使工作負載能夠高效地運作在 Kubernetes 之上。

1、分布式訓練

KubeDL 支援了主流的機器學習分布式訓練架構（TensorFlow / PyTorch / MPI / XGBoost / Mars 等），其中 Mars 是阿裡巴巴計算平台開源的基于張量的大規模資料計算架構，能夠分布式地加速 numpy，pandas 等資料處理架構的效率，幫助 Mars 作業以更 native 的方式內建進雲原生大資料生态中。

我們将各類訓練作業生命周期管理中的共同部分進行抽象，成為一層通用的運作時庫，被各分布式訓練作業控制器複用，同時使用者也可以在此基礎上快速擴充出自定義的 workload 控制器并複用現有的能力。借助聲明式 API 與 Kubernetes 網絡/存儲模型，KubeDL 能夠進行計算資源的申請/回收，各 Job Role 之間的服務發現與通信，運作時的 Fail-over 等，算法模型的開發者隻需聲明好此次訓練依賴的 Job Role 及各自的副本數，計算資源/異構資源數量等，然後送出任務。另外，我們針對訓練領域的痛點也做了諸多的特性設計來提升訓練的效率與體驗：

• 不同的訓練架構往往包含不同的 Job Role，如 TensorFlow 中的 PS/Chief/Worker 和 PyTorch 中的 Master/Worker，Role 與 Role 之間往往隐含着依賴關系，如 Worker 依賴 Master 啟動之後才能正常開始計算，錯亂的啟動順序不僅容易造成資源長時間空轉，甚至可能引發 Job 直接失敗。KubeDL 設計了基于 DAG（Direct Acyclic Graph）的排程編排控制流，很好地解決了 Role 之間的啟動依賴順序，并能夠靈活擴充。

• 大模型的訓練時長往往受制于計算節點間的通信效率，RDMA 等高性能網絡技術的應用将極大地提升了資料的傳輸速度，但這些定制網絡往往需要計算節點使用 Hostnetwork 進行互相通信，同時有些時候由于環境限制無法提供基于 Service 模式的服務發現機制。這就需要作業管理引擎能夠支援 Host 網絡模式下的服務發現機制，處理好各計算節點的網絡端口配置設定，并與各訓練架構的特性結合來處理節點 Fail-over 後的網絡連通性，KubeDL 支援了 Host 網絡模式下的高性能分布式訓練。

• Gang Scheduling 是分布式訓練作業排程場景中的常見需求，組成單個訓練作業的一簇 Pod 往往要求同時被排程，避免在叢集容量緊張時因作業間的資源競争出現活鎖，但 Kubernetes 強大的可擴充性也使得不同的排程器實作了不同的 Gang Scheduling 協定，如 YuniKorn, KubeBatch 等。為了避免與具體排程器實作的耦合，适應不同使用者環境的差異，KubeDL 将 Gang Scheduling 的協定實作插件化，按需啟用對應的插件即可與排程器互相協作，實作作業的批量排程。

• Job 是一次性的，但在實際的生産應用中我們經常會遇到反複訓練/定時訓練的場景，如每日拉取某一時間區間的離線表并進行資料清洗以及模型的 Re-train，KubeDL 提供了一類單獨的工作負載—Cron 來處理定時的訓練請求，并支援任意類型的訓練作業（如 TFJob，PyTorchJob 等），使用者可以送出 cron tab 風格的定時指令及作業模闆，并在 Cron 資源的狀态中追蹤訓練作業的曆史及目前進行中的作業。

針對海量離線作業中繼資料需要長時間儲存（Job CRD 被删除後中繼資料即從 etcd 銷毀）的訴求，KubeDL 還内置了中繼資料的持久化，實時監聽 Job/Pod/Events 等資源對象的變化，轉化成對應的 Databse Schema Object 并持久化到存儲後端中。存儲後端的設計也是插件化的，使用者可以根據自己的線上環境來實作存儲插件并在部署時 enable。在 KubeDL 中 Job/Pod 預設支援了 Mysql 的存儲協定，以及将 Events 收集到阿裡雲 SLS 服務中。

同時我們還提供了管控套件：KubeDL-Dashboard，使用者不需要去了解 Kubernetes 的衆多 API 并在各種 kubectl 指令中掙紮，即可界面化地上手簡單易用的機器學習作業。持久化的中繼資料也可以直接被 Dashboard 消費使用。Dashboard 提供了簡單的作業送出、作業管理、事件/日志檢視、叢集資源視圖等功能，以極低的學習門檻幫助機器學習使用者上手實驗。

2、推理服務規格調優

GPU 虛拟化與分時複用技術的發展和成熟，讓我們有機會在一塊 GPU 上同時運作多個推理服務，顯著降低成本。然而如何為推理服務選擇合适的 GPU 資源規格，尤其是不可壓縮的顯存資源，成為一個關鍵難題。一方面，頻繁的模型疊代讓算法工程師無暇去精确估計每個模型的資源需求，流量的動态變化也讓資源評估變得不準确，是以他們傾向于配置較多的 GPU 資源備援，在穩定性和效率之間選擇犧牲後者，造成大量浪費；另一方面，由于 Tensorflow 等機器學習架構傾向于占滿所有空閑的顯存，站在叢集管理者的角度，根據顯存的曆史用量來估計推理業務的資源需求也非常不準确。在 KubeDL-Morphling 這個元件中我們實作了推理服務的自動規格調優，通過主動壓測的方式，對服務在不同資源配置下進行性能畫像，最終給出最合适的容器規格推薦。畫像過程高度智能化：為了避免窮舉方式的規格點采樣，我們采用貝葉斯優化作為畫像采樣算法的内部核心驅動，通過不斷細化拟合函數，以低采樣率（<20%）的壓測開銷，給出接近最優的容器規格推薦結果。

3、模型管理與推理服務

模型是訓練的産物，是計算與算法結合後的濃縮精華，通常收集與維護模型的方式是托管在雲存儲上，通過組織檔案系統的方式來實作統一管理。這樣的管理方式依賴于嚴格的流程規範與權限控制，沒有從系統層面實作模型管理的不可變，而容器鏡像的誕生解決的就是 RootFS 的建構-分發-不可變等問題，KubeDL 将兩者進行結合，實作了基于鏡像的模型管理。訓練成功結束後，通過 Job Spec 中指定的 ModelVersion 會自動觸發模型鏡像的建構。使用者可以在 ModelVersion.Spec 中約定模型的存儲路徑，目标的鏡像 Registry 等基本資訊，将每次的訓練輸出 Push 到對應的鏡像倉庫。

同時鏡像作為訓練的輸出，以及推理服務的輸入，很好地串聯起了兩個階段，也借此實作了分布式訓練->模型建構與管理->推理服務部署的完整機器學習流水線。KubeDL 提供了 Inference 資源對象提供推理服務的部署與運作時控制，一個完整的 Inference 服務可以由單個或多個 Predictor 組成，每個 Predictor 對應前序訓練輸出的模型，模型會被自動拉取并挂載到主容器 Volume 中。當多個不同模型版本的 Predictor 并存時，可以根據配置設定的權重進行流量的分發與控制，達到 A/B Test 的對照實驗效果，後續我們還會在 Batching 批量推理和 AutoScale 上針對推理服務場景做更多的探索。

KubeDL 分布式訓練在公有雲上的實踐

• PAI-DLC

随着雲計算的深入人心以及越來越多的業務都用雲原生的方式進行，阿裡雲計算平台 PAI 機器學習團隊推出了 DLC（Deep Learning Cloud）這一深度學習平台産品。DLC 采用全新的雲原生架構，底層采用 Kubernetes 作為資源底座支援，而訓練部分全面采用 KubeDL 進行管理，是 KubeDL 在深度學習雲計算場景中的大規模實踐。

DLC 在阿裡集團内部内廣泛支撐了衆多的業務，包括淘系安全部達摩院的圖像視訊、自然語言、語音、多模态了解、自動駕駛等衆多業務部門的深度學習計算需求。在服務于深度學習驅動的前沿業務生産中，PAI 團隊在架構和平台建設方面積累了許多的經驗，沉澱了相容社群（eg，TensorFlow/PyTorch）并且具有鮮明特色的大規模工業界實踐過的架構平台能力，如萬億規模參數的M6模型的訓練、工業級圖神經網絡系統 Graph-Learn、極緻資源管理和複用能力等等。

如今，PAI-DLC 的能力也在全面擁抱公有雲，為開發者和企業提供的雲原生一站式的深度學習訓練平台，一個靈活、穩定、易用和高性能的機器學習訓練環境，以及全面支援支援多種社群和 PAI 深度優化的算法架構，高性能且穩定的運作超大規模分布式深度學習任務，為開發者和企業降本增效。

公有雲的 DLC 作為阿裡巴巴集團機器學習平台最佳實踐的透出，在産品細節、架構優化、平台服務等方面都吸取了工程實踐中的寶貴的經驗。除此之外，DLC 産品在設計之初就充分考量了公有雲場景中的獨特屬性，提供了競價執行個體、自動 Fail-Over、彈性擴縮等功能，為客戶努力降低 AI 算力成本。

進一步的，DLC 也與 PAI 的其他公有雲産品相結合，比如說服務于算法工程師模組化的 DSW、服務于企業級 AI 全流程的、自動化的 AutoML、線上推理服務 EAS 等，打造全流程的 AI 标杆性産品。