File size: 9,493 Bytes
01171b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 | from abc import abstractmethod
from typing import Any, Dict, Optional, TypeVar
import zmq
from typing_extensions import Protocol
from ladder.components.inference.api.core.requests import RequestType, parse_request_type
from ladder.components.inference.api.core.serialization import deserialize, serialize
from ladder.core.common.logger import get_logger
from ladder.core.databib.config import Config, Configurable
LOGGER = get_logger(__name__)
class CallbackProtocol(Protocol):
def __call__(self, message: Dict) -> Dict:
...
class RequestReplyServerConfig(Config):
pass
RequestReplyServerConfigT = TypeVar('RequestReplyServerConfigT', bound=RequestReplyServerConfig)
class RequestReplyServer(Configurable[RequestReplyServerConfigT]):
"""
Abstract base class for a request-reply server that listens
for requests, dispatches them to the provided callback, and
sends the output back to the client.
Methods:
start(on_request_cb: CallbackProtocol): Begins listening for requests
"""
@abstractmethod
def start(self, on_request_cb: CallbackProtocol) -> None:
pass
class RequestReplyClientConfig(Config):
pass
RequestReplyClientConfigT = TypeVar('RequestReplyClientConfigT', bound=RequestReplyClientConfig)
class RequestReplyClient(Configurable[RequestReplyClientConfigT]):
"""
Abstract base class for a request-reply client that sends
requests to the server and receives the output from the server.
Methods:
send(request: Dict) -> Optional[Dict]: Send a request to the server.
"""
@abstractmethod
def send(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
pass
class ZmqRequestReplyServerConfig(RequestReplyServerConfig):
port: int # Port number to listen on
recv_timeout_ms: int = 5000 # Timeout in milliseconds for calls to socket.recv
class ZmqRequestReplyServer(RequestReplyServer[ZmqRequestReplyServerConfig]):
"""
ZeroMQ-based implementation of RequestReplyServer.
Methods:
start(on_request_cb: CallbackProtocol): Begins listening for requests.
Notes:
Designed to be used in conjunction with the ZmqRequestReplyClient class.
"""
def __init__(self, config: ZmqRequestReplyServerConfig):
super().__init__(config)
self._context = zmq.Context()
self._init_socket()
def _init_socket(self):
self.socket = self._context.socket(zmq.REP)
self.socket.bind(f"tcp://*:{self.config.port}")
self.socket.setsockopt(zmq.SNDHWM, 5)
self.socket.setsockopt(zmq.LINGER, 0)
self.socket.setsockopt(zmq.RCVTIMEO, self.config.recv_timeout_ms)
def start(self, on_request_cb: CallbackProtocol) -> None:
"""
Starts listening for requests over the socket, blocking on the current thread.
When a request is received, the on_request_cb function is called with the request
as an argument, and the output is sent back as a reply to the client.
Args:
on_request_cb (CallbackProtocol): Function to be called when a request is received.
Notes: this function is blocking
"""
LOGGER.info("Starting ZMQ request reply server")
while True:
LOGGER.debug("Waiting for new request")
try:
message = self.socket.recv()
except zmq.Again:
LOGGER.warn("ZMQ receive timeout. Continuing.")
continue
except zmq.ZMQError as e:
LOGGER.error(f"A ZMQ error occured while receiving. Error: {e}")
raise e
LOGGER.debug(f"Received new request")
LOGGER.debug(f"Processing new request")
message = deserialize(message)
res = on_request_cb(message)
res = serialize(res)
LOGGER.debug(f"Request processed")
LOGGER.debug(f"Sending reply")
try:
self.socket.send(res)
except zmq.ZMQError as e:
LOGGER.error(f"A ZMQ error occured while sending. Error: {e}")
raise e
LOGGER.debug(f"Reply sent")
# If the request is an exit request, we stop the loop and
# let the process exit gracefully.
if parse_request_type(message) == RequestType.EXIT:
LOGGER.info("Received EXIT request. Stopping server.")
break
def __del__(self):
if self.socket:
self.socket.close()
self._context.term()
class ZmqRequestReplyClientConfig(RequestReplyClientConfig):
ip: str # IP address of the server
port: int # Port number of the server
recv_timeout_ms: int = 15000 # Timeout in milliseconds for calls to socket.recv
send_num_retries: int = 3 # Number of times to retry sending a message
# If True, recv timeouts are raised as ``zmq.Again`` to the caller after the broken
# REQ socket has been reset. Use when a timeout indicates the worker should fail and
# be respawned by its supervisor (e.g. a remote inference client). Default False
# keeps the "log and return None" behaviour.
propagate_timeout: bool = False
class ZmqRequestReplyClient(RequestReplyClient[ZmqRequestReplyClientConfig]):
"""
ZeroMQ-based implementation of RequestReplyClient.
It communicates with the server using a ZeroMQ in REQ mode.
Methods:
send(request: dict) -> Optional[Dict]: Send a request to the server and
receive a reply.
Notes:
This class is designed to be used in conjunction with the
ZmqRequestReplyServer class.
"""
def __init__(self, config: ZmqRequestReplyClientConfig):
super().__init__(config)
self._closed = False
self.context = zmq.Context()
self.socket = None
self._init_socket()
def _init_socket(self):
"""
Reset the socket connection, this is needed when REQ is in a broken state.
"""
if self.socket:
LOGGER.info(f"Closing and deleting existing socket")
self.socket.close()
del self.socket
self.socket = self.context.socket(zmq.REQ)
self.socket.setsockopt(zmq.RCVTIMEO, self.config.recv_timeout_ms)
# this option is needed, among other things, to prevent context.term()
# from handing if socket.recv() is interrupted e.g. by a KeyboardInterrupt
self.socket.setsockopt(zmq.LINGER, 0)
LOGGER.info(f"Connecting to {self.config.ip}:{self.config.port}")
self.socket.connect(f"tcp://{self.config.ip}:{self.config.port}")
LOGGER.info(f"Connected to {self.config.ip}:{self.config.port}")
def send(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
if self.socket is None or self.socket.closed:
LOGGER.error(f"The socket is closed, request not sent.")
return None
try:
LOGGER.debug(f"Sending request")
sent = False
request_serialized = serialize(request)
for i in range(self.config.send_num_retries):
try:
self.socket.send(request_serialized)
sent = True
LOGGER.debug(f"Request sent.")
break
except zmq.ZMQError as e:
LOGGER.error(
f"Failed to send message to {self.config.ip}:{self.config.port}. Error: {e}."
)
LOGGER.debug(f"Attempt {i} failed. Retrying.")
if not sent:
LOGGER.warn("Number of retries exceeded. Aborting.")
return None
# If the request is an exit request, we don't wait for a reply.
if request.get("type") == RequestType.EXIT:
return None
LOGGER.debug(f"Waiting for reply")
try:
message = self.socket.recv()
except zmq.Again:
# REQ socket is now stuck in "waiting for reply", so any subsequent send
# would fail with ZMQError. Reset before deciding how to surface the timeout.
self._init_socket()
if self.config.propagate_timeout:
raise
LOGGER.warn("ZMQ receive timeout. Aborting.")
return None
except zmq.ZMQError as e:
LOGGER.error(f"A ZMQ error occured while receiving. Error: {e}.")
raise e
return deserialize(message)
except KeyboardInterrupt:
# Ctrl+C might be sent mid-episode. In this case, if the exception is raised inside the try
# block, the socket will fall into a broken state. Here we reinitialize it to keep the
# communication working.
self._init_socket()
raise
def close(self) -> None:
"""Idempotently tear down the socket and ZMQ context. Safe to call multiple times."""
if self._closed:
return
if self.socket:
try:
self.socket.disconnect(f"tcp://{self.config.ip}:{self.config.port}")
except zmq.ZMQError:
# Socket already disconnected (e.g. _init_socket re-created it after a timeout).
pass
self.socket.close()
self.context.term()
self._closed = True
def __del__(self):
self.close()
|