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))