如何在K8S上玩轉TensorFlow？

女主宣言

該文章出自於ADDOPS團隊，是關於如何在K8S上玩轉tensorflow的主題，該文章深入淺出的給我們介紹了當前tensorflow的現狀和架構特點等，然後介紹了讓tensorflow如何基於k8s快速落地，讓大家都能簡單的上手tensorflow，整體文章脈絡清晰，內容適度，所以希望能給大家帶來啟發。

PS：豐富的一線技術、多元化的表現形式，盡在「HULK一線技術雜談」，點關注哦！

前言

Tensorflow作為深度學習領域逐漸成熟的項目，以其支持多種開發語言，支持多種異構平台，提供強大的演算法模型，被越來越多的開發者使用。但在使用的過程中，尤其是GPU集群的時候，我們或多或少將面臨以下問題：

資源隔離。Tensorflow（以下簡稱tf）中並沒有租戶的概念，何如在集群中建立租戶的概念，做到資源的有效隔離成為比較重要的問題；

缺乏GPU調度。tf通過指定GPU的編號來實現GPU的調度，這樣容易造成集群的GPU負載不均衡；

進程遺留問題。tf的分散式模式ps伺服器會出現tf進程遺留問題；

訓練的數據分發以及訓練模型保存，都需要人工介入；

訓練日誌保存、查看不方便；

因此，我們需要一個集群調度和管理系統，可以解決GPU調度、資源隔離、統一的作業管理和跟蹤等問題。

目前，社區中有多種開源項目可以解決類似的問題，比如yarn，kubernetes。yarn是hadoop生態中的資源管理系統，而kubernetes（以下簡稱k8s）作為Google開源的容器集群管理系統，在tf1.6版本加入GPU管理後，已經成為很好的tf任務的統一調度和管理系統。

下文是我們公司在tensorflow on kubernetes方面的實踐經驗。

設計目標

我們將tensorflow引入k8s，可以利用其本身的機制解決資源隔離，GPU調度以及進程遺留的問題。除此之外，我們還需要面臨下面問題的挑戰：

支持單機和分散式的tensorflow任務；

分散式的tf程序不再需要手動配置clusterspec信息，只需要指定worker和ps的數目，能自動生成clusterspec信息；

訓練數據、訓練模型以及日誌不會因為容器銷毀而丟失，可以統一保存；

為了解決上面的問題，我們開發了tensorflow on kubernetes系統。

架構

tensorflow on kubernetes包含三個主要的部分，分別是client、task和autospec模塊。client模塊負責接收用戶創建任務的請求，並將任務發送給task模塊。task模塊根據任務的類型（單機模式和分散式模式）來確定接下來的流程：

如果type選擇的是single（單機模式），對應的是tf中的單機任務，則按照按照用戶提交的配額來啟動container並完成最終的任務；

如果type選擇的是distribute（分散式模式），對應的是tf的分散式任務，則按照分散式模式來執行任務。需要注意的是，在分散式模式中會涉及到生成clusterspec信息，autospec模塊負責自動生成clusterspec信息，減少人工干預。

下面是tensorflow on kubernetes的架構圖：

接下來將對三個模塊進行重點介紹。

client模塊

tshell

在容器中執行任務的時候，我們可以通過三種方式獲取執行任務的代碼和訓練需要的數據：

將代碼和數據做成新的鏡像；

將代碼和數據通過卷的形式掛載到容器中；

從存儲系統中獲取代碼和數據；

前兩種方式不太適合用戶經常修改代碼的場景，最後一種場景可以解決修改代碼的問題，但是它也有下拉代碼和數據需要時間的缺點。綜合考慮後，我們採取第三種方式。

我們做了一個tshell客戶端，方便用戶將代碼和程序進行打包和上傳。比如給自己的任務起名字叫cifar10-multigpu，將代碼打包放到code下面，將訓練數據放到data下面。

提交任務

在提交任務的時候，需要指定提前預估一下執行任務需要的配額：cpu核數、內存大小以及gpu個數（默認不提供），當然也可以按照我們提供的初始配額來調度任務。

比如，按照下面格式來將配額信息、s3地址信息以及執行模式填好後，執行send_task.py我們就可以提交一次任務。

task模塊

單機模式

對於單機模式，task模塊的任務比較簡單，直接調用python的client介面啟動container。container裡面主要做了兩件事情，initcontainer負責從s3中下載事先上傳好的文件，container負責啟動tf任務，最後將日誌和模型文件上傳到s3里，完成一次tf單機任務。

分散式模式

對於分散式模式，情況要稍微複雜些。下面先簡單介紹一下tensforlow分散式框架。tensorflow的分散式並行基於gRPC框架，client負責建立Session，將計算圖的任務下發到TF cluster上。

TF cluster通過tf.train.ClusterSpec函數建立一個cluster，每個cluster包含若干個job。 job由好多個task組成，task分為兩種，一種是PS（Parameter server），即參數伺服器，用來保存共享的參數，還有一種是worker，負責計算任務。

我們在執行分散式任務的時候，需要指定clusterspec信息，如下面的任務，執行該任務需要一個ps和兩個worker，我們先需要手動配置ps和worker，才能開始任務。這樣必然會帶來麻煩。如何解決clusterspec，成為了一個必須要解決的問題。

所以在提交分散式任務的時候，task需要autospec模塊的幫助，收集container的ip後，才能真正啟動任務。所以分散式模式要做兩件事情：

按照yaml文件啟動container；

通知am模塊收集此次任務container的信息後生成clusterspec；

Autospec模塊

tf分散式模式的node按照角色分為ps（負責收集參數信息）和worker，ps負責收集參數信息，worker執行任務，並定期將參數發送給worker。

要執行分散式任務，涉及到生成clusterspec信息，模型的情況下，clusterspec信息是通過手動配置，這種方式比較麻煩，而且不能實現自動化，我們引入autospec模型很好的解決此類問題。

Autospec模塊只有一個用途，就是在執行分散式任務時，從container中收集ip和port信息後，生成clusterspec，發送給相應的container。下面是autospec模塊的工作流程圖：

Container設計

tf任務比較符合k8s中kind為job的任務，每次執行完成以後這個容器會被銷毀。我們利用了此特徵，將container都設置為job類型。

k8s中設計了一種hook：poststart負責在容器啟動之前做一些初始化的工作，而prestop負責在容器銷毀之前做一些備份之類的工作。

我們利用了此特點，在poststart做一些數據獲取的工作，而在prestop階段負責將訓練產生的模型和日誌進行保存。我們接下來分單機和分散式兩種模式來說明Container的設計思想。

總結

至此，我們已經介紹了tensorflow on kubernetes的主要流程。還有許多需要完善的地方，比如：web端提交任務以及查看運行狀況和作業的日誌；支持GPU的親和性等等，總之，這只是我們前期的探索，後面還有許多東西需要完善。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！