diff --git a/tfx/types/artifact_utils.py b/tfx/types/artifact_utils.py index 358400cbc4..8ef5539e21 100644 --- a/tfx/types/artifact_utils.py +++ b/tfx/types/artifact_utils.py @@ -169,7 +169,11 @@ def get_artifact_type_class( # We need to compare `.name` and `.properties` (and not the entire proto # directly), because the proto `.id` field will be populated when the type # is read from MLMD. - if (artifact_type.name == candidate_type.name and + artifact_type_name = artifact_type.name + if artifact_type_name.startswith('tfx.'): + artifact_type_name = artifact_type_name[4:] + + if (artifact_type_name == candidate_type.name and artifact_type.properties == candidate_type.properties): return cls diff --git a/tfx/types/artifact_utils_test.py b/tfx/types/artifact_utils_test.py index b4faa6299a..50033fc549 100644 --- a/tfx/types/artifact_utils_test.py +++ b/tfx/types/artifact_utils_test.py @@ -133,6 +133,16 @@ def testArtifactTypeRoundTrip(self): self.assertIs(_MyArtifact, artifact_utils.get_artifact_type_class(mlmd_artifact_type)) + # Test that artifact types prefixed with 'tfx.' (as produced by KFP v2) + # are correctly resolved to their native Python classes. + mlmd_artifact_type_prefixed = metadata_store_pb2.ArtifactType() + mlmd_artifact_type_prefixed.name = 'tfx.Examples' + mlmd_artifact_type_prefixed.properties['span'] = metadata_store_pb2.INT + mlmd_artifact_type_prefixed.properties['version'] = metadata_store_pb2.INT + mlmd_artifact_type_prefixed.properties['split_names'] = metadata_store_pb2.STRING + self.assertIs(standard_artifacts.Examples, + artifact_utils.get_artifact_type_class(mlmd_artifact_type_prefixed)) + def testValueArtifactTypeRoundTrip(self): mlmd_artifact_type = standard_artifacts.String.annotate_as( # pylint: disable=protected-access system_artifacts.Dataset)._get_artifact_type()