def __init__(self): self._scheduler = Scheduler() self._task_to_operation_index = {} self._operation_name_to_task_index = {}
class TaskManager(object): ''' Manages a number of Tasks and allows you to get the ones ready to be run. ''' def __init__(self): self._scheduler = Scheduler() self._task_to_operation_index = {} self._operation_name_to_task_index = {} @property def tasks(self): return self._task_to_operation_index.keys() @property def num_tasks(self): return len(self.tasks) def remove_all_tasks(self): self._task_to_operation_index = {} self._operation_name_to_task_index = {} def remove_task(self, task): if task in self._task_to_operation_index.keys(): operation = self._task_to_operation_index[task] del self._operation_name_to_task_index[operation.name] del self._task_to_operation_index[task] else: raise TaskError('Tried to remove a task that is not under management.') def add_root_task(self, root_task): removed_ops = self._scheduler.set_root_outputs(root_task.provided_ids) removed_tasks = [self._operation_name_to_task_index[op.name] for op in removed_ops] for task in removed_tasks: self.remove_task(task) return removed_tasks def add_task(self, new_task): # 1. make operation # 2. add operation into scheduler # 3. add task and operation to indecies # find a unique name for this operation. operation_name = new_task.plugin.name[:1] # TODO make shorter if operation_name in self._operation_name_to_task_index.keys(): operation_name += '_2' while operation_name in self._operation_name_to_task_index.keys(): num = int(operation_name.split('_')[-1]) operation_name = operation_name.replace('_%d' % num, '_%d' % (num+1)) operation = Operation(new_task.required_ids, new_task.provided_ids, name=operation_name) self._scheduler.add_operation(operation) self._task_to_operation_index[new_task] = operation self._operation_name_to_task_index[operation_name] = new_task def get_ready_tasks(self): ''' Return a list of runnable tasks. ''' potentials = self._scheduler.get_ready_operations() ready_tasks = [] for op in potentials: task = self._operation_name_to_task_index[op.name] if task.is_ready: ready_tasks.append(task) return ready_tasks def get_plot_dict(self): return self._scheduler.get_plot_dict() def checkout_task(self, task): if task not in self.tasks: raise TaskError('Task not under management: %s' % str(task)) else: operation = self._task_to_operation_index[task] self._scheduler.start_operation(operation) return task.checkout() def complete_task(self, task, result=None): if task not in self.tasks: raise TaskError('Task not under management: %s' % str(task)) else: operation = self._task_to_operation_index[task] self._scheduler.finish_operation(operation) self.remove_task(task) if result is not None: return task.complete(result) else: return task.skip() def __str__(self): str_list = ['TaskManager:'] for task in self.tasks: str_list.append(' %s' % str(task)) return '\n'.join(str_list)