【网络】Unified Communication X (UCX)|统一多通信后端

UCX 的意义

随着DPU的普及、各类DSA芯片的广泛使用，如何在这之上抽象出统一的内存访问语义和统一的通信方式确实是一个很有意思、也很有价值的课题。类似的：TensorPipe、华为内部的UMDK等等

UCX 通信接口简介

UCX 的全称是 Unified Communication X。正如它名字所展示的，UCX 旨在提供一个统一的抽象通信接口，能够适配任何通信设备，并支持各种应用的需求。

下图是 UCX 官方提供的架构图：

可以看到，UCX 整体分为两层：上层的 UCP 接口和底层的 UCT 接口。

底层的 UCT 适配了各种通信设备：从单机的共享内存，到常用的 TCP Socket，以及数据中心常用的 RDMA 协议，甚至新兴的 GPU 上的通信，都有很好的支持。

上层的 UCP 则是在 UCT 不同设备的基础上，封装了更抽象的通信接口，以方便应用使用。具体来说有以下几类：

Active Message：最底层的接口，提供类似 RPC 的语义。每条 Active Message 会触发接收端进行一些操作。
RMA / Atomic：是对远程直接内存访问（RDMA）的抽象。通信双方可以直接读写远端的内存，但是需要有额外的内存注册过程。
Tag Matching：常用于高性能计算 MPI 程序中。每条消息都会附带一个 64 位整数作为 tag，接收方每次可以指定接收哪种 tag 的消息。
Stream：对字节流（TCP）的抽象。

一般来说，和底层通信设备模型最匹配的接口具有最高的性能，其它不匹配的接口都会有一次软件转换过程。另一方面，同一种 UCP 接口发送不同大小的消息可能也会使用不同的 UCT 方法。例如在 RDMA 网络中，由于内存注册也有不小的开销，因此对于小消息来说，拷贝到预注册好的缓冲区再发送的性能更高。这些策略默认是由 UCX 自己决定的，用户也可以通过设置环境变量的方式手动修改。

在我们的系统中，使用了 UCP Tag 接口并基于此实现了轻量级的 RPC。在 RPC 场景下，Tag 可以用于区分不同上下文的消息：每个链接双方首先随机生成一个 tag 作为请求的标识，对于每次请求再随机生成一个 tag 作为回复的标识。此外 Tag 接口还支持 IO Vector，即将不连续的多个内存段合并成一个消息发送。这个特性可以用来将用户提供的数据缓冲区和 RPC 请求打包在一起，一定程度上避免数据拷贝。

UCX 编程模型简介

UCX 采用了以异步 IO 为核心的编程模型。其中 UCP 层定义的核心对象有以下四种：

Context：全局资源的上下文，管理所有通信设备。一般每个进程创建一个即可。
Worker：任务的管理调度中心，以轮询方式执行任务。一般每个线程创建一个，会映射为网卡上的一个队列。
Listener：类似 TCP Listener，用来在 worker 之间创建连接。
Endpoint：表示一个已经建立的连接。在此之上提供了各种类型的通信接口。

它们之间的所属关系如下图所示：

建立连接

UCX 中双方首先要建立连接，拿到一个 Endpoint 之后才能进行通信。建立连接一般要通过 Listener，过程和 TCP 比较类似：

通信双方 A/B 首先建立各自的 Context 和 Worker，其中一方 A 在 Worker 上创建 Listener 监听连接请求，Listener 的地址会绑定到本机的一个端口上。用户需要通过某种方法将这个地址传递给另一方 B。B 拿到地址后在 Worker 上发起 connect 操作，此时 A 会收到新连接请求，它可以选择接受或拒绝。如果接受则需要在 Worker 上 accept 这个请求，将其转换为 Endpoint。之后 B 会收到 A 的回复，connect 操作完成，返回一个 Endpoint。此后双方就可以通过这对 Endpoint 进行通信了。

内存注册

对于常规的通信接口，用户可以直接在 Endpoint 上发起请求。但对于 RMA（远程内存访问）操作，需要被访问的一方首先在自己的 Context 上注册内存，同时指定访问权限，获得一个 Mem handle。然后将这个本地 handle 转化为其他节点可以访问的一串 token，称为 remote key（rkey）。最后想办法把 rkey 传给远端。远端拿着这个 rkey 进行远程内存访问操作。

异步任务处理（重点）

为了发挥最高的性能，整个 UCX 通信接口是全异步的。所谓异步指的是 IO 操作的执行不会阻塞当前线程，一次操作的发起和完成是独立的两个步骤。如此一来 CPU 就可以同时发起很多 IO 请求，并且在它们执行的过程中可以做别的事情。

不过接下来问题来了：程序如何知道一个异步任务是否完成了？常见的有两种做法：主动轮询，被动通知。前者还是需要占用 CPU 资源，所以一般都采用通知机制。在 C 这种传统过程式语言中，异步完成的通知一般通过 回调函数（callback）实现：每次发起异步操作时，用户都需要传入一个函数指针作为参数。当任务完成时，后台的运行时框架会调用这个函数来通知用户。下面是 UCX 中一个异步接收接口的定义：

ucs_status_ptr_t ucp_tag_recv_nb (
  ucp_worker_h worker,
  void ∗ buffer,
  size_t count,
  ucp_datatype_t datatype,
  ucp_tag_t tag,
  ucp_tag_t tag_mask,
  ucp_tag_recv_callback_t cb  // <-- 回调函数
);

// 回调函数接口的定义
typedef void(∗ ucp_tag_recv_callback_t) (
  void ∗request, 
  ucs_status_t status,        // 执行结果，错误码
  ucp_tag_recv_info_t ∗info   // 更多信息，如收到的消息长度等
);

这个接口的语义是：发起一个异步 Tag-Matching 接收操作，并立即返回。当真的收到 tag 匹配的消息时，UCX 后台会处理这个消息，将其放到用户提供的 buffer 中，最后调用用户传入的 callback，通知用户任务的执行结果。

这里有一个很重要的问题是：上面提到的“后台处理”到底是什么时候执行的？答案是 UCX 并不会自己创建后台线程去执行它们，所有异步任务的后续处理和回调都是在 worker.progress() 函数中，也就是用户主动向 worker 轮询的过程中完成的。这个函数的语义是：“看看你手头要处理的事情，有哪些是能做的？尽力去推动一下，做完的通知我。” 换句话说，Worker 正在处理的所有任务组成了一个状态机，progress 函数的作用就是用新事件推动整个状态机的演进。后面我们会看到，对应到 async Rust 世界中，所有异步 IO 任务组成了最基础的 Future，worker 对应 Runtime，而 progress 及其中的回调函数则充当了 Reactor 的角色。

回到传统的 C 语言，在这里异步 IO 的最大难点是编程复杂性：多个并发任务在同一个线程上交替执行，只能通过回调函数来描述下一步做什么，会使得原本连续的执行逻辑被打散到多个回调函数中。本来局部变量就可以维护的状态，到这里就需要额外的结构体来在多个回调函数之间传递。随着异步操作数量的增加，代码的维护难度将会迅速上升。下面的伪代码展示了在 UCX 中如何通过异步回调函数来实现最简单的 echo 服务：

// 由于 C 语言语法的限制，这段代码需要从下往上读

// 这里存放所有需要跨越函数的状态变量
struct CallbackContext {
  ucp_endpoint_h ep;
  void *buf;
} ctx;

void send_cb(void ∗request, ucs_status_t status) {
  //【4】发送完毕
  ucp_request_free(request);
  exit(0);
}

void recv_cb(void ∗request, ucs_status_t status, ucp_tag_recv_info_t ∗info) {
  //【3】收到消息，发起发送请求
  ucp_tag_send_nb(ctx->ep, ctx->buf, info->length, ..., send_cb);
  ucp_request_free(request);
}

int main() {
  // 省略 UCX 初始化部分
  //【0】初始化任务状态
  ctx->ep = ep;
  ctx->buf = malloc(0x1000);
  //【1】发起异步接收请求
  ucp_tag_recv_nb(worker, ctx->buf, 0x1000, ..., recv_cb);
  //【2】不断轮询，驱动后续任务完成
  while(true) {
    ucp_worker_progress(worker);
  }
}

作为对比，假如 UCX 提供的是同步接口，那么同样的逻辑只需要以下几行就够了：

int main() {
  // 省略 UCX 初始化部分
  void *buf = malloc(0x1000);
  int len;
  ucp_tag_recv(worker, buf, 0x1000, &len, ...);
  ucp_tag_send(ep, buf, len, ...);
  return 0;
}

面对传统异步编程带来的“回调地狱”，主流编程语言经过了十几年的持续探索，终于殊途同归，纷纷引入了控制异步的终极解决方案—— async-await 协程。它的杀手锏就是能让开发者用同步的风格编写异步的逻辑。经过我们的封装过后，在 Rust 中用 async 协程编写同样的逻辑是长这样的：

async fn main() {
  // 省略 UCX 初始化部分
  let mut buf = vec![0u8; 0x1000];
  let len = worker.tag_recv(&mut buf, ...).await.unwrap();
  ep.tag_send(&buf[..len], ...).await.unwrap();
}

（摘抄自：《async Rust 封装 UCX 通信库》 https://zhuanlan.zhihu.com/p/397199431）

mellonx文档《Unified Communication X (UCX)》：https://www.hpcadvisorycouncil.com/events/2019/APAC-AI-HPC/uploads/2018/07/UCX_Accelerate_HPC_Application.pdf

posted on 2022-10-04 01:21 bdy 阅读(988) 评论(0) 编辑收藏举报

刷新页面返回顶部

UESTC