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()
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() except BaseException, error: description = traceback.format_exc(error) logger.warning("error during task %s:\n%s", task.key, description) 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