class MonitorEvents(SockJSConnection):
def _zmq_msg(self, msg):
#logging.debug(msg)
try:
msg_obj = json.loads(msg[0])
logging.debug(msg_obj)
if self.monitor != 'All':
if 'mon_id' in msg_obj and msg_obj['mon_id'] == self.monitor:
self.send(msg_obj)
else:
self.send(msg_obj)
except Exception as ex:
logging.error(ex)
def on_open(self, info):
logging.debug("Monitor ticker open: "+self.monitor)
zmq_socket = zmq.Context.instance().socket(zmq.SUB)
zmq_socket.connect(zmq_local_endpoint)
zmq_socket.setsockopt(zmq.SUBSCRIBE, '')
self.stream = ZMQStream(zmq_socket)
self.stream.on_recv(self._zmq_msg)
def on_close(self):
logging.debug("Monitor ticker close: "+self.monitor)
self.stream.stop_on_recv()
class LRUQueue(object):
def __init__(self, backend_socket, frontend_socket, clients, workers):
self.avaliable_workers = 0
self.workers = []
self.worker_num = workers
self.client_num = clients
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
worker_addr, empty, client_addr = msg[:3]
assert self.avaliable_workers < self.worker_num
self.avaliable_workers += 1
self.workers.append(worker_addr)
assert empty == ''
if client_addr != 'READY':
empty, reply = msg[3:]
assert empty == ''
self.frontend.send_multipart([client_addr, '', reply])
self.client_num -= 1
if 0 == self.client_num:
self.loop.add_timeout(time.time() + 1, self.loop.stop)
if self.avaliable_workers == 1:
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
client_addr, empty, request = msg
assert empty == ''
self.avaliable_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart([worker_id, '', client_addr, '', request])
if self.avaliable_workers == 0:
self.frontend.stop_on_recv()
class LRUQueue(object):
def __init__(self, backend_socket, frontend_socket, clients, workers):
self.avaliable_workers = 0
self.workers = []
self.worker_num = workers
self.client_num = clients
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
worker_addr, empty, client_addr = msg[:3]
assert self.avaliable_workers < self.worker_num
self.avaliable_workers += 1
self.workers.append(worker_addr)
assert empty == ""
if client_addr != "READY":
empty, reply = msg[3:]
assert empty == ""
self.frontend.send_multipart([client_addr, "", reply])
self.client_num -= 1
if 0 == self.client_num:
self.loop.add_timeout(time.time() + 1, self.loop.stop)
if self.avaliable_workers == 1:
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
client_addr, empty, request = msg
assert empty == ""
self.avaliable_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart([worker_id, "", client_addr, "", request])
if self.avaliable_workers == 0:
self.frontend.stop_on_recv()
class ManagerControlled(object):
def __init__(self, *args, **kwargs):
self.context = Context.instance()
self.loop = IOLoop.instance()
self.control_socket = self.context.socket(SUB)
self.control_socket.setsockopt(LINGER, 0) # discard unsent messages on close
self.control_socket.setsockopt(SUBSCRIBE, '')
self.control_socket.connect('tcp://{}:{}'.format(MANAGER_PUB_ADDRESS, MANAGER_PUB_PORT))
self.control_stream = ZMQStream(self.control_socket, self.loop)
self.control_stream.on_recv_stream(self.control_handler)
def control_handler(self, stream, message_list):
for message in message_list:
try:
notification, data = message.split()
except ValueError:
notification = message
if notification == NOTIFICATION_PROCESS_STOP:
self.stop()
def stop(self):
self.control_stream.stop_on_recv()
self.control_stream.close()
self.control_socket.close()
class AsyncCircusClient(object):
def __init__(self,
context=None,
endpoint=DEFAULT_ENDPOINT_DEALER,
timeout=5.0,
ssh_server=None,
ssh_keyfile=None):
self._init_context(context)
self.endpoint = endpoint
self._id = b(uuid.uuid4().hex)
self.socket = self.context.socket(zmq.DEALER)
self.socket.setsockopt(zmq.IDENTITY, self._id)
self.socket.setsockopt(zmq.LINGER, 0)
get_connection(self.socket, endpoint, ssh_server, ssh_keyfile)
self._timeout = timeout
self.timeout = timeout * 1000
self.stream = ZMQStream(self.socket, tornado.ioloop.IOLoop.instance())
def _init_context(self, context):
self.context = context or zmq.Context.instance()
def stop(self):
self.stream.stop_on_recv()
# only supported by libzmq >= 3
if hasattr(self.socket, 'disconnect'):
self.socket.disconnect(self.endpoint)
self.stream.close()
def send_message(self, command, **props):
return self.call(make_message(command, **props))
@tornado.gen.coroutine
def call(self, cmd):
if isinstance(cmd, string_types):
raise DeprecationWarning('call() takes a mapping')
call_id = uuid.uuid4().hex
cmd['id'] = call_id
try:
cmd = json.dumps(cmd)
except ValueError as e:
raise CallError(str(e))
try:
yield tornado.gen.Task(self.stream.send, cmd)
except zmq.ZMQError as e:
raise CallError(str(e))
while True:
messages = yield tornado.gen.Task(self.stream.on_recv)
for message in messages:
try:
res = json.loads(message)
if res.get('id') != call_id:
# we got the wrong message
continue
raise tornado.gen.Return(res)
except ValueError as e:
raise CallError(str(e))
class AsyncCircusClient(object):
def __init__(self, context=None, endpoint=DEFAULT_ENDPOINT_DEALER,
timeout=5.0, ssh_server=None, ssh_keyfile=None):
self._init_context(context)
self.endpoint = endpoint
self._id = b(uuid.uuid4().hex)
self.socket = self.context.socket(zmq.DEALER)
self.socket.setsockopt(zmq.IDENTITY, self._id)
self.socket.setsockopt(zmq.LINGER, 0)
get_connection(self.socket, endpoint, ssh_server, ssh_keyfile)
self._timeout = timeout
self.timeout = timeout * 1000
self.stream = ZMQStream(self.socket, tornado.ioloop.IOLoop.instance())
def _init_context(self, context):
self.context = context or zmq.Context.instance()
def stop(self):
self.stream.stop_on_recv()
# only supported by libzmq >= 3
if hasattr(self.socket, 'disconnect'):
self.socket.disconnect(self.endpoint)
self.stream.close()
@tornado.gen.coroutine
def send_message(self, command, **props):
res = yield self.call(make_message(command, **props))
raise tornado.gen.Return(res)
@tornado.gen.coroutine
def call(self, cmd):
if isinstance(cmd, string_types):
raise DeprecationWarning('call() takes a mapping')
call_id = uuid.uuid4().hex
cmd['id'] = call_id
try:
cmd = json.dumps(cmd)
except ValueError as e:
raise CallError(str(e))
try:
yield tornado.gen.Task(self.stream.send, cmd)
except zmq.ZMQError as e:
raise CallError(str(e))
while True:
messages = yield tornado.gen.Task(self.stream.on_recv)
for message in messages:
try:
res = json.loads(message)
if res.get('id') != call_id:
# we got the wrong message
continue
raise tornado.gen.Return(res)
except ValueError as e:
raise CallError(str(e))
class AsynchronousStatsConsumer(object):
def __init__(self,
topics,
loop,
callback,
context=None,
endpoint=DEFAULT_ENDPOINT_SUB,
ssh_server=None,
timeout=1.):
self.topics = topics
self.keep_context = context is not None
self.context = context or zmq.Context()
self.endpoint = endpoint
self.pubsub_socket = self.context.socket(zmq.SUB)
get_connection(self.pubsub_socket, self.endpoint, ssh_server)
for topic in self.topics:
self.pubsub_socket.setsockopt(zmq.SUBSCRIBE, topic)
self.stream = ZMQStream(self.pubsub_socket, loop)
self.stream.on_recv(self.process_message)
self.callback = callback
self.timeout = timeout
# Connection counter
self.count = 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
""" On context manager exit, destroy the zmq context """
self.stop()
def process_message(self, msg):
topic, stat = msg
topic = topic.split('.')
if len(topic) == 3:
__, watcher, subtopic = topic
self.callback(watcher, subtopic, json.loads(stat), self.endpoint)
elif len(topic) == 2:
__, watcher = topic
self.callback(watcher, None, json.loads(stat), self.endpoint)
def stop(self):
self.stream.stop_on_recv()
if self.keep_context:
return
try:
self.context.destroy(0)
except zmq.ZMQError as e:
if e.errno == errno.EINTR:
pass
else:
raise
class _Messenger(object):
def __init__(self, in_sock, out_sock, context, io_loop=None):
self._context = context
self._io_loop = io_loop or IOLoop.instance()
self._create_socket(in_sock, out_sock)
self._in_stream = ZMQStream(self._in_socket, io_loop)
self._out_stream = ZMQStream(self._out_socket, io_loop)
self._callbacks = defaultdict(list)
def _create_socket(self, in_sock, out_sock):
raise NotImplementedError()
def start(self):
self._in_stream.on_recv(self._on_receive)
def stop(self):
self._in_stream.stop_on_recv()
# self._publish(CTRL_MSG_WORKER, None, CTRL_MSG_WORKER_QUIT_ACK)
#
def close(self):
self._in_stream.close()
self._in_socket.close()
self._out_stream.close()
self._out_socket.close()
def _on_receive(self, zmq_msg):
msg = CtrlMessage.deserialize(zmq_msg)
if msg.topic in self._callbacks:
for callback in self._callbacks[msg.topic]:
callback(msg)
# if msg.data == CTRL_MSG_WORKER_QUIT:
# self.stop()
def add_callback(self, topic, callback):
self._callbacks[topic].append(callback)
def remove_callback(self, topic, callback):
if topic in self._callbacks and callback in self._callbacks[topic]:
self._callbacks[topic].remove(callback)
def publish(self, topic, identity, data):
msg = CtrlMessage(topic, identity, data)
self._out_stream.send_multipart(msg.serialize())
class AsynchronousStatsConsumer(object):
def __init__(self, topics, loop, callback, context=None,
endpoint=DEFAULT_ENDPOINT_SUB, ssh_server=None, timeout=1.):
self.topics = topics
self.keep_context = context is not None
self.context = context or zmq.Context()
self.endpoint = endpoint
self.pubsub_socket = self.context.socket(zmq.SUB)
get_connection(self.pubsub_socket, self.endpoint, ssh_server)
for topic in self.topics:
self.pubsub_socket.setsockopt_string(zmq.SUBSCRIBE, topic)
self.stream = ZMQStream(self.pubsub_socket, loop)
self.stream.on_recv(self.process_message)
self.callback = callback
self.timeout = timeout
# Connection counter
self.count = 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
""" On context manager exit, destroy the zmq context """
self.stop()
def process_message(self, msg):
topic, stat = msg
stat = stat.decode('utf-8')
topic = topic.decode('utf-8').split('.')
if len(topic) == 3:
__, watcher, subtopic = topic
self.callback(watcher, subtopic, json.loads(stat), self.endpoint)
elif len(topic) == 2:
__, watcher = topic
self.callback(watcher, None, json.loads(stat), self.endpoint)
def stop(self):
self.stream.stop_on_recv()
if self.keep_context:
return
try:
self.context.destroy(0)
except zmq.ZMQError as e:
if e.errno == errno.EINTR:
pass
else:
raise
class LRUQueue(object):
"""LRUQueue class using ZMQStream/IOLoop for event dispatching"""
def __init__(self, backend_socket, frontend_socket):
self.available_workers = 0
self.workers = []
self.client_nbr = NBR_CLIENTS
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
# Queue worker address for LRU routing
worker_addr, empty, client_addr = msg[:3]
assert self.available_workers < NBR_WORKERS
self.available_workers += 1
self.workers.append(worker_addr)
assert empty == b""
if client_addr != b"READY":
empty, reply = msg[3:]
assert empty == b""
self.frontend.send_multipart([client_addr, b'', reply])
self.client_nbr -= 1
if self.client_nbr == 0:
self.loop.add_timeout(time.time() + 1, self.loop.stop)
if self.available_workers == 1:
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
client_addr, empty, request = msg
assert empty == b""
self.available_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart([worker_id, b'', client_addr, b'', request])
if self.available_workers == 0:
self.frontend.stop_on_recv()
class ZMQCameraPubSub(object):
def __init__(self, callback):
self.callback = callback
self.name = "".join(
random.choice(string.ascii_lowercase + string.digits)
for x in range(6))
def connect(self, stream):
self.context = zmq.Context()
self.subscriber = self.context.socket(zmq.SUB)
# self.subscriber.setsockopt(zmq.RCVHWM, 1)
# self.subscriber.setsockopt(zmq.RCVBUF, 1*1024)
self.subscriber.setsockopt(zmq.LINGER, 0)
self.subscriber.connect(stream)
self.subscriber = ZMQStream(self.subscriber)
self.subscriber.on_recv(self.callback, copy=False)
# self.request.linger = 0
self.subscriber.setsockopt(zmq.SUBSCRIBE, b"")
self.subscriber.setsockopt(zmq.SUBSCRIBE, self.name.encode('ascii'))
def close(self):
if self.subscriber:
self.subscriber.stop_on_recv()
self.subscriber.close()
self.subscriber = None
def subscribe(self, to, topic=''):
subscribeto = to
if len(topic) > 0:
subscribeto = f"{subscribeto}.{topic}"
subscribeto = subscribeto.encode('ascii')
self.subscriber.setsockopt(zmq.SUBSCRIBE, subscribeto)
def unsubscribe(self, to, topic=''):
subscribetopic = to
if len(topic) > 0:
subscribetopic = f"{subscribetopic}.{topic}"
subscribetopic = subscribetopic.encode('ascii')
self.subscriber.setsockopt(zmq.UNSUBSCRIBE, subscribetopic)
class Downloader(object):
def __init__(self, in_socket, out_socket, io_loop):
self._in_socket = in_socket
self._in_socket.setsockopt(zmq.HWM, 10)
self._out_socket = out_socket
self._io_loop = io_loop
self._client = AsyncHTTPClient(self._io_loop,
max_clients=10,
max_simultaneous_connections=1)
self._stream = ZMQStream(self._in_socket, self._io_loop)
self._stream.on_recv(self._receive)
def _receive(self, msg):
"""
Msg is a URL we should download or 'EXIT'.
"""
if msg[0] == "EXIT":
print "stopping downloader"
self._stream.flush()
self._stream.stop_on_recv()
self._out_socket.send_unicode(msg[0])
else:
self._download_this(msg)
def _download_this(self, url):
print url[0]
req = HTTPRequest(url[0].encode("utf-8"))
self._client.fetch(req, self._handle_response)
def _handle_response(self, response):
if response.error:
print "Error: %s", response.error
else:
# simply send the response body to the ougoing ZMQ socket
self._out_socket.send_multipart(
[response.request.url,
str(response.request_time)])
class ZQueue(object):
def __init__(self, frontend_socket, backend_socket):
self.queue = WorkerQueue()
self.frontend = ZMQStream(frontend_socket)
self.backend = ZMQStream(backend_socket)
self.liveness = HEARTBEAT_LIVENESS
self.heartbeat = HEARTBEAT_INTERVAL
self.interval = INTERVAL_INIT
self.loop = IOLoop.instance()
self.hearbeats = 0
self.time = self.interval * self.heartbeat
self.heartbeat_at = time.time() + self.heartbeat * HEARTBEAT_LIVENESS
self.callback = None
self.timed_out = False
self.frontend.on_recv(self.handle_frontend)
self.backend.on_recv(self.handle_backend)
self.period = PeriodicCallback(self.purge,HEARTBEAT_INTERVAL*1000)
self.period.start()
def handle_frontend(self,msg):
m = msg[:]
if len(m) == 1:
times_str('Received heartbeat')
if self.timed_out:
self.loop.add_timeout(time.time()+HEARTBEAT_INTERVAL, self.send_heartbeat)
self.timed_out = False
self.loop.remove_timeout(self.callback)
elif len(m) == 3:
times_str('Received: '+str(m))
address, worker = self.queue.getLRU()
worker.working = True
m.insert(0,address)
self.backend.send_multipart(m)
self.heartbeat_at = time.time() + HEARTBEAT_INTERVAL * HEARTBEAT_LIVENESS
def handle_backend(self,msg):
m = msg[:]
address = m[0]
times_str('Backend Received: {}'.format(m))
self.queue.ready(WorkerModel(address))
self.backend.send_multipart([address,PPP_HEARTBEAT])
mm = m[1:]
if len(mm) == 1:
if mm[0] == PPP_HEARTBEAT:
self.hearbeats += 1
times_str('Got hearbeat {}'.format(self.hearbeats))
else:
times_str('Sending it back..')
self.frontend.send_multipart(mm)
if not self.queue.empty():
self.frontend.on_recv(self.handle_frontend)
def purge(self):
self.queue.purge()
if self.queue.empty():
self.frontend.stop_on_recv()
def send_heartbeat(self):
if time.time() > self.heartbeat_at:
self.time *= 2 if self.time < INTERVAL_MAX else 1
times_str('Timed out.. Retrying in {} seconds..'.format(self.time))
self.callback = self.loop.add_timeout(time.time()+self.time*1, self.send_heartbeat)
self.timed_out = True
return
self.time = self.interval * self.heartbeat
times_str('sending heartbeat..')
self.frontend.send(PPP_HEARTBEAT)
self.loop.add_timeout(time.time()+self.heartbeat, self.send_heartbeat)
def run(self):
try:
self.frontend.send(PPP_READY)
self.loop.start()
except KeyboardInterrupt:
times_str('ctrlc')
class ZmqStreamlet(Zmqlet):
"""A :class:`ZmqStreamlet` object can send/receive data to/from ZeroMQ stream and invoke callback function. It
has three sockets for input, output and control.
.. warning::
Starting from v0.3.6, :class:`ZmqStreamlet` replaces :class:`Zmqlet` as one of the key components in :class:`jina.peapods.runtime.BasePea`.
It requires :mod:`tornado` and :mod:`uvloop` to be installed.
"""
def __init__(
self,
*args,
**kwargs,
):
super().__init__(**kwargs)
def _register_pollin(self):
"""Register :attr:`in_sock`, :attr:`ctrl_sock` and :attr:`out_sock` in poller."""
with ImportExtensions(required=True):
import tornado.ioloop
get_or_reuse_loop()
self.io_loop = tornado.ioloop.IOLoop.current()
self.in_sock = ZMQStream(self.in_sock, self.io_loop)
self.out_sock = ZMQStream(self.out_sock, self.io_loop)
self.ctrl_sock = ZMQStream(self.ctrl_sock, self.io_loop)
if self.in_connect_sock is not None:
self.in_connect_sock = ZMQStream(self.in_connect_sock, self.io_loop)
self.in_sock.stop_on_recv()
def _get_dynamic_out_socket(self, target_pod):
return super()._get_dynamic_out_socket(target_pod, as_streaming=True)
def close(self, flush: bool = True, *args, **kwargs):
"""Close all sockets and shutdown the ZMQ context associated to this `Zmqlet`.
.. note::
This method is idempotent.
:param flush: flag indicating if `sockets` need to be flushed before close is done
:param args: Extra positional arguments
:param kwargs: Extra key-value arguments
"""
# if Address already in use `self.in_sock_type` is not set
if (
not self.is_closed
and hasattr(self, 'in_sock_type')
and self.in_sock_type == zmq.DEALER
):
try:
if self._active:
self._send_cancel_to_router(raise_exception=True)
except zmq.error.ZMQError:
self.logger.debug(
f'The dealer {self.name} can not unsubscribe from the router. '
f'In case the router is down this is expected.'
)
self._active = (
False # Important to avoid sending idle back while flushing in socket
)
if not self.is_closed:
# wait until the close signal is received
time.sleep(0.01)
if flush:
for s in self.opened_socks:
events = s.flush()
self.logger.debug(f'Handled #{events} during flush of socket')
super().close()
if hasattr(self, 'io_loop'):
try:
self.io_loop.stop()
# Replace handle events function, to skip
# None event after sockets are closed.
if hasattr(self.in_sock, '_handle_events'):
self.in_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.out_sock, '_handle_events'):
self.out_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.ctrl_sock, '_handle_events'):
self.ctrl_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.in_connect_sock, '_handle_events'):
self.in_connect_sock._handle_events = (
lambda *args, **kwargs: None
)
except AttributeError as e:
self.logger.error(f'failed to stop. {e!r}')
def pause_pollin(self):
"""Remove :attr:`in_sock` from the poller """
self.in_sock.stop_on_recv()
self.is_polling_paused = True
def resume_pollin(self):
"""Put :attr:`in_sock` back to the poller """
if self.is_polling_paused:
self.in_sock.on_recv(self._in_sock_callback)
self.is_polling_paused = False
def start(self, callback: Callable[['Message'], None]):
"""
Open all sockets and start the ZMQ context associated to this `Zmqlet`.
:param callback: callback function to receive the protobuf message
"""
def _callback(msg, sock_type):
msg = _parse_from_frames(sock_type, msg)
self.bytes_recv += msg.size
self.msg_recv += 1
callback(msg)
self._in_sock_callback = lambda x: _callback(x, self.in_sock_type)
self.in_sock.on_recv(self._in_sock_callback)
self.ctrl_sock.on_recv(lambda x: _callback(x, self.ctrl_sock_type))
if self.out_sock_type == zmq.ROUTER and not self.args.dynamic_routing_out:
self.out_sock.on_recv(lambda x: _callback(x, self.out_sock_type))
if self.in_connect_sock is not None:
self.in_connect_sock.on_recv(
lambda x: _callback(x, SocketType.ROUTER_CONNECT)
)
self.io_loop.start()
self.io_loop.clear_current()
self.io_loop.close(all_fds=True)
class ZMQNodePubSub(object):
def __init__(self, node, request_callback, subscribe_callback):
self.callback = request_callback
self.subscribe_callback = subscribe_callback
self.node = node
def connect(self):
# print("Node PUbsub connect", flush=True)
self.context = zmq.Context()
self.socket = self.context.socket(zmq.DEALER)
# self.socket.identity = b'tada123'
# print(f'ConnecT to Router {self.node.router_address}')
self.socket.connect(self.node.router_address)
self.stream = ZMQStream(self.socket)
self.stream.setsockopt(zmq.LINGER, 0)
self.stream.on_recv(self.callback)
self.subscriber = self.context.socket(zmq.SUB)
self.subscriber.connect(self.node.publisher_address)
self.subscriber = ZMQStream(self.subscriber)
self.subscriber.setsockopt(zmq.LINGER, 0)
# self.subscriber.setsockopt( ZMQ:IMMEDIATE)
self.subscriber.on_recv(self.subscribe_callback)
def subscribe(self, to, topic=''):
subscribeto = to
if len(topic) > 0:
subscribeto = f"{subscribeto}.{topic}"
subscribeto = self.node.identity + b'.' + subscribeto.encode('ascii')
self.subscriber.setsockopt(zmq.SUBSCRIBE, subscribeto)
def unsubscribe(self, to, topic=''):
subscribetopic = to
if len(topic) > 0:
subscribetopic = f"{subscribetopic}.{topic}"
subscribetopic = subscribetopic.encode('ascii')
self.subscriber.setsockopt(zmq.UNSUBSCRIBE, subscribetopic)
def close(self):
if self.subscriber:
self.subscriber.stop_on_recv()
self.subscriber.close()
self.subscriber = None
if self.stream:
self.stream.stop_on_recv()
self.stream.close()
self.stream = None
def make_request(self, target, action, msg=None):
kind, *cmds = action.split(".")
method = action
if len(cmds) > 0:
method = cmds[0]
wrapped = {"data": msg}
# send_to_socket = [self.node.identity, method.encode('ascii'), json.dumps(wrapped).encode('ascii')]
# print(f'Sending {send_to_socket}')
# self.stream.send_multipart(send_to_socket)
# self.socket.send_multipart([method.encode('ascii'), json.dumps(wrapped).encode('ascii'), target.encode('ascii')])
return self.node.run_action(method, wrapped, target)
class PubSub(BasePubSub):
"""
This class manages application PUB/SUB logic.
"""
NAME = 'ZeroMQ'
def __init__(self, application):
super(PubSub, self).__init__(application)
self.sub_stream = None
self.pub_stream = None
self.zmq_context = None
self.zmq_pub_sub_proxy = None
self.zmq_xpub = None
self.zmq_xsub = None
self.zmq_pub_port = None
self.zmq_sub_address = None
def initialize(self):
self.zmq_context = zmq.Context()
options = self.application.settings['options']
self.zmq_pub_sub_proxy = options.zmq_pub_sub_proxy
# create PUB socket to publish instance events into it
publish_socket = self.zmq_context.socket(zmq.PUB)
# do not try to send messages after closing
publish_socket.setsockopt(zmq.LINGER, 0)
if self.zmq_pub_sub_proxy:
# application started with XPUB/XSUB proxy
self.zmq_xsub = options.zmq_xsub
publish_socket.connect(self.zmq_xsub)
else:
# application started without XPUB/XSUB proxy
if options.zmq_pub_port_shift:
# calculate zmq pub port number
zmq_pub_port = options.port - options.zmq_pub_port_shift
else:
zmq_pub_port = options.zmq_pub_port
self.zmq_pub_port = zmq_pub_port
publish_socket.bind(
"tcp://%s:%s" %
(options.zmq_pub_listen, str(self.zmq_pub_port)))
# wrap pub socket into ZeroMQ stream
self.pub_stream = ZMQStream(publish_socket)
# create SUB socket listening to all events from all app instances
subscribe_socket = self.zmq_context.socket(zmq.SUB)
if self.zmq_pub_sub_proxy:
# application started with XPUB/XSUB proxy
self.zmq_xpub = options.zmq_xpub
subscribe_socket.connect(self.zmq_xpub)
else:
# application started without XPUB/XSUB proxy
self.zmq_sub_address = options.zmq_sub_address
for address in self.zmq_sub_address:
subscribe_socket.connect(address)
subscribe_socket.setsockopt_string(zmq.SUBSCRIBE,
six.u(CONTROL_CHANNEL))
subscribe_socket.setsockopt_string(zmq.SUBSCRIBE, six.u(ADMIN_CHANNEL))
def listen_socket():
# wrap sub socket into ZeroMQ stream and set its on_recv callback
self.sub_stream = ZMQStream(subscribe_socket)
self.sub_stream.on_recv(self.dispatch_published_message)
tornado.ioloop.IOLoop.instance().add_callback(listen_socket)
if self.zmq_pub_sub_proxy:
logger.info("ZeroMQ XPUB: {0}, XSUB: {1}".format(
self.zmq_xpub, self.zmq_xsub))
else:
logger.info("ZeroMQ PUB - {0}; subscribed to {1}".format(
self.zmq_pub_port, self.zmq_sub_address))
def publish(self, channel, message, method=None):
"""
Publish message into channel of stream.
"""
method = method or self.DEFAULT_PUBLISH_METHOD
message["message_type"] = method
message = json_encode(message)
to_publish = [utf8(channel), utf8(message)]
self.pub_stream.send_multipart(to_publish)
@coroutine
def dispatch_published_message(self, multipart_message):
"""
Got message, decide what is it and dispatch into right
application handler.
"""
channel = multipart_message[0]
if six.PY3:
channel = channel.decode()
message_data = json_decode(multipart_message[1])
if channel == CONTROL_CHANNEL:
yield self.handle_control_message(message_data)
elif channel == ADMIN_CHANNEL:
yield self.handle_admin_message(message_data)
else:
yield self.handle_channel_message(channel, message_data)
def subscribe_key(self, subscription_key):
self.sub_stream.setsockopt_string(zmq.SUBSCRIBE,
six.u(subscription_key))
def unsubscribe_key(self, subscription_key):
self.sub_stream.setsockopt_string(zmq.UNSUBSCRIBE,
six.u(subscription_key))
def clean(self):
"""
Properly close ZeroMQ sockets.
"""
if hasattr(self, 'pub_stream') and self.pub_stream:
self.pub_stream.close()
if hasattr(self, 'sub_stream') and self.sub_stream:
self.sub_stream.stop_on_recv()
self.sub_stream.close()
class ZmqMgmt(object):
"""
A :class:`ZMQStream` object handling the management sockets.
"""
def __init__(self, subscriber, publisher, **kwargs):
"""
Initialize the management interface.
The `subscriber` socket is the socket used by the Master to send
commands to the workers. The publisher socket is used to send commands
to the Master.
You have to set the `zmq.SUBSCRIBE` socket option yourself!
"""
self._io_loop = kwargs.get('io_loop', IOLoop.instance())
self._subscriber = subscriber
self._in_stream = ZMQStream(self._subscriber, self._io_loop)
self._publisher = publisher
self._out_stream = ZMQStream(self._publisher, self._io_loop)
self._callbacks = dict()
def _receive(self, raw_msg):
"""
Main method for receiving management messages.
`message` is a multipart message where `message[0]` contains the topic,
`message[1]` is 0 and `message[1]` contains the actual message.
"""
msg = MgmtMessage(raw_msg)
if msg.topic in self._callbacks:
for callback in self._callbacks[msg.topic]:
if callable(callback):
callback(msg)
if ZMQ_SPYDER_MGMT_WORKER_QUIT == msg.data:
self.stop()
def start(self):
"""
Start the MGMT interface.
"""
self._in_stream.on_recv(self._receive)
def stop(self):
"""
Stop the MGMT interface.
"""
self._in_stream.stop_on_recv()
self.publish(topic=ZMQ_SPYDER_MGMT_WORKER, identity=None,
data=ZMQ_SPYDER_MGMT_WORKER_QUIT_ACK)
def close(self):
"""
Close all open sockets.
"""
self._in_stream.close()
self._subscriber.close()
self._out_stream.close()
self._publisher.close()
def add_callback(self, topic, callback):
"""
Add a callback to the specified topic.
"""
if not callable(callback):
raise ValueError('callback must be callable')
if topic not in self._callbacks:
self._callbacks[topic] = []
self._callbacks[topic].append(callback)
def remove_callback(self, topic, callback):
"""
Remove a callback from the specified topic.
"""
if topic in self._callbacks and callback in self._callbacks[topic]:
self._callbacks[topic].remove(callback)
def publish(self, topic=None, identity=None, data=None):
"""
Publish a message to the intended audience.
"""
assert topic is not None
assert data is not None
msg = MgmtMessage(topic=topic, identity=identity, data=data)
self._out_stream.send_multipart(msg.serialize())
class LRUQueue(object):
"""LRUQueue class using ZMQStream/IOLoop for event dispatching"""
def __init__(self, backend_socket, frontend_socket):
self.available_workers = 0
self.workers = []
self.client_nbr = NBR_CLIENTS
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
# Queue worker address for LRU routing
worker_addr, empty, client_addr = msg[:3]
assert self.available_workers < NBR_WORKERS
# add worker back to the list of workers
self.available_workers += 1
self.workers.append(worker_addr)
# Second frame is empty
assert empty == b""
# Third frame is READY or else a client reply address
# If client reply, send rest back to frontend
if client_addr != b"READY":
empty, reply = msg[3:]
# Following frame is empty
assert empty == b""
self.frontend.send_multipart([client_addr, b'', reply])
self.client_nbr -= 1
if self.client_nbr == 0:
# Exit after N messages
self.loop.add_timeout(time.time()+1, self.loop.stop)
if self.available_workers == 1:
# on first recv, start accepting frontend messages
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
# Now get next client request, route to LRU worker
# Client request is [address][empty][request]
client_addr, empty, request = msg
assert empty == b""
# Dequeue and drop the next worker address
self.available_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart([worker_id, b'', client_addr, b'', request])
if self.available_workers == 0:
# stop receiving until workers become available again
self.frontend.stop_on_recv()
class Master(object):
"""
Broker for asynchronous interaction,
But I would like to call it Master!!!
"""
def __init__(self, url_worker, url_client, batch_size,
estimator_update_callable):
context = zmq.Context()
frontend = context.socket(zmq.ROUTER)
frontend.bind(url_client)
backend = context.socket(zmq.ROUTER)
backend.bind(url_worker)
self.available_workers = 0
self.workers = []
self.batch_size = batch_size
self.estimator_update = estimator_update_callable
self.backend = ZMQStream(backend)
self.frontend = ZMQStream(frontend)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
# Queue worker address for LRU routing
worker_addr, empty, client_addr = msg[:3]
# add worker back to the list of workers
self.available_workers += 1
self.workers.append(worker_addr)
# Third frame is READY or else a client reply address
# If client reply, send rest back to frontend
if client_addr != b"READY":
empty, reply = msg[3:]
self.frontend.send_multipart([client_addr, b'', reply])
if self.available_workers == 1:
# on first recv, start accepting frontend messages
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
# Now get next client request, route to LRU worker
# Client request is [address][empty][request]
client_addr, empty, request = msg
request = msgpack.loads(request)
if request[0] == 'reset':
state = request[1]
msg = [b'', client_addr, b'', msgpack.dumps([state])]
self.worker_send(msg)
elif request[0] == 'step':
t = Transition(*request[1:])
self.update(t)
if t.done:
self.frontend.send_multipart([client_addr, b'', b'reset'])
else:
msg = [b'', client_addr, b'', msgpack.dumps([t.next_state])]
self.worker_send(msg)
def worker_send(self, msg):
# Dequeue and drop the next worker address
self.available_workers -= 1
worker_id = self.workers.pop(0)
self.backend.send_multipart([worker_id] + msg)
if self.available_workers == 0:
# stop receiving until workers become available again
self.frontend.stop_on_recv()
class PubSub(BasePubSub):
"""
This class manages application PUB/SUB logic.
"""
NAME = 'ZeroMQ'
def __init__(self, application):
super(PubSub, self).__init__(application)
self.sub_stream = None
self.pub_stream = None
self.zmq_context = None
self.zmq_pub_sub_proxy = None
self.zmq_xpub = None
self.zmq_xsub = None
self.zmq_pub_port = None
self.zmq_sub_address = None
def initialize(self):
self.zmq_context = zmq.Context()
options = self.application.settings['options']
self.zmq_pub_sub_proxy = options.zmq_pub_sub_proxy
# create PUB socket to publish instance events into it
publish_socket = self.zmq_context.socket(zmq.PUB)
# do not try to send messages after closing
publish_socket.setsockopt(zmq.LINGER, 0)
if self.zmq_pub_sub_proxy:
# application started with XPUB/XSUB proxy
self.zmq_xsub = options.zmq_xsub
publish_socket.connect(self.zmq_xsub)
else:
# application started without XPUB/XSUB proxy
if options.zmq_pub_port_shift:
# calculate zmq pub port number
zmq_pub_port = options.port - options.zmq_pub_port_shift
else:
zmq_pub_port = options.zmq_pub_port
self.zmq_pub_port = zmq_pub_port
publish_socket.bind(
"tcp://%s:%s" % (options.zmq_pub_listen, str(self.zmq_pub_port))
)
# wrap pub socket into ZeroMQ stream
self.pub_stream = ZMQStream(publish_socket)
# create SUB socket listening to all events from all app instances
subscribe_socket = self.zmq_context.socket(zmq.SUB)
if self.zmq_pub_sub_proxy:
# application started with XPUB/XSUB proxy
self.zmq_xpub = options.zmq_xpub
subscribe_socket.connect(self.zmq_xpub)
else:
# application started without XPUB/XSUB proxy
self.zmq_sub_address = options.zmq_sub_address
for address in self.zmq_sub_address:
subscribe_socket.connect(address)
subscribe_socket.setsockopt_string(
zmq.SUBSCRIBE,
six.u(CONTROL_CHANNEL)
)
subscribe_socket.setsockopt_string(
zmq.SUBSCRIBE, six.u(ADMIN_CHANNEL)
)
def listen_socket():
# wrap sub socket into ZeroMQ stream and set its on_recv callback
self.sub_stream = ZMQStream(subscribe_socket)
self.sub_stream.on_recv(self.dispatch_published_message)
tornado.ioloop.IOLoop.instance().add_callback(
listen_socket
)
if self.zmq_pub_sub_proxy:
logger.info(
"ZeroMQ XPUB: {0}, XSUB: {1}".format(self.zmq_xpub, self.zmq_xsub)
)
else:
logger.info(
"ZeroMQ PUB - {0}; subscribed to {1}".format(self.zmq_pub_port, self.zmq_sub_address)
)
def publish(self, channel, message, method=None):
"""
Publish message into channel of stream.
"""
method = method or self.DEFAULT_PUBLISH_METHOD
message["message_type"] = method
message = json_encode(message)
to_publish = [utf8(channel), utf8(message)]
self.pub_stream.send_multipart(to_publish)
@coroutine
def dispatch_published_message(self, multipart_message):
"""
Got message, decide what is it and dispatch into right
application handler.
"""
channel = multipart_message[0]
if six.PY3:
channel = channel.decode()
message_data = json_decode(multipart_message[1])
if channel == CONTROL_CHANNEL:
yield self.handle_control_message(message_data)
elif channel == ADMIN_CHANNEL:
yield self.handle_admin_message(message_data)
else:
yield self.handle_channel_message(channel, message_data)
def subscribe_key(self, subscription_key):
self.sub_stream.setsockopt_string(
zmq.SUBSCRIBE, six.u(subscription_key)
)
def unsubscribe_key(self, subscription_key):
self.sub_stream.setsockopt_string(
zmq.UNSUBSCRIBE, six.u(subscription_key)
)
def clean(self):
"""
Properly close ZeroMQ sockets.
"""
if hasattr(self, 'pub_stream') and self.pub_stream:
self.pub_stream.close()
if hasattr(self, 'sub_stream') and self.sub_stream:
self.sub_stream.stop_on_recv()
self.sub_stream.close()
class ZmqStreamlet(Zmqlet):
"""A :class:`ZmqStreamlet` object can send/receive data to/from ZeroMQ stream and invoke callback function. It
has three sockets for input, output and control.
.. warning::
Starting from v0.3.6, :class:`ZmqStreamlet` replaces :class:`Zmqlet` as one of the key components in :class:`jina.peapods.pea.BasePea`.
It requires :mod:`tornado` and :mod:`uvloop` to be installed.
"""
def register_pollin(self):
use_uvloop()
import asyncio
asyncio.set_event_loop(asyncio.new_event_loop())
try:
import tornado.ioloop
self.io_loop = tornado.ioloop.IOLoop.current()
except (ModuleNotFoundError, ImportError):
self.logger.error(
'Since v0.3.6 Jina requires "tornado" as a base dependency, '
'we use its I/O event loop for non-blocking sockets. '
'Please try reinstall via "pip install -U jina" to include this dependency'
)
raise
self.in_sock = ZMQStream(self.in_sock, self.io_loop)
self.out_sock = ZMQStream(self.out_sock, self.io_loop)
self.ctrl_sock = ZMQStream(self.ctrl_sock, self.io_loop)
self.in_sock.stop_on_recv()
def close(self):
"""Close all sockets and shutdown the ZMQ context associated to this `Zmqlet`. """
if not self.is_closed:
# wait until the close signal is received
time.sleep(.01)
for s in self.opened_socks:
s.flush()
super().close()
try:
self.io_loop.stop()
# Replace handle events function, to skip
# None event after sockets are closed.
if hasattr(self.in_sock, '_handle_events'):
self.in_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.out_sock, '_handle_events'):
self.out_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.ctrl_sock, '_handle_events'):
self.ctrl_sock._handle_events = lambda *args, **kwargs: None
except AttributeError as e:
self.logger.error(f'failed to stop. {e}')
def pause_pollin(self):
"""Remove :attr:`in_sock` from the poller """
self.in_sock.stop_on_recv()
def resume_pollin(self):
"""Put :attr:`in_sock` back to the poller """
self.in_sock.on_recv(self._in_sock_callback)
def start(self, callback: Callable[['ProtoMessage'], 'ProtoMessage']):
def _callback(msg, sock_type):
msg = _parse_from_frames(sock_type, msg)
self.bytes_recv += msg.size
self.msg_recv += 1
msg = callback(msg)
if msg:
self.send_message(msg)
self._in_sock_callback = lambda x: _callback(x, self.in_sock_type)
self.in_sock.on_recv(self._in_sock_callback)
self.ctrl_sock.on_recv(lambda x: _callback(x, self.ctrl_sock_type))
if self.out_sock_type == zmq.ROUTER:
self.out_sock.on_recv(lambda x: _callback(x, self.out_sock_type))
self.io_loop.start()
self.io_loop.clear_current()
self.io_loop.close(all_fds=True)
class Socket:
"""Wrapper class for zmq.Socket
This class utilizes a tornado event loop to support using ZmqStream for sending
and receiving messages. Additionally it stages sockets with specific configurations
to allow request/reply sockets to be more robust to timeout conditions and failure states.
Attributes:
logger: Logger instance for all socket activity
loop: Tornado event loop instance
address: Assigned address of the zmq.Socket
protocol: Assigned zmq socket type
ctx: ZMQ context instance
stream: ZmqStream instance
zmq_socket: Underlying zmq.Socket object
"""
def __init__(self, loop, protocol):
"""Constructor for Socket class
Args:
loop: Tornado event loop
protocol: Assigned protocol for the zmq.Socket
"""
self.logger = logging.getLogger("Socket")
self.loop = loop
self.protocol = protocol
self.server = True if (protocol == zmq.PUB
or protocol == zmq.REP) else False
self.ctx = zmq.Context().instance()
self.stream = None
self.zmq_socket = None
self.address = None
self.port = None
self.create_socket(protocol)
def create_socket(self, protocol):
"""Helper function for creating a zmq.Socket of various types with various options
Assumes that request sockets need extra configuration options to prevent erroneous states
when two requests are sent before a reply is received. REQ_RELAXED will drop the first
request and reset the underlying socket automatically allowing the second request to be
processed. Additionally, these options ensure that an event loop can exit even if a send is
pending but hasn't been sent yet.
Args:
protocol: zmq socket type
"""
if protocol == zmq.REQ:
# make sure that replies back to req's are coordinated with header data
self.ctx.setsockopt(zmq.REQ_CORRELATE, 1)
# allow req socket to internall try to reconnect if two sends are sent in a row
self.ctx.setsockopt(zmq.REQ_RELAXED, 1)
# timeout for trying to send
self.ctx.setsockopt(zmq.SNDTIMEO, 1000)
# ensure that the socket doesn't block on close
self.ctx.setsockopt(zmq.LINGER, 0)
self.zmq_socket = self.ctx.socket(protocol)
def connect(self, address, port):
"""Connect the socket to a local or remote address:port
Args:
address: Decimal separated string (eg, 127.0.0.1) where service is bound
port: int associated with service port
"""
self.zmq_socket.connect("tcp://{}:{}".format(address, port))
self.address = address # 127.0.0.1
self.port = port # 10001
self.start_stream()
return (self.address, self.port)
def disconnect(self):
"""Stop the zmqStream and disconnect the underlying socket
"""
# TODO(pickledgator): Figure out why this fails with error: Socket operation on non-socket
self.stop_stream()
try:
self.zmq_socket.disconnect("tcp://{}:{}".format(
self.address, self.port))
except Exception as e:
pass
def get_local_ip(self):
"""Identifies the ip address of the local node
Returns:
String: Decimal separated string (eg, 127.0.0.1)
"""
# TODO(pickledgator): Check robustness of this strategy and switch to netifaces if needed
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
s.connect(("10.255.255.255", 1))
ip = s.getsockname()[0]
except:
ip = "127.0.0.1"
finally:
s.close()
return ip
def set_filter(self, filter_string=""):
"""Helper function to enable zmq.SUBSCRIBER filters
In zmq, the filter (string or bytes) can be used to screen multi-part messages such
that only messages with the first element of the array match an associated string.
Currently, we configure the subscriber sockets to receive all messages from any publisher.
TODO(pickledgator): Add support for filter->separate callbacks?
Args:
filter_string: Setup socket to only allow multi-part messages whose first element
match this string
"""
# "" is a wildcard to accept all messages
self.zmq_socket.setsockopt_string(zmq.SUBSCRIBE, filter_string)
def bind(self):
"""Bind the underlying zmq socket to an ip on the local machine at a random available port
Also kicks off the zmqStream after binding.
Returns:
(String, int): Tuple containing the address string and the port chosen
"""
# TODO(pickledgator): Find specific range that has the most availability
ip = self.get_local_ip()
port = self.zmq_socket.bind_to_random_port("tcp://{}".format(ip),
min_port=10001,
max_port=20000,
max_tries=100)
self.address = ip
self.port = port
self.start_stream()
return (self.address, self.port)
def unbind(self):
"""Reverse the bind of the underlying zmq socket and stop the zmqStream
"""
# TODO(pickledgator): Figure out why this fails with error: Socket operation on non-socket
self.stop_stream()
try:
self.zmq_socket.unbind("tcp://{}:{}".format(
self.address, self.port))
except Exception as e:
pass
def send(self, message):
"""Identifies the correct underlying zmq send method based on the type of message
Args:
message: Message to be sent (string or bytes)
"""
self.logger.debug("Sending message: {}".format(message))
if type(message) == str:
# assumes string
self.stream.send_string(message)
else:
# assumes bytes
self.stream.send(message)
def start_stream(self):
if not self.stream:
self.stream = ZMQStream(self.zmq_socket, self.loop)
def stop_stream(self):
if self.stream:
if not self.stream.closed():
# TODO(pickledgator): does this block if no recv?
self.stream.stop_on_recv()
self.stream.close()
self.stream = None
def cycle_socket(self):
self.close()
self.create_socket(self.protocol)
# self.connect(self.address, self.port)
def receive(self, handler, timeout_ms=None, timeout_callback=None):
def msg_handler(handler, timeout, message):
# this callback receives a message list, with one element, so just pass the contents to the
# application handler
handler(message[0].decode("utf-8"))
# if we received the message, then we need to cancel the watchdog timeout from
# the last receive call
if timeout:
self.loop.remove_timeout(timeout)
def handle_timeout(timeout_callback):
if timeout_callback:
timeout_callback()
# the event that we hit a timeout, cycle the socket so that it doesn't
# get stuck in a weird state
self.cycle_socket()
if self.stream:
if timeout_ms:
# if we want to detect when recv fails, setup a timeout that cleans up the
# socket (RequestClients)
timeout = self.loop.call_later(
timeout_ms / 1000.0,
functools.partial(handle_timeout, timeout_callback))
# always set the handler, in case it changed
self.stream.on_recv(
functools.partial(msg_handler, handler, timeout))
else:
# handle cases when we dont want to put a timeout on the recv function (subscribers)
self.stream.on_recv(
functools.partial(msg_handler, handler, None))
else:
self.logger.error("Stream is not open")
def close(self):
try:
self.disconnect()
except:
pass
self.zmq_socket.close()
class LRUQueue(object):
"""LRUQueue class using ZMQStream/IOLoop for event dispatching"""
def __init__(self, backend_socket, frontend_socket):
self.available_workers = 0
self.workers = []
self.client_nbr = NBR_CLIENTS
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
# Queue worker address for LRU routing
worker_addr, empty, client_addr = msg[:3]
assert self.available_workers < NBR_WORKERS
# add worker back to the list of workers
self.available_workers += 1
self.workers.append(worker_addr)
# Second frame is empty
assert empty == b""
# Third frame is READY or else a client reply address
# If client reply, send rest back to frontend
if client_addr != b"READY":
empty, reply = msg[3:]
# Following frame is empty
assert empty == b""
self.frontend.send_multipart([client_addr, b'', reply])
self.client_nbr -= 1
if self.client_nbr == 0:
# Exit after N messages
self.loop.add_timeout(time.time() + 1, self.loop.stop)
if self.available_workers == 1:
# on first recv, start accepting frontend messages
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
# Now get next client request, route to LRU worker
# Client request is [address][empty][request]
client_addr, empty, request = msg
assert empty == b""
# Dequeue and drop the next worker address
self.available_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart([worker_id, b'', client_addr, b'', request])
if self.available_workers == 0:
# stop receiving until workers become available again
self.frontend.stop_on_recv()
class CloneServer(object):
# Our server is defined by these properties
ctx = None # Context wrapper
kvmap = None # Key-value store
bstar = None # Binary Star
sequence = 0 # How many updates so far
port = None # Main port we're working on
peer = None # Main port of our peer
publisher = None # Publish updates and hugz
collector = None # Collect updates from clients
subscriber = None # Get updates from peer
pending = None # Pending updates from client
primary = False # True if we're primary
master = False # True if we're master
slave = False # True if we're slave
def __init__(self, primary=True, ports=(5556,5566)):
self.primary = primary
if primary:
self.port, self.peer = ports
frontend = "tcp://*:5003"
backend = "tcp://localhost:5004"
self.kvmap = {}
else:
self.peer, self.port = ports
frontend = "tcp://*:5004"
backend = "tcp://localhost:5003"
self.ctx = zmq.Context.instance()
self.pending = []
self.bstar = BinaryStar(primary, frontend, backend)
self.bstar.register_voter("tcp://*:%i" % self.port, zmq.ROUTER, self.handle_snapshot)
# Set up our clone server sockets
self.publisher = self.ctx.socket(zmq.PUB)
self.collector = self.ctx.socket(zmq.SUB)
self.collector.setsockopt(zmq.SUBSCRIBE, b'')
self.publisher.bind("tcp://*:%d" % (self.port + 1))
self.collector.bind("tcp://*:%d" % (self.port + 2))
# Set up our own clone client interface to peer
self.subscriber = self.ctx.socket(zmq.SUB)
self.subscriber.setsockopt(zmq.SUBSCRIBE, b'')
self.subscriber.connect("tcp://localhost:%d" % (self.peer + 1))
# Register state change handlers
self.bstar.master_callback = self.become_master
self.bstar.slave_callback = self.become_slave
# Wrap sockets in ZMQStreams for IOLoop handlers
self.publisher = ZMQStream(self.publisher)
self.subscriber = ZMQStream(self.subscriber)
self.collector = ZMQStream(self.collector)
# Register our handlers with reactor
self.collector.on_recv(self.handle_collect)
self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)
self.hugz_callback = PeriodicCallback(self.send_hugz, 1000)
# basic log formatting:
logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
level=logging.INFO)
def start(self):
# start periodic callbacks
self.flush_callback.start()
self.hugz_callback.start()
# Run bstar reactor until process interrupted
try:
self.bstar.start()
except KeyboardInterrupt:
pass
def handle_snapshot(self, socket, msg):
"""snapshot requests"""
if msg[1] != "ICANHAZ?" or len(msg) != 3:
logging.error("E: bad request, aborting")
dump(msg)
self.bstar.loop.stop()
return
identity, request = msg[:2]
if len(msg) >= 3:
subtree = msg[2]
# Send state snapshot to client
route = Route(socket, identity, subtree)
# For each entry in kvmap, send kvmsg to client
for k,v in self.kvmap.items():
send_single(k,v,route)
# Now send END message with sequence number
logging.info("I: Sending state shapshot=%d" % self.sequence)
socket.send(identity, zmq.SNDMORE)
kvmsg = KVMsg(self.sequence)
kvmsg.key = "KTHXBAI"
kvmsg.body = subtree
kvmsg.send(socket)
def handle_collect(self, msg):
"""Collect updates from clients
If we're master, we apply these to the kvmap
If we're slave, or unsure, we queue them on our pending list
"""
kvmsg = KVMsg.from_msg(msg)
if self.master:
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.send(self.publisher)
ttl = kvmsg.get('ttl')
if ttl is not None:
kvmsg['ttl'] = time.time() + ttl
kvmsg.store(self.kvmap)
logging.info("I: publishing update=%d", self.sequence)
else:
# If we already got message from master, drop it, else
# hold on pending list
if not self.was_pending(kvmsg):
self.pending.append(kvmsg)
def was_pending(self, kvmsg):
"""If message was already on pending list, remove and return True.
Else return False.
"""
found = False
for idx, held in enumerate(self.pending):
if held.uuid == kvmsg.uuid:
found = True
break
if found:
self.pending.pop(idx)
return found
def flush_ttl(self):
"""Purge ephemeral values that have expired"""
if self.kvmap:
for key,kvmsg in self.kvmap.items():
self.flush_single(kvmsg)
def flush_single(self, kvmsg):
"""If key-value pair has expired, delete it and publish the fact
to listening clients."""
if kvmsg.get('ttl', 0) <= time.time():
kvmsg.body = ""
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.send(self.publisher)
del self.kvmap[kvmsg.key]
logging.info("I: publishing delete=%d", self.sequence)
def send_hugz(self):
"""Send hugz to anyone listening on the publisher socket"""
kvmsg = KVMsg(self.sequence)
kvmsg.key = "HUGZ"
kvmsg.body = ""
kvmsg.send(self.publisher)
# ---------------------------------------------------------------------
# State change handlers
def become_master(self):
"""We're becoming master
The backup server applies its pending list to its own hash table,
and then starts to process state snapshot requests.
"""
self.master = True
self.slave = False
# stop receiving subscriber updates while we are master
self.subscriber.stop_on_recv()
# Apply pending list to own kvmap
while self.pending:
kvmsg = self.pending.pop(0)
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.store(self.kvmap)
logging.info ("I: publishing pending=%d", self.sequence)
def become_slave(self):
"""We're becoming slave"""
# clear kvmap
self.kvmap = None
self.master = False
self.slave = True
self.subscriber.on_recv(self.handle_subscriber)
def handle_subscriber(self, msg):
"""Collect updates from peer (master)
We're always slave when we get these updates
"""
if self.master:
logging.warn("received subscriber message, but we are master %s", msg)
return
# Get state snapshot if necessary
if self.kvmap is None:
self.kvmap = {}
snapshot = self.ctx.socket(zmq.DEALER)
snapshot.linger = 0
snapshot.connect("tcp://localhost:%i" % self.peer)
logging.info ("I: asking for snapshot from: tcp://localhost:%d",
self.peer)
snapshot.send_multipart(["ICANHAZ?", ''])
while True:
try:
kvmsg = KVMsg.recv(snapshot)
except KeyboardInterrupt:
# Interrupted
self.bstar.loop.stop()
return
if kvmsg.key == "KTHXBAI":
self.sequence = kvmsg.sequence
break # Done
kvmsg.store(self.kvmap)
logging.info ("I: received snapshot=%d", self.sequence)
# Find and remove update off pending list
kvmsg = KVMsg.from_msg(msg)
# update integer ttl -> timestamp
ttl = kvmsg.get('ttl')
if ttl is not None:
kvmsg['ttl'] = time.time() + ttl
if kvmsg.key != "HUGZ":
if not self.was_pending(kvmsg):
# If master update came before client update, flip it
# around, store master update (with sequence) on pending
# list and use to clear client update when it comes later
self.pending.append(kvmsg)
# If update is more recent than our kvmap, apply it
if (kvmsg.sequence > self.sequence):
self.sequence = kvmsg.sequence
kvmsg.store(self.kvmap)
logging.info ("I: received update=%d", self.sequence)
class ZmqWorker(object, LoggingMixin):
"""
This is the ZMQ worker implementation.
The worker will register a :class:`ZMQStream` with the configured
:class:`zmq.Socket` and :class:`zmq.eventloop.ioloop.IOLoop` instance.
Upon `ZMQStream.on_recv` the configured `processors` will be executed
with the deserialized context and the result will be published through the
configured `zmq.socket`.
"""
def __init__(self, insocket, outsocket, mgmt, processing, log_handler,
log_level, io_loop=None):
"""
Initialize the `ZMQStream` with the `insocket` and `io_loop` and store
the `outsocket`.
`insocket` should be of the type `zmq.socket.PULL` `outsocket` should
be of the type `zmq.socket.PUB`
`mgmt` is an instance of `spyder.core.mgmt.ZmqMgmt` that handles
communication between master and worker processes.
"""
LoggingMixin.__init__(self, log_handler, log_level)
self._insocket = insocket
self._io_loop = io_loop or IOLoop.instance()
self._outsocket = outsocket
self._processing = processing
self._mgmt = mgmt
self._in_stream = ZMQStream(self._insocket, self._io_loop)
self._out_stream = ZMQStream(self._outsocket, self._io_loop)
def _quit(self, msg):
"""
The worker is quitting, stop receiving messages.
"""
if ZMQ_SPYDER_MGMT_WORKER_QUIT == msg.data:
self.stop()
def _receive(self, msg):
"""
We have a message!
`msg` is a serialized version of a `DataMessage`.
"""
message = DataMessage(msg)
try:
# this is the real work we want to do
curi = self._processing(message.curi)
message.curi = curi
except:
# catch any uncaught exception and only log it as CRITICAL
self._logger.critical(
"worker::Uncaught exception executing the worker for URL %s!" %
(message.curi.url,))
self._logger.critical("worker::%s" % (traceback.format_exc(),))
# finished, now send the result back to the master
self._out_stream.send_multipart(message.serialize())
def start(self):
"""
Start the worker.
"""
self._mgmt.add_callback(ZMQ_SPYDER_MGMT_WORKER, self._quit)
self._in_stream.on_recv(self._receive)
def stop(self):
"""
Stop the worker.
"""
# stop receiving
self._in_stream.stop_on_recv()
self._mgmt.remove_callback(ZMQ_SPYDER_MGMT_WORKER, self._quit)
# but work on anything we might already have
self._in_stream.flush()
self._out_stream.flush()
def close(self):
"""
Close all open sockets.
"""
self._in_stream.close()
self._insocket.close()
self._out_stream.close()
self._outsocket.close()
class TornadoRPCService(RPCServiceBase): #{
""" An asynchronous RPC service that takes requests over a ROUTER socket.
Using Tornado compatible IOLoop and ZMQStream from PyZMQ.
"""
def __init__(self, context=None, ioloop=None, **kwargs): #{
"""
Parameters
==========
ioloop : IOLoop
An existing IOLoop instance, if not passed, zmq.IOLoop.instance()
will be used.
context : Context
An existing Context instance, if not passed, zmq.Context.instance()
will be used.
serializer : Serializer
An instance of a Serializer subclass that will be used to serialize
and deserialize args, kwargs and the result.
"""
assert context is None or isinstance(context, zmq.Context)
self.context = context if context is not None else zmq.Context.instance()
self.ioloop = IOLoop.instance() if ioloop is None else ioloop
self._is_started = False
super(TornadoRPCService, self).__init__(**kwargs)
#}
def _create_socket(self): #{
super(TornadoRPCService, self)._create_socket()
socket = self.context.socket(zmq.ROUTER)
self.socket = ZMQStream(socket, self.ioloop)
#}
def _handle_request(self, msg_list): #{
"""
Handle an incoming request.
The request is received as a multipart message:
[<id>..<id>, b'|', req_id, proc_name, <ser_args>, <ser_kwargs>, <ignore>]
First, the service sends back a notification that the message was
indeed received:
[<id>..<id>, b'|', req_id, b'ACK', service_id]
Next, the actual reply depends on if the call was successful or not:
[<id>..<id>, b'|', req_id, b'OK', <serialized result>]
[<id>..<id>, b'|', req_id, b'FAIL', <JSON dict of ename, evalue, traceback>]
Here the (ename, evalue, traceback) are utf-8 encoded unicode.
"""
req = self._parse_request(msg_list)
if req is None:
return
self._send_ack(req)
ignore = req['ignore']
try:
# raise any parsing errors here
if req['error']:
raise req['error']
# call procedure
res = req['proc'](*req['args'], **req['kwargs'])
except Exception:
not ignore and self._send_fail(req)
else:
def send_future_result(fut):
try: res = fut.result()
except: not ignore and self._send_fail(req)
else: not ignore and self._send_ok(req, res)
if isinstance(res, Future):
self.ioloop.add_future(res, send_future_result)
else:
not ignore and self._send_ok(req, res)
#}
def start(self): #{
""" Start the RPC service (non-blocking) """
assert self._is_started == False, "already started"
# register IOLoop callback
self._is_started = True
self.socket.on_recv(self._handle_request)
#}
def stop(self): #{
""" Stop the RPC service (non-blocking) """
# register IOLoop callback
self.socket.stop_on_recv()
self._is_started = False
#}
def serve(self): #{
""" Serve RPC requests (blocking) """
if not self._is_started:
self.start()
return self.ioloop.start()
class Channel(object):
"""Mother of all channels. Defines the interface.
Callbacks:
The callbacks will receive the channel as first parameter and
the message as second parameter. The error callback will get
the stream where the error occured as second parameter.
Attributes:
* stream_in, stream_out : the streams for eventlopp handling
* serializer : the serializer used
"""
def __init__(self, socket_in, socket_out, serializer):
self.stream_in = ZMQStream(socket_in)
self.stream_out = ZMQStream(socket_out)
self.serializer = serializer
self._cb_receive = None
self._cb_send = None
self._cb_error = None
self._chan_id = id(self)
return
def on_receive(self, callback):
"""Set callback to invoke when a message was received.
"""
self.stream_in.stop_on_recv()
self._cb_receive = callback
if callback:
self.stream_in.on_recv(self._on_recv)
return
def on_send(self, callback):
"""Set callback to invoke when a message was sent.
"""
self.stream_out.stop_on_send()
self._cb_send = callback
if callback:
self.stream_out.on_send(self._on_send)
return
def on_error(self, callback):
"""Set callback to invoke when an error event occured.
"""
self.stream_in.stop_on_err()
self.stream_out.stop_on_err()
self._cb_error = callback
if callback:
self.stream_in.on_err(self._on_err_in)
self.stream_out.on_err(self._on_err_out)
return
def send(self, message):
"""Send given message.
"""
m = self.serializer.serialize(message)
if self.serializer.multipart:
self.stream_out.send_multipart(m)
else:
self.stream_out.send(m)
return
def _on_recv(self, msg):
"""Helper interfacing w/ streams.
"""
if self.serializer.multipart:
msg = self.serializer.deserialize(msg)
else:
msg = self.serializer.deserialize(msg[0])
self._cb_receive(self, msg)
return
def _on_send(self, sent, _):
"""Helper interfacing w/ streams.
"""
msg = sent[0]
if self.serializer.multipart:
msg = self.serializer.deserialize(msg)
else:
msg = self.serializer.deserialize(msg[0])
self._cb_send(self, msg)
return
def _on_err_in(self):
self._cb_error(self, self.stream_in)
return
def _on_err_out(self):
self._cb_error(self, self.stream_out)
return
def __hash__(self):
return self._chan_id
def __eq__(self, other):
if not isinstance(other, self.__class__):
return False
return self._chan_id == other._chan_id
def __neq__(self, other):
if not isinstance(other, self.__class__):
return True
return self._chan_id != other._chan_id
class ZmqStreamlet(Zmqlet):
"""A :class:`ZmqStreamlet` object can send/receive data to/from ZeroMQ stream and invoke callback function. It
has three sockets for input, output and control.
.. warning::
Starting from v0.3.6, :class:`ZmqStreamlet` replaces :class:`Zmqlet` as one of the key components in :class:`jina.peapods.runtime.BasePea`.
It requires :mod:`tornado` and :mod:`uvloop` to be installed.
"""
def _register_pollin(self):
"""Register :attr:`in_sock`, :attr:`ctrl_sock` and :attr:`out_sock` in poller."""
with ImportExtensions(required=True):
import tornado.ioloop
get_or_reuse_loop()
self.io_loop = tornado.ioloop.IOLoop.current()
self.in_sock = ZMQStream(self.in_sock, self.io_loop)
self.out_sock = ZMQStream(self.out_sock, self.io_loop)
self.ctrl_sock = ZMQStream(self.ctrl_sock, self.io_loop)
self.in_sock.stop_on_recv()
def close(self):
"""Close all sockets and shutdown the ZMQ context associated to this `Zmqlet`.
.. note::
This method is idempotent.
"""
if not self.is_closed:
# wait until the close signal is received
time.sleep(0.01)
for s in self.opened_socks:
s.flush()
super().close()
if hasattr(self, 'io_loop'):
try:
self.io_loop.stop()
# Replace handle events function, to skip
# None event after sockets are closed.
if hasattr(self.in_sock, '_handle_events'):
self.in_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.out_sock, '_handle_events'):
self.out_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.ctrl_sock, '_handle_events'):
self.ctrl_sock._handle_events = lambda *args, **kwargs: None
except AttributeError as e:
self.logger.error(f'failed to stop. {e!r}')
def pause_pollin(self):
"""Remove :attr:`in_sock` from the poller """
self.in_sock.stop_on_recv()
def resume_pollin(self):
"""Put :attr:`in_sock` back to the poller """
self.in_sock.on_recv(self._in_sock_callback)
def start(self, callback: Callable[['Message'], 'Message']):
"""
Open all sockets and start the ZMQ context associated to this `Zmqlet`.
:param callback: callback function to receive the protobuf message
"""
def _callback(msg, sock_type):
msg = _parse_from_frames(sock_type, msg)
self.bytes_recv += msg.size
self.msg_recv += 1
msg = callback(msg)
if msg:
self.send_message(msg)
self._in_sock_callback = lambda x: _callback(x, self.in_sock_type)
self.in_sock.on_recv(self._in_sock_callback)
self.ctrl_sock.on_recv(lambda x: _callback(x, self.ctrl_sock_type))
if self.out_sock_type == zmq.ROUTER:
self.out_sock.on_recv(lambda x: _callback(x, self.out_sock_type))
self.io_loop.start()
self.io_loop.clear_current()
self.io_loop.close(all_fds=True)
class LRUQueue(object):
"""LRUQueue class using ZMQStream/IOLoop for event dispatching"""
def __init__(self, backend_socket, frontend_socket, heartbeat_socket):
self.available_workers = 0
self.unique_workers = []
self.workers = []
self.task_queue = []
self.task_count = 0
self.worker_stat = {}
self.client_nbr = NBR_CLIENTS
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.heartbeat = ZMQStream(heartbeat_socket)
self.backend.on_recv(self.handle_backend)
self.heartbeat.on_recv(self.handle_heartbeat)
self.loop = IOLoop.instance()
self.mpq = Queue()
self.worker_status = {}
def readyWorker(self, worker):
"""Prepare a worker by popping in the queue and placing inside again"""
pass
def purge(self):
"""Look for & kill expired workers."""
pass
# TODO: need to handle heartbeats itself, if the time is not new
# for N seconds, remove it from the available worker list
# Just time probably.
def handle_heartbeat(self, msg):
# print("Received heartbeat: ", msg)
worker_id, _, _, status, last_beat = msg
curr_stat = {}
curr_stat['last_beat'] = last_beat
curr_stat['status'] = status
self.worker_status[worker_id] = curr_stat
self.heartbeat.send_multipart([msg[0], b'', b"Pong..."])
print("Curr time is: {}".format(str(int(time.time()))))
print(self.worker_status)
# Is this the best place to put this?
self.purge()
def handle_backend(self, msg):
# Queue worker address for LRU routing
worker_addr, empty, client_addr = msg[:3]
assert self.available_workers < NBR_WORKERS
# add worker back to the list of workers
self.available_workers += 1
print('Added to available workers (new total): ',
self.available_workers)
# self.workers.append(worker_addr)
self.workers.insert(0, worker_addr)
self.unique_workers.append(worker_addr)
# Second frame is empty
assert empty == b""
# Third frame is READY or else a client reply address
if client_addr == b'READY':
print("Received READY signal from {} BE".format(
worker_addr.decode("utf-8")))
# If client reply, send rest back to frontend
if client_addr != b"READY":
empty, reply = msg[3:]
print("some reply: {}".format(reply.decode('utf-8')))
# Following frame is empty
assert empty == b""
print("Received task, done by {}, to be sent to {}".format(
worker_addr.decode("utf-8"), client_addr.decode("utf-8")))
print("Checking task list which has {} tasks left".format(
len(self.task_queue)))
# Check if still some task is available
if self.task_queue:
task = self.task_queue.pop()
print("Number of workers: {}".format(len(self.workers)))
worker_id = self.workers.pop()
self.available_workers -= 1
print("Sending to {}".format(worker_id))
if worker_id.decode('utf-8') in self.worker_stat:
self.worker_stat[worker_id.decode('utf-8')] = int(
self.worker_stat[worker_id.decode('utf-8')]) + 1
else:
self.worker_stat[worker_id.decode('utf-8')] = 1
print("Stat of worker:",
self.worker_stat[worker_id.decode("utf-8")])
self.backend.send_multipart(
[worker_id, b'', client_addr, b'',
task.encode('ascii')])
self.task_count += 1
else:
print(
"No more tasks available. Sent a total of {} tasks to {} unique workers"
.format(str(self.task_count),
len(set(self.unique_workers))))
print("stats:", self.worker_stat)
# Receive something from backend and if the same length as the task queue
# Aggregate and then send a response to the client
# Only reply when broker is done collecting?
self.frontend.send_multipart([client_addr, b'', reply])
# Should change this? Or I should break down the data
# received on the frontend to multiple tasks for different workers.
# a queue of tasks.
if self.available_workers == 1:
# on first recv, start accepting frontend messages
self.frontend.on_recv(self.handle_frontend)
# If a worker is available, pop a task from the queue.
def distribute_tasks(self, client_addr):
for _ in range(len(self.workers) - 1):
self.available_workers -= 1
task = self.task_queue.pop()
worker_id = self.workers.pop()
print("Sending to {}".format(worker_id))
if worker_id.decode('utf-8') in self.worker_stat:
self.worker_stat[worker_id.decode('utf-8')] = int(
self.worker_stat[worker_id.decode('utf-8')]) + 1
else:
self.worker_stat[worker_id.decode('utf-8')] = 1
self.task_count += 1
self.backend.send_multipart(
[worker_id, b'', client_addr, b'',
task.encode('ascii')])
def handle_frontend(self, msg):
print("Entered handle_frontend()...")
# Now get next client request, route to LRU worker
# Client request is [address][empty][request]
client_addr, empty, request = msg
print(client_addr)
assert empty == b""
tc = TaskCreator.TaskCreator(request)
self.task_queue = tc.parse_request()
print("Task Queue:", self.task_queue)
self.distribute_tasks(client_addr)
if self.available_workers == 0:
print('Stopping reception until workers are available.')
# stop receiving until workers become available again
self.frontend.stop_on_recv()
