Tensorflow網路傳輸效能分析
0. 寫在前面
tensorflow分散式訓練時，grpc的 慢 一直都被很多人所詬病。在早期的版本中，由於實現的一些原因，的確存在一些效能問題(可以參見這個 ofollow,noindex">issue )。
但隨著專案的迭代，現在效能如何，就有些莫衷一是了。這裡通過對兩個專案master分支程式碼的一些測試，希望能探討下這些問題。
1. 直觀的看傳輸速率
這裡先用 一個測試程式 測試下tensor在兩個機器中的傳輸速率。測試使用的兩臺機器配置的都是萬兆乙太網的網絡卡：
[work@host benchtools]$ ethtool eth0 Settings for eth0: ... Speed: 10000Mb/s ...
在兩臺機器上分別跑測試程式的worker和ps：
[host1] python tensor_transfer_throughput.py --ps_hosts=host1:12222 --worker_hosts=host2:12222 --job=ps --task=0 [host2] python tensor_transfer_throughput.py --ps_hosts=host1:12222 --worker_hosts=host2:12222 --data_mb=100 --job=worker --task=0 --iters=100
測試程式乾的事情很簡單：在ps和worker上各建立一個相同大小的variable, 然後worker反覆將自己的variable assign給ps。在上面的測試中，我們將variable的大小設定為100M，傳輸次數為100。
測試結果在worker執行結束後可以看到：
[host2] python tensor_transfer_throughput.py --ps_hosts=host1:12222 --worker_hosts=host2:12222 --data_mb=100 --job=worker --task=0 --iters=100 .... transfer rate: 173.488801 MB/s
利用ifstat工具也可以看到網路的傳輸效能：
[hosts1]$ ./ifstat eth0eth1 KB/s inKB/s outKB/s inKB/s out 191.95176435.60.000.00 206.18170675.30.000.00 222.45220156.50.000.00 162.84169024.80.000.00 224.44211070.70.000.00
可以看到兩種測試的througput效果差不多。理論上來說ifstat可能會比worker的輸出稍微大一點，因為grpc要為每次傳輸額外新增一些header資訊。但和100MB的資料相比，應該可以忽略不計。
但無論是哪個結果，離理論值的1.25GBps(10Gbps)差距仍舊非常大。所以初步來看，網絡卡的利用率是比較低的。
2. 單獨測試grpc
為了驗證問題是不是出在grpc這裡，我利用 另一個測試程式 ，來測試grpc本身的傳輸效率。
程式不太複雜，要點包括：
- client和server端的功能要簡單，儘量減少額外操作所帶來的時間開銷：client只負責無腦傳送，server端也要直接丟棄收到的資料。
- 直接利用grpc的ByteBuffer，從而避免掉在傳送和接收時的memcpy。這點和tensorflow傳送tensor的流程也是一致的。
- server端可以建立多個completion queue, 從而可以指定多個worker執行緒。
- client利用非同步介面。可以指定傳輸併發度，也可以允許grpc建立多個channel。
- 可以指定傳送資料和響應資料塊的大小。
然後將程式部署到兩臺機器上開始測試。client每次向server傳送100M資料，共傳送1000條：
[host1] ./grpc_raw --job_type=server --server_threads=1 --message_size=10 [host2] ./grpc_raw --job_type=client --job_type=client --target_ip=host1 --total_message=1000 --message_size=104857600
利用ifstat看結果：
[work@host2 benchtools]$ ./ifstat eth0eth1 KB/s inKB/s outKB/s inKB/s out 162.05198529.90.000.00 128.67150799.50.000.00 196.09203136.00.000.00 169.20192864.80.000.00 130.67146532.70.000.00
可以看到和測tensor傳輸時類似，也是170MBps左右，離1.25GBps的理論值也差距較大。
3. 為什麼慢
為了進一步確定問題，我用 iperf 工具對網路的throughput做了單獨的測試：
[host1] ./iperf3 -s -i 5 [host2] ./iperf3 -c host1 -i 5 -t 1000
測試結果如下：
[host2]$ ./iperf3 -c host1 -i 5 -t 1000 ... [5]0.00-5.00sec983 MBytes1.65 Gbits/sec315452.49 MBytes [5]5.00-10.00sec839 MBytes1.41 Gbits/sec35645889 KBytes [5]10.00-15.00sec830 MBytes1.39 Gbits/sec35863954 KBytes ...
可以看到大概也就是1.4Gbps(175MBps)左右，和 grpc的測試結果差不多 。
為什麼會這樣呢？事實上，當提高socket數後，結果就會大大改觀，總的傳輸速率會達到9.3 Gbps左右，從而和理論值接近：
[host2]$ ./iperf3 -c host1 -i 5 -t 1000 -P 8 ... [5]40.00-45.00sec621 MBytes1.04 Gbits/sec99362.06 MBytes .... [ 19]40.00-45.00sec206 MBytes346 Mbits/sec92290.5 KBytes [SUM]40.00-45.00sec5.43 GBytes9.33 Gbits/sec33646
這裡我們可以看到的一個結論是： 單個socket可能(遠遠)無法用滿網絡卡的頻寬 。
那麼如果把grpc的socket數增加如何？遺憾的是，目前grpc還不支援這樣的特性。在grpc裡，通訊是用 channel 來進行抽象的。 哪怕你在兩個機器間建立多個channel, 他們在底層也是會共享socket的 。
4. 單個socket用不滿網絡卡？
當我通過測試得出這個結論時，我內心也是無法接受的。我嘗試了
- 手動調整擁塞視窗(事實上也沒有必要，因為TCP會自發的增大它；穩定後的擁塞視窗大小，也沒有達到Linux的上限)。
- 關閉Nagel演算法
傳輸速率仍然沒有變化。
後來在組裡boss的建議下，我換了兩臺機器做測試。發現對於不同的機器組合，單socket的傳輸效能是不同的。 也存在一些機器，他們的單socket效能是可以達到網絡卡理論上限的 。
對於這一問題，現在懷疑可能和網路佈局以及中間的交換機有關係。但具體的根源究竟是什麼，還無從得知。
5. 繼續測試
在我換了 單socket可以打滿頻寬 的兩臺機器後，我把1和2中的實驗使用相同的引數重新做了一遍。結論如下：
- grpc在單server單client的前提下，網絡卡傳輸的利用率還是非常高的。在我的實驗中大概能到9Gbps左右，比iperf的結果稍遜一點，目測也就是5%左右。這可能和grpc在資料傳輸時的一些資料結構的分配、處理有關，但整理來說grpc效能已經比較可觀了。
- 對於傳輸tensor的測試而言，傳輸速率大概能到5Gbps左右，是裸grpc的一多半。
這裡有兩個問題：
1. 為什麼 傳輸tensor 的吞吐要低於 裸的grpc傳輸 ，問題在哪裡？
2. 在我們最開始的兩個實驗中，由於單socket極限頻寬較低，這二者的傳輸效率類似。為什麼提高單socket的極限頻寬後，二者開始體現出差別來？
其實這兩個問題並不難解釋：
- 在傳輸tensor時，除了有效的傳輸資料外，還有master驅動worker執行、序列化、反序列化、資料assign等其他操作。而我們測試看到的throughput，是把這些操作都當成 有效傳輸 而平均化後的一個結果。
- 兩個機器間頻寬越高，額外操作的佔比就越大，對總throughput的影響就越大。
6. 驗證假設
為了驗證我們的假設，我們需要知道 tensorflow在傳輸tensor時，真正用於資料傳輸的時間是多少 ，從而可以根據資料量大致推算一下傳輸時的網路頻寬。
可以先用timeline看一下每一步所有op的耗時，以及RecvTensor這個op的耗時。
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) run_metadata = tf.RunMetadata() sess.run(add_op.op, options=run_options, run_metadata=run_metadata) trace = timeline.Timeline(step_stats=run_metadata.step_stats) trace_file = open('timeline.ctf.json', 'w') trace_file.write(trace.generate_chrome_trace_format())
結果(dur表示op的耗時，單位為us)：
{
"name": "RecvTensor",
...
"dur": 183311
},
....
{
"name": "Assign",
...
"dur": 19925
}
耗時主要在RecvTensor和Assign上，總耗時有200ms左右。對於100M資料而言，這個耗時也和觀察到的5Gbps的吞吐大致吻合。
但我們仍舊不能知道真正在 傳輸 的時候頻寬能不能有效的利用。timeline所能給出的最小粒度就是op，而"RecvTensor"這個op，我們可以看到耗時是180ms左右。這比grpc的傳輸吞吐還是要低出不少來的。
我們知道，在Tensorflow中，一個RecvTensor是要分成如下幾個步驟的：
1. RecvOp的AsyncCompute，通過rendezvous介面，最終呼叫到grpc這一層。
2. 發起RecvTensor的請求，包括獲取一個grpc_remote_worker的handle，以及準備RecvTensorRequest的protobuf，然後建立和rpc call相關的資料結構
3. 呼叫grpc的API，將資料推到網路引擎，傳送資料。
4. server端從rendezvous_manager中獲取tensor, 並且和其他的meta資訊包裝成ByteBuffer返回給客戶端。
5. 客戶端將收到的ByteBuffer反序列化成Tensor。
所以整個傳輸過程的慢，可能會慢在以下幾個地方：
1. 做準備工作時，一些執行緒排程或者加鎖操作帶來開銷。
2. server的序列化費時間。
3. grpc的網路引擎就是慢，比如說引入額外的資料拷貝之類的，導致ByteBuffer傳輸很慢。
4. client的反序列化費時間。
第三點其實不太可能，因為我們已經拿裸的grpc+ByteBuffer做過測試，其頻寬利用率是比較高的。當然，我們也可以在Tensorflow中通過更細緻的metrics來驗證下這一點。
因為沒法用timeline，只能通過改tensorflow程式碼來測試。為此，我簡單修改了 tensorflow的程式碼 ，來觀察傳輸和客戶端處理的耗時。測試的結論如下：
- 對於100M的tensor，grpc的傳輸的時間大概在100ms左右。大概的資料傳輸率應該有9Gbps左右，比較高效。
- server資料序列化的時間佔比很小。這點tensorflow的確做過專門處理：tensor的記憶體是作為ByteBuffer直接傳輸的，很大程度避免了記憶體拷貝。
- 客戶端的訊息反序列化會佔用一定時間，大概佔到了RecvTensor的1/4多一些。主要原因是grpc ByteBuffer中的Tensor資料不滿足Tensor的記憶體佈局要求，所以必須得通過記憶體拷貝來一次重新整理。
7. 擴充套件性
前面分析了grpc在傳輸效率方面的效能，接下來看下有關擴充套件性方面的問題。
首先明確下，當我們討論擴充套件性時，應該從如下兩個角度來衡量：
- server端未到網絡卡的瓶頸時，通過增加client，server端的throughput能隨著client的個數線性增加。
- server端達到網絡卡瓶頸後，隨著client個數的增加， server端的吞吐最好基本不會下降，而client端的latency則會線性的增加。
這裡的測試細節就不再展開了。通過對這兩個方面的測試，我發現grpc在這兩個層面基本表現也比較良好。
8. 總結
測試的結論大致有如下幾個：
- 在開發分散式程式時，機房間機器的拓撲結構需要注意下，可能會影響單socket的極限頻寬。如果存在此類問題，多socket的rpc是一個可能可行的方案。
- grpc在大資料包的傳輸上，頻寬利用率和擴充套件性都還不錯。
- 對於tensorflow的RecvTensor，收到資料後的後續處理，會佔據一部分計算資源，對總體的網絡卡頻寬會存在影響。
幾個需要繼續調研的方面有：
- grpc在高併發處理小資料包上latency表現如何，可以調研一下。對與tensorflow而言，這其實不太重要。但對於latency敏感的線上服務而言，還是非常重要的。
- 在tensor的send方這邊， tensor table 是用一個非常粗粒度的互斥鎖保護的，在RecvTensor請求較多時候懷疑可能會成為瓶頸(比如很多個worker的分散式訓練)。這點需要拿大的訓練場景測試一下。