Question: How to implement multiprocessing.queues.SimpleQueue using ZeroMQ
### 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 条评论