Python与高性能IPC：利用RDMA（远程直接数据存取）实现跨节点零拷贝通信

Python 与高性能 IPC：利用 RDMA 实现跨节点零拷贝通信

大家好，今天我们来探讨一个重要的课题：如何利用 RDMA (Remote Direct Memory Access) 技术，结合 Python，实现跨节点的高性能零拷贝通信。在分布式计算、高性能计算以及大数据处理等领域，节点间的数据传输效率至关重要。传统的 TCP/IP 通信方式在数据传输过程中会涉及多次内核态和用户态的数据拷贝，带来显著的性能损耗。RDMA 技术允许网络适配器直接访问远程节点的内存，绕过操作系统内核，从而实现零拷贝的数据传输，显著降低延迟，提高带宽利用率。

1. RDMA 的基本原理

RDMA 是一种网络技术，允许计算机中的网卡直接访问另一台计算机的内存，而无需经过操作系统内核的参与。这意味着数据可以直接从一个应用程序的内存空间传输到另一个应用程序的内存空间，避免了 CPU 的参与和内存拷贝，从而显著降低了延迟和 CPU 负载。

关键概念：

• Queue Pair (QP): RDMA 通信的基本单元。每个 QP 由一个发送队列 (SQ) 和一个接收队列 (RQ) 组成。
• Memory Region (MR): RDMA 允许注册一段内存区域，使其可以被远程节点访问。
• Completion Queue (CQ): 用于通知操作完成的队列。
• Work Queue Entry (WQE): 描述一个 RDMA 操作的指令。

RDMA 的工作流程：

1. 注册内存区域 (MR): 应用程序首先需要注册一块内存区域，以便远程节点可以通过 RDMA 直接访问。注册时，会得到一个内存密钥 (rkey) 和一个虚拟地址 (VA)。
2. 发布信息： 本地节点需要将 MR 的 VA 和 rkey 信息传递给远程节点。这可以通过带外通信（例如 TCP/IP）进行。
3. 发起 RDMA 操作： 远程节点使用收到的 VA 和 rkey 信息，通过 RDMA 动词 (verbs) 发起读或写操作。
4. 数据传输： 网卡直接将数据从远程节点的内存复制到本地节点的内存，或者反之，无需 CPU 的参与。
5. 完成通知： RDMA 操作完成后，网卡会将一个完成事件放入 Completion Queue (CQ) 中，应用程序可以通过 CQ 知道操作已经完成。

RDMA 的优势：

• 零拷贝： 数据直接从一个应用程序的内存空间传输到另一个应用程序的内存空间，无需经过内核的参与和内存拷贝。
• 低延迟： 减少了 CPU 的参与，避免了上下文切换，从而降低了延迟。
• 高带宽： 允许网卡直接访问内存，可以充分利用网络带宽。
• CPU 卸载： 将数据传输的任务从 CPU 卸载到网卡，从而释放 CPU 资源。

2. RDMA 协议类型

常见的 RDMA 协议有三种：

• InfiniBand (IB): 一种高性能互连网络，专门为 RDMA 设计。
• RoCE (RDMA over Converged Ethernet): 在以太网上实现 RDMA，有两种版本：RoCEv1 (基于以太网链路层) 和 RoCEv2 (基于 UDP/IP)。
• iWARP (Internet Wide Area RDMA Protocol): 一种基于 TCP/IP 的 RDMA 协议。

协议	传输层	优点	缺点
InfiniBand	InfiniBand	非常低的延迟和高带宽，专为高性能计算设计	需要专门的 InfiniBand 硬件，成本较高
RoCEv1	以太网链路层	可以使用现有的以太网基础设施，延迟较低	RoCEv1 不支持路由，需要在同一个子网内
RoCEv2	UDP/IP	可以跨子网路由，兼容性好	延迟相对 RoCEv1 较高，受 UDP/IP 协议的限制
iWARP	TCP/IP	可以跨广域网使用，兼容性最好	延迟相对较高，受 TCP/IP 协议的限制

3. Python 与 RDMA：Pyverbs 库

虽然 Python 本身没有内置 RDMA 支持，但我们可以通过第三方库 pyverbs 来实现 Python 与 RDMA 的交互。pyverbs 是一个 Python 绑定库，它封装了 Linux 内核提供的 RDMA verbs API。

安装 pyverbs:

pip install pyverbs

基本用法示例：

以下代码示例展示了如何使用 pyverbs 创建一个简单的 RDMA 连接，并进行数据传输。

import pyverbs.device as pvdevice
import pyverbs.verbs as pvverbs
import socket
import struct
import threading

class RDMAClient:
    def __init__(self, ip, port):
        self.ip = ip
        self.port = port
        self.pd = None  # Protection Domain
        self.cq = None  # Completion Queue
        self.qp = None  # Queue Pair
        self.mr = None  # Memory Region
        self.sock = None # Socket for exchanging QPN and MR details

    def connect(self):
        # 1. Discover RDMA devices
        dev_list = pvdevice.get_device_list()
        if not dev_list:
            raise Exception("No RDMA devices found")
        self.device = dev_list[0]  # Use the first device

        # 2. Create context
        self.ctx = pvverbs.Context(self.device)

        # 3. Allocate protection domain
        self.pd = pvverbs.Pd(self.ctx)

        # 4. Create completion queue
        self.cq = pvverbs.Cq(self.ctx, 10, None, None, 0) # cqe=10

        # 5. Create queue pair
        qp_init_attr = pvverbs.QPInitAttr(qp_type=pvverbs.IBV_QPT_RC,
                                        sq_sig_all=0,
                                        send_cq=self.cq,
                                        recv_cq=self.cq,
                                        cap=pvverbs.QPCap(max_send_wr=10, max_recv_wr=10, max_send_sge=1, max_recv_sge=1))
        self.qp = pvverbs.Qp(self.pd, qp_init_attr)

        # 6. Establish TCP connection to exchange QP number and MR details
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.ip, self.port))
        print(f"Connected to server at {self.ip}:{self.port}")

        # 7. Get QP number and MR details from server
        data = self.sock.recv(1024) # Adjust buffer size if needed
        server_qpn, server_va, server_rkey = struct.unpack("!III", data)
        self.server_qpn = server_qpn
        self.server_va = server_va
        self.server_rkey = server_rkey
        print(f"Received server QPN: {self.server_qpn}, VA: {self.server_va}, RKEY: {self.server_rkey}")

        # 8. Modify QP state to establish RDMA connection
        self.modify_qp_to_init()
        self.modify_qp_to_rtr(self.server_qpn)
        self.modify_qp_to_rts()

    def modify_qp_to_init(self):
      attr = pvverbs.QpAttr(qp_state=pvverbs.IBV_QPS_INIT,
                           pkey_index=0,
                           port_num=1, # Assuming using port 1
                           qp_access_flags=pvverbs.IBV_ACCESS_REMOTE_WRITE | pvverbs.IBV_ACCESS_REMOTE_READ)
      self.qp.modify_qp(attr, pvverbs.IBV_QP_STATE | pvverbs.IBV_QP_PKEY_INDEX | pvverbs.IBV_QP_PORT | pvverbs.IBV_QP_ACCESS_FLAGS)

    def modify_qp_to_rtr(self, dest_qp_num):
      attr = pvverbs.QpAttr(qp_state=pvverbs.IBV_QPS_RTR,
                           dest_qp_num=dest_qp_num,
                           rq_psn=0,
                           max_dest_rd_atomic=1,
                           min_rnr_timer=12,
                           path_mtu=pvverbs.IBV_MTU_2048, # Adjust MTU as needed
                           dlid=1, # Assuming DLID is 1 (needs to be discovered properly)
                           sl=0,  # Service Level
                           src_qp_num=self.qp.qp_num)
      self.qp.modify_qp(attr, pvverbs.IBV_QP_STATE | pvverbs.IBV_QP_AV | pvverbs.IBV_QP_PATH_MTU | pvverbs.IBV_QP_DEST_QPN | pvverbs.IBV_QP_RQ_PSN | pvverbs.IBV_QP_MAX_DEST_RD_ATOMIC | pvverbs.IBV_QP_MIN_RNR_TIMER)

    def modify_qp_to_rts(self):
      attr = pvverbs.QpAttr(qp_state=pvverbs.IBV_QPS_RTS,
                           sq_psn=0,
                           max_rd_atomic=1,
                           retry_cnt=7,
                           rnr_retry=7,
                           timeout=14)
      self.qp.modify_qp(attr, pvverbs.IBV_QP_STATE | pvverbs.IBV_QP_SQ_PSN | pvverbs.IBV_QP_MAX_RD_ATOMIC | pvverbs.IBV_QP_RETRY_CNT | pvverbs.IBV_QP_RNR_RETRY | pvverbs.IBV_QP_TIMEOUT)

    def register_memory(self, buffer):
        self.mr = pvverbs.Mr(self.pd, buffer, pvverbs.IBV_ACCESS_LOCAL_WRITE | pvverbs.IBV_ACCESS_REMOTE_READ | pvverbs.IBV_ACCESS_REMOTE_WRITE)
        return self.mr.buf, self.mr.lkey, self.mr.rkey

    def write(self, buffer, remote_va, remote_rkey):
        mr_buf, lkey, _ = self.register_memory(buffer)
        sge = pvverbs.Sge(addr=mr_buf, length=len(buffer), lkey=lkey)
        wr = pvverbs.Wr(wr_id=1, sg_list=[sge], opcode=pvverbs.IBV_WR_RDMA_WRITE,
                        send_flags=pvverbs.IBV_SEND_SIGNALED,
                        rdma=pvverbs.Rdma(remote_addr=remote_va, rkey=remote_rkey))
        self.qp.post_send(wr)
        wc = self.poll_completion()
        if wc.status != pvverbs.IBV_WC_STATUS_SUCCESS:
            raise Exception(f"RDMA write failed with status: {wc.status}")

    def read(self, buffer, remote_va, remote_rkey):
        mr_buf, lkey, _ = self.register_memory(buffer)
        sge = pvverbs.Sge(addr=mr_buf, length=len(buffer), lkey=lkey)
        wr = pvverbs.Wr(wr_id=1, sg_list=[sge], opcode=pvverbs.IBV_WR_RDMA_READ,
                        send_flags=pvverbs.IBV_SEND_SIGNALED,
                        rdma=pvverbs.Rdma(remote_addr=remote_va, rkey=remote_rkey))
        self.qp.post_send(wr)
        wc = self.poll_completion()
        if wc.status != pvverbs.IBV_WC_STATUS_SUCCESS:
            raise Exception(f"RDMA read failed with status: {wc.status}")

    def poll_completion(self):
        while True:
            poll_result = self.cq.poll(1)
            if poll_result:
                return poll_result[0]

    def disconnect(self):
        if self.sock:
            self.sock.close()
        if self.mr:
            self.mr.dereg_mr()
        if self.qp:
            self.qp.destroy()
        if self.cq:
            self.cq.destroy()
        if self.pd:
            self.pd.dealloc()
        if self.ctx:
            self.ctx.close()

class RDMAServer:
    def __init__(self, ip, port):
        self.ip = ip
        self.port = port
        self.pd = None  # Protection Domain
        self.cq = None  # Completion Queue
        self.qp = None  # Queue Pair
        self.mr = None  # Memory Region
        self.sock = None # Socket for accepting client connection

    def start(self):
        # 1. Discover RDMA devices
        dev_list = pvdevice.get_device_list()
        if not dev_list:
            raise Exception("No RDMA devices found")
        self.device = dev_list[0]  # Use the first device

        # 2. Create context
        self.ctx = pvverbs.Context(self.device)

        # 3. Allocate protection domain
        self.pd = pvverbs.Pd(self.ctx)

        # 4. Create completion queue
        self.cq = pvverbs.Cq(self.ctx, 10, None, None, 0) # cqe=10

        # 5. Create queue pair
        qp_init_attr = pvverbs.QPInitAttr(qp_type=pvverbs.IBV_QPT_RC,
                                        sq_sig_all=0,
                                        send_cq=self.cq,
                                        recv_cq=self.cq,
                                        cap=pvverbs.QPCap(max_send_wr=10, max_recv_wr=10, max_send_sge=1, max_recv_sge=1))
        self.qp = pvverbs.Qp(self.pd, qp_init_attr)

        # 6. Set up TCP socket to listen for client connections
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind((self.ip, self.port))
        self.sock.listen(1)
        print(f"Server listening on {self.ip}:{self.port}")

        # 7. Accept client connection
        self.conn, addr = self.sock.accept()
        print(f"Accepted connection from {addr}")

        # 8. Modify QP state to establish RDMA connection
        self.modify_qp_to_init()
        self.modify_qp_to_rtr()
        self.modify_qp_to_rts()

        # 9. Send QP number and MR details to client
        self.send_qp_and_mr_info()

    def modify_qp_to_init(self):
      attr = pvverbs.QpAttr(qp_state=pvverbs.IBV_QPS_INIT,
                           pkey_index=0,
                           port_num=1, # Assuming using port 1
                           qp_access_flags=pvverbs.IBV_ACCESS_REMOTE_WRITE | pvverbs.IBV_ACCESS_REMOTE_READ)
      self.qp.modify_qp(attr, pvverbs.IBV_QP_STATE | pvverbs.IBV_QP_PKEY_INDEX | pvverbs.IBV_QP_PORT | pvverbs.IBV_QP_ACCESS_FLAGS)

    def modify_qp_to_rtr(self):
      attr = pvverbs.QpAttr(qp_state=pvverbs.IBV_QPS_RTR,
                           dest_qp_num=100, # Dummy value - will be overwritten by client's QPN
                           rq_psn=0,
                           max_dest_rd_atomic=1,
                           min_rnr_timer=12,
                           path_mtu=pvverbs.IBV_MTU_2048, # Adjust MTU as needed
                           dlid=1, # Assuming DLID is 1 (needs to be discovered properly)
                           sl=0,  # Service Level
                           src_qp_num=self.qp.qp_num)
      self.qp.modify_qp(attr, pvverbs.IBV_QP_STATE | pvverbs.IBV_QP_AV | pvverbs.IBV_QP_PATH_MTU | pvverbs.IBV_QP_DEST_QPN | pvverbs.IBV_QP_RQ_PSN | pvverbs.IBV_QP_MAX_DEST_RD_ATOMIC | pvverbs.IBV_QP_MIN_RNR_TIMER)

    def modify_qp_to_rts(self):
      attr = pvverbs.QpAttr(qp_state=pvverbs.IBV_QPS_RTS,
                           sq_psn=0,
                           max_rd_atomic=1,
                           retry_cnt=7,
                           rnr_retry=7,
                           timeout=14)
      self.qp.modify_qp(attr, pvverbs.IBV_QP_STATE | pvverbs.IBV_QP_SQ_PSN | pvverbs.IBV_QP_MAX_RD_ATOMIC | pvverbs.IBV_QP_RETRY_CNT | pvverbs.IBV_QP_RNR_RETRY | pvverbs.IBV_QP_TIMEOUT)

    def register_memory(self, buffer):
        self.mr = pvverbs.Mr(self.pd, buffer, pvverbs.IBV_ACCESS_LOCAL_WRITE | pvverbs.IBV_ACCESS_REMOTE_READ | pvverbs.IBV_ACCESS_REMOTE_WRITE)
        return self.mr.buf, self.mr.lkey, self.mr.rkey

    def send_qp_and_mr_info(self):
        # Server sends its QP number, VA, and RKEY to the client
        va, _, rkey = self.register_memory(bytearray(1024)) # Allocate a buffer
        qpn = self.qp.qp_num
        data = struct.pack("!III", qpn, va, rkey)
        self.conn.sendall(data)
        print(f"Sent QPN: {qpn}, VA: {va}, RKEY: {rkey} to client")

    def disconnect(self):
        if self.conn:
            self.conn.close()
        if self.sock:
            self.sock.close()
        if self.mr:
            self.mr.dereg_mr()
        if self.qp:
            self.qp.destroy()
        if self.cq:
            self.cq.destroy()
        if self.pd:
            self.pd.dealloc()
        if self.ctx:
            self.ctx.close()

def main():
    server_ip = "192.168.1.100"  # Replace with server IP
    server_port = 12345
    client_ip = "192.168.1.101" # Replace with client IP
    client_port = 12345

    # Server setup
    server = RDMAServer(server_ip, server_port)
    server_thread = threading.Thread(target=server.start)
    server_thread.daemon = True  # Allow the main thread to exit even if the server thread is running
    server_thread.start()

    # Client setup
    client = RDMAClient(server_ip, client_port) # Client connects to server IP
    client.connect()

    # Example: Client writes to server's memory
    message = b"Hello from RDMA client!"
    client.write(message, client.server_va, client.server_rkey) # Write to server's registered memory
    print("Client wrote message to server's memory.")

    # Example: Client reads from server's memory
    read_buffer = bytearray(len(message))
    client.read(read_buffer, client.server_va, client.server_rkey)
    print(f"Client read message from server: {read_buffer.decode()}")

    # Cleanup
    client.disconnect()
    server.disconnect()

if __name__ == "__main__":
    main()

代码解释：

1. 设备发现和上下文创建: pvdevice.get_device_list() 获取 RDMA 设备列表，然后创建上下文 pvverbs.Context。
2. 保护域 (Protection Domain): pvverbs.Pd 用于管理内存访问权限。
3. 完成队列 (Completion Queue): pvverbs.Cq 用于异步操作完成的通知。
4. 队列对 (Queue Pair): pvverbs.Qp 是 RDMA 通信的基本单元。pvverbs.QPInitAttr 用于初始化 QP 的属性，例如 QP 类型、最大发送/接收工作请求数量等。pvverbs.IBV_QPT_RC 指定 QP 类型为可靠连接 (Reliable Connected)。
5. 内存注册 (Memory Registration): pvverbs.Mr 用于注册内存区域，使其可以被远程节点访问。pvverbs.IBV_ACCESS_LOCAL_WRITE | pvverbs.IBV_ACCESS_REMOTE_READ | pvverbs.IBV_ACCESS_REMOTE_WRITE 指定内存区域的访问权限。
6. RDMA 写操作: pvverbs.Sge (Scatter/Gather Element) 描述要传输的数据。pvverbs.Wr (Work Request) 描述一个 RDMA 操作。pvverbs.IBV_WR_RDMA_WRITE 指定操作类型为 RDMA 写。pvverbs.Rdma 描述远程内存地址和密钥。qp.post_send() 将工作请求放入发送队列。
7. RDMA 读操作: 类似于 RDMA 写操作，但 opcode 为 pvverbs.IBV_WR_RDMA_READ。
8. 完成事件轮询: cq.poll() 用于轮询完成队列，检查操作是否完成。
9. TCP 连接建立: 使用socket建立TCP连接，用于交换客户端和服务端的QPN（Queue Pair Number），VA（Virtual Address）和RKEY（Remote Key）。
10. QP状态转换: QP需要经过INIT, RTR (Ready To Receive), RTS (Ready To Send) 三个状态才能进行RDMA通信。 modify_qp_to_init(), modify_qp_to_rtr(), modify_qp_to_rts() 分别用于将QP状态转换为这三个状态。 其中RTR状态需要指定目标QP的QPN，DLID等信息。

注意：

• 此示例代码仅为演示 RDMA 的基本用法，实际应用中需要根据具体场景进行优化和调整。
• 需要配置 RDMA 网络环境，例如安装 InfiniBand 驱动程序，配置 IP 地址等。
• 需要确保客户端和服务器端都运行在支持 RDMA 的硬件上，并且网络配置正确。
• DLID (Destination Local Identifier) 需要根据实际网络环境进行配置。可以使用 ibstat 命令查看。
• 错误处理和资源释放非常重要，需要仔细检查代码，确保在出现错误时能够正确释放资源。

4. RDMA 在 Python 中的应用场景

RDMA 技术可以广泛应用于各种需要高性能数据传输的场景。以下是一些典型的应用场景：

• 分布式数据库： RDMA 可以用于在数据库节点之间快速复制数据，提高数据同步的效率。
• 机器学习： RDMA 可以用于在训练集群中的节点之间传输大规模数据集，加速模型训练。
• 高性能计算： RDMA 可以用于在计算节点之间传输中间结果，提高计算效率。
• 存储系统： RDMA 可以用于在存储节点和计算节点之间传输数据，提高存储系统的性能。
• 金融交易： RDMA 可以用于在交易服务器之间快速传输交易数据，降低交易延迟。

5. 优化 RDMA 性能的策略

为了充分利用 RDMA 的优势，需要采取一些优化策略：

• 选择合适的 RDMA 协议： 根据网络环境和应用需求选择合适的 RDMA 协议。例如，如果需要在同一个子网内进行通信，可以选择 RoCEv1；如果需要跨子网通信，可以选择 RoCEv2 或 iWARP。
• 调整 MTU (Maximum Transmission Unit)： 选择合适的 MTU 可以提高数据传输效率。通常情况下，选择较大的 MTU 可以减少网络开销。
• 使用内存池： 避免频繁地分配和释放内存，可以使用内存池来管理内存，提高内存利用率。
• 批量处理： 将多个 RDMA 操作合并成一个批量操作，可以减少 CPU 的参与和网络开销。
• 重叠计算和通信： 在进行数据传输的同时进行计算，可以提高整体效率。
• 避免内存拷贝： 尽量使用零拷贝技术，避免不必要的内存拷贝。
• 优化数据布局： 将数据按照网络传输的顺序进行布局，可以提高数据传输效率。

6. Python RDMA 开发的挑战

使用 Python 进行 RDMA 开发也面临一些挑战：

• 学习曲线： RDMA 技术相对复杂，需要理解其基本原理和 API。
• 调试难度： RDMA 程序的调试相对困难，需要使用专门的工具和技术。
• 库的成熟度： pyverbs 虽然提供了基本的 RDMA 功能，但其功能和性能可能不如 C/C++ 版本的 RDMA 库。
• GIL (Global Interpreter Lock)： Python 的 GIL 可能会限制多线程 RDMA 程序的性能。可以使用多进程来绕过 GIL 的限制。
• 内存管理： Python 的垃圾回收机制可能会影响 RDMA 程序的性能。需要仔细管理内存，避免频繁的垃圾回收。

7. 未来发展趋势

RDMA 技术在高性能计算和数据中心领域具有广阔的应用前景。未来发展趋势包括：

• RDMA over TCP/IP： iWARP 协议的进一步发展，使得 RDMA 可以更方便地应用于现有的 TCP/IP 网络。
• NVMe over Fabrics (NVMe-oF)： 将 NVMe 存储协议与 RDMA 技术结合，实现高性能的远程存储访问。
• Composable Infrastructure： RDMA 可以用于构建可组合的基础设施，实现资源的动态分配和管理。
• AI/ML 加速： RDMA 可以用于加速 AI/ML 模型的训练和推理，提高 AI/ML 应用的性能。

8. 总结：利用 RDMA 加速 Python 应用

RDMA 技术为 Python 提供了实现高性能跨节点通信的途径。通过 pyverbs 库，我们可以利用 RDMA 的零拷贝、低延迟和高带宽特性，加速分布式计算、大数据处理等应用。虽然 Python RDMA 开发面临一些挑战，但通过合理的优化策略，我们可以充分发挥 RDMA 的优势，构建高性能的 Python 应用。掌握RDMA技术，可以为Python应用带来显著的性能提升，尤其是在需要大量数据传输的场景下。随着RDMA技术的不断发展，其在Python中的应用前景将更加广阔。

更多IT精英技术系列讲座，到智猿学院