def build_task_spec_for_exit_task(
task: pipeline_task.PipelineTask,
dependent_task: str,
pipeline_inputs: pipeline_spec_pb2.ComponentInputsSpec,
) -> pipeline_spec_pb2.PipelineTaskSpec:
"""Builds PipelineTaskSpec for an exit handler's exit task.
Args:
tasks: The exit handler's exit task to build task spec for.
dependent_task: The dependent task name for the exit task, i.e. the name
of the exit handler group.
pipeline_inputs: The pipeline level input definitions.
Returns:
A PipelineTaskSpec object representing the exit task.
"""
pipeline_task_spec = build_task_spec_for_task(
task=task,
parent_component_inputs=pipeline_inputs,
tasks_in_current_dag=[], # Does not matter for exit task
input_parameters_in_current_dag=pipeline_inputs.parameters.keys(),
input_artifacts_in_current_dag=[],
)
pipeline_task_spec.dependent_tasks.extend([dependent_task])
pipeline_task_spec.trigger_policy.strategy = (
pipeline_spec_pb2.PipelineTaskSpec.TriggerPolicy.TriggerStrategy.
ALL_UPSTREAM_TASKS_COMPLETED)
for input_name, input_spec in task.component_spec.inputs.items():
if type_utils.is_task_final_status_type(input_spec.type):
pipeline_task_spec.inputs.parameters[
input_name].task_final_status.producer_task = dependent_task
return pipeline_task_spec
def build_component_spec_for_task(
task: pipeline_task.PipelineTask,
is_exit_task: bool = False,
) -> pipeline_spec_pb2.ComponentSpec:
"""Builds ComponentSpec for a pipeline task.
Args:
task: The task to build a ComponentSpec for.
is_exit_task: Whether the task is used as exit task in Exit Handler.
Returns:
A ComponentSpec object for the task.
"""
component_spec = pipeline_spec_pb2.ComponentSpec()
component_spec.executor_label = component_utils.sanitize_executor_label(
task.name)
for input_name, input_spec in (task.component_spec.inputs or {}).items():
# Special handling for PipelineTaskFinalStatus first.
if type_utils.is_task_final_status_type(input_spec.type):
if not is_exit_task:
raise ValueError(
'PipelineTaskFinalStatus can only be used in an exit task.'
)
component_spec.input_definitions.parameters[
input_name].parameter_type = pipeline_spec_pb2.ParameterType.STRUCT
continue
# skip inputs not present, as a workaround to support optional inputs.
if input_name not in task.inputs and input_spec.default is None:
continue
if type_utils.is_parameter_type(input_spec.type):
component_spec.input_definitions.parameters[
input_name].parameter_type = type_utils.get_parameter_type(
input_spec.type)
if input_spec.default is not None:
component_spec.input_definitions.parameters[
input_name].default_value.CopyFrom(
_to_protobuf_value(input_spec.default))
else:
component_spec.input_definitions.artifacts[
input_name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(input_spec.type))
for output_name, output_spec in (task.component_spec.outputs
or {}).items():
if type_utils.is_parameter_type(output_spec.type):
component_spec.output_definitions.parameters[
output_name].parameter_type = type_utils.get_parameter_type(
output_spec.type)
else:
component_spec.output_definitions.artifacts[
output_name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(output_spec.type))
return component_spec
def _create_pipeline_from_component_spec(
self,
component_spec: structures.ComponentSpec,
) -> pipeline_spec_pb2.PipelineSpec:
"""Creates a pipeline instance and constructs the pipeline spec for a
primitive component.
Args:
component_spec: The ComponentSpec to convert to PipelineSpec.
Returns:
A PipelineSpec proto representing the compiled component.
"""
args_dict = {}
for arg_name, input_spec in component_spec.inputs.items():
arg_type = input_spec.type
if not type_utils.is_parameter_type(
arg_type) or type_utils.is_task_final_status_type(arg_type):
raise TypeError(
builder.make_invalid_input_type_error_msg(
arg_name, arg_type))
args_dict[arg_name] = dsl.PipelineParameterChannel(
name=arg_name, channel_type=arg_type)
task = pipeline_task.PipelineTask(component_spec, args_dict)
# instead of constructing a pipeline with pipeline_context.Pipeline,
# just build the single task group
group = tasks_group.TasksGroup(
group_type=tasks_group.TasksGroupType.PIPELINE)
group.tasks.append(task)
pipeline_inputs = component_spec.inputs or {}
# Fill in the default values.
args_list_with_defaults = [
dsl.PipelineParameterChannel(
name=input_name,
channel_type=input_spec.type,
value=input_spec.default,
) for input_name, input_spec in pipeline_inputs.items()
]
group.name = uuid.uuid4().hex
return builder.create_pipeline_spec_for_component(
pipeline_name=component_spec.name,
pipeline_args=args_list_with_defaults,
task_group=group,
)
def __init__(self, component_spec: structures.ComponentSpec):
"""Init function for BaseComponent.
Args:
component_spec: The component definition.
"""
self.component_spec = component_spec
self.name = component_spec.name
# Arguments typed as PipelineTaskFinalStatus are special arguments that
# do not count as user inputs. Instead, they are reserved to for the
# (backend) system to pass a value.
self._component_inputs = {
input_name for input_name, input_spec in (
self.component_spec.inputs or {}).items()
if not type_utils.is_task_final_status_type(input_spec.type)
}
def __call__(self, *args, **kwargs) -> pipeline_task.PipelineTask:
"""Creates a PipelineTask object.
The arguments are generated on the fly based on component input
definitions.
"""
task_inputs = {}
if len(args) > 0:
raise TypeError(
'Components must be instantiated using keyword arguments. Positional '
f'parameters are not allowed (found {len(args)} such parameters for '
f'component "{self.name}").')
for k, v in kwargs.items():
if k not in self._component_inputs:
raise TypeError(
f'{self.name}() got an unexpected keyword argument "{k}".')
if k in task_inputs:
raise TypeError(
f'{self.name}() got multiple values for argument "{k}".')
task_inputs[k] = v
# Skip optional inputs and arguments typed as PipelineTaskFinalStatus.
missing_arguments = [
input_name for input_name, input_spec in (
self.component_spec.inputs or {}).items()
if input_name not in task_inputs and not input_spec._optional and
not type_utils.is_task_final_status_type(input_spec.type)
]
if missing_arguments:
argument_or_arguments = 'argument' if len(
missing_arguments) == 1 else 'arguments'
arguments = ', '.join(missing_arguments)
raise TypeError(
f'{self.name}() missing {len(missing_arguments)} required '
f'{argument_or_arguments}: {arguments}.')
return pipeline_task.create_pipeline_task(
component_spec=self.component_spec,
args=task_inputs,
)
def expand_command_part(arg) -> Union[str, List[str], None]:
if arg is None:
return None
if isinstance(arg, (str, int, float, bool)):
return str(arg)
elif isinstance(arg, (dict, list)):
return json.dumps(arg)
elif isinstance(arg, structures.InputValuePlaceholder):
input_name = arg.input_name
if not type_utils.is_parameter_type(
inputs_dict[input_name].type):
raise TypeError(
f'Input "{input_name}" with type '
f'"{inputs_dict[input_name].type}" cannot be paired with '
'InputValuePlaceholder.')
if input_name in args or type_utils.is_task_final_status_type(
inputs_dict[input_name].type):
return arg.to_placeholder()
else:
input_spec = inputs_dict[input_name]
if input_spec.default is not None:
return None
else:
raise ValueError(
f'No value provided for input: {input_name}.')
elif isinstance(arg, structures.InputUriPlaceholder):
input_name = arg.input_name
if type_utils.is_parameter_type(inputs_dict[input_name].type):
raise TypeError(
f'Input "{input_name}" with type '
f'"{inputs_dict[input_name].type}" cannot be paired with '
'InputUriPlaceholder.')
if input_name in args:
input_uri = arg.to_placeholder()
return input_uri
else:
input_spec = inputs_dict[input_name]
if input_spec.default is not None:
return None
else:
raise ValueError(
f'No value provided for input: {input_name}.')
elif isinstance(arg, structures.InputPathPlaceholder):
input_name = arg.input_name
if type_utils.is_parameter_type(inputs_dict[input_name].type):
raise TypeError(
f'Input "{input_name}" with type '
f'"{inputs_dict[input_name].type}" cannot be paired with '
'InputPathPlaceholder.')
if input_name in args:
input_path = arg.to_placeholder()
return input_path
else:
input_spec = inputs_dict[input_name]
if input_spec._optional:
return None
else:
raise ValueError(
f'No value provided for input: {input_name}.')
elif isinstance(arg, structures.OutputUriPlaceholder):
output_name = arg.output_name
if type_utils.is_parameter_type(
outputs_dict[output_name].type):
raise TypeError(
f'Onput "{output_name}" with type '
f'"{outputs_dict[output_name].type}" cannot be paired with '
'OutputUriPlaceholder.')
return arg.to_placeholder()
elif isinstance(arg, structures.OutputPathPlaceholder):
output_name = arg.output_name
if type_utils.is_parameter_type(
outputs_dict[output_name].type):
output_path = structures.OutputParameterPlaceholder(
arg.output_name).to_placeholder()
else:
output_path = arg.to_placeholder()
return output_path
elif isinstance(arg, structures.OutputParameterPlaceholder):
output_name = arg.output_name
if not type_utils.is_parameter_type(
outputs_dict[output_name].type):
raise TypeError(
f'Onput "{output_name}" with type '
f'"{outputs_dict[output_name].type}" cannot be paired with '
'OutputUriPlaceholder.')
return arg.to_placeholder()
elif isinstance(arg, structures.ConcatPlaceholder):
expanded_argument_strings = expand_argument_list(arg.items)
return ''.join(expanded_argument_strings)
elif isinstance(arg, structures.IfPresentPlaceholder):
if arg.if_structure.input_name in argument_values:
result_node = arg.if_structure.then
else:
result_node = arg.if_structure.otherwise
if result_node is None:
return []
if isinstance(result_node, list):
expanded_result = expand_argument_list(result_node)
else:
expanded_result = expand_command_part(result_node)
return expanded_result
else:
raise TypeError(f'Unrecognized argument type: {arg}')
def to_pipeline_spec(self) -> pipeline_spec_pb2.PipelineSpec:
"""Creates a pipeline instance and constructs the pipeline spec for a
single component.
Args:
component_spec: The ComponentSpec to convert to PipelineSpec.
Returns:
A PipelineSpec proto representing the compiled component.
"""
# import here to aviod circular module dependency
from kfp.compiler import pipeline_spec_builder as builder
from kfp.components import pipeline_task
from kfp.components import tasks_group
from kfp.components.types import type_utils
args_dict = {}
pipeline_inputs = self.inputs or {}
for arg_name, input_spec in pipeline_inputs.items():
arg_type = input_spec.type
if not type_utils.is_parameter_type(
arg_type) or type_utils.is_task_final_status_type(
arg_type):
raise TypeError(
builder.make_invalid_input_type_error_msg(
arg_name, arg_type))
args_dict[arg_name] = dsl.PipelineParameterChannel(
name=arg_name, channel_type=arg_type)
task = pipeline_task.PipelineTask(self, args_dict)
# instead of constructing a pipeline with pipeline_context.Pipeline,
# just build the single task group
group = tasks_group.TasksGroup(
group_type=tasks_group.TasksGroupType.PIPELINE)
group.tasks.append(task)
# Fill in the default values.
args_list_with_defaults = [
dsl.PipelineParameterChannel(
name=input_name,
channel_type=input_spec.type,
value=input_spec.default,
) for input_name, input_spec in pipeline_inputs.items()
]
group.name = uuid.uuid4().hex
pipeline_name = self.name
pipeline_args = args_list_with_defaults
task_group = group
builder.validate_pipeline_name(pipeline_name)
pipeline_spec = pipeline_spec_pb2.PipelineSpec()
pipeline_spec.pipeline_info.name = pipeline_name
pipeline_spec.sdk_version = f'kfp-{kfp.__version__}'
# Schema version 2.1.0 is required for kfp-pipeline-spec>0.1.13
pipeline_spec.schema_version = '2.1.0'
pipeline_spec.root.CopyFrom(
builder.build_component_spec_for_group(
pipeline_channels=pipeline_args,
is_root_group=True,
))
deployment_config = pipeline_spec_pb2.PipelineDeploymentConfig()
root_group = task_group
task_name_to_parent_groups, group_name_to_parent_groups = builder.get_parent_groups(
root_group)
def get_inputs(task_group: tasks_group.TasksGroup,
task_name_to_parent_groups):
inputs = collections.defaultdict(set)
if len(task_group.tasks) != 1:
raise ValueError(
f'Error compiling component. Expected one task in task group, got {len(task_group.tasks)}.'
)
only_task = task_group.tasks[0]
if only_task.channel_inputs:
for group_name in task_name_to_parent_groups[only_task.name]:
inputs[group_name].add(
(only_task.channel_inputs[-1], None))
return inputs
inputs = get_inputs(task_group, task_name_to_parent_groups)
builder.build_spec_by_group(
pipeline_spec=pipeline_spec,
deployment_config=deployment_config,
group=root_group,
inputs=inputs,
dependencies={}, # no dependencies for single-component pipeline
rootgroup_name=root_group.name,
task_name_to_parent_groups=task_name_to_parent_groups,
group_name_to_parent_groups=group_name_to_parent_groups,
name_to_for_loop_group=
{}, # no for loop for single-component pipeline
)
return pipeline_spec
def test_is_task_final_statu_type(self, given_type, expected_result):
self.assertEqual(expected_result,
type_utils.is_task_final_status_type(given_type))