機器學習，特別是深度學習，在蘑菇街這樣的電商平臺有大量實際業務的落地場景，比如搜索推薦、圖像算法、交易風控反作弊等等。隨着業務的快速發展，之前已有的基於 Yarn 的調度平臺已經無法滿足大規模機器學習的計算需求，因此我們在 2018 年和算法工程團隊一起建設了基於 Kubeflow 和 Kubernetes 的分佈式機器學習平臺，並深入到業務層面進行分佈式改造，並且從Kubernetes、Tensorflow和業務等多個層面進行了一系列的性能調優。
背景
隨着機器學習和人工智能的迅猛發展，業界出現了許多開源的機器學習平臺。由於機器學習與大數據天然的緊密結合，基於 Hadoop Yarn 的分佈式任務調度仍是業界主流，但是隨着容器化的發展，Docker + Kubernetes 的雲原生組合，也展現出了很強的生命力。
| 公司 | 平臺名稱 | 集羣管理和調度 |
|---|---|---|
| 百度 | PaddlePaddle | Docker+Kubernetes |
| 騰訊 | Angel | Docker+Yarn |
| 阿里 | X-DeepLearning | Docker+Yarn |
| 360 | XLearning | Yarn |
| 京東 | 登月平臺 | Docker+Kubernetes |
| 才雲科技 | Clever | Docker+Kubernetes |
表1. 互聯網業界機器學習平臺架構對比
痛點
在建設分佈式訓練平臺的過程中，我們和機器學習的各個業務方，包括搜索推薦、圖像算法、交易風控反作弊等，進行了深入溝通，調研他們的痛點。從中我們發現，算法業務方往往專注於模型和調參，而工程領域是他們相對薄弱的一個環節。建設一個強大的分佈式平臺，整合各個資源池，提供統一的機器學習框架，將能大大加快訓練速度，提升效率，帶來更多的可能性，此外還有助於提升資源利用率。
痛點一：對算力的需求越來越強烈
算力代表了生產力。深度學習在多個領域的出色表現，給業務帶來了更多的可能性，同時對算力提出了越來越高的要求。深度學習模型參數的規模陡增，迭代的時間變的更長，從之前的小時級別，變成天級別，甚至月級別。以商品推薦爲例，面對幾億的參數，近百億的樣本數量，即使採用 GPU 機器，也需要長達一個星期的訓練時間；而圖像業務擁有更多的參數和更復雜的模型，面對 TB 級的訓練樣本，單機場景下往往需要長達近一個月的訓練時間。
再者，機器學習具有試錯性非常強的特點，更快的訓練速度可以帶來更多的嘗試，從而發掘更多的可能性。Tensorflow 從 0.8 版本開始支持分佈式訓練，至今爲止，無論高階還是低階的 API，對分佈式訓練已經有了完善的支持。同時，Kubernetes 和 Hadoop 等具有完善的資源管理和調度功能，爲 Tensorflow 分佈式訓練奠定資源層面的基礎。
Tensorflow On Yarn 和 Tensorflow On Spark 是較早的解決方案，奇虎360的 Xlearning 也得到衆人的青睞。而基於 Kubernetes 的 kubeflow 解決了 Yarn 的痛點，展現出旺盛的生命力。
上述方案無一例外的將各個部門分散的機器納入到統一的資源池，並提供資源管理和調度功能，讓基於 Excel 表的資源管理和人肉調度成爲過去，讓用戶更專注於算法模型，而非基礎設施。在幾十個 worker 下，無論是 CPU 還是 GPU 的分佈式訓練，訓練速度都能得到近乎線性的提升，將單機的訓練時間從月級別縮短到一天以內，提升效率的同時也大大提升了資源利用率。
蘑菇街早期的業務方往往獨立維護各自團隊的GPU機器“小池子”，機器的資源分配和管理存在盲區，缺乏統一管理和調度。GPU的資源存在不均衡和資源利用率低下的問題。事實上，大數據團隊之前已經將 CPU 類型的訓練任務接入了 Xlearning，嚐到了分佈式訓練的甜頭，但也發現一些問題：
- 公司目前的 Yarn 不支持 GPU 資源管理，雖然近期版本已支持該特性，但存在穩定性風險。
- 缺乏資源隔離和限制，同節點的任務容易出現資源衝突。
- 監控信息不完善。在發生資源搶佔時，往往無法定位根本原因。
- 缺少彈性能力，目前資源池是靜態的，如果能借助公有云的彈性能力，在業務高峯期提供更大的算力，將能更快的滿足業務需求。
痛點二：人肉管理的成本很高
業務方反饋的第二大問題是人肉管理的成本問題。人肉化的管理主要包含了部署和訓練任務管理兩大方面。
人肉部署
一個典型的場景是：團隊內的成員共享一批機器，每次跑訓練任務前，用戶手動登陸機器，下載代碼，安裝對應的 python 包，過程繁瑣且容易出現安裝包版本的衝突。
由於不同的訓練任務對 python 的版本和依賴完全不同，比如有些模型使用 python 2.7，有些使用 python 3.3，有些使用 tensorflow 1.8，有些使用 tensorflow 1.11 等等，非常容易出現依賴包衝突的問題。雖然沙箱能在一定程度上解決這問題，但是也帶來了額外的管理負擔。
此外，不同GPU機型依賴的 Nvidia 驅動也不同，較新的卡，比如 V100 所依賴的版本更高。人肉部署還需要管理和維護多個不同的驅動版本。
訓練任務管理
人肉啓動訓練任務時，需要手動查看/評估資源的剩餘可用情況，手動指定PS和Worker的數量，管理配置並進行服務發現。這些都給業務方帶來了很大的負擔。
解決的思路
Docker和Kubernetes很好的解決了人肉管理的問題:
- 部署：藉助Docker良好的隔離性，各個容器的文件系統互不影響，將訓練任務和驅動程序製作成鏡像，避免了多次安裝的繁瑣。此外 Kubernetes 提供了服務發現功能，簡化了分佈式的部署。
- 訓練任務生命週期管理：Kubernetes 提供了生命週期管理的 API，用戶基於 API 即可一鍵式完成訓練任務的增刪改查，避免人工 ssh 的各種繁瑣操作，可以大幅提升用戶體驗和效率。
痛點三：監控的缺失
監控也是分佈式訓練重要的一環，它是性能調優的重要依據。比如在 PS-Worker 的訓練框架下，我們需要爲每個 PS/Worker 配置合適的 GPU/CPU/Memory，並設置合適的 PS 和 Worker 數量。如果某個參數配置不當，往往容易造成性能瓶頸，影響整體資源的利用率。比如當 PS 的網絡很大時，我們需要增加 PS 節點，並對大參數進行 partition；當 worker CPU 負載過高時，我們應該增加 Worker 的核數。
早期版本的 Hadoop 和 Yarn 並不支持GPU的資源可視化監控，而 Kubernetes 已擁有非常成熟監控方案 Prometheus，它能週期採集 CPU，內存，網絡和 GPU 的監控數據，即每個 PS/Worker 的資源使用率都能得到詳細的展示，爲優化資源配置提供了依據。事實上，我們也是通過監控信息爲不同的訓練任務分配合適的資源配置，使得在訓練速度和整體的吞吐率上達到一個較好的效果。
痛點四：資源隔離較弱
早期的機器學習平臺基於 Yarn 的 Angel 和 XLearning，由於 Yarn 缺乏對實例之間的資源隔離，我們在內存，網絡，磁盤等均遇到不同程度的問題。
由於 Yarn 沒有對任務的內存進行隔離，所以，業務方常常因對內存資源估計不準而導致 worker 的進程 OOM。由於所有的進程都共用宿主機的 IP，很容易造成端口衝突，此外磁盤 IO 和網絡 IO 也很容易被打爆。
| Yarn | Kubernetes | |
|---|---|---|
| 調度 | 以任務爲粒度進行調度 支持Capacity Schedule、Fair Schedule等多種調度器 搶佔式調度，支持任務排隊，支持基於Tag的調度 |
以Pod爲粒度進行調度 Label/Selector機制，支持多調度器 Kube-batch支持搶佔式調度 |
| 隔離性 | 隔離性偏弱，高版本支持cgroup 磁盤和網絡IO未隔離，多任務容易造成搶佔 |
隔離性強 docker可以使用獨立IP 底層通過Linux內核的cgroup/namespace和TC等機制，保證磁盤和網絡IO的隔離 |
| 多租戶配額管理 | 通過隊列機制實現多租戶的配額管理 | 支持基於namespace的配額管理，包括GPU的配額 |
| 開源社區 | 活躍度相對低，Apache基金會 | 活躍度相對高，CNCF基金會 |
| 可擴展性 | 可擴展性好，可以支持其他 非 Yarn的調度框架 | 可擴展性好，支持CRD和operator，能滿足內部業務的定製化需求 |
表2. Hadoop Yarn和Kubernetes的橫向對比
Kubeflow及核心組件
Kubeflow 是由 Google 等公司於 2017 年推出的基於 Kubernetes 的開源項目，它支持常見的機器學習框架。
Kubeflow 簡介
The Kubeflow project is dedicated to making deployments of machine learning (ML) workflows on Kubernetes simple, portable and scalable.
Kubeflow 旨在支持多種機器學習框架運行在 Kubernetes 之上，比如 Tensorflow, Pytorch, Caffe 等常見框架。它包含了 operator、pipeline、超參數調優、serving 等諸多模塊。它通過提供對應的 operator，基於 Kubernetes 的 Pod/headless Service 等基礎資源爲框架提供與之相配的更高層次的資源。比如 tf-operator 爲 Tensorflow 提供了job維度的生命週期管理能力，以滿足 Tensorflow 分佈式訓練的資源和拓撲需求，達到了一鍵式部署 Tensorflow 訓練任務的效果。
Kubeflow 包含如下 operator，分別對應主流的分佈式計算框架。蘑菇街主要採用了 kubeflow 中的 tf-operator，來實現對機器學習任務的統一管理。
圖1. Kubeflow所支持的主流分佈式計算框架
Distributed Tensorflow
蘑菇街業務方使用的計算框架主要是 Tensorflow，因此有必要先介紹一下 Tensorflow 的分佈式訓練，Tensorflow 支持如下三種分佈式策略：
- MirroredStrategy：適用於單機多卡的訓練場景，功能有限，不在本文討論範圍內。
- ParameterServerStrategy：用於多機多卡場景，主要分爲 worker 節點和 PS 節點，其中模型參數全部存儲在 PS 節點，worker 在每個 step 計算完梯度後向 PS 更新梯度，蘑菇街當前使用這種方案。
- CollectiveAllReduceStrategy：用於多機多卡場景，通過 all-reduce 的方式融合梯度，只需要 worker 節點，不需要 PS 節點，從另外一個角度說，該節點既充當 worker 角色，又充當 PS 角色。該方案是帶寬優化的，具有性能好，可擴展性強的特點，是 Tensorflow 推薦的未來方向。
以 ParameterServerStrategy 爲例，一個分佈式訓練集羣至少需要兩種類型的節點：PS 和 worker。由於在訓練中需要一個 worker 節點來評估效果和保存 checkpoint，因此單獨把該節點作爲 chief(或者叫 master) 節點。通常情況下，一個集羣需要多個 worker 節點，多個 PS 節點，一個 chief 節點。所有 worker 節點的 CPU/內存/GPU等資源配置完全相同，所有 PS 節點的 CPU/內存等資源配置也相同。從資源拓撲角度出發，如果能夠提供一種 Kubernetes 資源，用戶可以基於該資源定義 PS/worker/chief 的數量和規格，用戶就可以一鍵式創建分佈式集羣，大大簡化了分佈式集羣的部署和配置。tf-operator 定義了 TFJob 資源，用戶可以藉助 tf-operator 在 Kubernetes 上一鍵拉起分佈式訓練集羣。
從 Tensorflow 分佈式訓練的使用方式出發，用戶在啓動每個節點的任務時，需要傳入集羣所有節點的網絡信息。這意味着分佈式訓練集羣的每個節點需要預先知道所有其它節點的網絡地址信息，即要求服務發現功能。tf-operator 基於 Kubernetes headless service，完美的提供了服務發現功能，無需用戶再手工指定PS/Worker的IP信息，極大的降低了用戶的部署成本。
圖2. Tensorflow分佈式訓練的ClusterSpec配置
落地實踐
圖3. 基於 Kubernetes 的機器學習基礎平臺總體架構
主要包含了以下幾部分：
Tensorflow生命週期管理(tf-operator)
在部署 tf-operator 之後，首先需要在 Kubernetes 中創建對應的 TFJob CRD，之後就可以創建 TFJob 資源了。
在如下的樣例中，我們定義了一個具有 10 個 worker，4 個 ps，1 個 chief 的分佈式訓練集羣。從 TFJob 參數不難發現，它對ParameterServerStrategy 和 CollectiveAllReduceStrategy 這兩種策略方式都支持，只是在 CollectiveAllReduceStrategy 場景下，無需要設置 PS 資源。
apiVersion: "kubeflow.org/v1alpha1"
kind: "TFJob"
metadata:
name: "example-job"
spec:
replicaSpecs:
- replicas: 1
tfReplicaType: CHIEF
template:
spec:
containers:
- image: gcr.io/tf-on-k8s-dogfood/chief_sample:latest
name: tensorflow
restartPolicy: OnFailure
- replicas: 10
tfReplicaType: WORKER
template:
spec:
containers:
- image: gcr.io/tf-on-k8s-dogfood/worker_sample:latest
name: tensorflow
restartPolicy: OnFailure
- replicas: 4
tfReplicaType: PS
template:
spec:
containers:
- image: gcr.io/tf-on-k8s-dogfood/ps_sample:latest
name: tensorflow
Tf-operator 啓動後，通過 list-watch 不斷的監聽 TFJob 資源相關事件，當收到創建 TFJob 事件時，tf-operator 依次創建 PS/Worker/Chief(Master) Replica 資源。以 PS Replica 爲例，根據 replicas 數量依次創建等同數量的 pod，併爲每個 pod 創建 headless service。此外，它還生成 TF_CONFIG 環境變量，這個環境變量記錄了所有 pod 的域名和端口，最終在創建 pod 時候注入到容器中。
圖4. tf-operator的配置示例
任務調度(kube-batch)
通過引入 kube-batch，滿足了業務方對批量任務調度的需求。沒有采用 Kubernetes 默認的調度器，主要是基於以下兩點考慮：
-
GANG scheduling：Kubernetes 默認的調度器是以 pod 爲粒度的，並不支持GANG scheduling。機器學習的訓練任務要求集羣保持一個整體，要麼所有的 pod 都能成功創建，要麼沒有一個 pod 能被創建。試想資源處於臨界狀態時，如果採用默認的調度器調度一個分佈式訓練集羣，則會導致部分節點因資源不足而無法成功調度，那些成功創建的 pod 空佔用資源，但是無法進行訓練。
-
任務排隊：Kubernetes 默認的調度器無法提供隊列調度的功能，而不會像 Hadoop 任務那樣進行排隊。而面向機器學習/大數據/HPC的批量任務調度，往往要求系統能維護一個或多個隊列，進行任務排隊。當任務結束並且釋放資源後，後續的任務可以繼續執行。
Kube-batch 目前是Kubernetes SIGs旗下的一個孵化項目，是一個運行在 Kubernetes 上面向機器學習/大數據/HPC 的批調度器（batch scheduler），它支持以Pod Group爲單位進行資源調度，並支持preempt和priority。對於暫時無法滿足資源條件的Pod，在 Kubernetes 中會處於pending狀態，直到資源釋放從而繼續執行。
Kube-batch 的基本流程如下圖，它通過 list-watch 監聽 Pod, Queue, PodGroup 和 Node 等資源，在本地維護一份集羣資源的全局緩存，依次通過如下的策略（reclaim, allocate, preemption，predict） 完成資源的調度。
圖5. kube-batch基本流程
圖6. kube-batch工作流程
kube-batch 旨在通過配置策略，提供 gang scheduling 的調度能力，由於是基於 queue 的調度思想，kube-batch 對機器學習及大數據任務的調度具有很大的潛力。
圖7. 蘑菇街的使用姿勢
由於 kube-batch 對任務的 actions 有四種，可以根據自身的業務特點，指定只使用其中的一種或者多種：如果簡單的分配，使用allocation 即可；如果資源存在碎片，讓調度器能夠感知並重新分配，可以將 allocation 與 backfill 結合起來使用。如此既可滿足業務的調度需求，又可以在一定程度上提升調度性能。
認證和授權（Dex)
我們引入了由CoreOS研發的dex 組件，並嵌入到我們的 SDK 中，實現了與公司 LDAP 的對接。我們將 Kubernetes 中的 namespace 映射成不同的應用組，應用的負責人會自動授予（Rolebinding）管理員的角色（Role），同時針對一些特權用戶還會創建 clusterRole 及 clusterRolebinding，方便這些特權用戶訪問某些需要高級權限的 API（如查詢 node 信息等）。
圖8. Dex認證流程
圖9. Dex授權流程
多租戶的隔離性
平臺需要接入多個不同的業務方，自然需要考慮多租戶間的隔離性。我們利用 Kubernetes 提供的 resource quota，以 namespace 爲單位對資源進行分配。
$ kubectl get resourcequota --all-namespaces
NAMESPACE NAME AGE
default compute-resources 73d
tinyplus-docker-algo-anticheat compute-resources 74d
tinyplus-docker-algo-image compute-resources 74d
tinyplus-docker-algo-search compute-resources 74d
tinyplus-docker compute-resources 74d
圖10. 多租戶下的資源配額
性能調優
這裏介紹一下蘑菇街在分佈式機器學習調優的一些經驗，主要分爲 Kubernetes 層面、Tensorflow 層面和業務層面的一些調優。
Kubernetes 層面調優
CPUShare VS CPUSet
以商品推薦常用的 wide and deep 模型爲例，該模型採用 CPU 資源進行訓練，宿主機的規格爲 80C256G。
從直覺上說，訓練本質上是大量的矩陣運算，而矩陣運算在內存地址空間具有良好連續性的特點。若充分利用 CPU cache，將有助於提升 cache 的命中率，最終加快訓練速度。Kubernetes 從 1.8 開始支持 CPU Manager 特性，該特性支持對於 Guaranteed 類型的 pod，採用 CpuSet 爲 Pod 分配獨佔的 CPU core。在相同訓練任務下，對比 CpuSet 和 CpuShare 對訓練速度的影響，發現在 worker CPU 核數較少的情況下，CpuSet 的性能遠遠超過 CpuShare。
圖11. CpuSet和CpuShare性能對比（Y軸數值越大越好）
雖然 Guaranteed 類型的 Pod 犧牲了資源彈性，但是 CpuSet 帶來的性能收益要高於彈性帶來的收益。即使對 CpuShare 容器不設置 cpu limits，當跑滿容器的CPU資源時，相同 cpu requests 下，CpuSet 容器的性能依舊比 CpuShare 的性能高出 20% 以上。
在kubelet中開啓CPU Manaer特性的配置參數如下
--feature-gates=CPUManager=true
--cpu-manager-policy=static
--cpu-manager-reconcile-period=5s
--kube-reserved=cpu=500m
Pod Affinity
PS/Worker 訓練框架下，PS 節點作爲中心節點，網絡流量往往非常大，容易把帶寬跑滿。通過 Pod AntiAffinity 將所有 PS 節點儘可能打散到不同的宿主機上，從而分攤網絡流量。
podAntiAffinity:
preferredDuringSchedulingIgnoredDuringExecution:
- weight: 100
podAffinityTerm:
labelSelector:
matchExpressions:
- key: tf-replica-type
operator: NotIn
values:
- ps
topologyKey: kubernetes.io/hostname
設置合適的資源配比
不同模型的計算量和參數各不相同，因此針對每個模型都應該設置一個合適的 PS/Worker 的規格和數量。在監控完善的條件下，可以根據監控數據分析瓶頸，調整實例的規格和數量，使得整體的訓練速度和平臺吞吐量能達到較好的水平。
Kube-batch調度
由於kube-batch默認開啓“reclaim, allocate, backfill, preempt”四種actions，導致每次調度時輪詢週期較長，通過配置actions爲allocate一種可以提高30%的調度效率。
Tensorflow 層面調優
Tensorflow 的調優主要參考了官網文檔。
線程數
由於 sysconf 系統調用等隔離性的問題，容器觀察到的 CPU 信息往往是宿主機的。Tensorflow 默認的線程數爲 CPU 核數，如此情況下，Tensorflow 創建的線程數遠遠超過實際分配到的 CPU 核數。同樣以 wide and deep 模型爲例，通過保持與 cpu limit
一致的線程數，上下文切換降低約 40%，訓練速度提升約 5%。
config = tf.ConfigProto()
config.intra_op_parallelism_threads = cpu_limit_cores
config.inter_op_parallelism_threads = cpu_limit_cores
tf.Session(config=config)
業務層面調優
Partition
在某次訓練中發現 PS 流量節點的分佈不均勻，其中某個 PS 節點的流量非常大，而其它 PS 節點的流量相對較低。通過分析 timeline.json 發現某個 embedding 向量非常大，所以通過 partition，將該 tensor 分散到不同的 PS 節點上，從而避免了某個 PS 節點成爲瓶頸。
圖12. 通過partition將tensor打散到不同的PS節點
partitioner = tf.fixed_size_partitioner(ps_number, axis=0)
with tf.variable_scope("emb_layer", partitioner= partitioner) as scope:
...
Adam 優化器
由於 Adam 優化器會更新所有參數的梯度，所以在大 embedding 向量下，如果採用 adam 優化器，會大大增加計算量，嚴重影響訓練速度。因此建議採用 Lazy_Adam_Optimizer 或者 Adadelta 優化器。
總結和體會
目前上述基於Kubernetes的機器學習平臺已經在生產環境爲多個業務方提供了服務。建設這套平臺的過程，也是我們探索的過程。以下我們總結了一些還不盡如人意的地方，以及我們對未來的展望。
tf-operator
從落地的情況來看，tf-operator 的功能能滿足基本的要求，穩定性較高。但是在故障恢復稍有欠缺，對於 pod 級別的故障，依賴 kubelet 來恢復；對於 node 級別的故障，目前還不支持故障恢復。分佈訓練下，故障概率隨着 worker 數量和訓練時間的增加而增加。worker 作爲無狀態節點，故障恢復既是可行的，也是非常有必要的。
kube-batch
目前 kube-batch 功能薄弱，成熟度有待商榷。比如 kube-batch 只有 predict 功能，沒有 priority 功能。並且 predict 功能也非常薄弱，僅支持部分基礎的 filter，比如 PodMatchNodeSelector, PodFitsHostPorts 以及基本的資源調度等。特別是 PodAffinity 特性，對於 PS-Worker 架構非常有用，因爲 PS 節點的網絡流量非常大，所以需要 PS 節點之間反親和，將各個 PS 節點分散。此外，kube-batch 也不支持多任務之間依賴關係。
Kube-batch 的落地效果差強人意，社區的維護力度較低。除了功能薄弱以外，我們也碰上了諸多的問題，處於勉強可用狀態。我們建議可將 pod 羣維度的資源判斷功能放到上層，只有當空閒資源滿足創建整個分佈式訓練集羣時，再將請求發送給 Kubernetes，由 Kubernetes 默認的調度器完成調度，當然，這種方式也存在一些缺點。
GPU虛擬化
目前GPU相關的監控、隔離性以及更細粒度的調度，仍然是業界的痛點問題。Nvidia提供的GPU虛擬化方案需要收取高額的Lincense費用，目前業界還沒有免費的GPU虛擬化方案。
在實際的業務場景中，一些業務的GPU使用率並不高，但以虛擬機或者容器方式運行時往往會獨佔一塊GPU卡，雖然可以通過設置CUDA_VISIBLE_DEVICES來實現多個容器共享一塊GPU卡，但任務之間的隔離性是無法保證的。
另外，Kubernetes進行GPU資源分配時，默認還不感知GPU的Topology，而不同分配的策略，對訓練的性能也會產生很大的影響。YARN和Kubernetes開源社區都在努力增強這塊的調度能力。
Kubeflow展望
Kubeflow是目前基於Kubernetes的主流機器學習解決方案，它抽象了和機器學習相關的PS-Worker模型，實現了一套pipeline的工作流，支持超參數訓練和Jupyter notebooks集成等能力。
由於 Kubeflow 解決了基於 Yarn 的機器學習平臺的痛點，同時容器化越來越普及。基於 Kubernetes 這樣大規模的企業級容器平臺，將會是未來的方向，相信 Kubeflow 在 2019 年將會有更好的發展。