最理想的点到点通信库究竟是怎样的？

 

撰文 | 袁进辉

 

在之前的文章《对抗软件系统复杂性：恰当分层，不多不少》讨论分布式深度学习框架的网络传输需求时，我们划分了几个抽象层次，比较底层的一个抽象层次是点到点传输（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