【干貨】開源Kubeflow：在Kubernetes上運行機器學習

  這篇文章主要介紹了 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 腳本 
https://github.com/kubernetes/kubernetes/blob/master/hack/local-up-cluster.sh


	  使用 minikube 項目 https://github.com/kubernetes/minikube


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


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

# Initialize a ksonnet APP
APP_NAME=my-kubeflow
ks init ${APP_NAME} 
cd ${APP_NAME} 
# Install Kubeflow components
ks registry add kubeflow github.com/kubeflow/kubeflow/tree/master/kubeflow
ks pkg install kubeflow/core
ks pkg install kubeflow/tf-serving
ks pkg install kubeflow/tf-job

# Deploy Kubeflow
NAMESPACE=default 
kubectl 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 的狀態，但它仍然值得一試。