From 47014ca36582cb1a049e86eb352a7a5359582df1 Mon Sep 17 00:00:00 2001 From: Yaqi Date: Tue, 21 Dec 2021 15:45:04 -0800 Subject: [PATCH] feat(sdk): add metadata --- sdk/python/kfp/v2/compiler/pipeline_spec_builder.py | 3 ++- .../test_data/pipeline_with_importer.json | 8 +++++--- .../test_data/pipeline_with_importer.py | 3 ++- sdk/python/kfp/v2/components/importer_node.py | 9 ++++++--- 4 files changed, 15 insertions(+), 8 deletions(-) diff --git a/sdk/python/kfp/v2/compiler/pipeline_spec_builder.py b/sdk/python/kfp/v2/compiler/pipeline_spec_builder.py index 7cfdcd79137..9bff9152320 100644 --- a/sdk/python/kfp/v2/compiler/pipeline_spec_builder.py +++ b/sdk/python/kfp/v2/compiler/pipeline_spec_builder.py @@ -355,7 +355,8 @@ def build_importer_spec_for_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) + reimport=task.importer_spec.reimport, + metadata=task.importer_spec.metadata) if isinstance(task.importer_spec.artifact_uri, pipeline_channel.PipelineParameterChannel): importer_spec.artifact_uri.runtime_parameter = 'uri' diff --git a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.json b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.json index d29ee5f8bd9..37c9f59db83 100644 --- a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.json +++ b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.json @@ -163,6 +163,7 @@ "artifactUri": { "constant": "gs://ml-pipeline-playground/shakespeare1.txt" }, + "metadata": {}, "typeSchema": { "schemaTitle": "system.Artifact", "schemaVersion": "0.0.1" @@ -174,6 +175,7 @@ "artifactUri": { "runtimeParameter": "uri" }, + "metadata": {}, "reimport": true, "typeSchema": { "schemaTitle": "system.Artifact", @@ -192,7 +194,7 @@ "command": [ "sh", "-c", - "\nif ! [ -x \"$(command -v pip)\" ]; then\n python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location 'kfp==1.8.9' && \"$0\" \"$@\"\n", + "\nif ! [ -x \"$(command -v pip)\" ]; then\n python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location 'kfp==1.8.10' && \"$0\" \"$@\"\n", "sh", "-ec", "program_path=$(mktemp -d)\nprintf \"%s\" \"$0\" > \"$program_path/ephemeral_component.py\"\npython3 -m kfp.v2.components.executor_main --component_module_path \"$program_path/ephemeral_component.py\" \"$@\"\n", @@ -212,7 +214,7 @@ "command": [ "sh", "-c", - "\nif ! [ -x \"$(command -v pip)\" ]; then\n python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location 'kfp==1.8.9' && \"$0\" \"$@\"\n", + "\nif ! [ -x \"$(command -v pip)\" ]; then\n python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location 'kfp==1.8.10' && \"$0\" \"$@\"\n", "sh", "-ec", "program_path=$(mktemp -d)\nprintf \"%s\" \"$0\" > \"$program_path/ephemeral_component.py\"\npython3 -m kfp.v2.components.executor_main --component_module_path \"$program_path/ephemeral_component.py\" \"$@\"\n", @@ -312,5 +314,5 @@ } }, "schemaVersion": "2.1.0", - "sdkVersion": "kfp-1.8.9" + "sdkVersion": "kfp-1.8.10" } \ No newline at end of file diff --git a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.py b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.py index 3564ea2fe10..46d5a12ef6f 100644 --- a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.py +++ b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_importer.py @@ -50,7 +50,8 @@ def my_pipeline(dataset2: str = 'gs://ml-pipeline-playground/shakespeare2.txt'): importer1 = importer( artifact_uri='gs://ml-pipeline-playground/shakespeare1.txt', artifact_class=Dataset, - reimport=False) + reimport=False, + metadata={'key': 'value'}) train1 = train(dataset=importer1.output) with dsl.Condition(train1.outputs['scalar'] == '123'): diff --git a/sdk/python/kfp/v2/components/importer_node.py b/sdk/python/kfp/v2/components/importer_node.py index 1bd377d6b05..0cfefa9aa62 100644 --- a/sdk/python/kfp/v2/components/importer_node.py +++ b/sdk/python/kfp/v2/components/importer_node.py @@ -13,7 +13,7 @@ # limitations under the License. """Utility function for building Importer Node spec.""" -from typing import Union, Type +from typing import Any, Union, Optional, Type, Mapping from kfp.v2.components import pipeline_task from kfp.v2.components import pipeline_channel @@ -29,7 +29,8 @@ def importer(artifact_uri: Union[pipeline_channel.PipelineParameterChannel, str], artifact_class: Type[artifact_types.Artifact], - reimport: bool = False) -> pipeline_task.PipelineTask: + reimport: bool = False, + metadata: Optional[Mapping[str, Any]] = None) -> pipeline_task.PipelineTask: """dsl.importer for importing an existing artifact. Only for v2 pipeline. Args: @@ -37,6 +38,7 @@ def importer(artifact_uri: Union[pipeline_channel.PipelineParameterChannel, artifact_type_schema: The user specified artifact type schema of the artifact to be imported. reimport: Whether to reimport the artifact. Defaults to False. + metadata: Properties of the artifact. Returns: A PipelineTask instance. @@ -52,7 +54,8 @@ def importer(artifact_uri: Union[pipeline_channel.PipelineParameterChannel, artifact_uri=placeholders.input_parameter_placeholder( INPUT_KEY), type_schema=artifact_class.TYPE_NAME, - reimport=reimport)), + reimport=reimport, + metadata=metadata)), inputs={INPUT_KEY: structures.InputSpec(type='String')}, outputs={OUTPUT_KEY: structures.OutputSpec(type='Artifact')}, )