[源碼解析] 深度學習分布式訓練框架 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機器學習的核心加速器