TensorFlow on Kubernetes的架構與實踐

這兩年，Kubernetes在企業中的DevOps、微服務領域取得了出色的成績，從2017年開始，將Kubernetes應用到HPC、AI等領域也成了技術熱點。這裡我給大家分享一下Kubernetes在AI中的落地經驗，內容包括TensorFlow on Kubernetes的架構與實踐，以及線上經驗和坑。

都是玩容器的老司機，都知道Kubernetes這兩年非常火，截止目前在GitHub上31K+ stars，然而相比於TensorFlow，也就只能說是一般般了。TensorFlow才兩年多，在GitHub上已經有86K+ stars，這是個什麼概念呢？要知道，Linux Kernel這麼多年才積累54K+ stars，當然，它們各自都是所在領域的霸主，這種對比只當閑談。

這兩年，Kubernetes在各個企業中的DevOps、微服務方向取得了出色的成績，從2017年開始，越來越多的企業也開始探索將Kubernetes應用到HPC、AI等領域。隨著公司AI業務的迅猛增長，vivo在2017年9月也開始基於Kubernetes強大的分散式能力，探索與TensorFlow等ML框架深度整合，提高資料中心資源利用率，加快演演算法迭代速度。

Kubernetes在AI中的應用與在DevOps中部署App相比，最大的差別在於容器的規模以及容器生命週期。在我們的實踐中，目前叢集伺服器規模很小的情況下，每天要排程近10W的容器，有很多容器可能只運行了十幾分鐘甚至幾分鐘，而且計劃在2018年，叢集規模還要翻十倍。在DevOps場景，應用釋出頻率再高，我相信一年下來能排程10W容器的企業並不多。下麵我將聊一下TensorFlow on Kubernetes的架構及在vivo的實踐。

分散式TensorFlow

TensorFlow是一個使用資料流圖進行數值計算的開源軟體庫。圖中的節點代表數學運算，而圖中的邊則代表在這些節點之間傳遞的多維陣列（張量）。這種靈活的架構可讓您使用一個API將計算工作部署到桌面裝置、伺服器或者移動裝置中的一個或多個CPU或GPU。 關於TensorFlow的基礎概念，我就不多介紹了。

單機TensorFlow

下麵是一個單機式TensorFlow訓練示意圖，透過Client提交Session，定義這個worker要用哪個CPU/GPU做什麼事。

分散式TensorFlow

2016年4月TensorFlow釋出了0.8版本宣佈支援分散式計算，我們稱之為Distributed TensorFlow。這是非常重要的一個特性，因為在AI的世界裡，訓練的資料量和模型引數通常會非常大。比如Google Brain實驗室今年發表的論文《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》中提到一個680億個Parameters的模型，如果只能單機訓練，那耗時難於接受。透過Distributed TensorFlow，可以利用大量伺服器構建分散式TensorFlow叢集來提高訓練效率，減少訓練時間。

透過TensorFlow Replcation機制，使用者可以將SubGraph分佈到不同的伺服器中進行分散式計算。TensorFlow的副本機制又分為兩種，In-graph和Between-graph。

In-graph Replication簡單來講，就是透過單個client session定義這個TensorFlow叢集的所有task的工作。

與之相對地，Between-graph Replication就是每個worker都有獨立的client來定義自己的工作。

下麵是抽象出來的分散式TensorFlow Framework如下：

我們先來瞭解裡面的幾個概念：

Cluster

一個TensorFlow Cluster有一個或多個jobs組成，每個job又由一個或多個tasks構成。Cluster的定義是透過tf.train.ClusterSpec來定義的。比如，定義一個由3個worker和2個PS的TensorFlow Cluster的ClusterSpec如下：

tf.train.ClusterSpec({
    "worker": [
        "worker0.example.com:2222",  //主機名也可以使用IP
        "worker1.example.com:2222",
        "worker2.example.com:2222"
    ],
    "ps": [
        "ps0.example.com:2222",
        "ps1.example.com:2222"
    ]})

Client

Client用來build一個TensorFlow Graph，並構建一個tensorflow::Session用來與叢集通訊。一個Client可以與多個TensorFlow Server互動，一個Server能服務多個Client。

Job

一個Job由tasks list組成，Job分PS和Worker兩種型別。PS即parameter server，用來儲存和更新variables的，而Worker可以認為是無狀態的，用來作為計算任務的。Workers中，一般都會選擇一個chief worker（通常是worker0），用來做訓練狀態的checkpoint，如果有worker故障，那麼可以從最新checkpoint中restore。

Task

每個Task對應一個TensorFlow Server，對應一個單獨的行程。一個Task屬於某個Job，透過一個index來標記它在對應Job的tasks中的位置。每個TensorFlow均實現了Master service和Worker service。Master service用來與叢集內的worker services進行gRPC互動。Worker service則是用local device來計算Subgraph。

關於Distributed TensorFlow的更多內容，請參考官方內容：www.tensorflow.org/deplopy/distributed。

分散式TensorFlow的缺陷

分散式TensorFlow能利用資料中心所有伺服器構成的資源池，讓大量PS和Worker能分佈在不同的伺服器進行引數儲存和訓練，這無疑是TensorFlow能否在企業落地的關鍵點。然而，這還不夠，它還存在一下先天不足：

• 訓練時TensorFlow各個Task資源無法隔離，很有可能會導致任務間因資源搶佔互相影響。

• 缺乏排程能力，需要使用者手動配置和管理任務的計算資源。

• 叢集規模大時，訓練任務的管理很麻煩，要跟蹤和管理每個任務的狀態，需要在上層做大量開發。

• 使用者要檢視各個Task的訓練日誌需要找出對應的伺服器，並ssh過去，非常不方便。

• TensorFlow原生支援的後端檔案系統只支援：標準Posix檔案系統（比如NFS）、HDFS、GCS、memory-mapped-file。大多數企業中資料都是存在大資料平臺，因此以HDFS為主。然而，HDFS的Read效能並不是很好。

• 當你試著去建立一個大規模TensorFlow叢集時，發現並不輕鬆。

TensorFlow on Kubernetes架構與原理

TensorFlow的這些不足，正好是Kubernetes的強項：

• 提供ResourceQuota、LimitRanger等多種資源管理機制，能做到任務之間很好的資源隔離。

• 支援任務的計算資源的配置和排程。

• 訓練任務以容器方式執行，Kubernetes提供全套的容器PLEG介面，因此任務狀態的管理很方便。

• 輕鬆對接EFK/ELK等日誌方案，使用者能方便的檢視任務日誌。

• 支援Read效能更優秀的分散式儲存（GlusterFS），但目前我們也還沒對接GlusterFS，有計劃但沒人力。

• 透過宣告式檔案實現輕鬆快捷的建立一個大規模TensorFlow叢集。

TensorFlow on Kubernetes架構

TensorFlow on Kubernetes原理

在我們的TensorFlow on Kubernetes方案中，主要用到以下的Kubernetes物件：

Kubernetes Job

我們用Kubernetes Job來部署TensorFlow Worker，Worker訓練正常完成退出，就不會再重啟容器了。註意Job中的Pod Template restartPolicy只能為Never或者OnFailure，不能為Always，這裡我們設定restartPolicy為OnFailure，worker一旦異常退出，都會自動重啟。但是要註意，要保證worker重啟後訓練能從checkpoint restore，不然worker重啟後又從step 0開始，可能跑了幾天的訓練就白費了。如果你使用TensorFlow高階API寫的演演算法，預設都實現了這點，但是如果你是使用底層Core API，一定要註意自己實現。

kind: Job
apiVersion: batch/v1
metadata:
  name: {{ name }}-{{ task_type }}-{{ i }}
  namespace: {{ name }}
spec:
  template:
    metadata:
      labels:
        name: {{ name }}
        job: {{ task_type }}
        task: "{{ i }}"
    spec:
      imagePullSecrets:
      - name: harborsecret
      containers:
      - name: {{ name }}-{{ task_type }}-{{ i }}
        image: {{ image }}
        resources:
          requests:
            memory: "4Gi"
            cpu: "500m"
        ports:
        - containerPort: 2222
        command: ["/bin/sh", "-c", "export CLASSPATH=.:/usr/lib/jvm/java-1.8.0/lib/tools.jar:$(/usr/lib/hadoop-2.6.1/bin/hadoop classpath --glob); wget -r -nH  -np --cut-dir=1 -R 'index.html*,*gif'  {{ script }}; cd ./{{ name }}; sh ./run.sh {{ ps_hosts() }} {{ worker_hosts() }} {{ task_type }} {{ i }} {{ ps_replicas }} {{ worker_replicas }}"]
      restartPolicy: OnFailure

Kubernetes Deployment

TensorFlow PS用Kubernetes Deployment來部署。為什麼不像worker一樣，也使用Job來部署呢？其實也未嘗不可，但是考慮到PS行程並不會等所有worker訓練完成時自動退出（一直掛起），所以使用Job部署沒什麼意義。

kind: Deployment
apiVersion: extensions/v1beta1
metadata:
  name: {{ name }}-{{ task_type }}-{{ i }}
  namespace: {{ name }}
spec:
  replicas: 1
  template:
    metadata:
      labels:
        name: {{ name }}
        job: {{ task_type }}
        task: "{{ i }}"
    spec:
      imagePullSecrets:
      - name: harborsecret
      containers:
      - name: {{ name }}-{{ task_type }}-{{ i }}
        image: {{ image }}
        resources:
          requests:
            memory: "4Gi"
            cpu: "500m"
        ports:
        - containerPort: 2222
        command: ["/bin/sh", "-c","export CLASSPATH=.:/usr/lib/jvm/java-1.8.0/lib/tools.jar:$(/usr/lib/hadoop-2.6.1/bin/hadoop classpath --glob); wget -r -nH  -np --cut-dir=1 -R 'index.html*,*gif'  {{ script }}; cd ./{{ name }}; sh ./run.sh {{ ps_hosts() }} {{ worker_hosts() }} {{ task_type }} {{ i }} {{ ps_replicas }} {{ worker_replicas }}"]
      restartPolicy: Always

關於TensorFlow PS行程掛起的問題，請參考：https://github.com/tensorflow/tensorflow/issues/4713。我們是這麼解決的，開發了一個模組，watch每個TensorFlow叢集的所有worker狀態，當所有worker對應Job都Completed時，就會自動去刪除PS對應的Deployment，從而kill PS行程釋放資源。

Kubernetes Headless Service

Headless Service通常用來解決Kubernetes裡面部署的應用叢集之間的內部通訊。在這裡，我們也是這麼用的，我們會為每個TensorFlow對應的Job和Deployment物件都建立一個Headless Service作為Worker和PS的通訊代理。

kind: Service
apiVersion: v1
metadata:
  name: {{ name }}-{{ task_type }}-{{ i }}
  namespace: {{ name }}
spec:
  clusterIP: None
  selector:
    name: {{ name }}
    job: {{ task_type }}
    task: "{{ i }}"
  ports:
  - port: {{ port }}
    targetPort: 2222

用Headless Service的好處，就是在KubeDNS中，Service Name的域名解析直接對應到PodIp，而沒有service VIP這一層，這就不依賴於kube-proxy去建立iptables規則了。少了kube-proxy的iptables這一層，帶來的效能的提升。

在TensorFlow場景中，這是不可小覷的，因為一個TensorFlow Task都會建立一個service，幾萬個service是很正常的事，如果使用Normal Service，iptables規則就幾十萬上百萬條了，增刪一條iptabels規則耗時幾個小時甚至幾天，叢集早已奔潰。關於kube-proxy iptables樣式的效能測試資料，請參考華為PaaS團隊的相關分享。

KubeDNS Autoscaler

前面提到，每個TensorFlow Task都會建立一個service，都會在KubeDNS中有一條對應的解析規則，但service數量太多的時候，我們發現有些Worker的域名解析失敗機率很大，十幾次才能成功解析一次。這樣會影響TensorFlow叢集內各個task的session建立，可能導致TensorFlow叢集起不來。

為瞭解決這個問題，我們引入了Kubernetes的孵化專案kubernetes-incubator/cluster-proportional-autoscaler來對KubeDNS進行動態伸縮。

TensorFlow on Kubernetes實踐

基於上面的方案，我們開發一個TaaS平臺，已經實現了基本的功能，包括演演算法管理、訓練叢集的建立和管理、模型的管理、模型上線（TensorFlow Serving）、一鍵建立TensorBoard服務、任務資源監控、叢集資源監控、定時訓練管理、任務日誌線上檢視和批次打包下載等等，這部分內容可以參考之前在DockOne上分享的文章《vivo基於Kubernetes構建企業級TaaS平臺實踐》。

這隻是剛開始，我正在做下麵的特性：

• 支援基於訓練優先順序的任務搶佔式排程：使用者在TaaS上建立TensorFlow訓練專案時，可以指定專案的優先順序為生產（Production）、迭代（Iteration）、調研（PTR），預設為迭代。優先順序從高到低依次為Production –> Iteration –> PTR。但叢集資源不足時，按照任務優先順序進行搶佔式排程。

• 提供像Yarn形式的資源分配檢視，讓使用者對自己的所有訓練專案的資源佔用情況變得清晰。

• 訓練和預測的混合部署，提供資料中心資源利用率。

• ……

經驗和坑

整個過程中，遇到了很多坑，有TensorFlow的，也有Kubernetes的，不過問題最多的還是我們用的CNI網路外掛Contiv Netplugin，每次大問題基本都是這個網路外掛造成的。Kubernetes是問題最少的，它的穩定性比我預期還要好。

• Contiv Netplugin的問題，在DevOps環境中還是穩定的，在大規模高併發的AI場景，問題就層出不窮了，產生大量垃圾IP和Openflow流表，直接把Node都成NotReady了，具體的不多說，因為據我瞭解，現在用這個外掛的公司已經很少了，想瞭解的私下找我。

• 在我們的方案中，一個TensorFlow訓練叢集就對應一個Kubernetes Namespace，專案初期我們並沒有對及時清理垃圾Namespace，到後來集群裡上萬Namespace的時候，整個Kubernetes叢集的相關API效能非常差了，導致TaaS的使用者體驗非常差。

• TensorFlow gRPC效能差，上千個Worker的訓練叢集，機率性的出現這樣的報錯：grpc_chttp2_stream request on server; last grpc_chttp2_stream id=231, new grpc_chttp2_stream id=227，只能嘗試透過升級TensorFlow來升級gRPC了。目前我們透過增加單個Worker的計算負載來減少Worker數量的方法，減少gRPC壓力。除此之外，規模大一點TensorFlow還有各種問題，甚至還有自己的OOM機制，防不勝防。不過，高速成長的東西，我們要給它足夠的容忍。

• 還有TensorFlow OOM的問題等等。

Q&A;

Q：Worker為什麼不直接用Pod，而用的 Job？

A：Kubernetes中是不建議直接使用Pod的，建議透過一個控制器（RS、RC、Deploy、StatefulSets、Job等）來建立和管理Pod。

Q：我看訓練的時候並沒有指定資料集和訓練引數等，是都放到訓練指令碼內了嗎，還有訓練集是放到Gluster，掛載到容器內嗎，還是換存到本地？

A：目前訓練資料和引數都是在腳本里使用者自己搞定，正在把這些放到Portal，提供“命令列引數自定義”的能力。目前我們的訓練資料是從HDFS叢集直接走網路讀取，Kubernetes本身也沒有HDFS的Volume Plugin。

Q：請問PS的個數是使用者指定還是根據Worker的數量來指派？

A：PS和Worker數都是透過使用者指定，但實際上都是使用者根據自己的訓練資料大小來計算的。

Q：分散式訓練的時候，Training data是如何分發到各個Worker節點的？Tensorflow API可以做到按節點數自動分發嗎？

A：各個Worker都是從HDFS叢集讀取訓練資料載入到記憶體的。資料的分配都是使用者自己在指令碼種實現的。

基於Kubernetes的DevOps實踐培訓

本次培訓包含：Kubernetes核心概念；Kubernetes叢集的安裝配置、運維管理、架構規劃；Kubernetes元件、監控、網路；針對於Kubernetes API介面的二次開發；DevOps基本理念；微服務架構；微服務的容器化等，點選識別下方二維碼加微信好友瞭解具體培訓內容。

本週四正式開課，點選閱讀原文連結即可報名。