歡迎關注公眾號“Tim在路上”
在Spark3.2中引入了領英設計的一種新的shuffle方案，今天我們先來了解下其大致的設計原理，之后會再分析其具體的代碼實現。
當我們在Yarn上部署Spark時，通常會使用ESS來管理shuffle數據（具體可見什么是ESS的文章）。我們先來回顧下基于ESS進行shuffle的過程。
- 每個 Spark Executor 在啟動后都會和位于同一個節點上面的 Spark External Shuffle Service （ESS）進行注冊。此類注冊允許 Spark ESS 了解來自每個注冊 Executor 的本地 Map 任務產生的物化 Shuffle 數據的位置。請注意，Spark ESS 的實例在 Spark Executor 的外部，并且可以在多個 Spark 應用程序中共享。
- Shuffle Map Stage 中的每個任務都會處理部分數據。在 Map 任務結束時，它會產生 2 個文件，一個用來存儲 Shuffle 數據，另一個用來索引前者的 Shuffle 塊。為了這樣做， Map 任務會根據分區鍵的散列值對所有轉換的記錄進行排序。在此過程中，如果無法在內存中對整個數據進行排序，則 Map 任務會溢出中間數據到磁盤。一旦排序，將生成 Shuffle 數據文件，其中屬于相同 Shuffle 分區的所有記錄都會被組合到一起，放到一個 Shuffle 塊中。還會生成匹配的 Shuffle 索引文件，用來記錄塊邊界的偏移量。
- 當下一個 Stage 的 Reduce 任務開始運行時，它們會查詢 Spark 的Driver 以獲取輸入的 Shuffle 塊的位置。一旦此信息變為可用，每個Reduce 任務將會建立和對應的 Spark ESS 實例的連接，以便獲取其輸入數據。 Spark ESS 在接收到這樣的請求時，會利用 Shuffle 索引文件來跳到 Shuffle 數據文件中對應塊數據，從磁盤讀取它，并將其發送回 Reduce 任務。
然而在實踐中任然存在些問題，使得spark任務的穩定性不高。
- Spark ESS 每個 FETCH 請求只會讀取一個 Shuffle 塊，因此Shuffle 塊的平均大小決定了每次盤讀的平均數據量，如果存在大量小 Shuffle 塊導致磁盤 I/O 低效。
- Reduce 任務在建立與遠程 Spark ESS 的連接時出現失敗的情況，它會立即失敗整個的 Shuffle Reduce Stage，導致前面的 Stage 重試，來重新生成拉取不到的 Shuffle 數據。
- 如果 Shuffle 塊在 Reduce 任務中本地可用，則任務可以直接從磁盤讀取，繞過 Shuffle 服務，這有助于減少 Shuffle 期間的 RPC 連接數。但是 Spark 當前的 Shuffle 機制會導致 Reduce 任務的數據本地性很少，因為它們的任務輸入數據分散在所有的 Map 任務中。
Push-based shuffle架構流程
PBS主要結構和流程：
- Spark driver組件，協調整體的shuffle操作;
- map任務的shuffle writer過程完成后，增加了一個額外的操作push-merge，將數據復制一份推到遠程shuffle服務上；
- magnet shuffle service是一個強化版的ESS。將隸屬于同一個shuffle partition的block，會在遠程傳輸到magnet 后被merge到一個文件中；
- reduce任務從magnet shuffle service 接收合并好的shuffle數據，不同reduce任務可以共享shuffle數據來提升shuffle傳輸效率。
幾個重要的特性：
Push-Merge Shuffle - Magnet采用 Push-Merge Shuffle 機制，其中 Mapper 生成的 Shuffle 數據被推送到遠程的 Magnet Shuffle Service，從而實現每個 shuffle 分區都能被合并。這允許Magnet將小的 Shuffle 塊的隨機讀取轉化成 MB 大小塊的順序讀取。此外，此推送操作與 Mapper 分離，這樣的話，如果操作失敗，也不會增加 Map Task 的運行時間或者導致 Map Task 失敗。
最給力的兜底方法 - Magnet不需要塊 push 操作完成的那么完美。通過執行Push-Merge Shuffle，Magnet有效地復制了 shuffle 數據。Magnet允許 reducer 獲取合并的和未合并的 shuffle 數據都作為任務輸入。這使得Magnet能夠容忍塊 push 操作的部分完成。
靈活的部署策略 - Magnet 通過在頂層構建的方式集成了 Spark 原生的 shuffle。這使得Magnet可以部署在具有相同位置的計算和存儲節點的 on-prem 集群中與disaggrecated存儲層的cloud-based的集群中。在前一種情況下，隨著每次 Reduce Task 的大部分都合并在一個位置，Magnet利用這種本地性來調度 Reduce Task 并實現更好的 Reducer 數據本地性。在后一種情況下，代替數據本地性，Magnet可以選擇較少負載的遠程 shuffle 服務，從而更好的優化了負載均衡。
緩解落后/數據傾斜 - Magnet可以處理落后和數據傾斜。由于Magnet可以容忍塊 push 操作的部分完成，因此可以通過停止慢速 push 操作或跳過 push 大/傾斜的 block 塊來緩解落后和數據傾斜。
為遠程 push 準備 block 塊
push-merge 的根本目的是減少reduce側的隨機IO, 在Magnet上把小文件block合并后, 將隨機IO轉變為順序IO。reduce task可以讀取連續存儲的、大小在MB級別的文件。
為了解決map端的小文件問題，提高磁盤 I/O 效率，我們需要增加每次 I/O 操作的數據量。這里提出了采用合并屬于同一個 Shuffle 分區的 Shuffle block 塊，以創建更大的數據塊的方式。
下面我們來詳細解釋下：
首先，push-merge的基本單位是chunk，map task輸出block后，首先要將block以算法的方式分配到chunk中去。
這里的算法的簡單思想就是將block塊合并為chunk，當chunk的長度超過超限之后又push到magent上的過程。具體的實現在方法ShuffleBlockPusher.prepareBlockPushRequests方法中：
for (reduceId <- 0 until numPartitions) {
val blockSize = partitionLengths(reduceId)
if (blockSize > 0) {
// [1] 通過以下公式，更新一下merge service機器編號，把chunk發送到下一臺機器上
val mergerId = math.min(math.floor(reduceId * 1.0 / numPartitions * numMergers),
numMergers - 1).asInstanceOf[Int]
// [2] 當chunk長度沒有超過限制maxBlockSizeToPush，將block append到chunk中，更新chunk長度
// service, and does not go beyond existing limitations.
if (currentReqSize + blockSize <= maxBlockBatchSize
&& blocks.size < maxBlocksInFlightPerAddress
&& mergerId == currentMergerId && blockSize <= maxBlockSizeToPush) {
// Add current block to current batch
currentReqSize += blockSize.toInt
// [3] 當chunk長度超過限制，將chunk推到編號為currentMergerId的Magnet機器上，之后寫入新的block進去（重新初始化)
} else {
if (blocks.nonEmpty) {
// Convert the previous batch into a PushRequest
requests +=PushRequest(mergerLocs(currentMergerId), blocks.toSeq,
createRequestBuffer(transportConf, dataFile, currentReqOffset, currentReqSize))
blocks = new ArrayBuffer[(BlockId, Int)]
}
// Start a new batch
currentReqSize = 0
// Set currentReqOffset to -1 so we are able to distinguish between the initial value
// of currentReqOffset and when we are about to start a new batch
currentReqOffset = -1
currentMergerId = mergerId
}
// push的blocks長度都是小于maxBlockSizeToPush
// Only push blocks under the size limit
if (blockSize <= maxBlockSizeToPush) {
val blockSizeInt = blockSize.toInt
blocks += ((ShufflePushBlockId(shuffleId, shuffleMergeId, partitionId,
reduceId), blockSizeInt))
// Only update currentReqOffset if the current block is the first in the request
if (currentReqOffset == -1) {
currentReqOffset = offset
}
if (currentReqSize == 0) {
currentReqSize += blockSizeInt
}
}
}
offset += blockSize
}
可見這里的算法的流程為：
- [1] 通過math.min(math.floor(reduceId * 1.0 / numPartitions * numMergers),numMergers - 1)公式，更新一下merge service機器編號，把chunk發送到下一臺機器上；
- [2] 當chunk長度沒有超過限制maxBlockSizeToPush，將block append到chunk中，更新chunk長度
- [3] 當chunk長度超過限制，將chunk推到編號為currentMergerId的Magnet機器上，之后寫入新的block進去（重新初始化)
同時需要注意這里push的blocks的大小都是小于maxBlockSizeToPush，這里用于跳過數據傾斜的分區塊。
// Randomize the orders of the PushRequest, so different mappers pushing blocks at the same
// time won't be pushing the same ranges of shuffle partitions.
pushRequests ++= Utils.randomize(requests)
另外為了避免順序的構造push chunk，導致Magnet上資源的熱點和嚴重的爭用沖突。在完成準備將shuffle data轉換為push request后，將chunk按照編號進行了隨機化處理，來避免所有map task按照相同次序push chunk。
合并Magnet shuffle service的 block 塊
在 Magnet shuffle service一側，對于正在主動合并的每個 Shuffle 分區，Magnet shuffle service會生成合并的 Shuffle 文件，用來添加所有接收的相應 block 塊。
它還為每個主動合并的 Shuffle 分區維護一些元數據。
這份元數據的唯一鍵由 applicationID，shuffle ID和 shuffle partition ID混合組成，并且放到一個 ConcurrentHashMap 中。
Magnet機器上需要維護一些重要的元信息，如上圖所示，包括：
- bitmap: 存儲已merge的mapper id，防止重復merge；
- position offset: 如果本次block沒有正常merge，可以恢復到上一個block的位置；
- currentMapId：標識當前正在append的block，保證不同mapper 的block能依次append。
當Magnet shuffle service接收到 block 塊時，在嘗試添加到對應的 shuffle 合并文件之前，它首先要檢索相應的 Shuffle 分區元數據。元數據可以幫助Magnet shuffle service正確處理一些潛在的異常場景。
例如，bitmap可幫助Magnet shuffle service識別任何潛在的重復塊，因此沒有多余的數據會被寫入 Shuffle 合并文件中。
currentMapId用于保證當前正在append的block，即使Magnet shuffle service可以從不同的 Map 任務中接收同一個 shuffle 分區的多個 block 塊，只有當currentMapId的 block 塊完整地添加到 Shuffle 合并文件中，下一次寫入才可開始。
并且，在遇到足以損壞整個 shuffle 合并文件的故障之前，可以將 block 塊部分地添加到 Shuffle 合并文件中。當發生這種情況時，position offset會有助于將 Shuffle 合并文件帶回到健康狀態。下一個 block 塊會從位置偏移量處開始添加，這可以有效地覆蓋損壞的部分。如果損壞的 block 塊是最后一個的話，block 合并操作結束之后將截斷損壞的部分。通過追蹤這份元數據，Magnet shuffle service可以在 block 塊合并操作期間適宜地去處理重復，沖撞和故障的情況。
提升 Shuffle 的可靠性
magnet shuffle服務通過Best-effort的方式來解決海量連接可靠性低的問題。在該體系上，所有連接異常都是non-fatal的，可以理解為每個環節上的連接斷開或異常，都有一個對應的備選和兜底方案：
- 如果Map task輸出的Block沒有成功Push到magnet上，并且反復重試仍然失敗，則reduce task直接從ESS上拉取原始block數據；
- 如果magnet上的block因為重復或者沖突等原因，沒有正常完成merge的過程，則reduce task直接拉取未完成merge的block；
- 如果reduce拉取已經merge好的block失敗，則會直接拉取merge前的原始block。
對于一個有著 M 個 Map 任務和 R 個 Reduce 任務的 Shuffle 來說，Spark Driver 會收集 M 個 MapStatus和 R 個 MergeStatus。
這些元數據會告訴 Spark Driver 每個未合并的 Shuffle block 塊和已合并的 Shuffle 文件的位置和大小，還有哪些 block 塊會合并到每一個 Shuffle 合并文件中。
因此，Spark Driver 可以完整的看到，怎樣去獲取每個 Reduce 任務已合并的Shuffle 文件和未合并的 Shuffle 塊。當 Reduce 任務沒能獲取到 Shuffle 合并 block 塊時，元數據便會能夠回過頭來獲取原始的未合并的 block 塊。
Magnet 盡最大可能有效地維護了兩份 Shuffle 數據的副本。
靈活的部署策略
Magnet允許 Spark原生地去管理 Shuffle 的各個方面，包括存儲 Shuffle 數據，提供容錯能力，還有可以追蹤 Shuffle 數據的位置元數據信息。
在這種情況下，Spark 不依賴于外部的系統進行 Shuffle。
這允許靈活地將Magnet部署在計算/存儲同一節點的 on-prem 集群和具有disaggregated storage layer的cloud-based的集群。對于計算和存儲同一個節點的on prem數據中心，Shuffle Reduce 任務的數據本地性可以帶來很多好處。
其中包括提高 I/O 效率，并且由于繞過網絡傳輸減少了 Shuffle 獲取失敗的情況。
通過利用 Spark 的位置感知任務調度并且基于 Spark Executor 的位置信息選擇 Magnet shuffle service來 push Shuffle block 塊，實現 Shuffle 數據本地性似乎微不足道。
動態分配的功能使得 Spark 在一段時間內如果沒有任務運行，則釋放空閑的 Executor，并且如果任務再次待辦，則可以稍后重新啟動 Executor。
這使得 Spark 應用程序在多租戶集群中資源更加富裕。
通過 Spark 動態分配，當 Driver 在 Shuffle Map Stage 的開頭選擇Magnet shuffle service列表時，由于 Executor 在前一個 Stage 的結尾會釋放，活躍的 Spark Executor 的數量可能小于需求的數量。如果我們選擇基于 Spark Executor 位置信息的 Magnet shuffle service，我們最終可能比需求的 Shuffle 服務更少。
為了解決這個問題，我們選擇在活躍 Spark Executor 之外位置的Magnet shuffle service，并通過基于所選Magnet shuffle service位置信息的動態分配機制來啟動 Spark Executor。這樣的話，我們基于Magnet shuffle service的位置信息來啟動 Spark Executor，而不會去基于 Spark Executor 的位置信息來選擇 Magnet shuffle service。由于Magnet和 Spark 原生的 Shuffle 集成，因此可以進行這種優化。
對于cloud-based的集群部署，計算和存儲節點通常是分開的。
在這樣的部署中， Shuffle 中間數據可以通過快速網絡連接在disaggregated storage中物化。
Shuffle Reduce 任務的數據本地性在這種設置中不再重要。然而，Magnet仍然適合這種cloud-based的部署。Magnet shuffle service在計算節點上運行，在 disaggregated storage 節點上面存儲合并的 shuffle 文件。通過讀取更大的數據 chunk 塊而不是橫跨網絡的細碎的 shuffle block 塊，Magnet有助于更好地利用可用網絡帶寬。
此外，Spark Executor 在選擇Magnet shuffle service的時候可以選擇優化更好的負載均衡而不是數據本地性。Spark Driver 可以查詢可用Magnet shuffle service的負載，以便選擇負載低的。在我們的Magnet實現中，我們允許通過靈活的政策來選擇Magnet shuffle service的位置。
因此，我們可以選擇根據集群的部署模式要么優化數據本地性要么優化負載均衡，或者兩者都有也行。
處理落后和數據傾斜
- 解決Task Straggler問題
當所有的 Map 任務在 Shuffle Map Stage 結尾完成的時候，Shuffle block 塊推送操作可能還沒有完全完成。
此時有一批 Map 任務剛剛開始推送 block 塊，也可能有落后者做不到足夠快地推送 block 塊。不同于 Reduce 任務中的落后者，我們在 Shuffle Map Stage 結尾經歷的任何延遲都將直接影響作業的運行時間。
為了緩解這樣的落后，Magnet允許 Spark Driver 設置期望等待 block 塊推送/合并操作的時間上限。
magnet服務設置了push-merge超時時間，如果block沒有在超時時間內完成push-merge，magnet服務會停止繼續接受block，提前讓reduce task開始執行；而未完成push-merge的block，根據上面中提到的Best-effort方案，reduce task會從MapStatus中獲取狀態與位置信息，直接拉取沒有merge的block數據。
然而，它確保Magnet可以提供 push/merge shuffle 的大部分益處，同時將落后者的負面影響限制在 Spark 應用程序的運行時間內。
- 解決數據傾斜
在Spark shuffle過程中，如果某個partition的shuffle數據量遠高于其他partition，則會出現數據傾斜（data skew）問題。 data skew 不是magnet特有的問題，而是在Spark上已經有成熟解決方案，即 AQE。
magnet需要適配Spark 的adaptive execution特性，同時防止一個magnet服務上因data skew而導致有 100GB / 1TB級別的數據需要merge。為此，針對上文的算法可以看出，push的blocks的大小都是小于maxBlockSizeToPush，通過限制 size超過閾值的block被并入到chunk中；如果超過閾值，則會利用上節中的Best-effort方案，直接拉取未完成merge的block數據。而普通的、未有data skew情況的block，則會走正常的push-merge流程。
push-based shuffle 配置
服務器端配置(yarn-site.xml)
# 默認的push based shuffle是關閉的。如果需要開啟請設置為：org.apache.spark.network.shuffle.RemoteBlockPushResolver。
spark.shuffle.push.server.mergedShuffleFileManagerImpl=org.apache.spark.network.shuffle.NoOpMergedShuffleFileManager
# 在push-based shuffle期間將合并的shuffle文件劃分為多個塊時最小的大小，默認為2m。
spark.shuffle.push.server.minChunkSizeInMergedShuffleFile=2m
# 緩存大小，可以存儲合并的索引文件
spark.shuffle.push.server.mergedIndexCacheSize=100m
客戶端配置
接下來我們將從源碼的角度進行進一步的分析。
