blob: 0debbb701e60351a6e056f87c77f2e2fcbda91a3 [file] [log] [blame]
# 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 u'{name}(id={self._id}, name={self.name}, graph={self._graph!r})'.format( # pylint: disable=redundant-keyword-arg
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(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
"""
task_ids = topological_sort(self._graph)
if reverse:
task_ids = reversed(tuple(task_ids))
for task_id in task_ids:
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(u'Task id: {0}'.format(dependent_task.id))
for _, dependency_id in self._graph.out_edges(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(u'Task id: {0}'.format(dependency_task.id))
for dependent_id, _ in self._graph.in_edges(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(u'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