撰文 | 袁進輝
在之前的文章《對抗軟件系統復雜性：恰當分層，不多不少》討論分布式深度學習框架的網絡傳輸需求時，我們劃分了幾個抽象層次，比較底層的一個抽象層次是點到點傳輸（point-to-point）。
在本文中，我們討論一下，一個最理想的點到點通信庫應該是什么樣？如果現在還沒有這樣的庫，我們何不一起做一個這方面的開源項目？
1
什么是點到點通信？
什么是點到點通信？維基百科上的定義：In telecommunications, a point-to-point connection refers to a communications connection between two communication endpoints or nodes. 簡言之，就是一對一的傳輸，只有一個發送方和只有一個接收方。
點到點傳輸為什么重要？因為點到點傳輸是用來構建上層任何復雜傳輸模式的基本單元。
譬如分布式深度學習訓練中常用的ring all-reduce或者tree all-reduce就是基于最基本的點對點傳輸功能組合拼裝起來的；點對點傳輸庫還可以經過封裝變成對用戶更加友好易用的接口，譬如各種遠程過程調用（remote procedure call, RPC)的庫也是基于點到點傳輸實現的。
提前說明，本文只介紹 cpu to cpu 的傳輸，實際項目中更多的是gpu to gpu 的傳輸，會更復雜一點，其中最簡單的是 GPUDirect RDMA，和 CPU 上的 RDMA 編程一致，但是僅支持數據中心級別GPU，否則應該是 gpu-cpu-net-cpu-gpu 的模式。
2
什么是點到點通信庫？
事實上，操作系統層面提供的網絡編程API就是點到點的，譬如套接字（Socket), RDMA 底層庫本身就是點到點的API。
為什么還需要一個庫？主要目的是在不損失性能的情況下更易用，更通用，隱藏多樣性的底層編程接口，對上層應用暴露一致的API，實現一致的編程體驗，譬如無論是TCP/IP的套接字，還是RDMA網絡，它們本身的編程接口不一樣，但我們希望上層應用程序編寫的程序是通過一致的接口來調用底層的傳輸能力。
ZeroMQ (https://zeromq.org/) 是一個應用范圍很廣的點到點通信庫（當然它也支持了一些多方通信的功能），使得Socket編程更簡單，性能也很高，使得編寫高性能網絡應用程序更簡單。
3
為什么要造一個新的點到點通信庫？
已有的點對點傳輸庫，都有各種各樣的問題。ZeroMQ不支持RDMA，在深度學習場景下不適合。
OneFlow中有一個模塊叫CommNet，同時支持Socket和RDMA，接口和實現都令人滿意，不過不夠獨立，與OneFlow系統耦合比較深，不方便被其它項目使用。Facebook為PyTorch項目搭建了TensorPipe，支持套接字和RDMA，從接口定義到實現都非常接近我的想象，但仍有不滿意的地方，希望你讀完整篇文章之后會理解這一點。
4
理想中的點到點通信庫有哪些特征？
從底層傳輸機制、上層應用需求以及已有點到點通信庫的經驗中可以提煉出這三點：
- 編程簡單，易于滿足各種上層應用，包括封裝成RPC使用，在OneFlow這樣的深度學習框架中使用，甚至被用在HPC和深度學習中常見的集群通信原語中（all-reduce, broadcast等）；
- 高性能：表現為零拷貝、低延時、高吞吐；
- 底層支持TCP/IP套接字和RDMA傳輸。
為了滿足這些需求，這個通信庫在技術上要實現這四點：
- 面向消息的編程模型；
- 非阻塞的接口；
- 零拷貝；
- 對小消息和大消息都友好。
下面我們更詳細的討論一下，為什么這些比較關鍵。
5
面向消息的編程模型
無論是Socket還是ZeroMQ 都把點到點通信的通路抽象成一個管道（pipe），發送方通過如下的send函數向管道中寫數據，接收方通過recv函數從管道中讀取數據（在函數輸入參數里我們特意省略了發送方和接收方的endpoint地址，譬如Socket的文件描述符）。
int64_t send(void* in_buf, int64_t size);
int64_t recv(void* out_buf, int64_t size);
通信庫并不關心傳輸的具體內容，統一視為字節序列（也就是序列化和反序列化是上層應用的責任），通信庫只關心傳輸量的大小（字節數size）。假設通信雙方預先已知傳輸量大小，即size，那么發送方會預先分配size大小的in_buf，并把將要發送的內容放到in_buf里，并調用send函數；接收方同樣會預先分配size大小的out_buf，并調用recv函數接收數據。注意，這里我們假設輸入參數中的緩沖區in_buf和out_buf都是用戶管理的。
為簡化用戶的編程，接口應該是面向完整消息的，而不是面向字節流的。也就是不管傳輸多大的數據，send和recv函數返回時應“一次性”把任務完成，這樣用戶每次有傳輸需求，只需要調用一次函數，而不關心底層是不是把數據分成多段傳輸。
ZeroMQ符合這個語義，Socket編程中的阻塞模式也符合這個語義，在阻塞模式的Socket編程中，直到數據傳輸完畢函數才會返回。但是，非阻塞模式的Socket編程不符合這個語義，在操作系統無法滿足一次性把數據傳輸完成時會先完成一部分并返回真正傳輸的字節數，用戶可能需要在后面再次調用send和recv進行傳輸。
6
非阻塞的調用模式
以Socket編程為例，在阻塞模式下，只有當數據真正完成傳輸時函數才會返回，但傳輸時間決定于傳輸量和傳輸帶寬，可能需要等待較長一段時間，在等待傳輸完成的這段時間內，調用send和recv的線程只能休眠，不能處理其它事情，為了提高系統的吞吐量，可能得啟動和管理很多線程。
而非阻塞模式下，在調用send和recv時，假如系統不能一次性完成傳輸任務，也會把用戶空間的一段數據拷貝到內核空間（盡管這個拷貝執行時間非常短，不過我們需要注意它的存在），并返回這次傳輸的數據量，提示用戶“并沒有完全傳完，請在合適的時間再次調用繼續傳輸”。
以上兩種模式對用戶來說都不夠友好，最好的方式是，傳輸庫作為面向上層需求的一個服務，上層應用把任務交給傳輸庫就立刻返回，當傳輸完成時再通知上層應用即可。為此目的，API 可以調整成：
void send(void* in_buf, int64_t size, Callback done);
void recv(void* out_buf, int64_t size, Callback done);
也就是每個函數都增加一個輸入的回調函數，send和recv會立刻返回，當數據傳輸全部完成時就執行用戶自定義的回調函數done。
當然，有了這個非阻塞的編程接口是非常容易做一點點工作就把阻塞模式支持起來的。
7
零拷貝
在上面的討論中，我們假設了 in_buf 和 out_buf 的內存是被上層應用管理的，譬如在調用send之前分配了in_buf，send函數返回后，in_buf就可以釋放了。但是，請注意，非阻塞模式下，即使send返回了，數據也可能還沒有發送過去，因此通信庫必須在send函數內部申請一段內存，并把in_buf的數據拷貝到這段由通信庫管理的內存上，這樣通信庫可以一直使用這段由自己管理的內存，直到真正把數據傳輸過去再釋放。
但上述方案有一些缺點，譬如每一次傳輸數據時，通信庫都要額外分配與用戶傳進來的緩沖區同等大小的內存，分配內存需要花費時間，把數據從應用程序的緩沖區拷貝到通信庫管理的緩沖區上也需要時間，還增加了內存使用量。
更理想的方式是：雖然in_buf是上層應用分配的，但在調用send函數那一刻，該緩沖區的內存的所有權就轉移給了通信庫，在send函數返回后并不能立即釋放in_buf，因為send發送過程中直接使用in_buf，當發送真正完成時，才能在回調函數done里釋放in_buf的內存。
同樣，即使out_buf是通信庫分配的，在recv輸入的回調函數done執行那一刻，out_buf的所有權也被轉移給上層應用，而不是把out_buf再拷貝到一個應用管理的緩沖區上去。
上文我們討論了一些比較通用的需求，下面我們需要把一些細節補全。
8
通信兩端如何協商傳輸量？
此前，我們假設發送方和接收方都已經知道了傳輸數據量的大小，也就是參數size的數值，這個假設不太實際，但還不算離譜。
首先，每次有傳輸需求，雖然傳輸量不盡相同，發送方是一定知道傳輸量的大小的，而接收方不一定知道。其次，每次傳輸真正的數據之前，發送方可以先把size數值發過去，這樣接收方就知道真實要傳輸的數據大小了，就可以提前把內存分配好。
需要注意的是，雙方在傳輸真正的數據之前需要先溝通傳輸量的大小，也就是size的值，這個size的值也是通過send/recv來傳送的，這個size值的大小是固定的，雙方不需要溝通，這有點類似一個bootstrap的過程。
假設從A向B要發送一次數據，我們都要至少調用3次send/recv對來完成，如下圖所示：
第一次由A到B，分別調用send和recv，A把size傳送給B，B在收到之后根據size為out_buf分配內存 （alloc) 。當B分配好內存之后，第二次通信是從B到A，B向A發送一個please start的信號，這個信號很短且是固定長度，不需要A和B雙方協商分配內存。當B收到please start的信號后，第三次通信就可以開始了，從A到B傳輸真正的數據。
上述方案有什么問題呢？
首先，每次通信都需要調用send和recv三次，即使本來傳輸的數據size就很小，也必須承受三次通信的延遲。
其次，send和recv必須配對使用，發送方和接收方必須按相同的節奏來調用才行，譬如發送方調用了send，接收方沒有調用recv，并不能成功，或者發送方調用了兩次send,但接收方只調用了一次recv，第二次也會失敗。但是，什么時候有傳輸需求是由發送方決定的，接收方是被動的，它并不知道什么時候需要調用recv，上面的規范使用起來并不好。
怎么辦呢？
對第一個問題，可以對短消息和長數據設計兩種傳輸模式，對于長度小于某個閾值的數據傳輸不需要雙方協商就直接發送，發送方可以假定接收方一定能成功接收，而且發送方也假設接收方一定提前調用了recv來和send配對。傳輸長數據時必須通過如上三次調用才能完成。
對第二個問題，通信庫總是提前為不知何時從何地發送過來的短消息需求做好準備，也就是提前準備了固定數量的recv調用。這一點不太好理解，熟悉Grpc異步編程或RDMA編程的朋友應該對這個比較熟悉，每個通信進程在啟動時就提前準備若干PostRecvRequest，而且每和別處的send配對一次，就消費掉一個RecvRequest，并及時補充一個新的RecvRequest。
最后，可能有的朋友對傳輸長數據時為什么接收方需要提前知道size大小不解。這主要是為了提前分配好內存，確保數據傳輸可以成功，并且在傳輸過程中不需要再分配內存，也可以實現零拷貝。
否則，假設不提前分配好內存，就需要在傳輸過程中不斷根據實際需求去分配內存，有可能分配不成功，就需要因為內存資源不夠的原因打斷傳輸過程，當然，也實現不了零拷貝。
9
API設計
有了以上討論，看上去只需要send/recv接口就能滿足所有需求了，它可以滿足傳輸短消息和長數據的需求。
不過，除了這個API，發送方和接收方還有一些復雜的邏輯來處理，接收方總要提前準備好一些RecvRequest，以及傳輸長數據時，發送方和接收方都需要來回協商幾次。從設計底層庫的角度來說，我們希望盡可能簡化用戶使用時的負擔，把和需求無關的細節隱藏起來。這樣看，只有send/recv還不夠。
對于短消息，我們希望發送方可以直接發送，通信庫來保證在接收方有準備好recv調用，這個recv不需要用戶來顯式調用，也就是，在短消息場景下，recv這個API是不必要的。用戶只需要為通信庫提供一個收到短消息之后的回調函數即可，每當接收方收到一個短消息，就調用相應的回調函數來處理這個短消息即可。
如果業務需要多種類型的短消息，那么可以對短消息分類，并為每種不同的短消息類型提供相應類型的回調函數即可。
對于長數據傳輸的第二次和第三次通信，接收方需要調用一次send和一次recv，發送方需要調用一次send，但這些調用細節應該對用戶透明。所有這些操作可以由通信庫底層來完成，用戶編程接口可以合并成一個單邊操作read由接收方調用，而發送方的應用程序不需要做任何操作，當然數據傳輸完成之后需要調用用戶指定的callback函數來處理接收到的數據。
也就是點到點通信庫的最小API可以是如下的形式：
void send(void* in_buf, int64_t size);
void read(void* out_buf, int64_t size, Callback done);
注意，在實際實現中，read接口實際上還需要一個標志發送端數據位置的token，通過這個token才能遠程讀取到正確的數據。
10
OneFlow CommNet的設計
CommNet 滿足 OneFlow 的功能需要兩個最重要的抽象，Eager Message 和 RMA Read。目前的實現中，Massage用于傳輸ActorMsg，RMA Read用于傳輸regst的實際內容。
Eager Message的設定：
- 點對點，每個消息對應一個發送端、對應一個接收端
- 發送端發送一個消息，接收端在未來接收到對應消息
- 發送端直接向接收端發送消息，無需事先協商
- 接收端無條件接受消息
- 發送端可以假設發送一定會成功，接收端未來一定可以收到消息
- 接收端通過輪詢或者注冊回調的方式處理消息
- 有連接或者無連接抽象，無連接抽象中，發送端使用接收端標識作為發送參數，有連接抽象中，發送端需事先與接收端建立連接，并使用連接標識作為發送參數
- 同一個線程向同一個接收端或者同一個連接發送的不同消息，需保證接收端接收到的順序與發送的順序一致
- 消息本身為固定大小或者動態大小的數據塊，無需關心上層協議
- 一般為處理小塊數據而設計
- 關鍵指標一般是延遲與吞吐率
Remote Memory Access (RMA) Read的設定：
- 點對點，每次操作對應一個本地端與遠端
- 本地端發起操作，操作的結果為將遠端地址空間里面的一段數據讀取掉本地內存空間
- 遠端需要事先生成訪問令牌（token），本地端必須通過令牌才能訪問遠端在生成令牌時注冊的地址范圍內的數據。操作發起前，本地端和遠端需通過其他任何方式交換訪問令牌
- 一次訪問本地端可以讀取訪問令牌對應的范圍內的任意范圍數據，同一位置的數據可以被讀取任意次數
- 讀取過程中，遠端不需要參與
- 本地端通過輪詢或者注冊回調的方式處理傳輸完成事件
- 本地端認為遠端內存一直可用
- 一般為處理大塊數據而設計
- 關鍵指標一般是帶寬/吞吐率
11
討論
為什么要把第二次和第三次通信抽象成一個read單邊操作，為什么不讓發送方顯式調用send或者write呢？這個調用是沒有必要的，它的執行時機應該是被接收方決定的，而且應該是自動執行的，沒有必要暴露給上層應用接口。
實際上，熟悉RDMA編程的朋友，應該很熟悉在RDMA里提供了send，沒有recv接口，同時提供了Write和Read這樣的單邊操作，我們上述討論表明作為點到點通信庫只需要Read這一種單邊操作就可以了。參考RDMA編程接口的設計，可以進一步驗證我們提議的編程API的合理性。
事實上，研究MPI的學者中已經有人提出了類似的接口設計，例如為解決現有MPI接口的不足，一批研究下一代MPI的學者就在一篇題為《Towards millions of communicating threads（https://snir.cs.illinois.edu/listed/C101.pdf）》的文章中提出了類似的設計，在這篇文章中，短消息的傳輸需求被命名為eager-protocol，而長數據的傳輸需要雙方協商，被稱之為rendezvous protocol （沒錯，TensorFlow的分布式設計中也有這個概念），特別感謝閆嘉昆告訴我這篇文章。
以上的討論都是從上層應用的需求角度出發來設計API，當然，API的設計也需要考慮底層實現，譬如面向Socket的epoll編程模型就和RDMA編程模型不同，我們的通信庫需要支持這些不同的傳輸機制，API的設計也要兼顧使用不同傳輸機制時編程的難度。
我們了解到，RDMA本身已經提供了send和read的單邊操作，使用RDMA來支持本文提議的API應該比較自然，不過當我們在未來的文章展開進一步的細節內容時，還是能發現一些復雜之處，譬如RDMA的傳輸需要鎖頁內存，對于變長數據傳輸，每次都在線分配鎖頁內存的開銷比較高，怎么解決這個問題并不簡單。epoll則沒有完全對應的概念，那么使用epoll實現這個通信庫，就可能需要額外更多的工作。
在后續文章中，我們會進一步討論使用RDMA和epoll實現這個通信庫的方法。
題圖源自TheDigitalArtist, Pixabay
歡迎下載體驗OneFlow新一代開源深度學習框架：
https://github.com/Oneflow-Inc/oneflow
