ZeroMQ无法发送超过2GB的NumPy数组问题及解决方法

问题背景 #

使用ZeroMQ的REQ/REP模式发送NumPy数组时，当数组大小超过2GB（例如n=17000的二维随机数组），客户端发送阶段出现明显卡顿，尽管设备内存充足（32GB MacBook Pro）。

原因分析 #

ZeroMQ本身没有硬性的消息大小上限，但默认的单消息处理机制在应对超大内存块时，会因为以下原因导致卡顿：

• 内部缓冲区管理：超大消息会触发ZeroMQ复杂的内存页分配和拷贝逻辑，即使设置copy=False，仍会有额外的内核态/用户态内存交互开销。
• 系统层面限制：部分操作系统对单块内存映射或网络发送缓冲区有隐含的大小阈值，超过后会触发低效的分段处理。

解决方案：分片发送与拼接 #

最可靠的解决方式是将超大数组拆分为多个小块分片发送，服务端接收后再拼接还原。以下是修改后的代码实现：

修改后的客户端代码（client.py） #

import sys
import numpy as np
import zmq

class Client:
    """Client for sending/receiving messages."""

    def __init__(self, address="tcp://localhost:5555"):
        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect(address)
        self.socket = socket

    def send_array(self, array: np.ndarray):
        # 设置分片大小为1GB，可根据网络情况调整
        chunk_size = 1024 * 1024 * 1024  # 1GB per chunk
        total_bytes = array.nbytes
        # 拆分数组为二进制分片
        raw_data = array.tobytes()
        chunks = [raw_data[i:i+chunk_size] for i in range(0, total_bytes, chunk_size)]
        
        # 发送元数据（包含分片数量）
        md = {
            "dtype": str(array.dtype),
            "shape": array.shape,
            "chunk_count": len(chunks)
        }
        self.socket.send_json(md, zmq.SNDMORE)
        
        # 逐个发送分片
        for idx, chunk in enumerate(chunks):
            # 最后一个分片不需要SNDMORE标记
            flags = zmq.SNDMORE if idx < len(chunks) - 1 else 0
            self.socket.send(chunk, copy=False, flags=flags)

    def recv_message(self):
        reply = self.socket.recv_string()
        print("Received reply:", reply)

def main():
    # Create array
    n = 17000  # 超过2GB的数组
    x = np.random.rand(n, n)
    print(f"Array shape:           {x.shape}")
    print(f"First three elements:  {x[0, 0:3]}")
    print(f"Size of array data:    {x.nbytes} bytes, {x.nbytes / 1000**2:.2f} MB")

    # Create client and send array
    client = Client()
    client.send_array(x)
    client.recv_message()

if __name__ == "__main__":
    main()

修改后的服务端代码（server.py） #

from typing import Any
import zmq
import numpy as np

class Server:
    """Server for receiving/sending messages."""

    def __init__(self, address="tcp://localhost:5555"):
        context = zmq.Context()
        socket = context.socket(zmq.REP)
        socket.bind(address)
        self.socket = socket
        print("Server started, waiting for array...")

    def _recv_array(self):
        # 接收元数据
        md: Any = self.socket.recv_json()
        chunk_count = md["chunk_count"]
        
        # 接收所有分片
        chunks = []
        for _ in range(chunk_count):
            msg: Any = self.socket.recv(copy=False)
            chunks.append(msg.bytes)
        
        # 拼接分片并重构NumPy数组
        raw_data = b''.join(chunks)
        array = np.frombuffer(raw_data, dtype=md["dtype"])
        return array.reshape(md["shape"])

    def run(self):
        """Run the server."""
        while True:
            # Receive the NumPy array
            array = self._recv_array()
            print("Received array with shape:", array.shape)
            print(f"First three elements:  {array[0, 0:3]}")

            # Send a confirmation reply
            self.socket.send_string("Array received")

def main():
    server = Server()
    server.run()

if __name__ == "__main__":
    main()

额外优化建议 #

• 调整分片大小：根据网络带宽和延迟调整分片大小（比如千兆网络可设置为100MB-1GB），平衡发送效率和内存占用。
• 替换序列化方式：用msgpack替代json发送元数据，提升序列化/反序列化速度，尤其当元数据复杂时。
• 调整ZeroMQ套接字选项：设置socket.setsockopt(zmq.SNDHWM, 0)和socket.setsockopt(zmq.RCVHWM, 0)关闭高水位线限制，避免大消息被阻塞。

内容的提问来源于stack exchange，提问作者wigging