全新快取元件，大幅加速雲上飛槳分散式訓練作業

語言: CN / TW / HK
時間 2022-01-07 20:04:37
測試阻礙交付，如何破解這一難題？>>>

在Kubernetes的架構體系中，計算與儲存是分離的，這給資料密集型的深度學習作業帶來較高的網路IO開銷。為了解決該問題，我們基於JuiceFS在開源專案Paddle Operator中實現了樣本快取元件，大幅提升了雲上飛槳分散式訓練作業的執行效率。

*JuiceFS：

http://github.com/juicedata/juicefs

*Paddle Operator：

http://github.com/PaddleFlow/paddle-operator

 

背景介紹

由於雲端計算平臺具有高可擴充套件性、高可靠性、廉價性等特點，越來越多的機器學習任務執行在Kubernetes叢集上。因此我們開源了Paddle Operator專案，通過提供PaddleJob自定義資源，讓雲上使用者可以很方便地在Kubernetes叢集使用飛槳（PaddlePaddle）深度學習框架執行模型訓練作業。

然而，在深度學習整個pipeline中，樣本資料的準備工作也是非常重要的一環。目前雲上深度學習模型訓練的常規方案主要採用手動或指令碼的方式準備資料，這種方案比較繁瑣且會帶來諸多問題。比如將HDFS裡的資料複製到計算叢集本地，然而資料會不斷更新，需要定期的同步資料，這個過程的管理成本較高；或者將資料匯入到遠端物件儲存，通過製作PV和PVC來訪問樣本資料，從而模型訓練作業就需要訪問遠端儲存來獲取樣本資料，這就帶來較高的網路IO開銷。

*Paddle Operator：

http://github.com/PaddleFlow/paddle-operator

*PaddlePaddle：

http://www.paddlepaddle.org.cn/

 

為了方便雲上使用者管理樣本資料，加速雲上飛槳框架分散式訓練作業，我們與JuiceFS社群合作，聯合推出了面向飛槳框架的樣本快取與管理方案，該方案期望達到如下目標：

  • 資料集及其管理操作的自定義資源抽象。將樣本資料集及其管理操作抽象成Kubernetes的自定義資源，遮蔽資料操作的底層細節，減輕使用者心智負擔。使用者可以很方便地通過操作自定義資源物件來管理資料，包括資料同步、資料預熱、清理快取、淘汰歷史資料等，同時也支援定時任務。

  • 基於JuiceFS加速遠端資料訪問。JuiceFS是一款面向雲環境設計的高效能共享檔案系統，其在資料組織管理和訪問效能上進行了大量針對性的優化。基於JuiceFS實現樣本資料快取引擎，能夠提供高效的檔案訪問效能。

  • 充分利用本地儲存，快取加速模型訓練。要能夠充分利用計算叢集本地儲存，比如記憶體和磁碟，來快取熱點樣本資料集，並配合快取親和性排程，在使用者無感知的情況下，智慧地將作業排程到有快取的節點上。這樣就不用反覆訪問遠端儲存，從而加速模型訓練速度，一定程度上也能提升GPU資源的利用率。

  • 統一資料介面，支援多種儲存後端。樣本快取元件要能夠支援多種儲存後端，並且能提供統一的POSIX協議介面，使用者無需在模型開發和訓練階段使用不同的資料訪問介面，降低模型開發成本。同時樣本快取元件也要能夠支援從多個不同的儲存源匯入資料，適配使用者現有的資料儲存狀態。

     

 

面臨的挑戰

然而，在Kubernetes的架構體系中，計算與儲存是分離的，這種架構給上訴目標的實現帶來了些挑戰，主要體現在如下幾點：

  • Kubernetes 排程器是快取無感知的，也就是說kube-scheduler並沒有針對本地快取資料的排程策略，因此模型訓練作業未必能排程到有快取的節點，從而導致快取無法重用。如何實現快取親和性排程，協同編排訓練作業與快取資料，是我們面臨的首要問題。

  • 在資料並行的分散式訓練任務中，單機往往存放不下所有的樣本資料，因此樣本資料是要能夠以分割槽的形式分散快取在各計算節點上。然而，我們知道負責管理自定義資源的控制器（Controller Manager）不一定執行在快取節點上，如何通過自定義控制器來管理分散式的快取資料，也是實現該方案時要考慮的難點問題。

  • 除了前述兩點，如何結合飛槳框架合理地對樣本資料進行分發和預熱，提高本地快取命中率，減少作業訪問遠端資料的次數，從而提高作業執行效率，這還需要進一步的探索。

針對上述問題，我們在開源專案Paddle Operator中提供了樣本快取元件，較好地解決了這些挑戰，下文將詳細闡述我們的解決方案。

 

 

整體設計方案

 

上圖是Paddle Operator的整體架構，其構建在Kubernetes上，包含如下三個主要部分：

1.自定義API資源（Custom Resource）

Paddle Operator定義了三個CRD，使用者可編寫和修改對應的YAML檔案來管理訓練作業和樣本資料集。

  • PaddleJob是飛槳分散式訓練作業的抽象，它將Parameter Server（引數伺服器）和Collective（集合通訊）兩種分散式深度學習架構模式統一到一個CRD中，使用者通過建立PaddleJob可以很方便地在Kubernetes叢集執行分散式訓練作業。

  • SampleSet是樣本資料集的抽象，資料可以來自遠端物件儲存、HDFS或Ceph等分散式檔案系統，並且可以指定快取資料的分割槽數、使用的快取引擎、 多級快取目錄等配置。

  • SampleJob定義了些樣本資料集的管理作業，包括資料同步、資料預熱、清除快取、淘汰歷史舊資料等操作，支援使用者設定各個資料操作命令的引數， 同時還支援以定時任務的方式執行資料管理作業。

 

2.自定義控制器（Controller Manager）

控制器在 Kubernetes 的 Operator 框架中是用來監聽 API 物件的變化（比如建立、修改、刪除等），然後以此來決定實際要執行的具體工作。

  • PaddleJob Controller負責管理PaddleJob的生命週期，比如建立引數伺服器和訓練節點的Pod，並維護工作節點的副本數等。

  • SampleSet Controller負責管理SampleSet的生命週期，其中包括建立 PV/PVC等資源物件、建立快取執行時服務、給快取節點打標籤等工作。

  • SampleJob Controller負責管理SampleJob的生命週期，通過請求快取執行時服務的介面，觸發快取引擎非同步執行資料管理操作，並獲取執行結果。

 

3.快取引擎（Cache Engine）

快取引擎由快取執行時服務（Cache Runtime Server）和JuiceFS儲存外掛（JuiceFS CSI Driver）兩部分組成，提供了樣本資料儲存、快取、管理的功能。

  • Cache Runtime Server負責樣本資料的管理工作，接收來自SampleSet Controller和SampleJob Controller的資料操作請求，並呼叫JuiceFS客戶端完成相關操作執行。

  • JuiceFS CSI Driver是JuiceFS社群提供的CSI外掛，負責樣本資料的儲存與快取工作，將樣本資料快取到叢集本地並將資料掛載進PaddleJob的訓練節點。

*JuiceFS儲存外掛：

http://github.com/juicedata/juicefs-csi-driver

 

難點突破與優化

在上述的整體架構中，Cache Runtime Server是非常重要的一個元件，它由Kubernetes原生的API資源StatefulSet實現，在每個快取節點上都會執行該服務，其承擔了快取資料分割槽管理等工作，也是解決難點問題的突破口。下圖是使用者建立SampleSet後，Paddle Operator對PaddleJob完成快取親和性排程的大概流程。

 

當用戶建立SampleSet後，SampleSet Controller就會根據 SampleSet中的配置創建出PV和PVC。在PV和PVC完成繫結後，SampleSet Controller則會將PVC新增到Runtime Pod模板中，並創建出指定分割槽數的Cache Runtime Server。在Cache Runtime Server成功排程到相應節點後，SampleSet Controller則會對該節點做標記，且標記中還帶有快取的分割槽數。這樣等下次使用者提交PaddleJob時，Paddle Controller會自動地給Paddle Worker Pod新增nodeAffinity和PVC欄位，這樣排程器（Scheduler）就能將Paddle Worker排程到指定的快取分割槽節點上，這即實現了對模型訓練作業的快取親和性排程。

值得一提的是，在該排程方案中，Paddle框架的訓練節點能夠做到與快取分割槽一一對應的，這能夠最大程度上地利用本地快取的優勢。當然，該方案同時也支援對PaddleJob的擴縮容，當PaddleJob的副本數大於SampleSet的分割槽數時（這也是可以調整的），PaddleJob Controller並不會對多出來Paddle Worker做nodeAffinity限制，這些Paddle Worker還可以通過掛載的Volume訪問遠端儲存來獲取樣本資料。

解決了快取親和性排程的問題，我們還面臨著SampleSet/SampleJob Controller如何管理分散式樣本快取資料集，如何合理地做分散式預熱的挑戰。使用JuiceFS CSI儲存外掛可以解決樣本資料儲存與快取的問題，但由於Kubernetes CSI外掛提供的介面有限（只提供了與 Volume 掛載相關的介面），所以需要有額外的資料管理服務駐守在快取節點，即：Cache Runtime Server。

 

上圖是 Cache Runtime Server 內部工作流程示意圖，其中封裝了JuiceFS客戶端的可執行檔案。Runtime Server提供了三種類型的介面，它們的作用分別是：

1.上傳資料操作命令的引數

2.獲取資料操作命令的結果

3.獲取樣本資料集及其快取的狀態。

Server在接收到Controller上傳的命令引數後，會將引數寫到指定路徑，然後會觸發Worker程序非同步地執行相關操作命令，並將命令執行結果寫到結果路徑，然後Controller可以通過呼叫相關介面獲取資料管理作業的執行結果。

資料同步、清理快取、淘汰舊資料這三個操作比較容易實現，而資料預熱操作的實現相對會複雜些。因為，當樣本資料量比較大時，單機儲存無法快取所有資料，這時就要考慮對資料進行分割槽預熱和快取。為了最大化利用本地快取和儲存資源，我們期望對資料預熱的策略要與飛槳框架讀取樣本資料的介面保持一致。

因此，對於WarmupJob我們目前實現了兩種預熱策略:Sequence和Random，分別對應飛槳框架SequenceSampler和DistributedBatchSampler兩個資料取樣API。JuiceFS的Warmup命令支援通過--file引數指定需要預熱的檔案路徑，故將0號Runtime Server作為Master，負責給各個分割槽節點分發待預熱的資料，即可實現根據使用者指定的策略對樣本資料進行分散式預熱的功能。

至此，難點問題基本都得以解決，該方案將快取引擎與飛槳框架緊密結合，充分利用了本地快取來加速雲上飛槳的分散式訓練作業。

 

使用示例

下面我們通過訓練ResNet50模型的例子來簡要說明下如何使用Paddle Operator，您只需要準備兩個YAML配置檔案，即可輕鬆地在Kubernetes叢集上完成複雜的分散式深度學習任務。第一步是編寫SampleSet的YAML檔案，並指定遠端資料來源，如下:

apiVersion: batch.paddlepaddle.org/v1alpha1
kind: SampleSet
metadata:
  name: imagenet
  namespace: paddle-system
spec:
  # 分割槽數，一個Kubernetes節點表示一個分割槽
  partitions: 2
  source:
    uri: bos://paddleflow-public.hkg.bcebos.com/imagenet
  secretRef:
    name: imagenet

 

通過kubectl apply -f命令可以建立該樣本資料集，您還可以檢視樣本資料集及其快取的狀態：

 

等SampleSet的狀態為Ready後，您就可以通過編寫如下的PaddleJob來完成ResNet50模型的訓練了：

apiVersion: batch.paddlepaddle.org/v1
kind: PaddleJob
metadata:
  name: resnet
spec:
  # 指定要使用的 SampleSet
  sampleSetRef:
    name: imagenet
    mountPath: /data
  worker:
    replicas: 2
    template:
      spec:
        containers:
          - name: resnet
            image: registry.baidubce.com/paddle-operator/demo-resnet:v1
            command:
            - python
            args:
            - "-m"
            - "paddle.distributed.launch"
            - "./tools/train.py"
            - "-c"
            - "./config/ResNet50.yaml"
            resources:
              limits:
                nvidia.com/gpu: 2

 

更多的使用文件可以參考Paddle Operator。

 

效能測試

為了驗證Paddle Operator樣本快取加速方案的實際效果，我們選取了常規的ResNet50模型以及ImageNet資料集來進行效能測試，並使用了PaddleClas專案中提供的模型實現程式碼。具體的實驗配置如下：

 

基於以上配置，我們做了兩組實驗。

第一組實驗對比了樣本資料快取前後的效能差異，從而驗證了使用樣本快取元件來加速雲上飛槳訓練作業的必要性。

如上圖，在樣本資料快取到訓練節點前，1機1卡和1機2卡的訓練速度分別為211.26和224.23 images/s；在樣本資料得以快取到本地後，1機1卡和1機2卡的訓練速度分別為383.01和746.76 images/s。可以看出，在計算與儲存分離的Kubernetes叢集裡，由於頻寬有限（本實驗的頻寬為50Mbps），訓練作業的主要效能瓶頸在於遠端樣本資料IO上。頻寬的瓶頸並不能通過調大訓練作業的並行度來解決，並行度越高，算力浪費越為嚴重。

因此，使用樣本快取元件提前將資料預熱到訓練叢集本地，可以大幅加速雲上飛槳訓練作業的執行效率。在本組實驗中，1機1卡訓練效率提升了81.3%，1機2卡的訓練速度提升了233%。

此外，使用樣本快取元件預熱資料後，1機1卡383.01 images/s的訓練速度與直接在宿主機上的測試結果一致，也就是說，快取引擎本身基本上沒有帶來效能損耗。

第二組實驗對比了使用1機1卡時樣本資料預熱前後JuiceFS與BOS FS的效能差異，相比於BOS FS，JuiceFS在訪問遠端小檔案的場景下具有更優的效能表現。

如上圖，在樣本資料預熱到本地前，需要通過CSI外掛訪問遠端物件儲存（如BOS）中的樣本資料，使用BOS FS CSI外掛與JuiceFS CSI外掛的訓練速度分別是69.71和211.26 images/s；在樣本資料得以快取在本地後，使用BOS FS CSI外掛與JuiceFS CSI外掛的訓練速度分別是381.98和382.43 images/s，效能基本沒有差異。

由此可以看出，在訪問遠端小檔案的場景下，JuiceFS相比BOS FS有近3倍的效能提升。兩者的效能差異可能與JuiceFS的檔案儲存格式和BOS FS的實現有關，這有待進一步驗證。除了效能上的考量，JuiceFS提供的檔案系統掛載服務更加穩定，功能完善且易用，這也是優先選擇JuiceFS作為底層快取引擎的重要原因。

 

總結及展望

樣本資料的準備工作是深度學習pipeline中的重要一環，本文總結了在雲上管理樣本資料並加速其訪問效率所面臨的挑戰，在文章的第3和第4節我們給出了這些問題的詳細解決方案。進一步地，我們對Paddle Operator的樣本快取方案進行了效能測試，實驗結果表明該方案能大幅加速雲上飛槳的模型訓練作業。現該方案已在Paddle Operator專案中開源，歡迎大家來體驗並使用。

在後續的工作中，我們將繼續完善樣本資料管理功能，提供更多的資料管理策略，並進一步加強Operator與飛槳框架的協同優化工作。除了離線訓練的場景，我們也在探索Operator支援線上學習場景的解決方案，如您對這些工作也很感興趣，歡迎參與進來一起討論、開發。

點選進入獲得更多技術資訊~~

 
「」
「」