[原始碼解析] 深度學習分散式訓練框架 horovod (8) --- on spark
0x00 摘要
Horovod 是Uber於2017年釋出的一個易於使用的高效能的分散式訓練框架，在業界得到了廣泛應用。
本系列將通過原始碼分析來帶領大家瞭解 Horovod。接下來幾篇介紹 horovod 如何執行在 spark 之上。本文是第八篇，介紹 horovod on spark 的總體架構。
Horovod on spark 的目的就是讓 horovod 能跑到 spark 叢集上，從而把資料處理，模型訓練，模型評估這一個機器學習的迴圈都放在Spark技術棧之中。
本系列其他文章如下：
[原始碼解析] 深度學習分散式訓練框架 Horovod (1) --- 基礎知識
[原始碼解析] 深度學習分散式訓練框架 horovod (2) --- 從使用者角度切入
[原始碼解析] 深度學習分散式訓練框架 horovod (3) --- Horovodrun背後做了什麼
[原始碼解析] 深度學習分散式訓練框架 horovod (4) --- 網路基礎 & Driver
[原始碼解析] 深度學習分散式訓練框架 horovod (5) --- 融合框架
[原始碼解析] 深度學習分散式訓練框架 horovod (6) --- 後臺執行緒架構
[原始碼解析] 深度學習分散式訓練框架 horovod (7) --- DistributedOptimizer
0x01 Spark相關知識
1.1 為什麼整合 Spark
Spark是一個分散式通用計算框架，而以 tensorflow 為代表的深度學習框架是分散式模型訓練框架，這些框架更多專注用迭代來計算梯度。很多業內公司都是用spark來獲取/處理資料，然後把spark處理好的資料結果發給Tensorflow進行訓練。
目前我們已經知道，Horovod 可以把 Tensorflow等深度學習框架和MPI緊密結合起來，那麼為什麼要再把 spark 整合進來呢？整合的意義在哪裡？具體如下：
- MPI是一個較低層級的庫，專注於提供可移植的效能而忽略了程式設計師生產力（原語過於低階，開發程式碼量大）。Spark是一個更高階別的框架，更專注於程式設計師的生產力。Spark可以使開發者用單機序列程式的思維來開發分散式程式，這樣使用者可以更加專注於演算法本身，而不需將精力過多放在分散式邏輯上。
- 整合之後，可以讓整個特徵處理和訓練流程都統一在 spark 環境內，從而實現更好的分散式訓練和資料傳輸。
- MPI叢集的任務成功率並不高，如果某個任務失敗，往往需要重啟整個MPI叢集。因為 MPI的容錯性較差，所以希望能夠藉助spark的容錯機制。
Horovod 需要解決的核心問題是：如何將spark作為分散式tensorflow的底層調動機制，從而通過spark executor就可以把 tensorflow 的程式調動起來，這樣進行tensorflow訓練時就不需要手動地去組建網路。
因此能想到的其他問題是：
- Spark如何開始執行？當某一個 Executor 啟動後就可以執行？還是需要所有的 Executor 都準備好之後才能一起跑？
- 如何釋出 訓練程式碼？
- 如何在 Spark Executor 之上啟動使用者程式碼？
- MPI 在這個機制中起到什麼作用？
我們在隨後一一分析。
1.2 Spark 簡單架構
簡要來說，Spark分成幾個角色：
- Driver。這是一個程式，我們編寫好的Spark程式在spark-submit提交之後，就是由Driver程式執行。充當Driver的可能是Spark叢集的某個節點、比如就是你提交Spark程式的機器。
- Executor。也是一個程式，在一個Executor程式裡面會有多個task執行緒。這裡的Executor和task主要負責對RDD的partition進行平行計算，也就是執行我們在程式中指定的RDD運算元（map、flatMap、reduceByKey等）。
- Task。是一個執行緒，主要是負責實際的執行運算元任務。一個 task 對應一個執行緒，多個 task 可以並行的執行在 executor 之中。使用者程式碼經過Spark Driver 排程之後，被封裝成若干Task，Driver 再將這些Task資訊發給Executor執行，Task資訊包括程式碼邏輯以及資料資訊。Executor不直接執行使用者的程式碼。
1.3 Pyspark 原理
當我們用python編寫程式時，其實使用的是 Pyspark 介面。所以我們介紹一下 pyspark，可以和 Horovod 做比對。
1.3.1 架構修改
如果我們使用Java或者Scala開發Spark相關程式，Driver 和 Executor 執行任務的載體是Java虛擬機器（JVM）。但是 Python 使用的是 Python自己的虛擬機器，這就產生了一個問題，核心架構是基於JVM還是PVM。
為了保持核心架構一致性，Spark依然使用JVM作為核心，核心功能依然基於JVM，其中包括：申請計算資源，管理/分配task，driver與executor之間的通訊等等。在此核心架構外圍則封裝了一層python。
因此，PySpark 採用了 Python程式和JVM 程式分離的多程式架構，在 Driver和Executor 端都同時有 Python和JVM 兩個程式。
1.3.2 Driver端
如果使用者提交一個Python 指令碼，Spark Driver 會：
- 執行這個指令碼；
- 通過Python 啟動 JVM；
- 如果Python指令碼中呼叫了DataFrame或者RDD操作，則會通過Py4j呼叫到Java方法，將使用者的"Spark"操作對映到JVM之中。比如python呼叫
a.map(lambda x:(x,1))，則這個rdd的操作會對映到JVM之中被執行。
1.3.3 Executor端
在Executor則正好相反，因為Executor端執行的Task邏輯（序列化後的位元組碼）是由Driver發過來的，所以 Executor 本來是可以直接執行Task，並不需要藉助任何Py4j。但是因為Python指令碼中會存在使用者定義的python函式（或者Lambda表示式），所以Executor必須再啟動Python程式進行相關處理：
- 當Driver申請到Executor的資源之後，會啟動Executor 的 JVM 程式，如果沒有Task下發過來，則Executor只有JVM，沒有其他程式，即沒有下面提到的Python程式；
- Executor接到任務之後，會啟動一個task進行處理。如果不存pyspark.deamon後臺公共程式，則Executor會通過Java Process的方式啟動pyspark.deamon後臺公共程式，pyspark.deamon負責接收Task的相關請求。此deamon在每個Executor上只有一個。
- pyspark.deamon接收到請求之後，會為每一個Task單獨啟動一個Python子程式（pyspark worker）；
- RDD的載體依然在Executor之中，當有udf和lambda邏輯時，Executor會通過socket作為載體，同pyspark worker進行資料通訊，把資料不停的提供給 pyspark worker；
- 當pyspark worker執行之後會把結果通過socket返回給JVM；
1.3.4 流程
互動流程如下圖，實線是方法呼叫，虛線是返回結果。
架構圖如下：
0x02 機器學習 on Spark
2.1 機器學習的特點
機器學習演算法和計算機領域的其他演算法相比，有自己的一些獨特特點。例如：
- 迭代性。模型的更新並非一次完成，需要迴圈迭代多次；
- 容錯性。即使在每個迴圈中產生一些錯誤，模型最終的收斂也不會受到影響。這於傳統分散式系統形成鮮明對比，比如分散式檔案系統就無法接受任何資料塊的寫入錯誤。
- 引數收斂的非均勻性。模型中某些引數可能經過幾個迴圈便不再改變，而某些引數需要很長時間多次迭代才能收斂。
- 網路是瓶頸。頻繁更新模型引數需要消耗大量頻寬，而GPU速度越快，網路瓶頸就越成為問題所在。
以上這些特點決定了機器學習系統的設計和其他計算系統的設計有很大區別。和傳統分散式系統比較，機器學習系統在通訊，同步和容錯等方面都活動空間極大。因為大量資源都會浪費在通訊，等待，協調這些非計算任務上，所以導致分散式機器學習任務往往並不能隨著機器數量隨之的增加而能力也線性提升。
因此，在設計大規模機器學習系統（比如深度學習/邏輯迴歸/主題模型/矩陣分解等依賴於SGD或者L-BFGS最優化的演算法）時，需要解決一系列挑戰，比如提高並行度，減少同步等待延遲，容錯以及巨大頻寬（頻繁訪問修改模型引數時所需）等。
2.2 機器學習 on Spark
MPI 的主要缺點是：
- 原語過於低階。用MPI寫演算法，往往程式碼量比較大也比較複雜。
- 容錯機制差。如果某個任務失敗，往往需要重啟整個MPI叢集，而MPI叢集的任務成功率並不高。
- MPI本身也無法支撐大規模資料。
Spark在一定層度上解決了MPI的問題。
2.2.1 簡單模型
Spark訓練的一個最最簡陋的整體流程圖如下：
- Map 操作定義了資料分發和在工作節點的計算：
- 首先在map階段對資料進行分割，分發給每一個 Executor；
- 在 Executor 之中，利用隨機梯度等方法逼近最優解；
- 在 reduce 階段定義了模型引數的聚合過程。
- 最後 Executor 輸出一個模型；
+----------------+
| |
| Spark Driver |
| |
+----------------+
+
Map Stage | Reduce Stage
|
|
+--------------------+ |
| Spark Executor | |
| +----------+
+-------> | User function | | |
| | | | |
| +--------------------+ | |
| | |
| +--------------------+ | | +------------------+
| | Spark Executor | | +-> | Spark Executor |
+---+--+ | | | | | +-----+
| Data +----> | User function +------------> | +----> |model|
+---+--+ | | | | User function| +-----+
| +--------------------+ | +-> | |
| | | +------------------+
| +--------------------+ | |
| | Spark Executor | | |
| | | | |
+-------> | User function +----------+
| | |
+--------------------+ +
但是我們發現，這個工作流程只能迭代一次，完全不匹配機器學習需要迴圈迭代多次的特點，於是還需要修改這個架構。
2.2.2 升級模型
於是我們修改角色如下：
- Spark driver不但要負責協調整個Spark任務執行，還需要儲存最近所有梯度，並且負責對Executor傳來的梯度做更新。
- 而executor負責分散式地計算梯度向量，並且梯度提交給driver。
迭代過程也擴充如下：
- 每輪迭代中，executor負責分散式地計算梯度向量，然後將每個 executor 計算的梯度更新值 Aggregate 到 driver。
- 全域性梯度 儲存在driver上，driver根據每個梯度的最新值進行聚合，並且更新模型引數值 w。
- Driver 將 更新後的引數值 w 廣播到每個Executor。
最後 reduce 階段匯出模型。
Map Stage +----------------+
1 | 2 | 1
+----------------> | Spark Driver | <-------------------+
| | | |
| +--+------+---+--+ |
| | 3| ^ | |
| | | | | |
| 3 | | | | 3 |
| +----------------+ | | +-------------------+ |
| | | | | |
| v v |1 v |
+-----+----+---------+ +--------+-----------+ +----------+----+----+
| Spark Executor | | Spark Executor | | Spark Executor |
| | | | | |
| User function | | User function | | User function |
| | | | | |
+-------------+------+ +--------+-----------+ +--------+-----------+
| | |
+----------------------------------------------------------------------+
| | |
+-----------+ | +----------------+
| | |
Reduce Stage v v v
+-+------+------+--+
| Spark Executor |
| |
| |
| User function|
| |
+--------+---------+
|4
|
v
+--+--+
|model|
+-----+
我們突然發現，這居然是一個引數伺服器的架構了，即 Spark Driver 充當了引數伺服器的角色。這和 Horovod 的 ring-allreduce 的架構顯然不符合。另外，Spark採用的完全是BSP協議，即第二輪迭代必須等到第一輪迭代所有的機器完成，這也會拖慢我們的訓練過程。
2.3 機器學習 on Spark 的缺陷
所以，我們在深入之前，需要先說說Spark 如果用於機器學習，會有哪些缺陷：
-
規模依舊不足。Spark受限於模型大小和記憶體限制，只是中等規模機器學習框架。其瓶頸就是Driver。
- Spark框架以Driver為核心，Driver 負責具體任務排程和引數彙總；
- driver又是單機結構，難以擴充套件；
- 當模型規模超過Driver或者Executor所在機器記憶體的時候，Spark就無法正常執行；
-
本質仍不匹配。機器學習的核心是迭代和引數更新。Spark的核心概念是RDD。這兩者的特點不能很好匹配。
-
RDD具備一系列transformation和action介面。使用者使用這些介面完成成不同的演算法或應用。但這組介面是通用介面，無法靈活高效應用於特定領域問題。
-
RDD 並不能很好地支援機器學習中的迭代運算，另外節點之間通訊也低效。
因為大規模機器學習，其模型引數會非常巨大，如果使用 RDD 去容納所有更新的模型引數。需要在每次迭代中建立新的 RDD，這涉及到機器和磁碟間的頻繁資料交換，這會帶來大量額外開銷。
-
RDD難以滿足引數反覆迭代更新的需求。
RDD使用不可變性這個特點來規避分散式環境下的並行問題。此抽象可以簡化運算元複雜度，提供高效能分散式資料處理能力，非常適合資料分析領域。然而不可變性卻不適合引數反覆更新這個需求。
-
雖然 Spark 對於機器學習來說有各種缺陷，但是對於中等規模的學習確實非常有用，所以就有了 Horovod on spark。我們接下來就要看看 Horovod 是如何處理（緩解）這些問題的。大規模機器學習的目的就是解決"資料和偏差"規模非常大的時候所帶來的理論/工程問題。
0x03 整體架構
3.1 整體思路
Tensorflow是C++開發的，而python是機器學習世界的主宰。所以，如果Spark要和TensorFlow 進行整合，一般來說有以下三種方式：
- 通過Tensorflow Java API；
- 通過Tensorflow Python API；
- 通過JNI來呼叫Tensorflow C++ API；
但是 Horovod 的思路又比較別緻，可以認為是按照 Spark 的思路，在 Spark 之上又實現了一套自己的。即：
- Horovod 也有自己的 DriverService（可以認為其對應了 spark driver），或者說 Horovod job 自己就變成了 Spark driver，負責全域性初始化，啟動協調和後續任務分發；
- Horovod 也有自己的 TaskService（可以認為其對應了 spark Executor）；
- Horovod DriverService 用
horovod.spark._make_spark_thread建立了 Spark 叢集； - Horovod DriverService 然後在Spark 叢集上建立了
num_proc個 tasks（Horovod TaskService），這些 tasks 都註冊到 driver 之上，因此 driver 知道已經啟動的所有 task資訊（ip，port，路由，...），這些task 也把自己的 host hash（一個被 MPI 當作 host 的字串）傳送給Horovod DriverService ； - Horovod DriverService 會 通知 Horovod TaskService 啟動訓練；
- 每個 Horovod TaskService 在其所在的 Spark Executor之上，通過呼叫本地程式的方式 mpi 程式，在mpi程式之中又啟動Tensorflow或者Torch來訓練模型。這樣相當於：
- Spark變成容器進行計算資源的排程；
- Tensorflow或者Torch來訓練模型;
- mpi來在各個 Executor 之間做互動做 all-reduce，從而更新梯度等；
這樣就充分利用了已有的大資料體系的資料和計算特性。其實，絕大多數大規模機器學習的平臺/系統都可以看做這由這兩個角色構成 ：Model node（driver node）和 Data node（worker node）。每個角色都有自己一套計算邏輯。從 Horovod來說，Horovod DriverService 就是 driver node，Horovod TaskService就是 data node：
- 資料分佈在 n 個 data node節點上，data node 從 model node 接收任務和程式碼，然後進行計算，並且把計算結果傳送給模型節點。Horovod TaskService 就是完成如下操作，只是不需要傳送計算結果給Horovod DriverService；
- 模型（程式碼）分佈在 m 個model node節點上。在模型結點上進行模型更新，更新是依據"當前模型在資料節點計算/彙總結果 VS 理想模型" 這個偏差來完成。Horovod DriverService （系統中只有一個）就負責維護程式碼，把任務和程式碼發給Horovod TaskService，但是Horovod DriverService沒有更新模型的操作，轉而由Horovod TaskService 通過 Ring-Allreduce 自行完成。
大致如下，其中 SparkDriverService 對應了Horovod DriverService，SparkTaskService對應了Horovod TaskService：
+------------------------------+
| Horovod Main thread |
| |
| |
| SparkDriverService |
| |
| +----------------+ |
| | Spark Driver | |
| +----------------+ |
+------------------------------+
|
|
+--------------------------------------------------------+
| | |
| | |
v v v
+------------------------+ +----------------------+ +------------------------+
| Spark Executor | | Spark Executor | | Spark Executor |
| | | | | |
| +-------------------+ | | +------------------+ | | +-------------------+ |
| | SparkTaskService | | | | SparkTaskService | | | | SparkTaskService | |
| | | | | | | | | | | |
| | TensorFlow | | | | TensorFlow | | | | TensorFlow | |
| | | | | | | | | | | |
| | | | | | | | | | | |
| | MPI | | | | MPI | | | | MPI | |
| | + | | | | + | | | | + | |
| | | | | | | | | | | | | | |
| +-------------------+ | | +------------------+ | | +-------------------+ |
| | | | | | | | |
| | | | | | | | |
+------------------------+ +----------------------+ +------------------------+
| | |
| | |
| | |
+----------------------------+---------------------------+
手機如下：
3.2 具體分析
具體分析如下。
- 在 Horovod 的主程式中執行一個 SparkDriverService（對應 spark driver），或者說就是 Spark driver。
- 利用 _make_spark_thread 啟動 Spark Executor，從而建立了一個Spark叢集，然後 horovod 會等待所有Executor啟動結束;
- 在 spark 的 每個 Executor 上執行一個 SparkTaskService（對應 spark Executor）。
- MPI 需要得到 host 之間的路由資訊，所以 horovod 需要得到這些資訊：
- 回憶一下，在沒有 spark 的情況下，也需要獲取到這些 host 之間的路由資訊。因為 host 之間是一個環形，構成了 ring allreduce。
- 在 Hovorod on spark 狀態下，我們的訓練函式實際上是在 Spark Executor 中執行，為了進行 ring allreduce，所以現在需要知道 spark Executor 之間的路由，以及 driver & tasks 對應關係。
- SparkTaskService 把自己的地址和埠註冊到 SparkDriverService 之上。
- 這樣 SparkTaskService 通過 SparkDriverService 可以獲得自己和彼此的各種資訊。
- SparkTaskService 通過函式，也能夠知道 spark Executor 之間的路由，從而可以互相訪問。
- 從邏輯上來說， spark exector 自己本身的邏輯任務此時已經結束了，因為以後都是 SparkTaskService 自己獨立完成的動作，SparkTaskService 來負責從SparkDriverService接收訓練程式碼，啟動訓練；
- SparkDriverService 知道所有 SparkTaskService 啟動之後，會通知他們進入下一個階段，即等待任務。
- Horovod main thread 在通過SparkDriverService 知道所有 task 啟動之後，會 用
mpi_run來在這些 tasks 之中啟動 python function（通過 RPC）。- 通常，MPI 會通過 SSH 來連線 hosts，但是這種方式無法在 Spark Executor 之中啟動 Python function。
- 因此 MPI 使用 RPC 來啟動使用者程式碼，即使用
horovod.spark.driver.mpirun_rsh來連線每個 Executor，然後 "remote shell" 到這些 spark executors 之中。 horovod.spark.driver.mpirun_rsh是與每個 host hash 之中 最小 index 的 task進行通訊，這個 task 就執行 MPI 的orted命令。因此，每個 Executor 之中只會執行一個mpi orted程式，即使這個 executor 有多個 tasks。其他的 非orted程式 task會等待orted程式 task 結束。
- 在mpirun_rsh之中， SparkDriverService 給 SparkTaskService 傳送 RunCommandRequest，要求 Task 啟動訓練。
- SparkTaskService 在 spark Executor 內部將會使用
_run_command在 spark 之中啟動訓練job。具體如下：- mpi_run 實際上是在 每一個 Spark Executor 之上執行 mpi 程式。即，Horovod 呼叫 mpi_run （又利用到 mpirun_rsh.py）在每一個 spark executor 上啟動 orted，以啟動 MPI cluster。
- SparkTaskService 可以 從 SparkDriverService 得到訓練程式碼；
- orted 在每一個 executor 之上執行訓練程式碼，即 python function；
- 我們的訓練程式碼也是一個 mpirun 程式，即使執行了 tensor flow，也是一個mpi程式，因為一開始從 SparkTaskService 得到了地址和埠，所以可以彼此互動，實現 ring-allreduce。
備註：
Hovorod 期望所有的 task 都同時執行，因此 cluster 應該至少提供同樣個數的 core，每個 executor 可以有多個 core，因此一個 executor 可以處理多個 tasks，host 可以有多個 executor。
具體如下圖：
+--------------------------+ +---------------------------------+ +-------------------------+
| Horovod Main thread | | Spark Executor | | Spark Executor |
| | | | | |
| | | | | |
| +--------------------+ | 1 register | +----------------------+ | | +--------------------+ |
| | SparkDriverService +<---------------------------------+ SparkTaskService | | | | SparkTaskService | |
| | | | | | | | | | | |
| | | | 2 notify start | | | | | | | |
| | +--------------------------------> | | | | | | |
| | | | | | | | | | | |
| | | | | | | | | | | |
| | | | 3 RunCommandRequest | | | | | | | |
| | +---------------------------------------> orted mpirun_rsh| | | | | |
| | | | | | + | | | | | |
| | | | | | | 4 | | | | | |
| | | | | | | | | | | | |
| | | | | | v | | | | | |
| | | | | | task_exec | | | | | |
| | | | | | + | | | | | |
| | | | | | | 5 | | | | | |
| | | | + | | | | | | | |
| | | |6 set_local_rank_to_rank | v | | | | | |
| | +-------------------------+---------> SparkTaskClient | | | | | |
| | | | | | | | | | | |
| | | | | | +------------------+ | | | | +----------------+ | |
| | | | 7 code() | | | | | | | | | | | |
| | +----------------------------------------> 8 train() | | | | | | train() | | |
| | | | | | | | | | | | | | | |
| | | | | | | MPI <----------------------> MPI | | |
| | | | | | | | | | | | | | | |
| | | | | | +------------------+ | | | | +----------------+ | |
| +--------------------+ | | +----------------------+ | | +--------------------+ |
+--------------------------+ +---------------------------------+ +-------------------------+
手機如下：
3.3 Horovod on Spark 架構圖
在 Horovod 原始碼中，有一個架構圖。我們可以大致瞭解其架構。
但是因為這部分實在複雜，所以單憑這一幅圖很難了解其實現，所以我們需要做深入研究。
首先我們看看 Driver 的特點。
3.4 普通狀況 Driver
我們首先用普通Horovod驅動做個對比。
在沒有 spark 的情況下，假設有多個 hosts，需要獲取到這些 host 之間的路由資訊。因為 host 之間是一個環形，構成了 ring allreduce。
Tasks ping each other in a circular fashion to determine interfaces reachable within the cluster.
Driver 服務由 HorovodRunDriverService 提供，Task 服務由 HorovodRunTaskService 等提供。
其功能主要是維護各種 task 地址以及相應關係。具體各種 task 地址就是 Task 服務 來註冊的。
需要注意的是：HorovodRunDriverService 和 HorovodRunTaskService 都最終繼承了 network.BasicService，他們之間可以是異地執行互動。
3.5 Spark 相關的Driver
在 Hovorod on spark 狀態下，我們的訓練函式實際上是在 Spark Executor 中執行，因為面對的情況不同，所以我們對於 Driver 需求是不同的。之前記錄的是 host 之間的路由以及 driver & tasks 對應關係。現在需要知道 spark Executor 之間的路由，以及 driver & tasks 對應關係。
0x04 Spark 模式入口
4.1 示例程式碼
從原始碼中找到示例程式碼如下，可以看到，horovod.spark.run 是入口。
# Horovod: run training.
history, best_model_bytes = \
horovod.spark.run(train_fn, args=(model_bytes,), num_proc=args.num_proc,
stdout=sys.stdout, stderr=sys.stderr, verbose=2,
prefix_output_with_timestamp=True)[0]
4.2 Horovod.spark.run 邏輯
fn 就是訓練函式，被使用者程式碼傳進來的，具體被賦值之後，在 SparkDriverService 之中儲存（具體是在其成員變數 _fn 之中），以後會使用。這樣就解決了程式碼釋出問題。
driver = driver_service.SparkDriverService(settings.num_proc, settings.num_proc,
fn, args, kwargs,
settings.key, settings.nics)
Horovod.spark.run 的邏輯是：
- 處理各種配置，比如timeout，nice...；
- 獲取 spark 資訊，比如從 pyspark 之中獲取SparkContext；
- 構建驅動 SparkDriverService（Spark driver service）；
- 利用 _make_spark_thread 來啟動 spark executor（以及在每一個 spark executor 之中啟動一個SparkTaskService），這樣就構建了 cluster；
- 利用 _notify_and_register_task_addresses 等待所有 spark task 都結束；
- 利用 _launch_job 啟動訓練；
- 利用 spark_thread.join 來收集訓練結果；
具體程式碼如下：
def run(fn, args=(), kwargs={}, num_proc=None, start_timeout=None,
use_mpi=None, use_gloo=None, extra_mpi_args=None,
env=None, stdout=None, stderr=None, verbose=1, nics=None,
prefix_output_with_timestamp=False):
# 處理各種配置，比如timeout，nice...
if start_timeout is None:
# Lookup default timeout from the environment variable.
start_timeout = int(os.getenv('HOROVOD_SPARK_START_TIMEOUT', '600'))
# nics needs to be a set
if nics and not isinstance(nics, set):
nics = set(nics)
tmout = timeout.Timeout(start_timeout, message)
settings = hvd_settings.Settings(verbose=verbose,
extra_mpi_args=extra_mpi_args,
key=secret.make_secret_key(),
start_timeout=tmout,
nics=nics,
run_func_mode=True,.....)
# 獲取 spark 資訊，比如從 pyspark 之中獲取SparkContext
spark_context = pyspark.SparkContext._active_spark_context
settings.num_proc = num_proc
result_queue = queue.Queue(1)
# 利用 _make_spark_thread 來啟動 spark executor（以及在每一個 spark executor 之中啟動一個SparkTaskService）
# start Spark driver service and launch settings.num_proc Spark tasks
spark_job_group = 'horovod.spark.run.%d' % job_id.next_job_id()
driver = driver_service.SparkDriverService(settings.num_proc, settings.num_proc,
fn, args, kwargs,
settings.key, settings.nics)
gloo_is_used = is_gloo_used(use_gloo=use_gloo, use_mpi=use_mpi, use_jsrun=False)
spark_thread = _make_spark_thread(spark_context, spark_job_group, driver,
result_queue, settings,
use_gloo=gloo_is_used, is_elastic=False)
try:
# 等待第一階段結束，即 等待所有 spark task 都結束
# wait for all tasks to register, notify them and initiate task-to-task address registration
_notify_and_register_task_addresses(driver, settings)
# Determine the index grouping based on host hashes.
# Barrel shift until index 0 is in the first host.
host_hashes = list(driver.task_host_hash_indices().keys())
host_hashes.sort()
while 0 not in driver.task_host_hash_indices()[host_hashes[0]]:
host_hashes = host_hashes[1:] + host_hashes[:1]
settings.hosts = ','.join('%s:%d' % (host_hash, len(driver.task_host_hash_indices()[host_hash]))
for host_hash in host_hashes)
# Run the job，啟動訓練
_launch_job(use_mpi, use_gloo, settings, driver, env, stdout, stderr)
except:
# Terminate Spark job.
spark_context.cancelJobGroup(spark_job_group)
# Re-raise exception.
raise
finally:
spark_thread.join()
driver.shutdown()
# Make sure Spark Job did not fail.
driver.check_for_spark_job_failure()
# get ranks from driver
indices_in_rank_order = _get_indices_in_rank_order(driver)
# If there's no exception, execution results are in this queue.
results = result_queue.get_nowait()
return [results[index] for index in indices_in_rank_order]
既然知道了總體程式碼，下一篇我們就介紹 Horovod on spark 如何啟動，敬請期待。
0x05 總結
至此，我們分析了 Horovod on spark 的總體架構，幾個相關問題回答如下：
-
如何將spark作為分散式tensorflow的底層調動機制，通過spark executor去把tensorflow 的程式調動起來，這樣在進行tensorflow訓練時就不需要手動地去組建網路。
- 答案是： Horovod 的思路又比較別緻，可以認為是按照 Spark 的思路，在 Spark 之上又實現了一套自己的。即：
- Horovod 也有自己的 DriverService（對應了 spark driver），或者說 Horovod job 自己就變成了 Spark driver，負責全域性的初始化，建立 Cluster，啟動工作 和 後續任務分發；
- Horovod 也有自己的 TaskService（對應了 spark Executor）；Horovod DriverService 在Spark cluster上建立了
num_proc個 tasks，這些 tasks 都註冊到 driver 之上； - Horovod 的 DriverService 會 通知 TaskService 啟動訓練；
- 答案是： Horovod 的思路又比較別緻，可以認為是按照 Spark 的思路，在 Spark 之上又實現了一套自己的。即：
-
MPI 如何在 Spark Executor 之上啟動使用者程式碼？
- 答案是：
- 通常MPI 會通過 SSH 來連線 hosts，但是這種方式無法在 Spark Executor 之中啟動 Python function。
- 因此 MPI 使用 RPC 來啟動使用者程式碼，即使用
horovod.spark.driver.mpirun_rsh來連線每個 Executor，然後 "remote shell" 到這些 executors 之中。
- 答案是：
-
如何釋出 訓練程式碼？
- 答案是：SparkTaskService 可以 從 SparkDriverService 得到訓練程式碼，因為是 python 指令碼，所以可以直接通過 RPC 傳輸過來；
-
Spark如何開始執行？當某一個 Executor 啟動後就可以執行？還是需要所有的 Executor 都 ready 才能一起跑？
- 答案是：Hovorod 期望所有的 task 都同時執行，因此 cluster 應該至少提供同樣個數的 core，每個 executor 可以有多個 core，因此一個 executor 可以處理多個 tasks，host 可以有多個 executor。
我們在一篇文章中會繼續深入 Horovd on Spark。
0xEE 個人資訊
★★★★★★關於生活和技術的思考★★★★★★
微信公眾賬號：羅西的思考
如果您想及時得到個人撰寫文章的訊息推送，或者想看看個人推薦的技術資料，敬請關注。
0xFF 參考
sona：Spark on Angel大規模分散式機器學習平臺介紹
Spark on Angel：Spark機器學習的核心加速器