Spark任務間消息傳遞方法研究

夏立斌，劉曉宇，孫 瑋，姜曉巍，孫功星

1.中國科學院 高能物理研究所，北京 100049

2.中國科學院大學，北京 100049

為了對自然現象進行更深層次的理解分析，諸多科學研究正面臨著高性能計算、大數據分析與機器學習帶來的挑戰，如粒子物理實驗、氣候預測、生物數據分析、天文圖像處理等。這些科學計算問題不僅數據量大，并且其計算任務通常是以數值線性代數（numerical linear algebra）為核心的多維參數擬合、線性方程組求解等問題，具有很高的靈活性，不易被并行實現。因此，完善的科學計算系統需兼顧大數據處理和高性能計算的特點和優勢，以保證計算任務的執行效率。

近年來內存計算（in-memory computing）技術是互聯網行業的研究熱點，隨著內存容量和性能的提升，可以將計算時需要的數據盡量放在內存中實現本地化計算，避免與網絡、磁盤等低速設備進行頻繁的I/O 操作而損耗性能，其擁有“一次加載，多次使用”的特點，可以大大提升迭代計算的執行效率。分布式處理框架Spark 借助該技術，利用存放在內存中的彈性分布式數據集（resilient distributed datasets，RDDs）和DataSet、DataFrame 構建了一套MapReduce-like 的編程模型，并提供諸多通用的數據操作算子，方便用戶簡單地實現復雜數據處理任務；同時允許將可重用的數據緩存到內存中，通過構建任務依賴關系構成的DAG（directed acyclic graph）對任務進行分發，以提升迭代和交互任務的性能[1]。

由于Spark 編程模型的限制，其針對的應用特點多為易并行（embarrassingly parallel）且相互獨立的任務，可以方便地進行大規模橫向擴展實現并行計算（圖1（a））。而科學計算和機器學習任務間往往彼此相關，表現出具有通信性的特征，如基礎線性代數子程序（basic linear algebra subprograms，BLAS）中的矩陣向量運算。Spark 由于任務間無法相互通信，在其機器學習庫Spark MLlib中實現的BLAS只能通過Shuffle方式進行任務間的數據交換，會產生大量中間數據，對內存、網絡等帶來極大的負擔，導致性能下降；在科學計算領域，通常使用MPI（message passing interface）進行任務間的消息傳遞（圖1（b）），其高效的點對點和集合通信機制，可以使得算法進行靈活高效的實現，并且對高帶寬、低延遲的RDMA（remote direct memory access）網絡，如InfiniBand、RoCE（RDMA over converged ethernet）等有良好的支持。

針對科學計算的特點，本文研究了Spark 內存計算和MPI消息傳遞模型的無縫融合解決方案，通過MPI的多種通信編程機制對Spark編程模型的表達能力進行擴充，以實現大數據分析和數值計算的高效率執行。本文的主要貢獻為：（1）利用MPI增加了Spark任務間的通信機制，使得Spark 能夠高效處理科學計算算法及迭代密集型應用。（2）通過更改Spark任務DAG的劃分方式，添加MPI Stage在Spark框架中對MPI任務進行統一調度執行。（3）實現MPI在分布式內存數據集之中執行，對任務進行函數化、算子化，減輕開發難度，減少數據移動開銷，同時提升了任務的容錯性。

1 相關研究情況介紹

近年來已有一些對高性能計算和大數據技術融合的研究。早期的研究主要利用高性能計算中優異的網絡通信性能對大數據處理框架進行優化，如通過擴展MPI 鍵值對的支持以實現Hadoop/Spark 中的通信原語[2]，利 用RDMA 提 升Spark Shuffle 的 性 能[3]。隨 后Anderson、Gittens 等人通過Linux 共享內存和TCP/IP sockets 的方式對Spark 和MPI 進行連接，并在內存中對兩種框架的數據格式副本進行轉換[4-5]。 Malitsky利用PMI Server 和Spark Driver 分別對Spark 應用程序中的Spark任務與MPI任務進行管理[6]。除了對Spark與MPI結合的研究嘗試以外，SparkNet[7]、TensorFlowOnSpark[8]、RayDP[9]等項目都在探索基于不同工作負載的計算框架之間的結合，以實現端到端的計算流程。

上述工作的設計思想可以總結為兩點：利用高性能計算對框架本身的通信模塊進行修改；對不同的計算框架進行連接。前者無法解決Spark編程模型帶來的局限性；后者需要用戶編寫兩套不同的代碼，在執行過程中存在數據格式轉換的開銷，且Spark 與MPI 計算任務在相互獨立的上下文中管理，本質上依然需要對兩套集群分別進行維護和資源分配。文獻[4-6]中的解決方案，均停留在Spark 與MPI 連接調用的層面上，未能從Spark的編程和計算模型的角度出發進行思考。在對統一資源管理方法的研究中，GERBIL[10]和MPICH-yarn[11]基于Yarn對MPI任務的申請分配方式做出了探索，但僅涉及到資源管理器的層面，使得MPI可以根據獲取到的資源信息分配進程，仍未解決Spark與MPI間的無縫融合。

在HPC 應用中廣泛使用的MPI 庫相比于Spark/Hadoop 有著極高的性能優勢，但無法在生產率和容錯（fault-tolerant）等方面得到保證。在最新的MPI-4.0[12]標準中，MPI 本身仍然不提供錯誤處理機制，一旦MPI 作業中的某個進程出現錯誤，需要對作業進行重新運行?，F有對容錯的研究，主要通過Checkpoint/Restart（CPR）、Global-Restart、MPI Stage 等方式對作業的運行狀態進行記錄，并盡可能地減少作業啟動時間的開銷[13]。在大數據領域，Spark 計算框架和Alluxio 分布式內存系統[14]的核心是以數據為中心，通過DAG對計算/數據的執行流程進行記錄，從而可以很方便地使作業從錯誤中恢復。

2 Spark消息通信設計與實現

2.1 編程和計算模型

2.1.1 編程模型

Spark的編程模型是基于分布式數據集的函數式編程思想，即將一系列算子作用在分布式數據集（RDD）之上，進行MapReduce-like 操作，然后根據依賴關系生成DAG 執行。對于任務之間可以完全并行的操作，能夠利用多線程技術提升性能，其提供的諸多算子和編程庫使得用戶能便捷地開發出并行應用。MPI 則是基于計算資源（如處理器、緩存）和緩沖區的編程模型，核心為進程間通信，通過使用OpenMP 可以實現混合編程（hybrid programming），由于用戶可以直接結合物理資源和算法特性進行細粒度的優化，因此被廣泛地應用到高性能計算領域中。具體的MPI與Spark編程模型和框架軟件生態的對比如表1所示。

表1 Spark與MPI對比Table 1 Comparison of Spark and MPI

為彌補Spark 編程模型表達能力的局限性，同時為MPI 添加了基于數據的DAG 執行能力，本文將兩種編程模型進行了融合，使得MPI可以充分利用大數據生態內的內存計算能力和分布式文件系統等組件，將大數據處理和高性能計算結合在一起。如代碼1所示，編程模型的核心為多個分區構成的分布式數據集。通過借鑒Spark編程模型中的map操作，使用mpimap將MPI函數映射到數據集的不同分區上進行計算。此時的數據分區單位，同時也是MPI進程單位，數據分區集合RDD也代表著同一個Communicator下的MPI進程集合?；谝陨暇幊谈拍?，使得Spark 的map 任務間可以通過MPI通信原語進行交互，彌補了Spark 編程模型中map 任務間不能交互的不足。并且每個mpimap操作產生的計算中間數據都會存放在內存中，可供后續算子復用。因此，編程時可以將復雜的MPI邏輯解耦，并與Spark計算中便捷的數據處理方法結合，在統一的應用中通過數據流的形式進行串聯。在實際使用的過程中，用戶可以根據性能需求靈活地實現算法，而不必局限于Spark 原生算子及MapReduce編程模型的限制。

基于上述編程模型，底層的計算框架可根據數據分區信息，采用數據本地性及網絡拓撲結構的算法進行任務（進程）分配。通過使用以數據為中心的模型，取代了傳統mpirun 命令行啟動MPI 進程的方式，加強數據和計算之間的聯系，降低了編程的復雜度，同時也方便用戶進行代碼復用。

2.1.2 消息通信DAG計算模型

為了將編程模型中提出的概念設計與實際任務的物理分配建立起有效的聯系，Spark利用不同RDD之間的依賴（Dependency）關系，生成出計算DAG 以供任務調度。如圖2 所示，Spark 默認提供的依賴關系包括NarrowDependency 和ShuffleDependency 兩種。為了減少重復計算，同時提升任務的并行度，Spark 通過ShuffleDependency劃分出Stage，同一個Stage中的任務作用在單獨的Partition 之上，不需要與其他任務的Partition數據進行交換，在不同的Stage 間通過Shuffle 進行數據交換。

在對Spark 中加入消息通信機制后，低效的Shuffle數據交換方式已被MPI 取代。為了實現二者的無縫融合，必須對原有的DAG 構建策略進行修改。本文添加了與MPI 相關的RDD、Dependency 和Stage，并對DAG 的劃分策略進行了修改。如圖3 所示，對RDD 進行mpimap 函數計算會產生新的MPIMapPartitionRDD，以及相對應的MPIDependency 依賴，并在DAG 生成的過程中劃分出MPI Stage。鑒于MPI 通信模型的要求，該Stage 中的任務在Barrier mode 下執行，即必須獲取到所有對應的進程資源，否則計算不會被執行。不同于Spark 自身實現的ShuffleMapStage，MPI Stage 在任務間利用消息傳遞進行通信，從而擴充Spark的局限性，以提升計算效率。

2.1.3 MPI計算容錯

根據以上規則生成出的DAG，每個作用在RDD 分區上的MPI 函數都會在一個Stage 中執行，并在函數執行成功之后產生一個新的RDD。由于對MPI程序進行了算子化操作，在一個作業中會有多個MPI函數被分配到相互獨立的Stage中，根據數據的依賴關系并行執行。上述的編程模型和計算模型，不僅僅對Spark 自身的表達能力進行了擴充，同時也使得MPI可以實現以數據為中心的內存計算。MPI 與Spark 的無縫融合，為使用HDFS，Alluxio等分布式文件系統也提供了便利。在對于容錯技術的研究上，圍繞MPI進行的相關工作都集中在Checkpoint和減少作業恢復的啟動時間上，對于某一個進程產生的錯誤是不可避免的。Spark、Alluxio 則是從計算/數據的DAG出發，根據中間計算執行流和中間數據緩存，保證作業從錯誤中快速恢復。因此，在對MPI函數算子化后，作業的執行過程變成了諸多作用在分布式內存數據集上的變換，相當于對一個大的MPI作業進行了分割，并可以對中間計算結果進行保存。當部分任務出現錯誤時，可以基于DAG 和RDD 進行恢復，而無需重啟整個作業。以上容錯的處理方式，尤其是對運行在大規模數據集上的作業，能夠減輕大量狀態恢復所需的計算和I/O開銷。

2.2 Blaze架構設計

根據2.1 節闡述的編程和計算模型，基于Spark 和OpenMPI 構建了一套計算系統Blaze，其總體設計和工作流程如下所示。

2.2.1 OpenMPI組件設計

基于MPI 標準的實現有多種，其中OpenMPI 的采用模塊-組件的架構設計（modular component architecture，MCA），如圖4 所示，能夠便捷地進行模塊插件開發，實現自定義功能[15]。為了實現MPI與Spark的融合，OpenMPI端的主要任務為接收Spark調度器中產生的進程分配信息，從而啟動MPI Executor（Process）執行具體任務。傳統的MPI 程序主要使用mpirun 命令行執行，即通過OpenMPI 的Run-Time 層（PRTE）產生資源分配信息，并實現進程啟動。因此，Blaze需要對上述執行過程進行修改以完成對Spark的結合。

Blaze 在設計上并未改變OpenMPI 執行的邏輯，而是充分利用MCA 架構的靈活性，編寫了自定義的功能模塊進行替換。按照OpenMPI 模塊調用的順序，主要包括進程生命周期管理模塊（PLM）、資源分配模塊（RAS）、進程映射模塊（RMAPS）、啟動模塊（ODLS）及存儲上述信息的數據庫。使用JNI的方式將PRTE集成在Master中提供服務，并通過對上述與進程啟動相關模塊的改寫，使得MPI資源分配與調度所需的信息完全由Spark提供。

具體的實現方法為：（1）PRTE根據PLM管理的作業生命周期維護了一個有限狀態機，在Blaze 中由Master進行管理。（2）Master 中記錄著系統整體的資源使用情況，以及提交的應用信息，RAS 被Master 中的作業資源申請功能所替代。（3）RMAPS 進程映射規則由Blaze 任務調度器根據2.1.2小節的DAG 計算模型產生，目前使用數據本地性算法進行調度。（4）ODLS 在OpenMPI 中負責本地進程的啟動，但在Blaze 中將此過程改為Worker根據調度信息產生Executor。

至此，相對于傳統MPI 進程啟動的關鍵組件已在Blaze 中實現。MPI 進程與Spark Executor 的概念建立起來了連接，MPI_Init 可直接利用Executor 中包含的Namespace 及Rank 信息完成進程間的感知，進而實現Spark任務間的消息通信。

2.2.2 Blaze處理流程

Blaze處理作業的總體過程如圖5所示，架構設計和執行邏輯在主體上與Spark保持一致。不同之處主要在于，本文2.1.2 小節設計的DAG 計算模型通過DAG Scheduler 進行了實現，并由Task Scheduler 根據數據本地性原則完成MPI 任務的分配；同時在Master 中添加了2.2.1 小節介紹的MPI RTE，通過與Spark 提供的信息進行交互以完成任務調度；最終的計算任務在MPI Executor中執行。

Blaze 完整的作業處理流程為，Master 接受客戶端提交的作業，并將其調度到Worker 節點中啟動App Driver，此處保存著應用程序運行時的上下文BlazeContext，代表與Spark 和MPI 之間的連接。用戶提交的應用程序代碼在Driver 中執行，調度器根據編程模型的規則生成DAG 計算圖，并按照依賴關系將Stage 轉換成TaskSet 按順序發送到Executor 執行。對于一個MPIStage，在任務調度的過程中，需要根據RDD Partition的存儲位置等信息做出決策，與Master中的MPI RTE通信，生成MPI Namespace 和Rank。每一個MPIStage 擁有一個Namespace，其生命周期由TasksetManager管理，使得被傳送到Executor 的Task 閉包能夠完成MPI 的初始化。因此由Worker產生的MPI Executor為單線程模式，即每個Executor分配一個CPU core。同時為了對科學計算中的C/C++ MPI程序進行支持，可將Spark執行中傳送的Closure替換為待執行的二進制文件或動態鏈接庫在MPI Executor中執行。

3 實驗測試及結果分析

在科學計算及機器學習領域，分布式稠密矩陣（dense matrix）和稀疏矩陣（sparse matrix）的運算，以及迭代算法是各種計算作業的基礎。例如在高能物理領域，分波分析（partial wave analysis，PWA）需要在高統計量樣本數據上進行數值擬合的計算，其核心過程為使用最大似然法估計待定參數，需要反復迭代以求得最優參數[16]；格點QCD（lattice quantum chromodynamics，LQCD）從第一性原理出發將連續的量子場離散化為時間-空間格點，計算熱點為大規模線性方程組的求解[17]。因此，本文對實際物理分析中的計算特征進行提取，選取分布式稠密矩陣乘法，線性方程組求解作為性能測試的用例，并與Spark的原生實現進行對比。

本文的實驗環境是一個8 節點組成的集群，包括1個BlazeMaster和8個Worker節點。各個節點的軟硬件情況如表2所示。

表2 測試集群軟硬件環境Table 2 Software and hardware environment of cluster

3.1 矩陣乘法測試

對于稠密矩陣的乘法測試，本文選取BLAS的第三層次矩陣-矩陣乘（matrix-matrix multiplication）進行實驗。Spark在其機器學習庫Spark MLlib中基于MapReduce實現了分布式分塊矩陣（BlockMatrix）的乘法運算，而本文則利用Blaze 中添加的消息傳遞機制，基于MPI和Spark 的RDD-based BlockMatrix 實 現 了Cannon 算法，與Spark MLlib中默認的稠密矩陣乘法進行對比。

Cannon 算法將待乘矩陣A和B按處理器數量p分成p×p的矩陣Aij和Bij，(0＜i,j ＜p-1)，分布在虛擬的方陣格點之上。在完成每輪結果Cij的計算后，對A矩陣沿著i方向循環左移，B矩陣沿著j循環上移，共執行p次循環移位后輸出結果C。常規的n×n矩陣乘法的復雜度為O(n3)，Cannon算法將計算復雜度降低到了O(n3/p)。但由于其在計算的過程中存在大量的分塊矩陣數據傳遞過程，無法在Spark 中高效實現，而在本文設計的Blaze中則可以借助MPI對Cannon算法進行實現，其中本地矩陣向量運算使用OpenBLAS處理。

實驗結果對比的基準為矩陣計算部分的性能，即矩陣已被加載到Spark 中的DenseMatrix 分布式內存數據后的計算過程。 圖6（a）是一組在20 160×20 160 大小上的矩陣乘法測試，并根據Cannon 算法要求選擇了n2數量的CPU 核心。結果顯示，基于消息傳遞方法的矩陣乘法相比于Spark MLlib 中的實現有57%～68%的性能提升，與C++實現的矩陣乘法性能十分接近，額外的開銷在于序列化反序列化過程及編程語言層面。同時隨著CPU 核心數目的增長，消息傳遞方法的加速曲線逐漸變緩，是由于進程數目增加導致的通信開銷不斷增長。圖6（b）則展示了隨著矩陣維數的增長在相同CPU核心數目（64 核）情況下的性能對比，消息傳遞相比于Spark MLlib 有50%～69%的提升。并且當矩陣維度增大到64 000時，Spark會由于OOM無法得到計算結果。

由上述測試可以看出，矩陣維度及CPU 核心數目等參數的變化都會對程序的執行效率造成影響，但在相同的軟硬件配置下，同樣基于BlockMatrix 和MLlib BLAS 實現的消息傳遞方式的矩陣乘法在性能測試中的表現均優于Spark MLlib。同時本文提出的Spark 消息傳遞方法，相比于文獻[5]中提出的Alchemist有10%～27%的性能提升，原因在于Alchemist 并未將MPI 融合在Spark中，而是充當MPI與Spark交互的連接器，因此存在額外的數據傳遞及格式轉換開銷。并且用戶需使用Alchemist 提供的API 編寫程序，在程序運行時調用MPI 實現的編程庫進行計算，相比于在Spark 中直接實現的任務間消息傳遞，在易用性、靈活性及性能表現上均存在不足。

3.2 迭代計算測試

共軛梯度下降法（conjugate gradient method）是一種在Krylov子空間中迭代的算法，被廣泛應用在線性系統的求解中，如Ax=b。其實現如算法1 所示，計算的熱點為矩陣-向量乘和向量-向量乘，當結果小于規定的殘差值或迭代到達最大輪次后輸出結果。

由于Spark Mllib 中未包含CG 算法，本文對Spark-CG進行了實現，其核心思想為，對迭代中需復用的分布式矩陣A進行Cache，而每輪迭代更新的參數通過Broadcast 發送到Worker 中使用，在每個任務內完成分布式矩陣向量乘法及向量內積運算。基于消息傳遞方式的實現則通過MPI_Allreduce 完成參數更新，并令所有的計算過程在同一個Stage中完成。

圖7 分別展示了在不同CPU 核心數量（32 000×32 000 矩陣）及不同維度矩陣（32 核）下通過CG 求解線性方程組的平均迭代耗時，結果顯示消息傳遞方式的實現相比于Spark 分別有51%～94%和58%～92%的性能提升，相比于Alchemist 有7%～17%的性能提升。在對Spark的測試過程中，隨著CPU核心數目的提升，并沒有帶來執行時間上的顯著縮短，原因在于隨著節點增加帶來的調度和網絡開銷以及不同任務的執行時間差異影響了整體性能，實際上每一個任務內的計算時間仍然在下降，與此同時，消息傳遞方式則表現出了與預期相符的性能表現。故采用消息傳遞的方式可以大量減少冗余任務的分發調度過程，同時能夠更高效地實現分布式環境下的矩陣向量運算。因此對于迭代密集型應用，均可選擇使用消息傳遞的方式替換Spark 中的Broadcast及Reduce等操作，從而獲得極大的性能提升。

4 結束語

本文針對Spark不能滿足科學計算場景下對編程模型靈活性的需求，以及計算性能低下的問題，深入研究了Spark 及OpenMPI 框架的設計思想，提出了一種在Spark 任務間進行消息傳遞的方法。通過修改Spark 與OpenMPI的資源管理和任務調度的運行邏輯，實現了大數據和高性能計算框架的融合，有效提升了Spark 執行數值計算任務的性能。因此，本文提出的Spark 任務間消息傳遞方法，可適用于大規模數據量環境下的科學計算與數據處理任務，如迭代密集型計算或具有任務通信特征的算法與應用等。

為了便于移植歷史程序，或出于性能考慮需使用基于C/C++編寫的MPI程序時，會存在大量的數據轉換和序列化/反序列化開銷。在Executor 內執行的MPI 計算任務也存在類似的問題，同時會面臨JVM 的內存壓力。因此，除計算引擎和調度系統之外，高效統一的數據格式以及內存管理系統也是需要關注的重點。RDD雖然提供了良好的分布式內存計算抽象，但無法滿足復雜數據類型的高效使用與管理。本文下一步的工作會專注于分布式內存管理相關技術的優化研究，以最終實現可以高效率執行多種作業任務的統一計算系統。