def get_artifact_type_class( artifact_type: metadata_store_pb2.ArtifactType) -> Type[Artifact]: """Get the artifact type class corresponding to an MLMD type proto.""" # Make sure this module path containing the standard Artifact subclass # definitions is imported. Modules containing custom artifact subclasses that # need to be deserialized should be imported by the entrypoint of the # application or container. from tfx.types import standard_artifacts # pylint: disable=g-import-not-at-top,import-outside-toplevel,unused-import,unused-variable # Enumerate the Artifact type ontology, separated into auto-generated and # natively-defined classes. artifact_classes = _get_subclasses(Artifact) native_artifact_classes = [] generated_artifact_classes = [] for cls in artifact_classes: if not cls.TYPE_NAME: # Skip abstract classes. continue if getattr(cls, '_AUTOGENERATED', False): generated_artifact_classes.append(cls) else: native_artifact_classes.append(cls) # Try to find an existing class for the artifact type, if it exists. Prefer # to use a native artifact class. for cls in itertools.chain(native_artifact_classes, generated_artifact_classes): candidate_type = cls._get_artifact_type() # pylint: disable=protected-access # 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.properties == candidate_type.properties): return cls # Generate a class for the artifact type on the fly. absl.logging.warning( ('Could not find matching artifact class for type %r (proto: %r); ' 'generating an ephemeral artifact class on-the-fly. If this is not ' 'intended, please make sure that the artifact class for this type can ' 'be imported within your container or environment where a component ' 'is executed to consume this type.') % (artifact_type.name, str(artifact_type))) new_artifact_class = _ArtifactType(mlmd_artifact_type=artifact_type) setattr(new_artifact_class, '_AUTOGENERATED', True) return new_artifact_class
class _MyArtifact(artifact.Artifact): TYPE_NAME = 'MyTypeName' PROPERTIES = { 'int1': artifact.Property(type=artifact.PropertyType.INT), 'int2': artifact.Property(type=artifact.PropertyType.INT), 'float1': artifact.Property(type=artifact.PropertyType.FLOAT), 'float2': artifact.Property(type=artifact.PropertyType.FLOAT), 'string1': artifact.Property(type=artifact.PropertyType.STRING), 'string2': artifact.Property(type=artifact.PropertyType.STRING), } _MyArtifact2 = artifact._ArtifactType( # pylint: disable=invalid-name name='MyTypeName2', properties={ 'int1': artifact.Property(type=artifact.PropertyType.INT), 'int2': artifact.Property(type=artifact.PropertyType.INT), 'float1': artifact.Property(type=artifact.PropertyType.FLOAT), 'float2': artifact.Property(type=artifact.PropertyType.FLOAT), 'string1': artifact.Property(type=artifact.PropertyType.STRING), 'string2': artifact.Property(type=artifact.PropertyType.STRING), }) _mlmd_artifact_type = metadata_store_pb2.ArtifactType() json_format.Parse( json.dumps({ 'name': 'MyTypeName3', 'properties': { 'int1': 'INT', 'int2': 'INT', 'float1': 'DOUBLE', 'float2': 'DOUBLE', 'string1': 'STRING',
'string1': artifact.Property(type=artifact.PropertyType.STRING), 'string2': artifact.Property(type=artifact.PropertyType.STRING), } _MyArtifact2 = artifact._ArtifactType( # pylint: disable=invalid-name name='MyTypeName2', properties={ 'int1': artifact.Property(type=artifact.PropertyType.INT), 'int2': artifact.Property(type=artifact.PropertyType.INT), 'float1': artifact.Property(type=artifact.PropertyType.FLOAT), 'float2': artifact.Property(type=artifact.PropertyType.FLOAT), 'string1': artifact.Property(type=artifact.PropertyType.STRING), 'string2': artifact.Property(type=artifact.PropertyType.STRING), 'jsonvalue_string': artifact.Property(type=artifact.PropertyType.JSON_VALUE), 'jsonvalue_dict': artifact.Property(type=artifact.PropertyType.JSON_VALUE), 'jsonvalue_int': artifact.Property(type=artifact.PropertyType.JSON_VALUE), 'jsonvalue_float': artifact.Property(type=artifact.PropertyType.JSON_VALUE), 'jsonvalue_list': artifact.Property(type=artifact.PropertyType.JSON_VALUE), 'jsonvalue_null': artifact.Property(type=artifact.PropertyType.JSON_VALUE), 'jsonvalue_empty': artifact.Property(type=artifact.PropertyType.JSON_VALUE), }) _mlmd_artifact_type = metadata_store_pb2.ArtifactType() json_format.Parse( json.dumps({