我們在協助某AI客戶排查一個UFS文件存儲的性能case時發現，其使用的Pytorch訓練IO性能和硬件的IO能力有很大的差距（後面內容有具體性能對比數據）。
讓我們感到困惑的是：UFS文件存儲，我們使用fio自測可以達到單實例最低10Gbps帶寬、IOPS也可達到2w以上。該AI客戶在高IOPS要求的AI單機小模型訓練場景下，或者之前使用MXNet、TensorFlow框架時，IO都能跑到UFS理論性能，甚至在大型分佈式訓練場景中，UFS也可以完全勝任。
於是我們開啓了和客戶的一次深度聯合排查。
初步嘗試優化
1、調整參數
基於上述情況，首先考慮是不是使用Pytorch的姿勢不對？參考網上提到經驗，客戶調整batch_size、Dataloader等參數。
Batch_size
默認batch_size爲256，根據內存和顯存配置嘗試更改batch_size大小，讓一次讀取數據更多，發現實際對效率沒有提升。通過分析是由於batch_size設置與數據讀取邏輯沒有直接關係，IO始終會保留單隊列與後端交互，不會降低網絡交互上的整體延時（因爲用的是UFS文件存儲，後面會講到爲什麼用）。
Pytorch Dataloader
Pytorch框架dataloader的worker負責數據的讀取和加載、分配。通過batch_sampler將batch數據分配給對應的worker，由worker從磁盤讀取數據並加載數據到內存，dataloader從內存中讀取相應batch做迭代訓練。這裏嘗試調整了worker_num參數爲CPU核數或倍數，發現提升有限，反而內存和CPU的開銷提升了不少，整體加重了訓練設備的負擔，通過 worker加載數據時的網絡開銷並不會降低，與本地SSD盤差距依然存在。
這個也不難理解，後面用strace排查的時候，看到CPU更多的時候在等待。
所以：從目前信息來看，調整Pytorch框架參數對性能幾乎沒有影響。
2、嘗試不同存儲產品
在客戶調整參數的同時，我們也使用了三種存儲做驗證，來看這裏是否存在性能差異、差異到底有多大。在三種存儲產品上放上同樣的數據集：
1、 單張平均大小20KB的小圖片，總量2w張。
2、 以目錄樹方式存到三種存儲下的相同路徑，使用Pytorch常用的標準讀圖接口CV2和PIL
測試結果，如下圖：
| 讀取方式 | SSHFS | 本地SSD | UFS文件存儲 |
| CV2 | 319.94張/s | 554.73張/s | 72.41張/s |
| PIL（image_open） | 435.93張/s | 3507.93張/s | 115.78張/s |
注：SSHFS基於X86物理機（32核/64G/480G SSD*6 raid10）搭建，網絡25Gbps
結論：通過對存儲性能實測， UFS文件存儲較本地盤、單機SSHFS性能差距較大。
爲什麼會選用這兩種存儲（SSHFS和本地SSD）做UFS性能對比？
當前主流存儲產品的選型上分爲兩類：自建SSHFS/NFS或採用第三方NAS服務（類似UFS產品），個別場景中也會將需要的數據下載到本地SSD盤做訓練。傳統SSD本地盤擁有極低的IO延時，一個IO請求處理基本會在us級別完成，針對越小的文件，IO性能越明顯。受限於單臺物理機配置，無法擴容，數據基本 “即用即棄”。而數據是否安全也只能依賴磁盤的穩定性，一旦發生故障，數據恢復難度大。但是鑑於本地盤的優勢，一般也會用作一些較小模型的訓練，單次訓練任務在較短時間即可完成，即使硬件故障或者數據丟失導致訓練中斷，對業務影響通常較小。
用戶通常會使用SSD物理機自建SSHFS/NFS共享文件存儲，數據IO會通過以太網絡，較本地盤網絡上的開銷從us級到ms級，但基本可以滿足大部分業務需求。但用戶需要在日常使用中同時維護硬件和軟件的穩定性，並且單臺物理機有存儲上限，如果部署多節點或分佈式文件系統也會導致更大運維精力投入。
我們把前面結論放到一起看：
1、 隱形結論：Tensorflow、Mxnet框架無問題。
2、 調整Pytorch框架參數對性能幾乎沒有影響。
3、Pytorch+UFS的場景下， UFS文件存儲較本地SSD盤、單機SSHFS性能差距大。
結合以上幾點信息並與用戶確認後的明確結論：UFS結合非Pytorch框架使用沒有性能瓶頸， Pytorch框架下用本地SSD盤沒有性能瓶頸，用SSHFS性能可接受。那原因就很明顯了，就是Pytorch+UFS文件存儲這個組合存在IO性能問題。
深入排查優化
看到這裏，大家可能會有個疑問：是不是不用UFS，用本地盤就解決了？
答案是不行，原因是訓練所需的數據總量很大，很容易超過了單機的物理介質容量，另外也出於數據安全考慮，存放單機有丟失風險，而UFS是三副本的分佈式存儲系統，並且UFS可以提供更彈性的IO性能。
根據以上的信息快速排查3個結論，基本上可以判斷出：Pytorch在讀UFS數據過程中，文件讀取邏輯或者UFS存儲IO耗時導致。於是我們通過strace觀察Pytorch讀取數據整體流程：
通過strace發現，CV2方式讀取UFS裏的文件（NFSV4協議）有很多次SEEK動作，即便是單個小文件的讀取也會“分片”讀取，從而導致了多次不必要的IO讀取動作，而最耗時的則是網絡，從而導致整體耗時成倍增長。這也是符合我們的猜測。
簡單介紹一下NFS協議特點：NAS所有的IO都需要經過以太網，一般局域網內延時在1ms以內。以NFS數據交互爲例，通過圖中可以看出，針對一次完整的小文件IO操作將涉及元數據查詢、數據傳輸等至少5次網絡交互，每次交互都會涉及到client與server集羣的一個TTL，其實這樣的交互邏輯會存在一個問題，當單文件越小、數量越大時則延時問題將越明顯，IO過程中有過多的時間消耗在網絡交互，這也是NAS類存儲在小文件場景下面臨的經典問題。
對於UFS的架構而言，爲了達到更高擴展性、更便利的維護性、更高的容災能力，採用接入層、索引層和數據層的分層架構模式，一次IO請求會先經過接入層做負載均衡，client端再訪問後端UFS索引層獲取到具體文件信息，最後訪問數據層獲取實際文件，對於KB級別的小文件，實際在網絡上的耗時比單機版NFS/SSHFS會更高。
從Pytorch框架下兩種讀圖接口來看：CV2讀取文件會“分片”進行，而PIL雖然不會“分片”讀取，但是基於UFS分佈式架構，一次IO會經過接入、索引、數據層，網絡耗時也佔比很高。我們存儲同事也實際測試過這2種方法的性能差異：通過strace發現，相比OpenCV的方式，PIL的數據讀取邏輯效率相對高一些。
優化方向一：如何降低與UFS交互頻次，從而降低整體存儲網絡延時
1、CV2：對單個文件而言，“分片讀取”變“一次讀取”
通過對Pytorch框架接口和模塊的調研，如果使用 OpenCV方式讀取文件可以用2個方法， cv2.imread和cv2.imdecode。
默認一般會用cv2.imread方式，讀取一個文件時會產生9次lseek和11次read，而對於圖片小文件來說多次lseek和read是沒有必要的。cv2.imdecode可以解決這個問題，它通過一次性將數據加載進內存，後續的圖片操作需要的IO轉化爲內存訪問即可。
兩者的在系統調用上的對比如下圖：
我們通過使用cv2.imdecode方式替換客戶默認使用的cv2.imread方式，單個文件的總操作耗時從12ms下降到6ms。但是內存無法cache住過大的數據集，不具備任意規模數據集下的訓練，但是整體讀取性能還是提升明顯。使用cv2版本的benchmark對一個小數據集進行加載測試後的各場景耗時如下(延遲的非線性下降是因爲其中包含GPU計算時間):
| 本地SSD | UFS imread方式 | UFS imdecode方式 |
| 95s | 150s | 100s左右 |
2、PIL：優化dataloader元數據性能，緩存文件句柄
通過PIL方式讀取單張圖片的方式，Pytorch處理的平均延遲爲7ms(不含IO時間)，單張圖片讀取(含IO和元數據耗時)平均延遲爲5-6ms，此性能水平還有優化空間。
由於訓練過程會進行很多個epoch的迭代，而每次迭代都會進行數據的讀取，這部分操作從多次訓練任務上來看是重複的，如果在訓練時由本地內存做一些緩存策略，對性能應該有提升。但直接緩存數據在集羣規模上升之後肯定是不現實的，我們初步只緩存各個訓練文件的句柄信息，以降低元數據訪問開銷。
我們修改了Pytorch的dataloader實現，通過本地內存cache住訓練需要使用的文件句柄，可以避免每次都嘗試做open操作。測試後發現1w張圖片通過100次迭代訓練後發現，單次迭代的耗時已經基本和本地SSD持平。但是當數據集過大，內存同樣無法cache住所有元數據，所以使用場景相對有限，依然不具備在大規模數據集下的訓練伸縮性。
| 存儲類型 | Real time（s） | User time（s） | Sys time（s） |
| 本地SSD | 25.10 | 475.89 | 3.66 |
| UFS | 25.90 | 491.62 | 3.73 |
3、UFS server端元數據預加載
以上client端的優化效果比較明顯，但是客戶業務側需要更改少量訓練代碼，最主要是client端無法滿足較大數據量的緩存，應用場景有限，我們繼續從server端優化，儘量降低整個鏈路上的交互頻次。
正常IO請求通過負載均衡到達索引層時，會先經過索引接入server，然後到索引數據server。考慮到訓練場景具有目錄訪問的空間局部性，我們決定增強元數據預取的功能。通過客戶請求的文件，引入該文件及相應目錄下所有文件的元數據，並預取到索引接入server，後續的請求將命中緩存，從而減少與索引數據server的交互，在IO請求到達索引層的第一步即可獲取到對應元數據，從而降低從索引數據server進行查詢的開銷。
經過這次優化之後，元數據操作的延遲較最初能夠下降一倍以上，在客戶端不做更改的情況下，讀取小文件性能已達到本地SSD盤的50%。看來單單優化server端還是無法滿足預期，通過執行Pytorch的benchmark程序，我們得到UFS和本地SSD盤在整個數據讀取耗時。
| 存儲類型 | 1w張耗時（s） | 2w張耗時（s） |
| 本地SSD | 48.86 | 97.06 |
| UFS | 97.98 | 195.82 |
此時很容易想到一個問題：非Pytorch框架在使用UFS做訓練集存儲時，爲什麼使用中沒有遇到IO性能瓶頸？
通過調研其他框架的邏輯發現：無論是MXNet的rec文件，Caffe的LMDB，還是TensorFlow的npy文件，都是在訓練前將大量圖片小文件轉化爲特定的數據集格式，所以使用UFS在存儲網絡交互更少，相對Pytorch直接讀取目錄小文件的方式，避免了大部分網絡上的耗時。這個區別在優化時給了我們很大的啓示，將目錄樹級別小文件轉化成一個特定的數據集存儲，在讀取數據做訓練時將IO發揮出最大性能優勢。
優化方向二：目錄級內的小文件轉換爲數據集，最大程度降到IO網絡耗時
基於其他訓練框架數據集的共性功能，我們UFS存儲團隊趕緊開工，幾天開發了針對Pytorch框架下的數據集轉換工具，將小文件數據集轉化爲UFS大文件數據集並對各個小文件信息建立索引記錄到index文件，通過index文件中索引偏移量可隨機讀取文件，而整個index文件在訓練任務啓動時一次性加載到本地內存，這樣就將大量小文件場景下的頻繁訪問元數據的開銷完全去除了，只剩下數據IO的開銷。該工具後續也可直接應用於其他AI類客戶的訓練業務。
工具的使用很簡單，只涉及到兩步：
- 使用UFS自研工具將Pytorch數據集以目錄形式存儲的小文件轉化爲一個大文件存儲到UFS上，生成date.ufs和index.ufs。
- 使用我方提供Folder類替換pytorch原有代碼中的torchvision.datasets.ImageFolder數據加載模塊（即替換數據集讀取方法），從而使用UFS上的大文件進行文件的隨機讀取。只需更改3行代碼即可。
- 20行：新增from my_dataloader import *
- 205行：train_dataset = datasets.ImageFolder改爲train_dataset = MyImageFolder
- 224行：datasets.ImageFolder改爲MyImageFolder
通過github上Pytorch測試demo對imagenet數據集進行5、10、20小時模擬訓練，分別讀取不同存儲中的數據，具體看下IO對整體訓練速度的影響。（數據單位：完成的epoch的個數）
| 數據讀取方式 | 第一次（5小時） | 第二次（10小時） | 第三次（20小時） |
| 本地SSD盤 | 5.226 | 10.384 | 20.873 |
| UFS目錄小文件 | 1.238 | 3.180 | 6.394 |
| UFS數據集 | 5.46 | 10.739 | 21.784 |
| SSHFS目錄小文件 | 3.721 | 7.398 | 14.797 |
測試條件：GPU服務器：P40*4物理機，48核256G，數據盤800G*6 SATA SSD RAID10SSHFS：X86物理機32核/64G，數據盤480G*6 SATA SSD RAID10Demo：https://github.com/pytorch/examples/tree/master/imagenet數據集：總大小148GB、圖片文件數量120w以上
通過實際結果可以看出：UFS數據集方式效率已經達到甚至超過本地SSD磁盤的效果。而UFS數據集轉化方式，客戶端內存中只有少量目錄結構元數據緩存，在100TB數據的體量下，元數據小於10MB，可以滿足任意數據規模，對於客戶業務上的硬件使用無影響。
UFS產品
針對Pytorch小文件訓練場景，UFS通過多次優化，吞吐性能已得到極大提升，並且在後續產品規劃中，我們也會結合現有RDMA網絡、SPDK等存儲相關技術進行持續優化。詳細請點擊訪問鏈接：https://docs.ucloud.cn/storage_cdn/ufs/overview。