ITADN

Question: How to implement multiprocessing.queues.SimpleQueue using ZeroMQ

#2125Openhaiminh2001 创建于 2025-09-08
H
haiminh2001commented
### This is a pyzmq bug - [x] This is a pyzmq-specific bug, not an issue of zmq socket behavior. Don't worry if you're not sure! We'll figure it out together. ### What pyzmq version? 27.0.2 ### What libzmq version? 4.3.5 ### Python version (and how it was installed) python3.10 ### OS ubuntu 20.04 ### What happened? I am using concurrent.futures.ProcessPoolExecutor to handle large data processing. I find it so slow and a big bottleneckt is sending data through multiprocessin.Pipe, which is used by multiprocessing.queues.SimpleQueue. Therefore, I am implementing an alternative using ZeroMQ, also add some out-of-band data transferring (which is also a bottleneck of standard Python multiprocessing). My code works fine with `max_workers` > 1, and I saw big improvement in my pipeline throughput. But when `max_workers` is set to 1, I can only run one task before it get stuck, Also with `max_workers` > 1, only one worker process can run. I think there is something to do with my exposing ZMQ Socket FD trick to the Executor. Can you please explain why ? I will be very appreciate, ### Code to reproduce bug ```python from multiprocessing.context import SpawnContext from multiprocessing.reduction import ForkingPickler from multiprocessing.context import assert_spawning import io import pickle import types import os import tempfile import uuid import zmq class OOBPickler(pickle.Pickler): dispatch_table = ForkingPickler(io.BytesIO()).dispatch_table.copy() def dumps(obj, protocol=None): buf = io.BytesIO() oob_buffers = [] pickler = OOBPickler(buf, protocol, buffer_callback=oob_buffers.append) pickler.dump(obj) return buf.getbuffer(), oob_buffers class SocketWrapper: def __init__(self, sock: zmq.Socket): self.sock = sock def fileno(self): return self.sock.FD def recv(self): bufs = self.sock.recv_multipart(copy=False) return pickle.loads(bufs[0], buffers=bufs[1:]) def __getattr__(self, name: str): # This method is only called for attributes that don't exist on SocketWrapper return getattr(self.sock, name) class SimpleQueue: """ A drop-in replacement for multiprocessing.SimpleQueue using ZeroMQ for high performance. """ def __init__(self, *, ctx=None): # 1. Initialize ZMQ context and sockets # The main object holds both the reader (PULL) and writer (PUSH) self._context = zmq.Context() # The PULL socket is the "server" or sink for all messages self._reader = SocketWrapper(self._context.socket(zmq.PULL)) ipc_dir = os.path.join(tempfile.gettempdir(), 'zmq_queues') os.makedirs(ipc_dir, exist_ok=True) self._address = f"ipc://{os.path.join(ipc_dir, uuid.uuid4().hex)}" self._reader.bind(self._address) # The writer connects to the reader. In the main object, this allows # putting and getting from the same queue instance. self._writer = self._context.socket(zmq.PUSH) self._writer.connect(self._address) def close(self): """Closes all sockets and terminates the context.""" # Set a linger period of 0 to discard pending messages immediately if self._writer and not self._writer.closed: self._writer.setsockopt(zmq.LINGER, 0) self._writer.close() if self._reader and not self._reader.closed: self._reader.setsockopt(zmq.LINGER, 0) self._reader.close() if self._context and not self._context.closed: self._context.term() def empty(self): """ Checks for any incoming messages without blocking. Returns True if the queue is likely empty. """ # Poll the reader socket with a timeout of 0 (non-blocking) return self._reader.poll(0, zmq.POLLIN) == 0 def get(self): """ Receives an item from the queue. Blocks until an item is available. """ # ZMQ sockets are not thread-safe, but this pattern is for # inter-process, not inter-thread communication. No lock needed. return self._reader.recv() def put(self, obj): """ Puts an item into the queue. Blocks if the high-water mark is reached. """ buf, oob_buffers = dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) self._writer.send_multipart([buf] + oob_buffers, copy=False) def __getstate__(self): """ Prepare the queue to be pickled and sent to another process. We only send the address, not the sockets themselves. """ assert_spawning(self) return (self._address,) def __setstate__(self, state): """ Recreate the queue in a new process. This instance will only have a "writer" (PUSH) socket. """ (self._address,) = state self._context = zmq.Context() self._writer = self._context.socket(zmq.PUSH) self._writer.connect(self._address) # The new process's queue object cannot read. self._reader = None __class_getitem__ = classmethod(types.GenericAlias) class ZMQ_QueueContext(SpawnContext): def SimpleQueue(self): return SimpleQueue(ctx=self.get_context()) --- import asyncio from concurrent.futures import ProcessPoolExecutor import multiprocessing as mp import numpy as np from tqdm.asyncio import tqdm_asyncio from context import ZMQ_QueueContext mp.set_start_method('spawn', force=True) ARR_SIZE = (3, 10240, 720) def bar(x, y): return x + y async def main(): x = np.random.rand(*ARR_SIZE) y = np.random.rand(*ARR_SIZE) with ProcessPoolExecutor(max_workers=1, mp_context=ZMQ_QueueContext()) as executor: loop = asyncio.get_running_loop() tasks = [] for i in range(20): tasks.append(loop.run_in_executor(executor, bar, x, y)) await tqdm_asyncio.gather(*tasks) if __name__ == "__main__": asyncio.run(main()) ``` ### Traceback, if applicable ```shell ``` ### More info _No response_
0 条评论