Kubeflow 安利：在 Kubernetes 上進行機器學習

02-17

這篇文章主要介紹了 Kubeflow 的使用，以及未來的計劃，面向人群為對在 Kubernetes 上運行機器學習負載感興趣的同學。

問題背景

Kubernetes 本來是一個用來管理無狀態應用的容器平台，但是在近兩年，有越來越多的公司用它來運行各種各樣的工作負載，尤其是機器學習煉丹。各種 AI 公司或者互聯網公司的 AI 部門都會嘗試在 Kubernetes 上運行 TensorFlow，Caffe，MXNet 等等分散式學習的任務，這為 Kubernetes 帶來了新的挑戰。

首先，分散式的機器學習任務一般會涉及參數伺服器（以下稱為 PS）和工作節點（以下成為 worker）兩種不同的工作類型。而且不同領域的學習任務對 PS 和 worker 有不同的需求，這體現在 Kubernetes 中就是配置難的問題。以 TensorFlow 為例，TensorFlow 的分散式學習任務通常會啟動多個 PS 和多個 worker，而且在 TensorFlow 提供的最佳實踐中，每個 worker 和 PS 要求傳入不同的命令行參數。舉例說明：

# On ps0.example.com:$ python trainer.py --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 --job_name=ps --task_index=0# On ps1.example.com:$ python trainer.py --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 --job_name=ps --task_index=1# On worker0.example.com:$ python trainer.py --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 --job_name=worker --task_index=0# On worker1.example.com:$ python trainer.py --ps_hosts=ps0.example.com:2222,ps1.example.com:2222 --worker_hosts=worker0.example.com:2222,worker1.example.com:2222 --job_name=worker --task_index=1

其中需要的參數有四個，一個是所有的 PS 的網路地址（主機名-埠），以及所有的 worker 的網路地址。另外是 job 的類型，分為 PS 與 worker 兩種。最後是任務的 index，從 0 開始遞增。因此在此例中，用戶需要寫至少四個 pod 的配置文件，以及四個 service 的配置文件，使得 PS 跟 worker 可以互相訪問，況且這只是一個機器學習任務。如果大規模地在 Kubernetes 上運行 TensorFlow 分散式任務，可以預見繁雜的配置將成為機器學習工程師們新的負擔。

其次，Kubernetes 默認的調度器對於機器學習任務的調度並不友好。如果說之前的問題只是在應用與部署階段比較麻煩，那調度引發的資源利用率低，或者機器學習任務效率下降的問題，就格外值得關注。機器學習任務對於計算和網路的要求相對較高，一般而言所有的 worker 都會使用 GPU 進行訓練，而且為了能夠得到一個較好的網路支持，儘可能地同一個機器學習任務的 PS 和 worker 放在同一台機器或者網路較好的相鄰機器上會降低訓練所需的時間。

Hello, Kubeflow

針對這些問題，Kubeflow 項目應運而生，它以 TensorFlow 作為第一個支持的框架，在 Kubernetes 上定義了一個新的資源類型：TFJob，即 TensorFlow Job 的縮寫。通過這樣一個資源類型，使用 TensorFlow 進行機器學習訓練的工程師們不再需要編寫繁雜的配置，只需要按照他們對業務的理解，確定 PS 與 worker 的個數以及數據與日誌的輸入輸出，就可以進行一次訓練任務。在本節中，我們將從零開始搭建一個 Kubernetes 集群，並且將 Kubeflow 運行在其上，最後利用其進行一次完整的學習任務運行。

首先，我們需要有一個正在運行的 Kubernetes 集群，而且集群的版本要大於等於 1.8。在這一步里，個人推薦以下兩種方式創建一個單節點的本地 Kubernetes 集群：

使用 Kubernetes 里的 local-up-cluster.sh 腳本
使用 minikube 項目

其中前者會在本地創建一個 native 的 Kubernetes 集群，而後者則會在本地的虛擬機里創建出 Kubernetes 集群。因為本文側重點不在此，因此整個過程不再贅述。

如果你已經成功地創建了一個 Kubernetes 集群，那麼接下來就是在這一集群上創建 Kubeflow 所有的組件，這一步需要用到 ksonnet，一個簡化應用在 Kubernetes 上的分發與部署的命令行工具，它會幫助你創建 Kubeflow 所需組件。在安裝了 ksonnet 後，接下來就是一片坦途了，只需要運行下面的命令，就可以完成 Kubeflow 的部署。

# Initialize a ksonnet APPAPP_NAME=my-kubeflowks init ${APP_NAME}cd ${APP_NAME}# Install Kubeflow componentsks registry add kubeflow github.com/kubeflow/kubeflow/tree/master/kubeflowks pkg install kubeflow/coreks pkg install kubeflow/tf-servingks pkg install kubeflow/tf-job# Deploy KubeflowNAMESPACE=defaultkubectl create namespace ${NAMESPACE}ks generate core kubeflow-core --name=kubeflow-core --namespace=${NAMESPACE}ks apply default -c kubeflow-core

Kubeflow 的部署會附帶一個 JupyterHub 但筆者並不知道如何使用它，因此下面的操作是用 Docker 打包訓練數據和代碼，用 kubectl 在 Kubernetes 上啟動一次訓練任務的。

示例代碼可見 tf_smoke.py，與正常的訓練代碼類似，只不過 clusterspec 的傳遞方式是遵循了 Cloud ML 的 TF_CONFIG 的方式。Kubeflow 已經根據這一訓練文件打好了一個 Docker 鏡像：gcr.io/tf-on-k8s-dogfood/tf_sample:dc944ff，在這裡直接使用就好：

kubectl create -f https://raw.githubusercontent.com/tensorflow/k8s/master/examples/tf_job.yaml

Kubeflow 實現介紹

本部分主要涉及對 Kubeflow 內部實現的介紹和未來可能的開發計劃，如果不感興趣可以就此打住

對分散式訓練任務的支持

為了解決配置困難的問題，Kubeflow 以 TensorFlow 作為第一個支持的框架，為其實現了一個在 Kubernetes 上的 operator：tensorflow/k8s。由於在 Kubernetes 上內置的資源類型，如 deployment，replicaset，或者是 pod 等，都很難能夠簡練而清晰地描述一個分散式機器學習的任務，因此我們利用 Kubernetes 的 Custom Resource Definition 特性，定義了一個新的資源類型：TFJob，即 TensorFlow Job 的縮寫。一個 TFJob 配置示例如下所示：

apiVersion: "kubeflow.org/v1alpha1"kind: "TFJob"metadata: name: "example-job"spec: replicaSpecs: - replicas: 1 tfReplicaType: MASTER template: spec: containers: - image: gcr.io/tf-on-k8s-dogfood/tf_sample:dc944ff name: tensorflow restartPolicy: OnFailure - replicas: 1 tfReplicaType: WORKER template: spec: containers: - image: gcr.io/tf-on-k8s-dogfood/tf_sample:dc944ff name: tensorflow restartPolicy: OnFailure - replicas: 2 tfReplicaType: PS template: spec: containers: - image: gcr.io/tf-on-k8s-dogfood/tf_sample:dc944ff name: tensorflow restartPolicy: OnFailure

其中每個欄位就不多介紹了，這裡主要是說一下實現。任何一個 PS 或者 worker，都由兩個資源組成，分別是 job 和 service。其中 job 負責創建出 PS 或者 worker 的 pod，而 service 負責將其暴露出來。這裡社區目前也在重新考慮選型，目前希望可以直接創建 pod 而非 job，而用 headless service 替代 service，因為 PS worker 不需要暴露給除了該分散式學習任務外的其他服務。

TFJob operator 的實現早期是從 etcd-operator 複製來的，因此整體的架構在最初是完全仿照其改寫而成。在最初的實現中，當有一個 TFJob 被創建時，在 operator 內都會有一個新的 goroutine，以輪詢的方式獲取 TFJob 的狀態，然後基於此狀態做出相應的操作，相當於是在 operator 內部維護了一個狀態機。這樣的方式會有一些缺點：

這樣的架構使得 operator 是有狀態的，使得狀態很難橫向擴展
維護基於 Phase 的狀態機是 Kubernetes 社區不推崇的一種方式

基於這些問題，operator 的架構正在往事件驅動重構，這部分工作由 @caicloud 在推進。重構之後，operator 會在 Kubernetes 的一些資源上註冊 informer 的事件回調，比較現在的狀態與理想狀態的不同而採取相應的操作。比如當有一個新的 TFJob 被創建時，理想狀態是所有對應的 PS， worker 都被創建好，而當下的狀態則是沒有任何 pod 和 service 被創建，此時 operator 會創建出對應的 PS，worker 的 pod 和 service，以達到理想狀態，這也是 Kubernetes 社區對於 operator/controller 的最佳實踐。

對分散式學習任務效率的關注

目前社區還停留在如何對 AI 工程師更友好，更好地維護上面提到的 operator 這一步，在效率方面考慮地較少。目前有利用 kube-arbitrator 來進行 gang scheduling 的探索，目前還沒有嘗試過因此不好評價。但是整體來說 Kubeflow 的性能提高還有很大的空間。

因為機器學習任務根據模型的不同，其輸入數據的規模，特徵，模型的大小等等都有很大不同。比如 CV 領域與推薦領域的學習模型就有完全不同的特點，因此 TensorFlow 的分散式模型提供了極強的靈活性。而對於 Kubernetes 而言，如何能夠在保持靈活性的基礎上，同時也保證任務在較高的性能下運行，同時集群的利用率也相對較高，是一個值得研究的問題。

對其他機器學習框架的支持

目前 Kubeflow 主要關注 TensorFlow，而其他機器學習框架的支持將於之後展開，目前有一些第三方實現的 operator，比如 MXNet operator，但是質量難以保證。

開發情況與未來展望

目前 Kubeflow 有來自 Google，Caicloud，RedHat 等公司的積极參与，短期的目標有這麼幾個：

operator 方面

使用 pod 替換 job tensorflow/k8s#325
使用 headless service 替換 service tensorflow/k8s#40
由 etcd operator 主動輪詢的方式改為事件驅動 tensorflow/k8s#314
分離對 TensorBoard 的支持 tensorflow/k8s#347
支持細粒度的任務狀態 tensorflow/k8s#333

模型服務方面

GPU 支持 kubeflow/kubeflow#64
監控支持 kubeflow/kubeflow#64
多框架支持下的統一 API 支持 kubeflow/kubeflow#102

UI 方面

為各個部件支持統一的 UI kubeflow/kubeflow#199

目前 Kubeflow 在 GitHub 上有 2400 多個 star，有 40 個左右的貢獻者。其長期的目標是成為 CNCF 的一個項目，目前實現仍存在很多問題，竊以為也並不是 production ready 的狀態，但它仍然值得一試。

關於作者

@gaocegege，上海交通大學軟體學院研究生在讀，Kubeflow Committer，尋找 2018 暑期實習。

License

This article is licensed under CC BY-NC-SA 3.0.
Please contact me for commercial use.