# Copyright 2021 The Kubeflow Authors
#
# Licensed 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.
"""Definition for TasksGroup."""
import copy
import enum
from typing import List, Optional, Union
import warnings
from kfp.dsl import for_loop
from kfp.dsl import pipeline_channel
from kfp.dsl import pipeline_context
from kfp.dsl import pipeline_task
class TasksGroupType(str, enum.Enum):
"""Types of TasksGroup."""
PIPELINE = 'pipeline'
CONDITION = 'condition'
CONDITION_BRANCHES = 'condition-branches'
FOR_LOOP = 'for-loop'
EXIT_HANDLER = 'exit-handler'
class TasksGroup:
"""Represents a logical group of tasks and groups of TasksGroups.
This class is the base class for groups of tasks, such as tasks
sharing an exit handler, a condition branch, or a loop. This class
is not supposed to be used by pipeline authors. It is useful for
implementing a compiler.
Attributes:
group_type: The type of the TasksGroup.
tasks: A list of all PipelineTasks in this group.
groups: A list of TasksGroups in this group.
display_name: The optional user given name of the group.
dependencies: A list of tasks or groups this group depends on.
is_root: If TasksGroup is root group.
"""
def __init__(
self,
group_type: TasksGroupType,
name: Optional[str] = None,
is_root: bool = False,
) -> None:
"""Create a new instance of TasksGroup.
Args:
group_type: The type of the group.
name: The name of the group. Used as display name in UI.
"""
self.group_type = group_type
self.tasks = []
self.groups = []
self.display_name = name
self.dependencies = []
self.is_root = is_root
# backref to parent, set when the pipeline is called in pipeline_context
self.parent_task_group: Optional[TasksGroup] = None
def __enter__(self):
if not pipeline_context.Pipeline.get_default_pipeline():
raise ValueError('Default pipeline not defined.')
self._make_name_unique()
pipeline_context.Pipeline.get_default_pipeline().push_tasks_group(self)
return self
def __exit__(self, *unused_args):
pipeline_context.Pipeline.get_default_pipeline().pop_tasks_group()
def _make_name_unique(self):
"""Generates a unique TasksGroup name in the pipeline."""
if not pipeline_context.Pipeline.get_default_pipeline():
raise ValueError('Default pipeline not defined.')
group_id = pipeline_context.Pipeline.get_default_pipeline(
).get_next_group_id()
self.name = f'{self.group_type.value}-{group_id}'
self.name = self.name.replace('_', '-')
def remove_task_recursive(self, task: pipeline_task.PipelineTask):
"""Removes a task from the group recursively."""
if self.tasks and task in self.tasks:
self.tasks.remove(task)
for group in self.groups or []:
group.remove_task_recursive(task)
[docs]class ExitHandler(TasksGroup):
"""A class for setting an exit handler task that is invoked upon exiting a
group of other tasks.
Args:
exit_task: The task that is invoked after exiting a group of other tasks.
name: The name of the exit handler group.
Example:
::
exit_task = ExitComponent(...)
with ExitHandler(exit_task):
task1 = my_component1(...)
task2 = my_component2(...)
"""
def __init__(
self,
exit_task: pipeline_task.PipelineTask,
name: Optional[str] = None,
) -> None:
"""Initializes a Condition task group."""
super().__init__(
group_type=TasksGroupType.EXIT_HANDLER,
name=name,
is_root=False,
)
if exit_task.dependent_tasks:
raise ValueError('exit_task cannot depend on any other tasks.')
# Removing exit_task form any group
pipeline_context.Pipeline.get_default_pipeline(
).remove_task_from_groups(exit_task)
# Set is_exit_handler since the compiler might be using this attribute.
exit_task.is_exit_handler = True
self.exit_task = exit_task
class ConditionBranches(TasksGroup):
_oneof_id = 0
def __init__(self) -> None:
super().__init__(
group_type=TasksGroupType.CONDITION_BRANCHES,
name=None,
is_root=False,
)
def get_oneof_id(self) -> int:
"""Incrementor for uniquely identifying a OneOf for the parent
ConditionBranches group.
This is analogous to incrementing a unique identifier for tasks
groups belonging to a pipeline.
"""
self._oneof_id += 1
return self._oneof_id
class _ConditionBase(TasksGroup):
"""Parent class for condition control flow context managers (Condition, If,
Elif, Else).
Args:
condition: A list of binary operations to be combined via conjunction.
name: The name of the condition group.
"""
def __init__(
self,
conditions: List[pipeline_channel.ConditionOperation],
name: Optional[str] = None,
) -> None:
super().__init__(
group_type=TasksGroupType.CONDITION,
name=name,
is_root=False,
)
self.conditions: List[pipeline_channel.ConditionOperation] = conditions
[docs]class If(_ConditionBase):
"""A class for creating a conditional control flow "if" block within a
pipeline.
Args:
condition: A comparative expression that evaluates to True or False. At least one of the operands must be an output from an upstream task or a pipeline parameter.
name: The name of the condition group.
Example:
::
task1 = my_component1(...)
with dsl.If(task1.output=='pizza', 'pizza-condition'):
task2 = my_component2(...)
"""
def __init__(
self,
condition,
name: Optional[str] = None,
) -> None:
super().__init__(
conditions=[condition],
name=name,
)
if isinstance(condition, bool):
raise ValueError(
f'Got constant boolean {condition} as a condition. This is likely because the provided condition evaluated immediately. At least one of the operands must be an output from an upstream task or a pipeline parameter.'
)
copied_condition = copy.copy(condition)
copied_condition.negate = True
self._negated_upstream_conditions = [copied_condition]
[docs]class Condition(If):
"""Deprecated.
Use dsl.If instead.
"""
def __enter__(self):
super().__enter__()
warnings.warn(
'dsl.Condition is deprecated. Please use dsl.If instead.',
category=DeprecationWarning,
stacklevel=2)
return self
[docs]class Elif(_ConditionBase):
"""A class for creating a conditional control flow "else if" block within a
pipeline. Can be used following an upstream dsl.If or dsl.Elif.
Args:
condition: A comparative expression that evaluates to True or False. At least one of the operands must be an output from an upstream task or a pipeline parameter.
name: The name of the condition group.
Example:
::
task1 = my_component1(...)
task2 = my_component2(...)
with dsl.If(task1.output=='pizza', 'pizza-condition'):
task3 = my_component3(...)
with dsl.Elif(task2.output=='pasta', 'pasta-condition'):
task4 = my_component4(...)
"""
def __init__(
self,
condition,
name: Optional[str] = None,
) -> None:
prev_cond = pipeline_context.Pipeline.get_default_pipeline(
).get_last_tasks_group()
if not isinstance(prev_cond, (Condition, If, Elif)):
# prefer pushing toward dsl.If rather than dsl.Condition for syntactic consistency with the if-elif-else keywords in Python
raise InvalidControlFlowException(
'dsl.Elif can only be used following an upstream dsl.If or dsl.Elif.'
)
if isinstance(condition, bool):
raise ValueError(
f'Got constant boolean {condition} as a condition. This is likely because the provided condition evaluated immediately. At least one of the operands must be an output from an upstream task or a pipeline parameter.'
)
copied_condition = copy.copy(condition)
copied_condition.negate = True
self._negated_upstream_conditions = _shallow_copy_list_of_binary_operations(
prev_cond._negated_upstream_conditions) + [copied_condition]
conditions = _shallow_copy_list_of_binary_operations(
prev_cond._negated_upstream_conditions)
conditions.append(condition)
super().__init__(
conditions=conditions,
name=name,
)
def __enter__(self):
if not pipeline_context.Pipeline.get_default_pipeline():
raise ValueError('Default pipeline not defined.')
pipeline = pipeline_context.Pipeline.get_default_pipeline()
maybe_make_and_insert_conditional_branches_group(pipeline)
self._make_name_unique()
pipeline.push_tasks_group(self)
return self
[docs]class Else(_ConditionBase):
"""A class for creating a conditional control flow "else" block within a
pipeline. Can be used following an upstream dsl.If or dsl.Elif.
Args:
name: The name of the condition group.
Example:
::
task1 = my_component1(...)
task2 = my_component2(...)
with dsl.If(task1.output=='pizza', 'pizza-condition'):
task3 = my_component3(...)
with dsl.Elif(task2.output=='pasta', 'pasta-condition'):
task4 = my_component4(...)
with dsl.Else():
my_component5(...)
"""
def __init__(
self,
name: Optional[str] = None,
) -> None:
prev_cond = pipeline_context.Pipeline.get_default_pipeline(
).get_last_tasks_group()
# if it immediately follows as TasksGroup, this is because it immediately
# follows Else in the user code and we wrap Else in a TasksGroup
if isinstance(prev_cond, ConditionBranches):
# prefer pushing toward dsl.If rather than dsl.Condition for syntactic consistency with the if-elif-else keywords in Python
raise InvalidControlFlowException(
'Cannot use dsl.Else following another dsl.Else. dsl.Else can only be used following an upstream dsl.If or dsl.Elif.'
)
if not isinstance(prev_cond, (Condition, If, Elif)):
# prefer pushing toward dsl.If rather than dsl.Condition for syntactic consistency with the if-elif-else keywords in Python
raise InvalidControlFlowException(
'dsl.Else can only be used following an upstream dsl.If or dsl.Elif.'
)
super().__init__(
conditions=prev_cond._negated_upstream_conditions,
name=name,
)
def __enter__(self):
if not pipeline_context.Pipeline.get_default_pipeline():
raise ValueError('Default pipeline not defined.')
pipeline = pipeline_context.Pipeline.get_default_pipeline()
maybe_make_and_insert_conditional_branches_group(pipeline)
self._make_name_unique()
pipeline.push_tasks_group(self)
return self
def __exit__(self, *unused_args):
pipeline = pipeline_context.Pipeline.get_default_pipeline()
pipeline.pop_tasks_group()
# since this is an else, also pop off the parent dag for conditional branches
# this parent TasksGroup is not a context manager, so we simulate its
# __exit__ call with this
pipeline.pop_tasks_group()
def maybe_make_and_insert_conditional_branches_group(
pipeline: 'pipeline_context.Pipeline') -> None:
already_has_pipeline_wrapper = isinstance(
pipeline.get_last_tasks_group(),
Elif,
)
if already_has_pipeline_wrapper:
return
condition_wrapper_group = ConditionBranches()
condition_wrapper_group._make_name_unique()
# swap outer and inner group ids so that numbering stays sequentially consistent with how such hypothetical code would be authored
def swap_group_ids(parent: TasksGroup, cond: TasksGroup):
parent_name, parent_id = parent.name.rsplit('-', 1)
cond_name, cond_id = cond.name.split('-')
cond.name = f'{cond_name}-{parent_id}'
parent.name = f'{parent_name}-{cond_id}'
# replace last pushed group (If or Elif) with condition group
last_pushed_group = pipeline.groups[-1].groups.pop()
swap_group_ids(condition_wrapper_group, last_pushed_group)
pipeline.push_tasks_group(condition_wrapper_group)
# then repush (__enter__) and pop (__exit__) the last pushed group
# before the wrapper to emulate re-entering and exiting its context
pipeline.push_tasks_group(last_pushed_group)
pipeline.pop_tasks_group()
class InvalidControlFlowException(Exception):
pass
def _shallow_copy_list_of_binary_operations(
operations: List[pipeline_channel.ConditionOperation]
) -> List[pipeline_channel.ConditionOperation]:
# shallow copy is sufficient to allow us to invert the negate flag of a ConditionOperation without affecting copies. deep copy not needed and would result in many copies of the full pipeline since PipelineChannels hold references to the pipeline.
return [copy.copy(operation) for operation in operations]
[docs]class ParallelFor(TasksGroup):
"""A class for creating parallelized for loop control flow over a static
set of items within a pipeline definition.
Args:
items: The items to loop over. It can be either a constant Python list or a list output from an upstream task.
name: The name of the for loop group.
parallelism: The maximum number of concurrent iterations that can be scheduled for execution. A value of 0 represents unconstrained parallelism (default is unconstrained).
Example:
::
with dsl.ParallelFor(
items=[{'a': 1, 'b': 10}, {'a': 2, 'b': 20}],
parallelism=1
) as item:
task1 = my_component(..., number=item.a)
task2 = my_component(..., number=item.b)
In the example, the group of tasks containing ``task1`` and ``task2`` would
be executed twice, once with case ``args=[{'a': 1, 'b': 10}]`` and once with
case ``args=[{'a': 2, 'b': 20}]``. The ``parallelism=1`` setting causes only
1 execution to be scheduled at a time.
"""
def __init__(
self,
items: Union[for_loop.ItemList, pipeline_channel.PipelineChannel],
name: Optional[str] = None,
parallelism: Optional[int] = None,
) -> None:
"""Initializes a for loop task group."""
parallelism = parallelism or 0
if parallelism < 0:
raise ValueError(
f'ParallelFor parallelism must be >= 0. Got: {parallelism}.')
super().__init__(
group_type=TasksGroupType.FOR_LOOP,
name=name,
is_root=False,
)
if isinstance(items, pipeline_channel.PipelineParameterChannel):
self.loop_argument = for_loop.LoopParameterArgument.from_pipeline_channel(
items)
self.items_is_pipeline_channel = True
elif isinstance(items, pipeline_channel.PipelineArtifactChannel):
self.loop_argument = for_loop.LoopArtifactArgument.from_pipeline_channel(
items)
self.items_is_pipeline_channel = True
else:
self.loop_argument = for_loop.LoopParameterArgument.from_raw_items(
raw_items=items,
name_code=pipeline_context.Pipeline.get_default_pipeline()
.get_next_group_id(),
)
self.items_is_pipeline_channel = False
# TODO: support artifact constants here.
self.parallelism_limit = parallelism
def __enter__(
self
) -> Union[for_loop.LoopParameterArgument, for_loop.LoopArtifactArgument]:
super().__enter__()
return self.loop_argument