小米雲深度學習平台的架構設計與實現

作者 | 陳迪豪

責編 | 何永燦

深度學習服務介紹

機器學習與人工智慧，相信大家已經耳熟能詳，隨著大規模標記數據的積累、神經網路演算法的成熟以及高性能通用GPU的推廣，深度學習逐漸成為計算機專家以及大數據科學家的研究重點。近年來，無論是圖像的分類、識別和檢測，還是語音生成、自然語言處理，甚至是AI下圍棋或者打遊戲都基於深度學習有了很大的突破。而隨著TensorFlow、Caffe等開源框架的發展，深度學習的門檻變得越來越低，甚至初中生都可以輕易實現一個圖像分類或者自動駕駛的神經網路模型，但目前最前沿的成果主要還是出自Google、微軟等巨頭企業。

Google不僅擁有優秀的人才儲備和大數據資源，其得天獨厚的基礎架構也極大推動了AI業務的發展，得益於內部的大規模集群調度系統Borg，開發者可以快速申請大量GPU資源進行模型訓練和上線模型服務，並且通過資源共享和自動調度保證整體資源利用率也很高。Google開源了TensorFlow深度學習框架，讓開發者可以在本地輕易地組合MLP、CNN和RNN等模塊實現複雜的神經網路模型，但TensorFlow只是一個數值計算庫，並不能解決資源隔離、任務調度等問題，將深度學習框架集成到基於雲計算的基礎架構上將是下一個關鍵任務。

除了Google、微軟，國內的百度也開源了PaddlePaddle分布式計算框架，並且官方集成了Kubernetes等容器調度系統，用戶可以基於PaddlePaddle框架實現神經網路模型，同時利用容器的隔離性和Kubernetes的資源共享、自動調度、故障恢復等特性，但平台不能支持更多深度學習框架介面。而亞馬遜和騰訊雲相繼推出了面向開發者的公有雲服務，可以同時支持多種主流的開源深度學習框架，阿里、金山和小米也即將推出基於GPU的雲深度學習服務，還有無數企業在默默地研發內部的機器學習平台和大數據服務。

面對如此眼花繚亂的雲服務和開源技術，架構師該如何考慮其中的技術細節，從用戶的角度又該如何選擇這些平台或者服務呢。我將介紹小米雲深度學習平台的架構設計與實現細節，希望能給AI領域的研發人員提供一些思考和啟示。

雲深度學習平台設計

雲深度學習平台，我定義為Cloud Machine Learning，就是基於雲計算的機器學習和深度學習平台。首先TensorFlow、MXNet是深度學習框架或者深度學習平台，但並不是雲深度學習平台，它們雖然可以組成一個分布式計算集群進行模型訓練，但需要用戶在計算伺服器上手動啟動和管理進程，並沒有雲計算中任務隔離、資源共享、自動調度、故障恢復以及按需計費等功能。因此我們需要區分深度學習類庫以及深度學習平台之間的關係，而這些類庫實現的隨機梯度下降和反向傳播等演算法卻是深度學習應用所必須的，這是一種全新的編程範式，需要我們已有的基礎架構去支持。

雲計算和大數據發展超過了整整十年，在業界催生非常多優秀的開源工具，如實現了類似AWS IaaS功能的OpenStack項目，還有Hadoop、Spark、Hive等大數據存儲和處理框架，以及近年很火的Docker、Kubernetes等容器項目，這些都是構建現代雲計算服務的基石。這些雲服務有共同的特點，例如我們使用HDFS進行數據存儲，用戶不需要手動申請物理資源就可以做到開箱即用，用戶數據保存在幾乎無限制的公共資源池中，並且通過租戶隔離保證數據安全，集群在節點故障或者水平擴容時自動觸發Failover且不會影響用戶業務。雖然Spark通過MLib介面提供部分機器學習演算法功能，但絕不能替代TensorFlow、Caffe等深度學習框架的作用，因此我們仍需要實現Cloud Machine Learning服務，並且確保實現雲服務的基本特性——我將其總結為下面幾條：

屏蔽硬體資源保證開箱即用

縮短業務環境部署和啟動時間

提供「無限」的存儲和計算能力

實現多租戶隔離保證數據安全

實現錯誤容忍和自動故障遷移

提高集群利用率和降低性能損耗

相比於MapReduce或者Spark任務，深度學習的模型訓練時間周期長，而且需要調優的超參數更多，平台設計還需要考慮以下幾點：

支持通用GPU等異構化硬體

支持主流的深度學習框架介面

支持無人值守的超參數自動調優

支持從模型訓練到上線的工作流

這是我個人對雲深度學習平台的需求理解，也是小米在實現cloud-ml服務時的基本設計原則。雖然涉及到高可用、分布式等頗具實現難度的問題，但藉助目前比較成熟的雲計算框架和開源技術，我們的架構和實現基本滿足了前面所有的需求，當然如果有更多需求和想法歡迎隨時交流。

雲深度學習平台架構

遵循前面的平台設計原則，我們的系統架構也愈加清晰明了，為了滿足小米內部的所有深度學習和機器學習需求，需要有一個多租戶、任務隔離、資源共享、支持多框架和GPU的通用服務平台。通過實現經典的MLP、CNN或RNN演算法並不能滿足業務快速發展的需求，因此我們需要支持TensorFlow等用戶自定義的模型結構，並且支持高性能GPU和分布式訓練是這個雲深度學習平台的必須功能，不僅僅是模型訓練，我們還希望集成模型服務等功能來最大化用戶的使用效益。

計算機領域有句名言「任何計算機問題都可以通過增加一個中間層來解決」。無論是AWS、OpenStack、Hadoop、Spark還是TCP/IP都是這樣做的，通過增加一個抽象層來屏蔽底層資源，對上層提供更易用或者更可靠的訪問介面。小米的cloud-ml平台也需要實現對底層物理資源的屏蔽，尤其是對GPU資源的抽象和調度，但我們不需要重新實現，因為社區已經有了很多成熟的分布式解決方案，如OpenStack、Yarn和Kubernetes。目前OpenStack和Yarn對GPU調度支持有所欠缺，虛擬機也存在啟動速度慢、性能overhead較大等問題，而容器方案中的Kubernetes和Mesos發展迅速，支持GPU調度等功能，是目前最值得推薦的架構選型之一。

目前小米cloud-ml平台的任務調度和物理機管理基於多節點的分布式Kubernetes集群，對於OpenStack、Yarn和Mesos我們也保留了實現介面，可以通過實現Mesos後端讓用戶的任務調度到Mesos集群進行訓練，最終返回給用戶一致的使用介面。目前Kubernetes最新穩定版是1.6，已經支持Nvidia GPU的調度和訪問，對於其他廠商GPU暫不支持但基本能滿足企業內部的需求，而且Pod、Deployment、Job、StatefulSet等功能日趨穩定，加上Docker、Prometheus、Harbor等生態項目的成熟，已經在大量生產環境驗證過，可以滿足通用PaaS或者Cloud Machine learning等定製服務平台的需求。

使用Kubernetes管理用戶的Docker容器，還解決了資源隔離的問題，保證不同深度學習訓練任務間的環境不會衝突，並且可以針對訓練任務和模型服務使用Job和Deployment等不同的介面，充分利用分布式容器編排系統的重調度和負載均衡功能。但是，Kubernetes並沒有完善的多租戶和Quota管理功能，難以與企業內部的許可權管理系統對接，這要求我們對Kubernetes API進行再一次「抽象」。我們通過API Server實現了內部的AKSK簽名和認證授權機制，在處理用戶請求時加入多租戶和Quota配額功能，並且對外提供簡單易用的RESTful API，進一步簡化了整個雲深度學習平台的使用流程，整體架構設計如圖1。

圖1 雲深度學習平台整體架構

通過實現API Server，我們對外提供了API、SDK、命令行以及Web控制台多種訪問方式，最大程度上滿足了用戶複雜多變的使用環境。集群內置了Docker鏡像倉庫服務，託管了我們支持的17個深度學習框架的容器鏡像，讓用戶不需要任何初始化命令就可以一鍵創建各框架的開發環境、訓練任務以及模型服務。多副本的API Server和Etcd集群，保證了整個集群所有組件的高可用，和Hadoop或者Spark一樣，我們的cloud-ml服務在任意一台伺服器經歷斷網、宕機、磁碟故障等暴力測試下都能自動Failover保證業務不受任何影響。

前面提到，我們通過抽象層定義了雲深度學習平台的介面，無論後端使用Kubernetes、Mesos、Yarn甚至是OpenStack、AWS都可以支持。通過容器的抽象可以定義任務的運行環境，目前已經支持17個主流的深度學習框架，用戶甚至可以在不改任何一行代碼的情況下定義自己的運行環境或者使用自己實現的深度學習框架。在靈活的架構下，我們還實現了分布式訓練、超參數自動調優、前置命令、NodeSelector、Bring Your Own Image和FUSE集成等功能，將在下面逐一介紹。

雲深度學習平台實現

前面提到我們後端使用Kubernetes編排系統，通過API Server實現授權認證和Quota配額功能。由於雲深度學習服務是一個計算服務，和我以前做過的分布式存儲服務有著本質的區別，計算服務離線運算時間較長，客戶端請求延時要求較低而且吞吐很小，因此我們的API服務在易用性和高性能上可以選擇前者，目前主流的Web伺服器都可以滿足需求。基於Web伺服器我們可以實現集成內部許可權管理系統的業務邏輯，小米生態雲提供了類似AWS的AKSK簽名認證機制，用戶註冊登錄後可以自行創建Access key和Secret key，請求時在客戶端進行AKSK的簽名後發送，這樣用戶不需要把賬號密碼或密鑰加到請求中，即使密鑰泄露也可以由用戶來禁用，請求時即使簽名被嗅探也只能重放當前的請求內容，是非常可靠的安全機制。除此之外，我們參考OpenStack項目的體系架構，實現了多租戶和Quota功能，通過認證和授權的請求需要經過Quota配額檢查，在高可用資料庫中持久化相應的數據，這樣平台管理員就可以動態修改每個租戶的Quota，而且用戶可以隨時查看自身的審計信息。

小米cloud-ml服務實現了深度學習模型的開發、訓練、調優、測試、部署和預測等完整功能，都是通過提交到後端的Kubernetes集群來實現，完整的功能介紹可以查看官方文檔http://docs.api.xiaomi.com/cloud-ml/。Kubernetes對外提供了RESTful API訪問介面，通過YAML或者JSON來描述不同的任務類型，不同編程語言實現的系統也可以使用社區開發的SDK來訪問。對於我們支持的多個深度學習框架，還有開發環境、訓練任務、模型服務等功能，都需要定製Docker鏡像，提交到Kubernetes時指定使用的容器鏡像、啟動命令等參數。通過對Kubernetes API的封裝，我們可以簡化Kubernetes的使用細節，保證了對Mesos、Yarn等後端支持的兼容性，同時避免了直接暴露Kubernetes API帶來的授權問題以及安全隱患。

除了可以啟動單個容器執行用戶的訓練代碼，小米cloud-ml平台也支持TensorFlow的分布式訓練，使用時只需要傳入ps和worker個數即可。考慮到對TensorFlow原生API的兼容性，我們並沒有定製修改TensorFlow代碼，用戶甚至可以在本地安裝開源的TensorFlow測試後再提交，同樣可以運行在雲平台上。但本地運行分布式TensorFlow需要在多台伺服器上手動起進程，同時要避免進程使用的埠與其他服務衝突，而且要考慮系統環境、內存不足、磁碟空間等問題，代碼更新和運維壓力成倍增加，Cloud Machine Learning下的分布式TensorFlow只需要在提交任務時多加兩個參數即可。有人覺得手動啟動分布式TensorFlow非常繁瑣，在雲端實現邏輯是否更加複雜？其實並不是，通過雲服務的控制節點，我們在啟動任務前就可以分配不會衝突的埠資源，啟動時通過容器隔離環境資源，而用戶不需要傳入Cluster spec等繁瑣的參數，我們遵循Google CloudML標準，會自動生成Cluster spec等信息通過環境變數加入到容器的啟動任務中。這樣無論是單機版訓練任務，還是幾個節點的分布式任務，甚至是上百節點的分布式訓練任務，cloud-ml平台都可以通過相同的鏡像和代碼來運行，只是啟動時傳入的環境變數不同，在不改變任何外部依賴的情況下優雅地實現了看似複雜的分布式訓練功能。

圖2 雲深度學習平台分布式訓練

看到這裡大家可能認為，小米的cloud-ml平台和Google的CloudML服務，都有點類似之前很火的PaaS（Platform as a Service）或者CaaS（Container as a Service）服務。確實如此，基於Kubernetes或者Mesos我們可以很容易實現一個通用的CaaS，用戶上傳應用代碼和Docker鏡像，由平台調度和運行，但不同的是Cloud Machine Learning簡化了與機器學習無關的功能。我們不需要用戶了解PaaS的所有功能，也不需要支持所有編程語言的運行環境，暴露提交任務、查看任務、刪除任務等更簡單的使用介面即可，而要支持不同規模的TensorFlow應用代碼，用戶需要以標準的Python打包方式上傳代碼。Python的標準打包方式獨立於TensorFlow或者小米cloud-ml平台，幸運的是目前Google CloudML也支持Python的標準打包方式，通過這種標準介面，我們甚至發現Google CloudML打包好的samples代碼甚至可以直接提交到小米cloud-ml平台上訓練。這是非常有意思的嘗試，意味著用戶可以使用原生的TensorFlow介面來實現自己的模型，在本地計算資源不充足的情況下可以提交到Google CloudML服務上訓練，同時可以一行代碼不用改直接提交到小米或者其他雲服務廠商中的雲平台上訓練。如果大家在實現內部的雲深度學習平台，不妨也參考下標準的Python打包方式，這樣用戶同一份代碼就可以兼容所有雲平台，避免廠商綁定。

除了訓練任務，Cloud Machine Learning平台最好也能集成模型服務、開發環境等功能。對於模型服務，TensorFlow社區開源了TensorFlow Serving項目，可以載入任意TensorFlow模型並且提供統一的訪問介面，而Caffe社區也提供了Web demo項目方便用戶使用。目前Kubernetes和Mesos都實現了類似Deployment的功能，通過製作TensorFlow Serving等服務的容器鏡像，我們可以很方便地為用戶快速啟動對應的模型服務。通過對Kubernetes API的封裝，我們在暴露給用戶API時也提供了replicas等參數，這樣用戶就可以直接通過Kubernetes API來創建多副本的Deployment實例，並且由Kubernetes來實現負載均衡等功能。除此之外，TensorFlow Serving本身還支持在線模型升級和同時載入多個模型版本等功能，我們在保證TensorFlow Serving容器正常運行的情況下，允許用戶更新分布式對象存儲中的模型文件就可以輕易地支持在線模型升級的功能。

對於比較小眾但有特定使用場景的深度學習框架，Cloud Macine Learning的開發環境、訓練任務和模型服務都支持Bring Your Own Image功能，也就是說用戶可以定製自己的Docker鏡像並在提交任務時指定使用。這種靈活的設置極大地降低了平台管理者的維護成本，我們不需要根據每個用戶的需求定製通用的Golden image，事實上也不可能有完美的鏡像可以滿足所有需求，用戶不同的模型可能有任意的Python或者非Python依賴，甚至是自己實現的私有深度學習框架也可以直接提交到Cloud Machine Learning平台上訓練。內測BYOI功能時，我們還驚喜地發現這個功能對於我們開發新的深度學習框架支持，以及提前測試鏡像升級有非常大的幫助，同時用戶自己實現的Caffe模型服務和XGBoost模型服務也可以完美支持。

當然Cloud Machine Learning平台還可以實現很多有意思的功能，例如通過對線上不同GPU機型打label，通過NodeSelector功能可以允許用戶選擇具體的GPU型號進行更細粒度的調度，這需要我們暴露更底層Kubernetes API實現，這在集群測試中也是非常有用的功能。而無論是基於GPU的訓練任務還是模型服務，我們都製作了對應的CUDA容器鏡像，通過Kubernetes調度到對應的GPU計算節點就可以訪問本地圖像處理硬體進行高性能運算了。小米cloud-ml還開放了前置命令和後置命令功能，允許用戶在啟動訓練任務前和訓練任務結束後執行自定義命令，對於不支持分布式存儲的深度學習框架，可以在前置命令中掛載S3 fuse和FDS fuse到本地目錄，或者初始化HDFS的Kerberos賬號，靈活的介面可以實現更多用戶自定義的功能。還有超參數自動調優功能，與Google CloudML類似，用戶可以提交時指定多組超參數配置，雲平台可以自動分配資源起多實例並行計算，為了支持讀取用戶自定義的指標數據，我們實現了類似TensorBoard的Python介面直接訪問TensorFlow event file數據，並通過命令行返回給用戶最優的超參數組合。最後還有TensorFlow Application Template功能，在Cloud Machine Learning平台上用戶可以將自己的模型代碼公開或者使用官方維護的開源TensorFlow應用，用戶提交任務時可以直接指定這些開源模板進行訓練，模板已經實現了MLP、CNN、RNN和LR等經典神經網路結構，還可以通過超參數來配置神經網路每一層的節點數和層數，而且可以支持任意稠密和稀疏的數據集，這樣不需要編寫代碼就可以在雲平台上訓練自己的數據快速生成AI模型了。

在前面的平台設計和平台架構後，要實現完整的雲深度學習服務並不困難，尤其是集成了Docker、Etcd、Kubernetes、TensorFlow等優秀開源項目，組件間通過API松耦合地交互，需要的重複工作主要是打通企業內部許可權系統和將用戶請求轉化成Kubernetes等後端請求而已，而支持標準的打包方式還可以讓業務代碼在任意雲平台上無縫遷移。

雲深度學習平台實踐

目前小米雲深度學習平台已經在內部各業務部門推廣使用，相比於直接使用物理機，雲服務擁有超高的資源利用率、快速的啟動時間、近乎「無限」的計算資源、自動的故障遷移、支持分布式訓練和超參數自動調優等優點，相信可以得到更好的推廣和應用。

除了完成上述的功能，我們在實踐時也聽取了用戶反饋進行改進。例如有內部用戶反饋，在雲端訓練的TensorFlow應用把event file也導出到分布式存儲中，使用TensorBoard需要先下載文件再從本地起服務相對麻煩，因此我們在原有基礎架構實現了TensorboardService功能，可以一鍵啟動TensorBoard服務，用戶只需要用瀏覽器就可以打開使用。

管理GPU資源和排查GPU調度問題也是相當繁瑣的，尤其是需要管理不同GPU設備和不同CUDA版本的異構集群，我們統一規範了CUDA的安裝方式，保證Kubernetes調度的容器可以正常訪問宿主機的GPU設備。當然對於GPU資源的調度和釋放，我們有完善的測試文檔可以保證每一個GPU都可以正常使用，根據測試需求實現的NodeSelector功能也幫忙我們更快地定位問題。

由於已經支持幾十個功能和十幾個深度學習框架，每次升級都可能影響已有服務的功能，因此我們會在多節點的分布式staging集群進行上線演習和測試，並且實現smoke test腳本進行完整的功能性測試。服務升級需要更新代碼，但是為了保證不影響線上業務，無論是Kubernetes還是我們實現的API Server都有多副本提供服務，通過高可用技術先遷移服務進行滾動升級，對於一些單機運行的腳本也通過Etcd實現了高可用的搶主機制，保證所有組件沒有單點故障。

大家可以通過前面提到的文檔地址和cloud-ml-sdk項目了解到更多細節，或者關注我微博（@tobe-陳迪豪）與我交流。

總結

本文介紹了實現企業級雲深度學習平台需要的概念和知識，基於小米cloud-ml服務探討了雲平台的設計、架構、實現以及實踐這四方面的內容，希望大家看完有所收穫。

作者簡介

陳迪豪，小米深度學習工程師，負責小米雲深度學習平台的架構和實現，目前專註於TensorFlow和Kubernetes社區。

點擊展開全文

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！