最理想的點到點通訊庫究竟是怎樣的?

語言: CN / TW / HK

撰文 | 袁進輝

在之前的文章《對抗軟體系統複雜性:恰當分層,不多不少》討論分散式深度學習框架的網路傳輸需求時,我們劃分了幾個抽象層次,比較底層的一個抽象層次是點到點傳輸(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