ZMQ服务器如何批量处理请求？

Question

我想在客户端/服务器 model 中使用 ZMQ，但是每个服务器都会等待，直到它收到 100 个请求，共同处理它们，然后将 100 个响应发送回正确的客户端。 这背后的原因是服务器在 GPU 上执行计算，该计算仅在批量执行时计算效率高。 如何用 ZMQ 做到这一点？

以下是我尝试过的，不出所料地引发了zmq.error.ZMQError: Operation cannot be accomplished in current state因为服务器试图按顺序接收多个请求，而没有中间交错recv_pyobj()调用与send_pyobj()响应。

import multiprocessing as mp
import numpy as np
import time
import zmq

def computation(inputs):
  time.sleep(1)  # Simulate constant GPU overhead.
  results = np.zeros((len(inputs), 8))
  return results

def server(port, batch=100):
  context = zmq.Context()
  socket = context.socket(zmq.REP)
  socket.bind(f'tcp://*:{port}')
  while True:
    inputs = np.empty((100, 64))
    for i in range(batch):
      inputs[i] = socket.recv_pyobj()
    results = computation(inputs)
    for i in range(batch):
      socket.send_pyobj(results[i])

def client(ports):
  context = zmq.Context()
  socket = context.socket(zmq.REQ)
  for port in ports:
    socket.connect(f'tcp://localhost:{port}')
  while True:
    input_ = np.zeros(64)
    socket.send_pyobj(input_)
    result = socket.recv_pyobj()

if __name__ == '__main__':
  num_clients = 10
  num_servers = 3
  ports = list(range(5550, 5550 + num_servers))
  for port in ports:
    mp.Process(target=server, args=(port,)).start()
  for _ in range(num_clients):
    mp.Process(target=client, args=(ports,)).start()

Answer 1

ROUTER套接字类型与recv_multipart()和send_multipart()相结合是我正在寻找的。 最有用的资源是官方指南中Advanced Request-Reply Patterns下的rtreq示例。 除此之外，我添加了msgpack序列化并减少了批处理大小，因为它必须小于客户端的数量。 问题片段的工作版本如下。

import msgpack
import multiprocessing as mp
import numpy as np
import time
import uuid
import zmq

def computation(inputs):
  time.sleep(1)  # Simulate constant GPU overhead.
  results = np.zeros((len(inputs), 8))
  return results

def server(port, batch=10):
  context = zmq.Context.instance()
  socket = context.socket(zmq.ROUTER)
  socket.bind(f'tcp://*:{port}')
  while True:
    inputs = np.empty((batch, 64))
    addresses = []
    for i in range(batch):
      address, empty, payload = socket.recv_multipart()
      inputs[i] = unpack(payload)
      addresses.append(address)
    print('Collected request batch.', flush=True)
    results = computation(inputs)
    for i, address in enumerate(addresses):
      payload = pack(results[i])
      socket.send_multipart([address, b'', payload])
    print('Send response batch.', flush=True)

def client(ports):
  context = zmq.Context.instance()
  socket = context.socket(zmq.REQ)
  socket.setsockopt(zmq.IDENTITY, uuid.uuid4().bytes)
  for port in ports:
    socket.connect(f'tcp://localhost:{port}')
  while True:
    input_ = np.zeros(64)
    socket.send(pack(input_))
    result = unpack(socket.recv())

def pack(array):
  return msgpack.packb((array.shape, str(array.dtype), array.tobytes()))

def unpack(buffer):
  shape, dtype, value = msgpack.unpackb(buffer)
  return np.frombuffer(value, dtype).reshape(shape)

if __name__ == '__main__':
  num_clients = 100
  num_servers = 3
  ports = list(range(5550, 5550 + num_servers))
  for port in ports:
    mp.Process(target=server, args=(port,)).start()
  for _ in range(num_clients):
    mp.Process(target=client, args=(ports,)).start()