如何优化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耗时的核心原因
- 内存拷贝开销:默认情况下,ZMQ会将网络数据从内核态拷贝到自身内部缓冲区,再拷贝到Python的
bytes对象,两次拷贝是主要耗时来源。同时每次调用都会malloc新的缓冲区存储消息,频繁的内存分配释放也会增加开销。 - GIL竞争:单线程server处理32个客户端的请求,Python的GIL会限制CPU利用率,导致
recv_multipart在Python层的阻塞等待耗时占比升高。 - 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内部缓冲区到Python
bytes的拷贝,显著降低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和系统内核参数,减少内核态到用户态的拷贝次数。例如:
同时修改Linux内核参数(需root权限):# 设置socket接收缓冲区为32MB socket.setsockopt(zmq.RCVBUF, 32 * 1024 * 1024)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
相关产品推荐
相关产品推荐

