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 build_importer_spec_for_task(
task: pipeline_task.PipelineTask
) -> pipeline_spec_pb2.PipelineDeploymentConfig.ImporterSpec:
"""Builds ImporterSpec for a pipeline task.
Args:
task: The task to build a ComponentSpec for.
Returns:
A ImporterSpec object for the task.
"""
type_schema = type_utils.get_artifact_type_schema(
task.importer_spec.type_schema)
importer_spec = pipeline_spec_pb2.PipelineDeploymentConfig.ImporterSpec(
type_schema=type_schema, reimport=task.importer_spec.reimport)
if task.importer_spec.metadata:
metadata_protobuf_struct = struct_pb2.Struct()
metadata_protobuf_struct.update(task.importer_spec.metadata)
importer_spec.metadata.CopyFrom(metadata_protobuf_struct)
if isinstance(task.importer_spec.artifact_uri,
pipeline_channel.PipelineParameterChannel):
importer_spec.artifact_uri.runtime_parameter = 'uri'
elif isinstance(task.importer_spec.artifact_uri, str):
importer_spec.artifact_uri.constant.string_value = task.importer_spec.artifact_uri
return importer_spec
def build_component_inputs_spec(
component_spec: pipeline_spec_pb2.ComponentSpec,
pipeline_params: List[_pipeline_param.PipelineParam],
is_root_component: bool,
) -> None:
"""Builds component inputs spec from pipeline params.
Args:
component_spec: The component spec to fill in its inputs spec.
pipeline_params: The list of pipeline params.
is_root_component: Whether the component is the root.
"""
for param in pipeline_params:
param_name = param.full_name
if _for_loop.LoopArguments.name_is_loop_argument(param_name):
param.param_type = param.param_type or 'String'
input_name = (param_name if is_root_component else
additional_input_name_for_pipelineparam(param_name))
if type_utils.is_parameter_type(param.param_type):
component_spec.input_definitions.parameters[
input_name].parameter_type = type_utils.get_parameter_type(
param.param_type)
elif input_name not in getattr(component_spec.input_definitions,
'parameters', []):
component_spec.input_definitions.artifacts[
input_name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(param.param_type))
def build_component_spec_from_structure(
component_spec: structures.ComponentSpec,
executor_label: str,
actual_inputs: List[str],
) -> pipeline_spec_pb2.ComponentSpec:
"""Builds an IR ComponentSpec instance from structures.ComponentSpec.
Args:
component_spec: The structure component spec.
executor_label: The executor label.
actual_inputs: The actual arugments passed to the task. This is used as a
short term workaround to support optional inputs in component spec IR.
Returns:
An instance of IR ComponentSpec.
"""
result = pipeline_spec_pb2.ComponentSpec()
result.executor_label = executor_label
for input_spec in component_spec.inputs or []:
# skip inputs not present
if input_spec.name not in actual_inputs:
continue
if type_utils.is_parameter_type(input_spec.type):
result.input_definitions.parameters[
input_spec.
name].parameter_type = type_utils.get_parameter_type(
input_spec.type)
else:
result.input_definitions.artifacts[
input_spec.name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(input_spec.type))
for output_spec in component_spec.outputs or []:
if type_utils.is_parameter_type(output_spec.type):
result.output_definitions.parameters[
output_spec.
name].parameter_type = type_utils.get_parameter_type(
output_spec.type)
else:
result.output_definitions.artifacts[
output_spec.name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(output_spec.type))
return result
def build_component_outputs_spec(
component_spec: pipeline_spec_pb2.ComponentSpec,
pipeline_params: List[_pipeline_param.PipelineParam],
) -> None:
"""Builds component outputs spec from pipeline params.
Args:
component_spec: The component spec to fill in its outputs spec.
pipeline_params: The list of pipeline params.
"""
for param in pipeline_params or []:
output_name = param.full_name
if type_utils.is_parameter_type(param.param_type):
component_spec.output_definitions.parameters[
output_name].parameter_type = type_utils.get_parameter_type(
param.param_type)
elif output_name not in getattr(component_spec.output_definitions,
'parameters', []):
component_spec.output_definitions.artifacts[
output_name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(param.param_type))
def build_component_spec_for_group(
pipeline_channels: List[pipeline_channel.PipelineChannel],
is_root_group: bool,
) -> pipeline_spec_pb2.ComponentSpec:
"""Builds ComponentSpec for a TasksGroup.
Args:
group: The group to build a ComponentSpec for.
pipeline_channels: The list of pipeline channels referenced by the group.
Returns:
A PipelineTaskSpec object representing the loop group.
"""
component_spec = pipeline_spec_pb2.ComponentSpec()
for channel in pipeline_channels:
input_name = (channel.name if is_root_group else
_additional_input_name_for_pipeline_channel(channel))
if isinstance(channel, pipeline_channel.PipelineArtifactChannel):
component_spec.input_definitions.artifacts[
input_name].artifact_type.CopyFrom(
type_utils.get_artifact_type_schema(channel.channel_type))
else:
# channel is one of PipelineParameterChannel, LoopArgument, or
# LoopArgumentVariable.
component_spec.input_definitions.parameters[
input_name].parameter_type = type_utils.get_parameter_type(
channel.channel_type)
if is_root_group:
_fill_in_component_input_default_value(
component_spec=component_spec,
input_name=input_name,
default_value=channel.value,
)
return component_spec
def test_get_artifact_type_schema(self, artifact_class_or_type_name,
expected_result):
self.assertEqual(
expected_result,
type_utils.get_artifact_type_schema(artifact_class_or_type_name))
