wigging
wigging

Reputation: 9190

Can't send a NumPy array larger than 2 GB with ZeroMQ

I'm using the Python code shown below to serialize and send a NumPy array from the client to the server using ZeroMQ. I noticed that when the NumPy array is larger than 2 GB the client seems to stall when sending the array. For example, in the client.py code shown below, if you use n = 17000 the client will stall after creating the array. I ran this code on a MacBook Pro laptop that has 32 GB of memory so there should be plenty of RAM available for the message. Is there a limit to the size of a NumPy array that I can send with ZeroMQ? If there is a limit, then how would I send an array that exceeds the size limit?

Client code that creates NumPy array (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):
        md = {"dtype": str(array.dtype), "shape": array.shape}
        self.socket.send_json(md, zmq.SNDMORE)  # send metadata
        self.socket.send(array, copy=False)     # send NumPy array data

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

def main():
    # Create array
    n = 16000  # 8000 is 500 MB, 11500 is 1 GB, 16000 is 2 GB, 17000 fails to send
    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} MB")
    print(f"Size of array object:  {sys.getsizeof(x)} bytes, {x.nbytes / 1000**2} MB")

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

if __name__ == "__main__":
    main()

Server code that receives NumPy array (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()               # receive metadata
        msg: Any = self.socket.recv(copy=False)         # receive NumPy array data
        array = np.frombuffer(msg, dtype=md["dtype"])   # reconstruct the NumPy array
        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()

Upvotes: 1

Views: 38

Answers (0)

Related Questions