test_ndarray_zmq.py

import zmq
import time
from multiprocessing import Process
import numpy as np

np.set_printoptions(precision=3)

def send_array(socket, A, flags=0, copy=True, track=False):
    """send a numpy array with metadata"""
    md = dict(dtype=str(A.dtype), shape=A.shape)
    socket.send_json(md, flags | zmq.SNDMORE)
    return socket.send(A, flags, copy=copy, track=track)


def recv_array(socket, flags=0, copy=True, track=False):
    """recv a numpy array"""
    md = socket.recv_json(flags=flags)
    msg = socket.recv(flags=flags, copy=copy, track=track)
    buf = buffer(msg)
    A = np.frombuffer(buf, dtype=md['dtype'])
    return A.reshape(md['shape'])


def server(port="5556"):
    context = zmq.Context()
    socket = context.socket(zmq.REP)
    socket.bind("tcp://*:%s" % port)
    print "Running server on port: ", port
    # serves only 5 request and dies
    for reqnum in range(5):
        # Wait for next request from client
        #message = socket.recv()
        message = recv_array(socket)
        print "Received request #%s: %s" % (reqnum, message)
        socket.send("World from %s" % port)


def client(ports=["5556"]):

    context = zmq.Context()
    print "Connecting to server with ports %s" % ports
    socket = context.socket(zmq.REQ)

    for port in ports:
        socket.connect("tcp://localhost:%s" % port)

    for request in range(20):
        print "Sending request ", request, "..."
        #socket.send("Hello")
        send_array(socket, np.random.randn(3))
        message = socket.recv()

        print "Received reply ", request, "[", message, "]"
        time.sleep(1)


if __name__ == "__main__":
    # Now we can run a few servers
    server_ports = range(5550, 5558, 2)

    for server_port in server_ports:
        Process(target=server, args=(server_port,)).start()

    # Now we can connect a client to all these servers
    Process(target=client, args=(server_ports,)).start()