def __poll_socket(self, socket: zmq.Socket) -> int:
waited_time = timedelta()
report_timeout_ms = self.REPORT_TIMEOUT / timedelta(milliseconds=1)
while socket.poll(report_timeout_ms) == 0:
waited_time += self.REPORT_TIMEOUT
logger.warning(f"Polled socket for {waited_time}")
if waited_time >= self.MAX_SOCKET_TIMEOUT:
raise ConnectionException(f"Polling timed out")
return socket.poll()
def zmq_import_poll(sock: zmq.Socket, timeout: float):
"""
Author:
Alexander Heilmeier
Description:
Handles incoming ZMQ messages and allows polling.
Inputs:
sock: ZMQ socket (see below how to create it)
timeout: [s] timeout for polling
How to create a ZMQ socket to import data?
import zmq
zmq_context = zmq.Context()
sock = zmq_context.socket(zmq.PUB)
sock.connect("tcp://%s:%s" % (ip, port))
sock.setsockopt_string(zmq.SUBSCRIBE, zmq_topic])
"""
# ------------------------------------------------------------------------------------------------------------------
# PREPARATION ----------------------------------------------------------------------------------------------------
# ------------------------------------------------------------------------------------------------------------------
# timeout conversion s -> ms
timeout *= 1000.0
# set standard return
data = None
# ------------------------------------------------------------------------------------------------------------------
# CHECK IF BUFFER ALREADY STORES VALUES ----------------------------------------------------------------------------
# ------------------------------------------------------------------------------------------------------------------
# non blocking socket to empty buffer
try:
while True:
sock.recv_string(flags=zmq.NOBLOCK)
data = sock.recv_pyobj(flags=zmq.NOBLOCK)
except zmq.Again:
pass
# ------------------------------------------------------------------------------------------------------------------
# WAIT FOR A SHORT TIME IF NO DATA WAS STORED IN THE BUFFER --------------------------------------------------------
# ------------------------------------------------------------------------------------------------------------------
if data is None:
# use poll to be able to set a timeout -> will return 1 if an event occurs, otherwise 0
events = sock.poll(timeout=timeout)
if events:
# non blocking socket to empty buffer
try:
while True:
sock.recv_string(flags=zmq.NOBLOCK)
data = sock.recv_pyobj(flags=zmq.NOBLOCK)
except zmq.Again:
pass
return data
def has_pending_messages(socket: zmq.Socket, timeout: int) -> bool:
"""
Checks if a socket has any pending messages
Parameters
----------
socket
Socket to check for messages
timeout
How long to wait for messages in milliseconds
Returns
-------
has_pending_messages
True if there are pending messages, false otherwise
"""
return socket.poll(timeout) == zmq.POLLIN
def _check_provider_job(socket: zmq.Socket, provider: ExecutionProvider, job_id: Any):
"""Poll for messages, checking that the provider's allocation is alive."""
while not socket.poll(1000, zmq.POLLIN):
if provider.status([job_id])[0].terminal:
raise RuntimeError("Provider job has terminated")
