You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何优化ZMQ的recv_multipart()以消除数据传输瓶颈?

优化ZMQ多部分消息接收性能:消除recv_multipart瓶颈

问题背景

我用多客户端向ZMQ服务器发送包含Numpy数组的字典,通过多部分消息打包避免了反序列化时的memcpy,吞吐量提升了一倍。但性能分析显示70%的耗时集中在recv_multipart()函数,推测该操作涉及网络数据到内存的拷贝。想明确以下几点:

  • 耗时是否来自malloc新缓冲区并拷贝消息?
  • ZMQ中能否复用接收缓冲区?
  • TCP协议是否存在无法进一步优化的根本性限制?

性能分析结果

Total Samples 30400
GIL: 73.00%, Active: 73.00%, Threads: 1

  %Own   %Total  OwnTime  TotalTime  Function (filename:line)
 70.00%  70.00%   203.7s    203.7s   recv_multipart (zmq/sugar/socket.py:808)
  1.00%   1.00%    3.01s     4.13s   recv_multipart (zmq/sugar/socket.py:807)
  0.00%   0.00%    2.62s     2.62s   <listcomp> (zmq_gbs_dict_seq.py:37)
  0.00%   0.00%    2.49s     2.49s   send (zmq/sugar/socket.py:696)
  0.00%   0.00%    1.32s     1.32s   unpack (zmq_gbs_dict_seq.py:35)
  0.00%   0.00%   0.690s     1.22s   __call__ (enum.py:717)
  0.00%  72.00%   0.520s    209.9s   server (zmq_gbs_dict_seq.py:82)
  1.00%   1.00%   0.500s    0.840s   inner (typing.py:341)
  0.00%   0.00%   0.500s     5.32s   server (zmq_gbs_dict_seq.py:83)
  0.00%   1.00%   0.400s     1.33s   recv_multipart (zmq/sugar/socket.py:812)
  1.00%   1.00%   0.360s     3.07s   send_multipart (zmq/sugar/socket.py:751)
  0.00%   0.00%   0.350s    0.350s   __new__ (enum.py:1106)
  0.00%   0.00%   0.300s    0.300s   __hash__ (typing.py:1352)
  0.00%   0.00%   0.270s    0.270s   <genexpr> (zmq_gbs_dict_seq.py:93)
  0.00%   0.00%   0.260s    0.260s   server (zmq_gbs_dict_seq.py:101)
  0.00%   0.00%   0.250s    0.660s   server (zmq_gbs_dict_seq.py:92)
  0.00%   0.00%   0.250s     3.04s   unpack (zmq_gbs_dict_seq.py:36)
  0.00%   0.00%   0.210s    0.210s   unpack (zmq_gbs_dict_seq.py:38)
  0.00%   0.00%   0.210s    0.210s   server (zmq_gbs_dict_seq.py:91)
  0.00%   0.00%   0.200s    0.200s   unpack (zmq_gbs_dict_seq.py:39)
  0.00%   1.00%   0.200s     4.04s   server (zmq_gbs_dict_seq.py:99)

测试代码

import multiprocessing
import pickle
import time

import numpy as np
import zmq


def client(port):
  socket = zmq.Context.instance().socket(zmq.DEALER)
  socket.set_hwm(0)
  socket.connect(f'tcp://localhost:{port}')
  data = {
      'foo': np.zeros((1024, 64, 64, 3), np.uint8),
      'bar': np.zeros((1024, 1024), np.float32),
      'baz': np.zeros((1024,), np.float32),
  }
  parts = pack(data)
  while True:
    socket.send_multipart(parts)
    msg = socket.recv()
    assert msg == b'done'
  socket.close()


def server(port):
  socket = zmq.Context.instance().socket(zmq.ROUTER)
  socket.set_hwm(0)
  socket.bind(f'tcp://*:{port}')
  time.sleep(3)
  print('Start')
  start = time.time()
  steps = 0
  nbytes = 0
  poller = zmq.Poller()
  poller.register(socket, zmq.POLLIN)
  while True:
    if poller.poll():
      addr, *parts = socket.recv_multipart(zmq.NOBLOCK)
      data = unpack(parts)
      steps += data['foo'].shape[0]
      nbytes += sum(v.nbytes for v in data.values())
      socket.send_multipart([addr, b'done'])
    duration = time.time() - start
    if duration > 1:
      fps = steps / duration
      gbs = (nbytes / 1024 / 1024 / 1024) / duration
      print(f'{fps/1e3:.2f}k fps {gbs:.2f} gb/s')
      start = time.time()
      steps = 0
      nbytes = 0
  socket.close()


def pack(data):
  dtypes, shapes, buffers = [], [], []
  items = sorted(data.items(), key=lambda x: x[0])
  keys, vals = zip(*items)
  dtypes = [v.dtype.name for v in vals]
  shapes = [v.shape for v in vals]
  buffers = [v.tobytes() for v in vals]
  meta = (keys, dtypes, shapes)
  parts = [pickle.dumps(meta), *buffers]
  return parts


def unpack(parts):
  meta, *buffers = parts
  keys, dtypes, shapes = pickle.loads(meta)
  vals = [
      np.frombuffer(b, d).reshape(s)
      for i, (d, s, b) in enumerate(zip(dtypes, shapes, buffers))]
  data = dict(zip(keys, vals))
  return data


def main():
  mp = multiprocessing.get_context('spawn')
  workers = []
  for _ in range(32):
    workers.append(mp.Process(target=client, args=(5555,)))
  workers.append(mp.Process(target=server, args=(5555,)))
  [x.start() for x in workers]
  [x.join() for x in workers]


if __name__ == '__main__':
  main()

问题分析与优化方案

一、recv_multipart耗时的核心原因

  1. 内存拷贝开销:默认情况下,ZMQ会将网络数据从内核态拷贝到自身内部缓冲区,再拷贝到Python的bytes对象,两次拷贝是主要耗时来源。同时每次调用都会malloc新的缓冲区存储消息,频繁的内存分配释放也会增加开销。
  2. GIL竞争:单线程server处理32个客户端的请求,Python的GIL会限制CPU利用率,导致recv_multipart在Python层的阻塞等待耗时占比升高。
  3. TCP协议栈开销:localhost通信虽然没有物理网络延迟,但TCP的Nagle算法、滑动窗口管理、内核态到用户态的拷贝仍然存在开销。

二、针对性优化方案

1. 复用ZMQ接收缓冲区(减少拷贝)

  • 使用copy=False参数接收消息:调用recv_multipart时传入copy=False,返回的是zmq.Frame对象,直接引用ZMQ内部缓冲区,避免将数据拷贝到Pythonbytes对象。修改server的接收逻辑:

    addr, *parts = socket.recv_multipart(zmq.NOBLOCK, copy=False)
    

    同时修改unpack函数,直接从Frame的缓冲区创建numpy数组:

    def unpack(parts):
      meta_frame, *buffer_frames = parts
      keys, dtypes, shapes = pickle.loads(meta_frame.bytes)
      vals = [
          np.frombuffer(frame.buffer, dtype=d).reshape(s)
          for d, s, frame in zip(dtypes, shapes, buffer_frames)]
      return dict(zip(keys, vals))
    

    这一步能消除ZMQ内部缓冲区到Pythonbytes的拷贝,显著降低recv_multipart的耗时。

  • 调整HWM值:将set_hwm(0)改为合理值(比如根据单条消息大小*并发数设置,例如socket.set_hwm(1024)),避免ZMQ因无限制接收频繁创建销毁缓冲区,减少malloc开销。

2. TCP与内核参数优化

  • 启用TCP_NODELAY:禁用Nagle算法,减少小数据包的延迟,适合高吞吐量场景。在client和server的socket初始化时添加:
    socket.setsockopt(zmq.TCP_NODELAY, 1)
    
  • 增大接收缓冲区:调整socket的RCVBUF和系统内核参数,减少内核态到用户态的拷贝次数。例如:
    # 设置socket接收缓冲区为32MB
    socket.setsockopt(zmq.RCVBUF, 32 * 1024 * 1024)
    
    同时修改Linux内核参数(需root权限):
    echo "net.core.rmem_max=67108864" >> /etc/sysctl.conf
    sysctl -p
    
  • 替换为IPC传输:如果是本地进程通信,用ipc://代替tcp://,IPC跳过TCP协议栈,开销远低于TCP。修改client和server的连接/绑定地址:
    # client
    socket.connect(f'ipc:///tmp/zmq_ipc_socket')
    # server
    socket.bind(f'ipc:///tmp/zmq_ipc_socket')
    

3. Python层优化

  • 多进程/异步处理:将单线程server改为多进程模式,每个进程处理一部分客户端请求,避免GIL竞争。或者使用zmq.asyncio异步API,提升CPU利用率。
  • 预分配numpy数组:如果消息格式固定,预分配numpy数组,直接将Frame缓冲区的数据写入预分配数组,避免每次创建新数组的开销:
    # 在server初始化时预分配
    pre_allocated = {
        'foo': np.zeros((1024, 64, 64, 3), np.uint8),
        'bar': np.zeros((1024, 1024), np.float32),
        'baz': np.zeros((1024,), np.float32),
    }
    
    # unpack时写入预分配数组
    def unpack(parts, pre_allocated):
      meta_frame, *buffer_frames = parts
      keys, dtypes, shapes = pickle.loads(meta_frame.bytes)
      for key, frame in zip(keys, buffer_frames):
        arr = pre_allocated[key]
        frame.buffer.readinto(arr)
      return pre_allocated
    
  • 替换序列化方式:用msgpack代替pickle,meta部分的序列化/反序列化速度更快,减少辅助开销。

三、根本性限制

  • TCP内核态拷贝:除非使用零拷贝技术(如Linux的sendfile、共享内存),否则TCP数据从内核skb到用户空间的拷贝无法完全消除。如果吞吐量要求极高,可以考虑:
    • 用UDP代替TCP(需自行实现丢包重传、顺序保证)
    • 结合共享内存传递数据,ZMQ仅传递控制信号,完全避免网络拷贝
  • ZMQ多部分消息开销:每个消息部分都有头部开销,若单条消息大小允许,可合并为单部分消息发送,减少头部处理开销。

内容的提问来源于stack exchange,提问作者danijar

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.07.17 14:54:56