support artifact types under google namespace

sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.json

@@ -46,6 +46,12 @@
                 "schemaTitle": "system.HTML",
                 "schemaVersion": "0.0.1"
               }
+            },
+            "input_i": {
+              "artifactType": {
+                "schemaTitle": "google.BQMLModel",
+                "schemaVersion": "0.0.1"
+              }
             }
           },
           "parameters": {
@@ -116,6 +122,12 @@
                 "schemaTitle": "system.HTML",
                 "schemaVersion": "0.0.1"
               }
+            },
+            "output_9": {
+              "artifactType": {
+                "schemaTitle": "google.BQMLModel",
+                "schemaVersion": "0.0.1"
+              }
             }
           },
           "parameters": {
@@ -223,6 +235,12 @@
                     "outputArtifactKey": "output_8",
                     "producerTask": "upstream"
                   }
+                },
+                "input_i": {
+                  "taskOutputArtifact": {
+                    "outputArtifactKey": "output_9",
+                    "producerTask": "upstream"
+                  }
                 }
               },
               "parameters": {
@@ -286,7 +304,7 @@
       }
     },
     "schemaVersion": "2.0.0",
-    "sdkVersion": "kfp-1.7.2"
+    "sdkVersion": "kfp-1.8.3"
   },
   "runtimeConfig": {
     "gcsOutputDirectory": "dummy_root",

sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.py

@@ -15,8 +15,7 @@
 import pathlib
 
 from kfp import components
-from kfp.v2 import dsl
-import kfp.v2.compiler as compiler
+from kfp.v2 import compiler, dsl
 
 component_op_1 = components.load_component_from_text("""
 name: upstream
@@ -34,6 +33,7 @@
 - {name: output_6, type: Some arbitrary type}
 - {name: output_7, type: {GcsPath: {data_type: TSV}}}
 - {name: output_8, type: HTML}
+- {name: output_9, type: google.BQMLModel}
 implementation:
   container:
     image: gcr.io/image
@@ -63,6 +63,7 @@
 - {name: input_f, type: Some arbitrary type}
 - {name: input_g, type: {GcsPath: {data_type: TSV}}}
 - {name: input_h, type: HTML}
+- {name: input_i, type: google.BQMLModel}
 implementation:
   container:
     image: gcr.io/image
@@ -95,6 +96,7 @@ def my_pipeline(input1: str, input3: str, input4: str = ''):
         input_f=component_1.outputs['output_6'],
         input_g=component_1.outputs['output_7'],
         input_h=component_1.outputs['output_8'],
+        input_i=component_1.outputs['output_9'],
     )
 
 

sdk/python/kfp/v2/components/types/type_utils.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 """Utilities for component I/O type mapping."""
 import inspect
+import re
 from typing import Dict, List, Optional, Type, Union
-from kfp.components import structures
-from kfp.components import type_annotation_utils
+
+from kfp.components import structures, type_annotation_utils
 from kfp.pipeline_spec import pipeline_spec_pb2
 from kfp.v2.components.types import artifact_types
 
-
 PARAMETER_TYPES = Union[str, int, float, bool, dict, list]
 
 # ComponentSpec I/O types to DSL ontology artifact classes mapping.
@@ -33,6 +33,9 @@
     'markdown': artifact_types.Markdown,
 }
 
+_GOOGLE_TYPES_PATTERN = r'^google.[A-Za-z]+$'
+_GOOGLE_TYPES_VERSION = '0.0.1'
+
 # ComponentSpec I/O types to (IR) PipelineTaskSpec I/O types mapping.
 # The keys are normalized (lowercased). These are types viewed as Parameters.
 # The values are the corresponding IR parameter primitive types.
@@ -88,6 +91,11 @@ def get_artifact_type_schema(
     type."""
     artifact_class = artifact_types.Artifact
     if isinstance(artifact_class_or_type_name, str):
+        if re.match(_GOOGLE_TYPES_PATTERN, artifact_class_or_type_name):
+            return pipeline_spec_pb2.ArtifactTypeSchema(
+                schema_title=artifact_class_or_type_name,
+                schema_version=_GOOGLE_TYPES_VERSION,
+            )
         artifact_class = _ARTIFACT_CLASSES_MAPPING.get(
             artifact_class_or_type_name.lower(), artifact_types.Artifact)
     elif inspect.isclass(artifact_class_or_type_name) and issubclass(

sdk/python/kfp/v2/components/types/type_utils_test.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from absl.testing import parameterized
-
 import sys
 import unittest
 from typing import Any, Dict, List
 
+from absl.testing import parameterized
 from kfp.components import structures
 from kfp.pipeline_spec import pipeline_spec_pb2 as pb
 from kfp.v2.components.types import artifact_types, type_utils
@@ -49,6 +48,14 @@ class _ArbitraryClass:
     pass
 
 
+class _VertexDummy(artifact_types.Artifact):
+    TYPE_NAME = 'google.VertexDummy'
+    VERSION = '0.0.2'
+
+    def __init__(self):
+        super().__init__(uri='uri', name='name', metadata={'dummy': '123'})
+
+
 class TypeUtilsTest(parameterized.TestCase):
 
     def test_is_parameter_type(self):
@@ -160,6 +167,27 @@ def test_is_parameter_type(self):
                 pb.ArtifactTypeSchema(
                     schema_title='system.Markdown', schema_version='0.0.1')
         },
+        {
+            'artifact_class_or_type_name':
+                'some-google-type',
+            'expected_result':
+                pb.ArtifactTypeSchema(
+                    schema_title='system.Artifact', schema_version='0.0.1')
+        },
+        {
+            'artifact_class_or_type_name':
+                'google.VertexModel',
+            'expected_result':
+                pb.ArtifactTypeSchema(
+                    schema_title='google.VertexModel', schema_version='0.0.1')
+        },
+        {
+            'artifact_class_or_type_name':
+                _VertexDummy,
+            'expected_result':
+                pb.ArtifactTypeSchema(
+                    schema_title='google.VertexDummy', schema_version='0.0.2')
+        },
     )
     def test_get_artifact_type_schema(self, artifact_class_or_type_name,
                                       expected_result):