背 景
隨著機器學習和人工智慧的迅猛發展，業界出現了許多開源的機器學習平臺。由於機器學習與大資料天然的緊密結合，基於 Hadoop Yarn 的分散式任務排程仍是業界主流，但是隨著容器化的發展，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 也很容易被打爆。
表 2. Hadoop Yarn 和 Kubernetes 的橫向對比
Kubeflow 及核心元件
Kubeflow 是由 Google 等公司於 2017 年推出的基於 Kubernetes 的開源專案，它支援常見的機器學習框架。
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 所支援的主流分散式計算框架
蘑菇街業務方使用的計算框架主要是 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 的機器學習基礎平臺總體架構
主要包含了以下幾部分：
在部署 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: tensorflowTf-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，滿足了業務方對批量任務排程的需求。沒有采用 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 結合起來使用。如此既可滿足業務的排程需求，又可以在一定程度上提升排程效能。
我們引入了由 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=500mPod 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 預設開啟“reclaim, allocate, backfill, preempt”四種 actions，導致每次排程時輪詢週期較長，通過配置 actions 為 allocate 一種可以提高 30% 的排程效率。
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 優化器會更新所有引數的梯度，所以在大 embedding 向量下，如果採用 adam 優化器，會大大增加計算量，嚴重影響訓練速度。因此建議採用 Lazy_Adam_Optimizer 或者 Adadelta 優化器。
總結和體會
目前上述基於 Kubernetes 的機器學習平臺已經在生產環境為多個業務方提供了服務。建設這套平臺的過程，也是我們探索的過程。以下我們總結了一些還不盡如人意的地方，以及我們對未來的展望。
從落地的情況來看，tf-operator 的功能能滿足基本的要求，穩定性較高。但是在故障恢復稍有欠缺，對於 pod 級別的故障，依賴 kubelet 來恢復；對於 node 級別的故障，目前還不支援故障恢復。分佈訓練下，故障概率隨著 worker 數量和訓練時間的增加而增加。worker 作為無狀態節點，故障恢復既是可行的，也是非常有必要的。
目前 kube-batch 功能薄弱，成熟度有待商榷。比如 kube-batch 只有 predict 功能，沒有 priority 功能。並且 predict 功能也非常薄弱，僅支援部分基礎的 filter，比如 PodMatchNodeSelector, PodFitsHostPorts 以及基本的資源排程等。特別是 PodAffinity 特性，對於 PS-Worker 架構非常有用，因為 PS 節點的網路流量非常大，所以需要 PS 節點之間反親和，將各個 PS 節點分散。此外，kube-batch 也不支援多工之間依賴關係。
Kube-batch 的落地效果差強人意，社群的維護力度較低。除了功能薄弱以外，我們也碰上了諸多的問題，處於勉強可用狀態。我們建議可將 pod 群維度的資源判斷功能放到上層，只有當空閒資源滿足建立整個分散式訓練叢集時，再將請求傳送給 Kubernetes，由 Kubernetes 預設的排程器完成排程，當然，這種方式也存在一些缺點。
目前 GPU 相關的監控、隔離性以及更細粒度的排程，仍然是業界的痛點問題。Nvidia 提供的 GPU 虛擬化方案需要收取高額的 Lincense 費用，目前業界還沒有免費的 GPU 虛擬化方案。
在實際的業務場景中，一些業務的 GPU 使用率並不高，但以虛擬機器或者容器方式執行時往往會獨佔一塊 GPU 卡，雖然可以通過設定 CUDA_VISIBLE_DEVICES 來實現多個容器共享一塊 GPU 卡，但任務之間的隔離性是無法保證的。
另外，Kubernetes 進行 GPU 資源分配時，預設還不感知 GPU 的 Topology，而不同分配的策略，對訓練的效能也會產生很大的影響。YARN 和 Kubernetes 開源社群都在努力增強這塊的排程能力。
Kubeflow 是目前基於 Kubernetes 的主流機器學習解決方案，它抽象了和機器學習相關的 PS-Worker 模型，實現了一套 pipeline 的工作流，支援超引數訓練和 Jupyter notebooks 整合等能力。
由於 Kubeflow 解決了基於 Yarn 的機器學習平臺的痛點，同時容器化越來越普及。基於 Kubernetes 這樣大規模的企業級容器平臺，將會是未來的方向，相信 Kubeflow 在 2019 年將會有更好的發展。
作者介紹
範德良，花名攬勝，畢業於浙江大學，專注於 OpenStack 和 Kubernetes，在私有云和公有云等基礎設施領域有多年建設經驗和維護經驗。目前正基於 Kubeflow 建設蘑菇街分散式機器學習平臺，服務演算法，影象，風控等業務。
周佳煊，花名楚河，前 HPE 軟體部高階技術顧問，主要服務於大型公有云廠商，並有多年基礎架構運維開發經驗，對 Cloud Native 領域有深刻理解，目前任職於蘑菇街 PaaS 平臺，結合騰訊 TKE 打造服務於蘑菇街全站業務的 PaaS 平臺。
張振華，花名郭嘉，蘑菇街 PaaS 平臺負責人，畢業於浙江大學軟體學院，10 餘年軟體研發和技術管理經驗，曾在英特爾、思科等公司工作。從零開始帶領團隊研發了蘑菇街的 IaaS 和 PaaS 平臺，在 Docker、Kubernetes、DevOps、Kubeflow、Serverless 等領域擁有豐富經驗。