aria/orchestrator/workflows/api/task_graph.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Task graph.
"""

from collections import Iterable

from networkx import DiGraph, topological_sort

from ....utils.uuid import generate_uuid
from . import task as api_task


class TaskNotInGraphError(Exception):
    """
    An error representing a scenario where a given task is not in the graph as expected.
    """
    pass


def _filter_out_empty_tasks(func=None):
    if func is None:
        return lambda f: _filter_out_empty_tasks(func=f)

    def _wrapper(task, *tasks, **kwargs):
        return func(*(t for t in (task,) + tuple(tasks) if t), **kwargs)
    return _wrapper


class TaskGraph(object):
    """
    Task graph builder.
    """

    def __init__(self, name):
        self.name = name
        self._id = generate_uuid(variant='uuid')
        self._graph = DiGraph()

    def __repr__(self):
        return '{name}(id={self._id}, name={self.name}, graph={self._graph!r})'.format(
            name=self.__class__.__name__, self=self)

    @property
    def id(self):
        """
        ID of the graph
        """
        return self._id

    # graph traversal methods

    @property
    def tasks(self):
        """
        Iterator over tasks in the graph.
        """
        for _, data in self._graph.nodes_iter(data=True):
            yield data['task']

    def topological_order(self, reverse=False):
        """
        Topological sort of the graph.

        :param reverse: whether to reverse the sort
        :return: list which represents the topological sort
        """
        for task_id in topological_sort(self._graph, reverse=reverse):
            yield self.get_task(task_id)

    def get_dependencies(self, dependent_task):
        """
        Iterates over the task's dependencies.

        :param dependent_task: task whose dependencies are requested
        :raises ~aria.orchestrator.workflows.api.task_graph.TaskNotInGraphError: if
         ``dependent_task`` is not in the graph
        """
        if not self.has_tasks(dependent_task):
            raise TaskNotInGraphError('Task id: {0}'.format(dependent_task.id))
        for _, dependency_id in self._graph.out_edges_iter(dependent_task.id):
            yield self.get_task(dependency_id)

    def get_dependents(self, dependency_task):
        """
        Iterates over the task's dependents.

        :param dependency_task: task whose dependents are requested
        :raises ~aria.orchestrator.workflows.api.task_graph.TaskNotInGraphError: if
         ``dependency_task`` is not in the graph
        """
        if not self.has_tasks(dependency_task):
            raise TaskNotInGraphError('Task id: {0}'.format(dependency_task.id))
        for dependent_id, _ in self._graph.in_edges_iter(dependency_task.id):
            yield self.get_task(dependent_id)

    # task methods

    def get_task(self, task_id):
        """
        Get a task instance that's been inserted to the graph by the task's ID.

        :param basestring task_id: task ID
        :raises ~aria.orchestrator.workflows.api.task_graph.TaskNotInGraphError: if no task found in
         the graph with the given ID
        """
        if not self._graph.has_node(task_id):
            raise TaskNotInGraphError('Task id: {0}'.format(task_id))
        data = self._graph.node[task_id]
        return data['task']

    @_filter_out_empty_tasks
    def add_tasks(self, *tasks):
        """
        Adds a task to the graph.

        :param task: task
        :return: list of added tasks
        :rtype: list
        """
        assert all([isinstance(task, (api_task.BaseTask, Iterable)) for task in tasks])
        return_tasks = []

        for task in tasks:
            if isinstance(task, Iterable):
                return_tasks += self.add_tasks(*task)
            elif not self.has_tasks(task):
                self._graph.add_node(task.id, task=task)
                return_tasks.append(task)

        return return_tasks

    @_filter_out_empty_tasks
    def remove_tasks(self, *tasks):
        """
        Removes the provided task from the graph.

        :param task: task
        :return: list of removed tasks
        :rtype: list
        """
        return_tasks = []

        for task in tasks:
            if isinstance(task, Iterable):
                return_tasks += self.remove_tasks(*task)
            elif self.has_tasks(task):
                self._graph.remove_node(task.id)
                return_tasks.append(task)

        return return_tasks

    @_filter_out_empty_tasks
    def has_tasks(self, *tasks):
        """
        Checks whether a task is in the graph.

        :param task: task
        :return: ``True`` if all tasks are in the graph, otherwise ``False``
        :rtype: list
        """
        assert all(isinstance(t, (api_task.BaseTask, Iterable)) for t in tasks)
        return_value = True

        for task in tasks:
            if isinstance(task, Iterable):
                return_value &= self.has_tasks(*task)
            else:
                return_value &= self._graph.has_node(task.id)

        return return_value

    def add_dependency(self, dependent, dependency):
        """
        Adds a dependency for one item (task, sequence or parallel) on another.

        The dependent will only be executed after the dependency terminates. If either of the items
        is either a sequence or a parallel, multiple dependencies may be added.

        :param dependent: dependent (task, sequence or parallel)
        :param dependency: dependency (task, sequence or parallel)
        :return: ``True`` if the dependency between the two hadn't already existed, otherwise
         ``False``
        :rtype: bool
        :raises ~aria.orchestrator.workflows.api.task_graph.TaskNotInGraphError: if either the
         dependent or dependency are tasks which are not in the graph
        """
        if not (self.has_tasks(dependent) and self.has_tasks(dependency)):
            raise TaskNotInGraphError()

        if self.has_dependency(dependent, dependency):
            return

        if isinstance(dependent, Iterable):
            for dependent_task in dependent:
                self.add_dependency(dependent_task, dependency)
        else:
            if isinstance(dependency, Iterable):
                for dependency_task in dependency:
                    self.add_dependency(dependent, dependency_task)
            else:
                self._graph.add_edge(dependent.id, dependency.id)

    def has_dependency(self, dependent, dependency):
        """
        Checks whether one item (task, sequence or parallel) depends on another.

        Note that if either of the items is either a sequence or a parallel, and some of the
        dependencies exist in the graph but not all of them, this method will return ``False``.

        :param dependent: dependent (task, sequence or parallel)
        :param dependency: dependency (task, sequence or parallel)
        :return: ``True`` if the dependency between the two exists, otherwise ``False``
        :rtype: bool
        :raises ~aria.orchestrator.workflows.api.task_graph.TaskNotInGraphError: if either the
         dependent or dependency are tasks which are not in the graph
        """
        if not (dependent and dependency):
            return False
        elif not (self.has_tasks(dependent) and self.has_tasks(dependency)):
            raise TaskNotInGraphError()

        return_value = True

        if isinstance(dependent, Iterable):
            for dependent_task in dependent:
                return_value &= self.has_dependency(dependent_task, dependency)
        else:
            if isinstance(dependency, Iterable):
                for dependency_task in dependency:
                    return_value &= self.has_dependency(dependent, dependency_task)
            else:
                return_value &= self._graph.has_edge(dependent.id, dependency.id)

        return return_value

    def remove_dependency(self, dependent, dependency):
        """
        Removes a dependency for one item (task, sequence or parallel) on another.

        Note that if either of the items is either a sequence or a parallel, and some of the
        dependencies exist in the graph but not all of them, this method will not remove any of the
        dependencies and return ``False``.

        :param dependent: dependent (task, sequence or parallel)
        :param dependency: dependency (task, sequence or parallel)
        :return: ``False`` if the dependency between the two hadn't existed, otherwise ``True``
        :rtype: bool
        :raises ~aria.orchestrator.workflows.api.task_graph.TaskNotInGraphError: if either the
         dependent or dependency are tasks which are not in the graph
        """
        if not (self.has_tasks(dependent) and self.has_tasks(dependency)):
            raise TaskNotInGraphError()

        if not self.has_dependency(dependent, dependency):
            return

        if isinstance(dependent, Iterable):
            for dependent_task in dependent:
                self.remove_dependency(dependent_task, dependency)
        elif isinstance(dependency, Iterable):
            for dependency_task in dependency:
                self.remove_dependency(dependent, dependency_task)
        else:
            self._graph.remove_edge(dependent.id, dependency.id)

    @_filter_out_empty_tasks
    def sequence(self, *tasks):
        """
        Creates and inserts a sequence into the graph, effectively each task i depends on i-1.

        :param tasks: iterable of dependencies
        :return: provided tasks
        """
        if tasks:
            self.add_tasks(*tasks)

            for i in xrange(1, len(tasks)):
                self.add_dependency(tasks[i], tasks[i-1])

        return tasks