feat(sdk): add metadata

kubeflow · Dec 21, 2021 · 47014ca · 47014ca
1 parent b37e575
commit 47014ca
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 8 deletions.
diff --git 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
@@ -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"
 }
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
@@ -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
@@ -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,14 +29,16 @@
 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:
       artifact_uri: The artifact uri to import from.
       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')},
     )