diff --git a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.json b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.json index ada9f5afe5c..cfda8706bed 100644 --- a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.json +++ b/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", diff --git a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.py b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.py index 111ce45c70b..0afcca594f7 100644 --- a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_various_io_types.py +++ b/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'], ) diff --git a/sdk/python/kfp/v2/components/types/type_utils.py b/sdk/python/kfp/v2/components/types/type_utils.py index b1f6e4c2d65..22b1a7ca164 100644 --- a/sdk/python/kfp/v2/components/types/type_utils.py +++ b/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( diff --git a/sdk/python/kfp/v2/components/types/type_utils_test.py b/sdk/python/kfp/v2/components/types/type_utils_test.py index 9e682120177..c94543c38be 100644 --- a/sdk/python/kfp/v2/components/types/type_utils_test.py +++ b/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):