class SubQueue(BaseQueue):
def __init__(self, connect, callback, ioloop=None):
super().__init__(zmq.PULL)
self.socket.connect(connect)
self.stream = ZMQStream(self.socket, ioloop)
self.stream.on_recv(callback)
self.stream.flush()
class RequestQueue(BaseQueue):
def __init__(self, connect, ioloop=None):
super().__init__(zmq.REQ)
self.socket.connect(connect)
self.socket.SNDTIMEO = 2000
self.socket.RCVTIMEO = 2000
self.stream = ZMQStream(self.socket, ioloop)
self.stream.flush()
def __getattr__(self, name):
return getattr(self.stream, name)
class ReplyQueue(BaseQueue):
def __init__(self, bind, callback, ioloop=None):
super().__init__(zmq.REP)
self.socket.bind(bind)
self.socket.SNDTIMEO = 1000
self.socket.RCVTIMEO = 1000
self.stream = ZMQStream(self.socket, ioloop)
self.stream.on_recv(callback)
self.stream.flush()
def __getattr__(self, name):
return getattr(self.stream, name)
class PubQueue(BaseQueue):
def __init__(self, bind, ioloop=None):
super().__init__(zmq.PUSH)
self.socket.bind(bind)
self.stream = ZMQStream(self.socket, ioloop)
self.stream.on_send(self.__on_send)
def __on_send(self, msq, status):
print(msq, status)
def send(self, data):
self.stream.send(data)
self.stream.flush()
def send_string(self, data):
self.stream.send_string(data)
self.stream.flush()
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 Client:
def __init__(self):
self._current_packet_id = 0
self._request_response_map = {}
self._new_response_condition = Condition()
self._packet_id_lock = Lock()
self._notification_callbacks = []
def connect_to_server(self, server):
self.context = zmq.Context()
self.socket = self.context.socket(zmq.DEALER)
tmpdir = tempfile.gettempdir()
connection_string = 'ipc://%s/%s_%s' % (tmpdir, 'zmq', server)
self.socket.connect(connection_string)
io_loop = ioloop.IOLoop(ioloop.ZMQPoller())
self.stream = ZMQStream(self.socket, io_loop)
# create partial function that has self as first argument
callback = partial(_on_recv, self)
self.stream.on_recv(callback)
self.event_loop = EventLoop(io_loop)
self.event_loop.start()
def disconnect(self):
self.stream.flush()
self.event_loop.stop()
self.socket.close()
def register_notification_callback(self, callback):
# check a valid function has been past
assert callable(callback)
self._notification_callbacks.append(callback)
def request_queue_list_update(self):
pass
def submit_job_request(self, request, timeout=None):
params = JsonRpc.jobrequest_to_json_params(request)
packet_id = self._next_packet_id()
jsonrpc = JsonRpc.generate_request(packet_id, 'submitJob', params)
self._send_request(packet_id, jsonrpc)
response = self._wait_for_response(packet_id, timeout)
# Timeout
if response == None:
return None
# if we an error occurred then throw an exception
if 'error' in response:
exception = JobRequestException(response['id'],
response['error']['code'],
response['error']['message'])
raise exception
# otherwise return the molequeue id
return response['result']['moleQueueId']
def cancel_job(self):
# TODO
pass
def lookup_job(self, molequeue_id, timeout=None):
params = {'moleQueueId': molequeue_id}
packet_id = self._next_packet_id()
jsonrpc = JsonRpc.generate_request(packet_id, 'lookupJob', params)
self._send_request(packet_id, jsonrpc)
response = self._wait_for_response(packet_id, timeout)
# Timeout
if response == None:
return None
# if we an error occurred then throw an exception
if 'error' in response:
exception = JobRequestInformationException(
response['id'], reponse['error']['data'],
reponse['error']['code'], reponse['error']['message'])
raise exception
jobrequest = JsonRpc.json_to_jobrequest(response)
return jobrequest
def _on_response(self, packet_id, msg):
if packet_id in self._request_response_map:
self._new_response_condition.acquire()
self._request_response_map[packet_id] = msg
# notify any threads waiting that their response may have arrived
self._new_response_condition.notify_all()
self._new_response_condition.release()
# TODO Convert raw JSON into a Python class
def _on_notification(self, msg):
for callback in self._notification_callbacks:
callback(msg)
def _next_packet_id(self):
with self._packet_id_lock:
self._current_packet_id += 1
next = self._current_packet_id
return next
def _send_request(self, packet_id, jsonrpc):
# add to request map so we know we are waiting on response for this packet
# id
self._request_response_map[packet_id] = None
self.stream.send(str(jsonrpc))
self.stream.flush()
def _wait_for_response(self, packet_id, timeout):
try:
start = time.time()
# wait for the response to come in
self._new_response_condition.acquire()
while self._request_response_map[packet_id] == None:
# need to set a wait time otherwise the wait can't be interrupted
# See http://bugs.python.org/issue8844
wait_time = sys.maxint
if timeout != None:
wait_time = timeout - (time.time() - start)
if wait_time <= 0:
break
self._new_response_condition.wait(wait_time)
response = self._request_response_map.pop(packet_id)
self._new_response_condition.release()
return response
except KeyboardInterrupt:
self.event_loop.stop()
raise
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 Client:
def __init__(self):
self._current_packet_id = 0
self._request_response_map = {}
self._new_response_condition = Condition()
self._packet_id_lock = Lock()
self._notification_callbacks = []
def connect_to_server(self, server):
self.context = zmq.Context()
self.socket = self.context.socket(zmq.DEALER)
tmpdir = tempfile.gettempdir()
connection_string = 'ipc://%s/%s_%s' % (tmpdir, 'zmq', server)
self.socket.connect(connection_string)
io_loop = ioloop.IOLoop(ioloop.ZMQPoller())
self.stream = ZMQStream(self.socket, io_loop)
# create partial function that has self as first argument
callback = partial(_on_recv, self)
self.stream.on_recv(callback)
self.event_loop = EventLoop(io_loop)
self.event_loop.start()
def disconnect(self):
self.stream.flush()
self.event_loop.stop()
self.socket.close()
def register_notification_callback(self, callback):
# check a valid function has been past
assert callable(callback)
self._notification_callbacks.append(callback)
def request_queue_list_update(self):
pass
def submit_job_request(self, request, timeout=None):
params = JsonRpc.jobrequest_to_json_params(request)
packet_id = self._next_packet_id()
jsonrpc = JsonRpc.generate_request(packet_id,
'submitJob',
params)
self._send_request(packet_id, jsonrpc)
response = self._wait_for_response(packet_id, timeout)
# Timeout
if response == None:
return None
# if we an error occurred then throw an exception
if 'error' in response:
exception = JobRequestException(response['id'],
response['error']['code'],
response['error']['message'])
raise exception
# otherwise return the molequeue id
return response['result']['moleQueueId']
def cancel_job(self):
# TODO
pass
def lookup_job(self, molequeue_id, timeout=None):
params = {'moleQueueId': molequeue_id}
packet_id = self._next_packet_id()
jsonrpc = JsonRpc.generate_request(packet_id,
'lookupJob',
params)
self._send_request(packet_id, jsonrpc)
response = self._wait_for_response(packet_id, timeout)
# Timeout
if response == None:
return None
# if we an error occurred then throw an exception
if 'error' in response:
exception = JobRequestInformationException(response['id'],
reponse['error']['data'],
reponse['error']['code'],
reponse['error']['message'])
raise exception
jobrequest = JsonRpc.json_to_jobrequest(response)
return jobrequest
def _on_response(self, packet_id, msg):
if packet_id in self._request_response_map:
self._new_response_condition.acquire()
self._request_response_map[packet_id] = msg
# notify any threads waiting that their response may have arrived
self._new_response_condition.notify_all()
self._new_response_condition.release()
# TODO Convert raw JSON into a Python class
def _on_notification(self, msg):
for callback in self._notification_callbacks:
callback(msg)
def _next_packet_id(self):
with self._packet_id_lock:
self._current_packet_id += 1
next = self._current_packet_id
return next
def _send_request(self, packet_id, jsonrpc):
# add to request map so we know we are waiting on response for this packet
# id
self._request_response_map[packet_id] = None
self.stream.send(str(jsonrpc))
self.stream.flush()
def _wait_for_response(self, packet_id, timeout):
try:
start = time.time()
# wait for the response to come in
self._new_response_condition.acquire()
while self._request_response_map[packet_id] == None:
# need to set a wait time otherwise the wait can't be interrupted
# See http://bugs.python.org/issue8844
wait_time = sys.maxint
if timeout != None:
wait_time = timeout - (time.time() - start)
if wait_time <= 0:
break;
self._new_response_condition.wait(wait_time)
response = self._request_response_map.pop(packet_id)
self._new_response_condition.release()
return response
except KeyboardInterrupt:
self.event_loop.stop()
raise
class Stream:
'''
This is the main class that interacts with the zmq library to send, and
receive messages.
'''
def __init__(self, socket, stream_info, path, on_recv=None, on_send=None, loop=None):
'''
Initializes instance of Stream.
@param socket - zmq socket that has alredy been bound
@param stream_info - this streams definition from the yaml config
@param path - path to the socket on the disk
@parma on_recv - callback that processes received messages
@parma on_send - callback that processes sent messages
@param loop - loop this socket will belong to. Default is global async loop.
'''
self._path = path
self._recv_type = stream_info.recv_type
self._send_type = stream_info.send_type
self._reply_type = stream_info.reply_type
self._on_recv = on_recv
self._on_send = on_send
self._stream = ZMQStream(socket, io_loop=loop)
if self._on_recv is not None:
self._stream.on_recv(self._recv_wrapper)
if self._on_send is not None:
self._stream.on_send(self._send_wrapper)
@property
def recv_type(self):
return self._recv_type
@property
def send_type(self):
return self._send_type
@property
def reply_type(self):
return self._reply_type
def on_send(self, callback):
'''
Set the callback to be invoked on every send command. on_send(None)
disables this callback.
@param callback - Callback must take exactly two arguments, which will
be the message being sent (always a list), and the
return result of socket.send_multipart(msg) -
MessageTracker or None.
'''
self._on_send = callback
if callback is None:
self._stream.on_send(None)
else:
self._stream.on_send(self._send_wrapper)
def on_recv(self, callback):
'''
Register a callback for when a message is ready to recv. There can be
only one callback registered at a time, so each call to on_recv
replaces previously registered callbacks. on_recv(None) disables recv
event polling.
@param callback - callback must take exactly one argument, which will
be a list, as returned by socket.recv_multipart()
if callback is None, recv callbacks are disabled.
'''
self._on_recv = callback
if callback is None:
self._stream.on_recv(None)
else:
self._stream.on_recv(self._recv_wrapper)
def flush(self, flag=3, limit=None):
'''
Flush pending messages.
This method safely handles all pending incoming and/or outgoing
messages, bypassing the inner loop, passing them to the registered
callbacks.
A limit can be specified, to prevent blocking under high load.
flush will return the first time ANY of these conditions are met:
No more events matching the flag are pending.
the total number of events handled reaches the limit.
@param flag - 0MQ poll flags. If flag|POLLIN, recv events will be
flushed. If flag|POLLOUT, send events will be flushed.
Both flags can be set at once, which is the default.
@param limit - None, or int. Optional. The maximum number of messages
to send or receive. Both send and receive count against
this limit
@returns int - count of events handled
'''
return self._stream.flush(flag, limit)
def send(self, msg, **kwds):
'''
Send the given message on this stream. The message type must match that
specified in the streams config, or a ValueError will be raised.
@param msg - Google protocol buffer msg to send over this stream
@param kwds - extra keywords that zmq's stream send accepts.
'''
if not isinstance(msg, self.send_type):
raise ValueError("Wrong message type being sent. {0} is not {1}"
.format(type(msg), type(self._send_type)))
data = msg.SerializeToString()
msgs = [SEPARATOR, RECV_CODE, data]
self._stream.send_multipart(msgs, **kwds)
def reply(self, ids, msg, **kwds):
'''
Reply with the given message on this stream. The reply will be routed
back to the initial socket, based on the given list of ids, if the
stream being replied to is the routing proxy. The message type must
match that specified in the streams config, or a ValueError will be
raised.
@param ids - a list of socket ids to reply to
@param msg - Google protocol buffer msg to send over this stream.
@param kwds - extra keywords that zmq's stream send accepts.
'''
if not isinstance(msg, self.reply_type):
raise ValueError("Wrong message type being replied. {0} is not {1}"
.format(type(msg), type(self._reply_type)))
data = msg.SerializeToString()
msgs = list(unroll_list([ids, SEPARATOR, REPLY_CODE, data]))
self._stream.send_multipart(msgs, **kwds)
def _get_msg_type(self, op_code):
'''
Given the op code determine what message type is expected.
@param op_code - op code indicating whether the message was sent as a
reply, or a regular send.
@returns message type appropriate for the op code
'''
return self.recv_type if (op_code == RECV_CODE) else self.reply_type
def _callback_wrapper(self, data, callback):
'''
Helper method to parse serialized messages that are being sent and
received for the respective callbacks.
@param data - data to be sent/received
@param callback - method to invoke
'''
envelope = []
msgs = []
it = iter(data)
# Get envelopes if any
for d in it:
if d == SEPARATOR:
break
else:
envelope.append(d)
# Determine message type to parse
msg_type = self._get_msg_type(next(it))
# Get actual GPB messages
for d in it:
msg = msg_type()
try:
msg.ParseFromString(d)
except:
log.exception("Unable to parse the protocol buffer message.")
continue
msgs.append(msg)
if msgs:
callback(envelope, msgs)
def _recv_wrapper(self, data):
self._callback_wrapper(data, self._on_recv)
def _send_wrapper(self, data, _):
self._callback_wrapper(data, self._on_send)
def close(self):
'''
Close this stream.
'''
self._stream.close()
