def manage(self):
"""Manage workers and tasks."""
import zmq
try:
poller = zmq.Poller()
poller.register(self._rep_socket, zmq.POLLIN)
while self._core.unfinished_count() > 0:
events = dict(poller.poll())
assert events.get(self._rep_socket) == zmq.POLLIN
message = condor.recv_pyobj_compressed(self._rep_socket)
(next_task, completed) = self._core.handle(message)
if next_task is None:
reply_message = None
else:
reply_message = condor.messages.TaskMessage(None, next_task, self._cache.fork())
condor.send_pyobj_compressed(self._rep_socket, reply_message)
if completed is not None:
yield completed
except KeyboardInterrupt:
# work around bizarre pyzmq SIGINT behavior
raise
def work_once(condor_id, req_socket, task_message):
"""Request and/or complete a single unit of work."""
# get an assignment
poller = zmq.Poller()
poller.register(req_socket, zmq.POLLIN)
if task_message is None:
condor.send_pyobj_compressed(
req_socket,
condor.messages.ApplyMessage(condor_id),
)
polled = poller.poll(timeout = 60 * 1000)
if dict(polled).get(req_socket, 0) & zmq.POLLIN:
task_message = condor.recv_pyobj_compressed(req_socket)
else:
task_message = None
if task_message is None:
logger.info("received null assignment or timed out; terminating")
return None
task = task_message.get_task()
# complete the assignment
try:
logger.info("starting work on task %s", task.key)
result = task()
except KeyboardInterrupt, error:
logger.warning("interruption during task %s", task.key)
condor.send_pyobj_compressed(
req_socket,
condor.messages.InterruptedMessage(condor_id, task.key),
)
req_socket.recv()
condor.send_pyobj_compressed(
req_socket,
condor.messages.ErrorMessage(condor_id, task.key, description),
)
req_socket.recv()
else:
logger.info("finished task %s", task.key)
condor.send_pyobj_compressed(
req_socket,
condor.messages.DoneMessage(condor_id, task.key, result),
)
return condor.recv_pyobj_compressed(req_socket)
return None
def work_loop(condor_id, req_socket):
"""Repeatedly request and complete units of work."""
task_message = None
while True:
try:
task_message = work_once(condor_id, req_socket, task_message)
except Exception:
raise
if task_message is None:
