Artifacts
zenml.artifacts
special
Artifacts are the data that power your experimentation and model training. It is actually steps that produce artifacts, which are then stored in the artifact store. Artifacts are written in the signature of a step like so:
def my_step(first_artifact: int, second_artifact: torch.nn.Module -> int:
# first_artifact is an integer
# second_artifact is a torch.nn.Module
return 1
Artifacts can be serialized and deserialized (i.e. written and read from the
Artifact Store) in various ways like TFRecords
or saved model
pickles, depending on what the step produces.The serialization and
deserialization logic of artifacts is defined by the appropriate Materializer.
base_artifact
The below code is copied from the TFX source repo with minor changes. All credits go to the TFX team for the core implementation
BaseArtifact (Artifact)
Base class for all ZenML artifacts.
Every implementation of an artifact needs to inherit this class.
While inheriting from this class there are a few things to consider:
- Upon creation, each artifact class needs to be given a unique TYPE_NAME.
- Your artifact can feature different properties under the parameter PROPERTIES which will be tracked throughout your pipeline runs.
Source code in zenml/artifacts/base_artifact.py
class BaseArtifact(Artifact):
"""Base class for all ZenML artifacts.
Every implementation of an artifact needs to inherit this class.
While inheriting from this class there are a few things to consider:
- Upon creation, each artifact class needs to be given a unique TYPE_NAME.
- Your artifact can feature different properties under the parameter
PROPERTIES which will be tracked throughout your pipeline runs.
"""
# TODO [ENG-172]: Write about the materializers
TYPE_NAME: str = "BaseArtifact" # type: ignore[assignment]
PROPERTIES: Dict[str, Property] = { # type: ignore[assignment]
MATERIALIZER_PROPERTY_KEY: MATERIALIZER_PROPERTY,
DATATYPE_PROPERTY_KEY: DATATYPE_PROPERTY,
}
_MLMD_ARTIFACT_TYPE: Any = None
def __init__(self, *args: Any, **kwargs: Any) -> None:
"""Init method for BaseArtifact"""
self.set_zenml_artifact_type()
super(BaseArtifact, self).__init__(*args, **kwargs)
@classmethod
def set_zenml_artifact_type(cls) -> None:
"""Set the type of the artifact."""
type_name = cls.TYPE_NAME
if not (type_name and isinstance(type_name, str)):
raise ValueError(
(
"The Artifact subclass %s must override the TYPE_NAME attribute "
"with a string type name identifier (got %r instead)."
)
% (cls, type_name)
)
artifact_type = metadata_store_pb2.ArtifactType()
artifact_type.name = type_name
if cls.PROPERTIES:
# Perform validation on PROPERTIES dictionary.
if not isinstance(cls.PROPERTIES, dict):
raise ValueError(
"Artifact subclass %s.PROPERTIES is not a dictionary." % cls
)
for key, value in cls.PROPERTIES.items():
if not (
isinstance(key, (str, bytes))
and isinstance(value, Property)
):
raise ValueError(
(
"Artifact subclass %s.PROPERTIES dictionary must have keys of "
"type string and values of type artifact.Property."
)
% cls
)
# Populate ML Metadata artifact properties dictionary.
for key, value in cls.PROPERTIES.items():
artifact_type.properties[
key
] = value.mlmd_type() # type: ignore[no-untyped-call]
cls._MLMD_ARTIFACT_TYPE = artifact_type
__init__(self, *args, **kwargs)
special
Init method for BaseArtifact
Source code in zenml/artifacts/base_artifact.py
def __init__(self, *args: Any, **kwargs: Any) -> None:
"""Init method for BaseArtifact"""
self.set_zenml_artifact_type()
super(BaseArtifact, self).__init__(*args, **kwargs)
set_zenml_artifact_type()
classmethod
Set the type of the artifact.
Source code in zenml/artifacts/base_artifact.py
@classmethod
def set_zenml_artifact_type(cls) -> None:
"""Set the type of the artifact."""
type_name = cls.TYPE_NAME
if not (type_name and isinstance(type_name, str)):
raise ValueError(
(
"The Artifact subclass %s must override the TYPE_NAME attribute "
"with a string type name identifier (got %r instead)."
)
% (cls, type_name)
)
artifact_type = metadata_store_pb2.ArtifactType()
artifact_type.name = type_name
if cls.PROPERTIES:
# Perform validation on PROPERTIES dictionary.
if not isinstance(cls.PROPERTIES, dict):
raise ValueError(
"Artifact subclass %s.PROPERTIES is not a dictionary." % cls
)
for key, value in cls.PROPERTIES.items():
if not (
isinstance(key, (str, bytes))
and isinstance(value, Property)
):
raise ValueError(
(
"Artifact subclass %s.PROPERTIES dictionary must have keys of "
"type string and values of type artifact.Property."
)
% cls
)
# Populate ML Metadata artifact properties dictionary.
for key, value in cls.PROPERTIES.items():
artifact_type.properties[
key
] = value.mlmd_type() # type: ignore[no-untyped-call]
cls._MLMD_ARTIFACT_TYPE = artifact_type
data_analysis_artifact
DataAnalysisArtifact (BaseArtifact)
Class for all ZenML data analysis artifacts.
This should act as a base class for all artifacts generated from processes such as data profiling, data drift analyses, model drift detection etc.
Source code in zenml/artifacts/data_analysis_artifact.py
class DataAnalysisArtifact(BaseArtifact):
"""Class for all ZenML data analysis artifacts.
This should act as a base class for all artifacts generated from
processes such as data profiling, data drift analyses, model drift
detection etc.
"""
TYPE_NAME = "DataAnalysisArtifact"
data_artifact
DataArtifact (BaseArtifact)
Class for all ZenML data artifacts.
Source code in zenml/artifacts/data_artifact.py
class DataArtifact(BaseArtifact):
"""Class for all ZenML data artifacts."""
TYPE_NAME = "DataArtifact"
model_artifact
ModelArtifact (BaseArtifact)
Class for all ZenML model artifacts.
Source code in zenml/artifacts/model_artifact.py
class ModelArtifact(BaseArtifact):
"""Class for all ZenML model artifacts."""
TYPE_NAME = "ModelArtifact"
schema_artifact
SchemaArtifact (BaseArtifact)
Class for all ZenML schema artifacts.
Source code in zenml/artifacts/schema_artifact.py
class SchemaArtifact(BaseArtifact):
"""Class for all ZenML schema artifacts."""
TYPE_NAME = "SchemaArtifact"
statistics_artifact
StatisticsArtifact (BaseArtifact)
Class for all ZenML statistics artifacts.
Source code in zenml/artifacts/statistics_artifact.py
class StatisticsArtifact(BaseArtifact):
"""Class for all ZenML statistics artifacts."""
TYPE_NAME = "StatisticsArtifact"
type_registry
ArtifactTypeRegistry
A registry to keep track of which datatypes map to which artifact types
Source code in zenml/artifacts/type_registry.py
class ArtifactTypeRegistry(object):
"""A registry to keep track of which datatypes map to which artifact
types"""
def __init__(self) -> None:
"""Initialization with an empty registry"""
self._artifact_types: Dict[
Type[Any], Tuple[Type["BaseArtifact"], ...]
] = {}
def register_integration(
self, key: Type[Any], type_: Iterable[Type["BaseArtifact"]]
) -> None:
"""Method to register an integration within the registry
Args:
key: any datatype
type_: the list of artifact type that the given datatypes is
associated with
"""
self._artifact_types[key] = tuple(type_)
def get_artifact_type(
self, key: Type[Any]
) -> Tuple[Type["BaseArtifact"], ...]:
"""Method to extract the list of artifact types given the data type
Args:
key: Indicates the type of object.
Returns:
A list of `Artifact` types that was registered for this key.
Raises:
StepInterfaceError: If the key (or any of its superclasses) is not
registered or the key has more than one superclass with
different default materializers/artifact types
"""
# Check whether the type is registered
if key in self._artifact_types:
return self._artifact_types[key]
else:
# If the type is not registered, check for superclasses
artifact_types_for_compatible_superclasses = {
artifact_types
for registered_type, artifact_types in self._artifact_types.items()
if issubclass(key, registered_type)
}
# Make sure that there is only a single list of artifact types
if len(artifact_types_for_compatible_superclasses) == 1:
return artifact_types_for_compatible_superclasses.pop()
elif len(artifact_types_for_compatible_superclasses) > 1:
raise StepInterfaceError(
f"Type {key} is subclassing more than one type and these "
f"types map to different materializers. These "
f"materializers feature a different list associated "
f"artifact types within the registry: "
f"{artifact_types_for_compatible_superclasses}. Please "
f"specify which of these artifact types you would like to "
f"use explicitly in your step."
)
raise StepInterfaceError(
f"Type {key} does not have a default `Materializer` thus it does "
f"not have any associated `ArtifactType`s! Please specify your "
f"own `Materializer`."
)
__init__(self)
special
Initialization with an empty registry
Source code in zenml/artifacts/type_registry.py
def __init__(self) -> None:
"""Initialization with an empty registry"""
self._artifact_types: Dict[
Type[Any], Tuple[Type["BaseArtifact"], ...]
] = {}
get_artifact_type(self, key)
Method to extract the list of artifact types given the data type
Parameters:
Name | Type | Description | Default |
---|---|---|---|
key |
Type[Any] |
Indicates the type of object. |
required |
Returns:
Type | Description |
---|---|
Tuple[Type[BaseArtifact], ...] |
A list of |
Exceptions:
Type | Description |
---|---|
StepInterfaceError |
If the key (or any of its superclasses) is not registered or the key has more than one superclass with different default materializers/artifact types |
Source code in zenml/artifacts/type_registry.py
def get_artifact_type(
self, key: Type[Any]
) -> Tuple[Type["BaseArtifact"], ...]:
"""Method to extract the list of artifact types given the data type
Args:
key: Indicates the type of object.
Returns:
A list of `Artifact` types that was registered for this key.
Raises:
StepInterfaceError: If the key (or any of its superclasses) is not
registered or the key has more than one superclass with
different default materializers/artifact types
"""
# Check whether the type is registered
if key in self._artifact_types:
return self._artifact_types[key]
else:
# If the type is not registered, check for superclasses
artifact_types_for_compatible_superclasses = {
artifact_types
for registered_type, artifact_types in self._artifact_types.items()
if issubclass(key, registered_type)
}
# Make sure that there is only a single list of artifact types
if len(artifact_types_for_compatible_superclasses) == 1:
return artifact_types_for_compatible_superclasses.pop()
elif len(artifact_types_for_compatible_superclasses) > 1:
raise StepInterfaceError(
f"Type {key} is subclassing more than one type and these "
f"types map to different materializers. These "
f"materializers feature a different list associated "
f"artifact types within the registry: "
f"{artifact_types_for_compatible_superclasses}. Please "
f"specify which of these artifact types you would like to "
f"use explicitly in your step."
)
raise StepInterfaceError(
f"Type {key} does not have a default `Materializer` thus it does "
f"not have any associated `ArtifactType`s! Please specify your "
f"own `Materializer`."
)
register_integration(self, key, type_)
Method to register an integration within the registry
Parameters:
Name | Type | Description | Default |
---|---|---|---|
key |
Type[Any] |
any datatype |
required |
type_ |
Iterable[Type[BaseArtifact]] |
the list of artifact type that the given datatypes is associated with |
required |
Source code in zenml/artifacts/type_registry.py
def register_integration(
self, key: Type[Any], type_: Iterable[Type["BaseArtifact"]]
) -> None:
"""Method to register an integration within the registry
Args:
key: any datatype
type_: the list of artifact type that the given datatypes is
associated with
"""
self._artifact_types[key] = tuple(type_)