Beispiel #1
0
 def __init__(self):
     self._scheduler = Scheduler()
     self._task_to_operation_index = {}
     self._operation_name_to_task_index = {}
Beispiel #2
0
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)