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