Materializers

zenml.materializers

Initialization of ZenML materializers.

Materializers are used to convert a ZenML artifact into a specific format. They are most often used to handle the input or output of ZenML steps, and can be extended by building on the BaseMaterializer class.

Attributes

__all__ = ['BuiltInContainerMaterializer', 'BuiltInMaterializer', 'BytesMaterializer', 'CloudpickleMaterializer', 'StructuredStringMaterializer', 'PydanticMaterializer', 'ServiceMaterializer', 'UUIDMaterializer'] module-attribute

Classes

BuiltInContainerMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle built-in container types (dict, list, set, tuple).

Define self.data_path and self.metadata_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/built_in_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path` and `self.metadata_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)
    self.metadata_path = os.path.join(self.uri, DEFAULT_METADATA_FILENAME)

Functions

extract_metadata(data: Any) -> Dict[str, MetadataType]

Extract metadata from the given built-in container object.

Parameters:

Name	Type	Description	Default
data	Any	The built-in container object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/built_in_materializer.py

def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given built-in container object.

    Args:
        data: The built-in container object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    if hasattr(data, "__len__"):
        return {"length": len(data)}
    return {}

load(data_type: Type[Any]) -> Any

Reads a materialized built-in container object.

If the data was serialized to JSON, deserialize it.

Otherwise, reconstruct all elements according to the metadata file: 1. Resolve the data type using find_type_by_str(), 2. Get the materializer via the default_materializer_registry, 3. Initialize the materializer with the desired path, 4. Use load() of that materializer to load the element.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Raises:

Type	Description
RuntimeError	If the data was not found.

Source code in src/zenml/materializers/built_in_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Reads a materialized built-in container object.

    If the data was serialized to JSON, deserialize it.

    Otherwise, reconstruct all elements according to the metadata file:
        1. Resolve the data type using `find_type_by_str()`,
        2. Get the materializer via the `default_materializer_registry`,
        3. Initialize the materializer with the desired path,
        4. Use `load()` of that materializer to load the element.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.

    Raises:
        RuntimeError: If the data was not found.
    """
    # If the data was not serialized, there must be metadata present.
    if not self.artifact_store.exists(
        self.data_path
    ) and not self.artifact_store.exists(self.metadata_path):
        raise RuntimeError(
            f"Materialization of type {data_type} failed. Expected either"
            f"{self.data_path} or {self.metadata_path} to exist."
        )

    # If the data was serialized as JSON, deserialize it.
    if self.artifact_store.exists(self.data_path):
        outputs = yaml_utils.read_json(self.data_path)

    # Otherwise, use the metadata to reconstruct the data as a list.
    else:
        metadata = yaml_utils.read_json(self.metadata_path)
        outputs = []

        # Backwards compatibility for zenml <= 0.37.0
        if isinstance(metadata, dict):
            for path_, type_str in zip(
                metadata["paths"], metadata["types"]
            ):
                type_ = find_type_by_str(type_str)
                materializer_class = materializer_registry[type_]
                materializer = materializer_class(uri=path_)
                element = materializer.load(type_)
                outputs.append(element)

        # New format for zenml > 0.37.0
        elif isinstance(metadata, list):
            for entry in metadata:
                path_ = entry["path"]
                type_ = source_utils.load(entry["type"])
                materializer_class = source_utils.load(
                    entry["materializer"]
                )
                materializer = materializer_class(uri=path_)
                element = materializer.load(type_)
                outputs.append(element)

        else:
            raise RuntimeError(f"Unknown metadata format: {metadata}.")

    # Cast the data to the correct type.
    if issubclass(data_type, dict) and not isinstance(outputs, dict):
        keys, values = outputs
        return data_type(zip(keys, values))
    if issubclass(data_type, tuple) and not isinstance(outputs, tuple):
        return data_type(outputs)
    if issubclass(data_type, set) and not isinstance(outputs, set):
        return data_type(outputs)
    return outputs

save(data: Any) -> None

Materialize a built-in container object.

If the object can be serialized to JSON, serialize it.

Otherwise, use the default_materializer_registry to find the correct materializer for each element and materialize each element into a subdirectory.

Tuples and sets are cast to list before materialization.

For non-serializable dicts, materialize keys/values as separate lists.

Parameters:

Name	Type	Description	Default
data	Any	The built-in container object to materialize.	required

Raises:

Type	Description
Exception	If any exception occurs, it is raised after cleanup.

Source code in src/zenml/materializers/built_in_materializer.py

def save(self, data: Any) -> None:
    """Materialize a built-in container object.

    If the object can be serialized to JSON, serialize it.

    Otherwise, use the `default_materializer_registry` to find the correct
    materializer for each element and materialize each element into a
    subdirectory.

    Tuples and sets are cast to list before materialization.

    For non-serializable dicts, materialize keys/values as separate lists.

    Args:
        data: The built-in container object to materialize.

    Raises:
        Exception: If any exception occurs, it is raised after cleanup.
    """
    # tuple and set: handle as list.
    if isinstance(data, tuple) or isinstance(data, set):
        data = list(data)

    # If the data is serializable, just write it into a single JSON file.
    if _is_serializable(data):
        yaml_utils.write_json(
            self.data_path,
            data,
            ensure_ascii=not ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS,
        )
        return

    # non-serializable dict: Handle as non-serializable list of lists.
    if isinstance(data, dict):
        data = [list(data.keys()), list(data.values())]

    # non-serializable list: Materialize each element into a subfolder.
    # Get path, type, and corresponding materializer for each element.
    metadata: List[Dict[str, str]] = []
    materializers: List[BaseMaterializer] = []
    try:
        for i, element in enumerate(data):
            element_path = os.path.join(self.uri, str(i))
            self.artifact_store.mkdir(element_path)
            type_ = type(element)
            materializer_class = materializer_registry[type_]
            materializer = materializer_class(uri=element_path)
            materializers.append(materializer)
            metadata.append(
                {
                    "path": element_path,
                    "type": source_utils.resolve(type_).import_path,
                    "materializer": source_utils.resolve(
                        materializer_class
                    ).import_path,
                }
            )
        # Write metadata as JSON.
        yaml_utils.write_json(self.metadata_path, metadata)
        # Materialize each element.
        for element, materializer in zip(data, materializers):
            materializer.validate_save_type_compatibility(type(element))
            materializer.save(element)
    # If an error occurs, delete all created files.
    except Exception as e:
        # Delete metadata
        if self.artifact_store.exists(self.metadata_path):
            self.artifact_store.remove(self.metadata_path)
        # Delete all elements that were already saved.
        for entry in metadata:
            self.artifact_store.rmtree(entry["path"])
        raise e

save_visualizations(data: Any) -> Dict[str, VisualizationType]

Save visualizations for the given data.

Parameters:

Name	Type	Description	Default
data	Any	The data to save visualizations for.	required

Returns:

Type	Description
Dict[str, VisualizationType]	A dictionary of visualization URIs and their types.

Source code in src/zenml/materializers/built_in_materializer.py

def save_visualizations(self, data: Any) -> Dict[str, "VisualizationType"]:
    """Save visualizations for the given data.

    Args:
        data: The data to save visualizations for.

    Returns:
        A dictionary of visualization URIs and their types.
    """
    # dict/list type objects are always saved as JSON files
    # doesn't work for non-serializable types as they
    # are saved as list of lists in different files
    if _is_serializable(data):
        return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
    return {}

BuiltInMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle JSON-serializable basic types (bool, float, int, str).

Define self.data_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/built_in_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)

Functions

extract_metadata(data: Union[bool, float, int, str]) -> Dict[str, MetadataType]

Extract metadata from the given built-in container object.

Parameters:

Name	Type	Description	Default
data	Union[bool, float, int, str]	The built-in container object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/built_in_materializer.py

def extract_metadata(
    self, data: Union[bool, float, int, str]
) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given built-in container object.

    Args:
        data: The built-in container object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    # For boolean and numbers, add the string representation as metadata.
    # We don't to this for strings because they can be arbitrarily long.
    if isinstance(data, (bool, float, int)):
        return {"string_representation": str(data)}

    return {}

load(data_type: Union[Type[bool], Type[float], Type[int], Type[str]]) -> Any

Reads basic primitive types from JSON.

Parameters:

Name	Type	Description	Default
data_type	Union[Type[bool], Type[float], Type[int], Type[str]]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Source code in src/zenml/materializers/built_in_materializer.py

def load(
    self, data_type: Union[Type[bool], Type[float], Type[int], Type[str]]
) -> Any:
    """Reads basic primitive types from JSON.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.
    """
    contents = yaml_utils.read_json(self.data_path)
    if type(contents) is not data_type:
        # TODO [ENG-142]: Raise error or try to coerce
        logger.debug(
            f"Contents {contents} was type {type(contents)} but expected "
            f"{data_type}"
        )
    return contents

save(data: Union[bool, float, int, str]) -> None

Serialize a basic type to JSON.

Parameters:

Name	Type	Description	Default
data	Union[bool, float, int, str]	The data to store.	required

Source code in src/zenml/materializers/built_in_materializer.py

def save(self, data: Union[bool, float, int, str]) -> None:
    """Serialize a basic type to JSON.

    Args:
        data: The data to store.
    """
    yaml_utils.write_json(
        self.data_path,
        data,
        ensure_ascii=not ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS,
    )

BytesMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle bytes data type, which is not JSON serializable.

Define self.data_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/built_in_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_BYTES_FILENAME)

Functions

load(data_type: Type[Any]) -> Any

Reads a bytes object from file.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Source code in src/zenml/materializers/built_in_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Reads a bytes object from file.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.
    """
    with self.artifact_store.open(self.data_path, "rb") as file_:
        return file_.read()

save(data: Any) -> None

Save a bytes object to file.

Parameters:

Name	Type	Description	Default
data	Any	The data to store.	required

Source code in src/zenml/materializers/built_in_materializer.py

def save(self, data: Any) -> None:
    """Save a bytes object to file.

    Args:
        data: The data to store.
    """
    with self.artifact_store.open(self.data_path, "wb") as file_:
        file_.write(data)

CloudpickleMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer using cloudpickle.

This materializer can materialize (almost) any object, but does so in a non-reproducble way since artifacts cannot be loaded from other Python versions. It is recommended to use this materializer only as a last resort.

That is also why it has SKIP_REGISTRATION set to True and is currently only used as a fallback materializer inside the materializer registry.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

load(data_type: Type[Any]) -> Any

Reads an artifact from a cloudpickle file.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The data type of the artifact.	required

Returns:

Type	Description
Any	The loaded artifact data.

Source code in src/zenml/materializers/cloudpickle_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Reads an artifact from a cloudpickle file.

    Args:
        data_type: The data type of the artifact.

    Returns:
        The loaded artifact data.
    """
    # validate python version
    source_python_version = self._load_python_version()
    current_python_version = Environment().python_version()
    if source_python_version != current_python_version:
        logger.warning(
            f"Your artifact was materialized under Python version "
            f"'{source_python_version}' but you are currently using "
            f"'{current_python_version}'. This might cause unexpected "
            "behavior since pickle is not reproducible across Python "
            "versions. Attempting to load anyway..."
        )

    # load data
    filepath = os.path.join(self.uri, DEFAULT_FILENAME)
    with self.artifact_store.open(filepath, "rb") as fid:
        data = cloudpickle.load(fid)
    return data

save(data: Any) -> None

Saves an artifact to a cloudpickle file.

Parameters:

Name	Type	Description	Default
data	Any	The data to save.	required

Source code in src/zenml/materializers/cloudpickle_materializer.py

def save(self, data: Any) -> None:
    """Saves an artifact to a cloudpickle file.

    Args:
        data: The data to save.
    """
    # Log a warning if this materializer was not explicitly specified for
    # the given data type.
    if type(self) is CloudpickleMaterializer:
        logger.warning(
            f"No materializer is registered for type `{type(data)}`, so "
            "the default Pickle materializer was used. Pickle is not "
            "production ready and should only be used for prototyping as "
            "the artifacts cannot be loaded when running with a different "
            "Python version. Please consider implementing a custom "
            f"materializer for type `{type(data)}` according to the "
            "instructions at https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types"
        )

    # save python version for validation on loading
    self._save_python_version()

    # save data
    filepath = os.path.join(self.uri, DEFAULT_FILENAME)
    with self.artifact_store.open(filepath, "wb") as fid:
        cloudpickle.dump(data, fid)

PydanticMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle Pydantic BaseModel objects.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

extract_metadata(data: BaseModel) -> Dict[str, MetadataType]

Extract metadata from the given BaseModel object.

Parameters:

Name	Type	Description	Default
data	BaseModel	The BaseModel object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/pydantic_materializer.py

def extract_metadata(self, data: BaseModel) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given BaseModel object.

    Args:
        data: The BaseModel object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    return {"schema": data.schema()}

load(data_type: Type[BaseModel]) -> Any

Reads BaseModel from JSON.

Parameters:

Name	Type	Description	Default
data_type	Type[BaseModel]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Source code in src/zenml/materializers/pydantic_materializer.py

def load(self, data_type: Type[BaseModel]) -> Any:
    """Reads BaseModel from JSON.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.
    """
    data_path = os.path.join(self.uri, DEFAULT_FILENAME)
    contents = yaml_utils.read_json(data_path)
    return data_type.model_validate_json(contents)

save(data: BaseModel) -> None

Serialize a BaseModel to JSON.

Parameters:

Name	Type	Description	Default
data	BaseModel	The data to store.	required

Source code in src/zenml/materializers/pydantic_materializer.py

def save(self, data: BaseModel) -> None:
    """Serialize a BaseModel to JSON.

    Args:
        data: The data to store.
    """
    data_path = os.path.join(self.uri, DEFAULT_FILENAME)
    yaml_utils.write_json(data_path, data.model_dump_json())

ServiceMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer to read/write service instances.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

extract_metadata(service: BaseService) -> Dict[str, MetadataType]

Extract metadata from the given service.

Parameters:

Name	Type	Description	Default
service	BaseService	The service to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/service_materializer.py

def extract_metadata(
    self, service: BaseService
) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given service.

    Args:
        service: The service to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    from zenml.metadata.metadata_types import Uri

    if prediction_url := service.get_prediction_url() or None:
        return {"uri": Uri(prediction_url)}
    return {}

load(data_type: Type[Any]) -> BaseService

Creates and returns a service.

This service is instantiated from the serialized service configuration and last known status information saved as artifact.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type of the data to read.	required

Returns:

Type	Description
BaseService	A ZenML service instance.

Source code in src/zenml/materializers/service_materializer.py

def load(self, data_type: Type[Any]) -> BaseService:
    """Creates and returns a service.

    This service is instantiated from the serialized service configuration
    and last known status information saved as artifact.

    Args:
        data_type: The type of the data to read.

    Returns:
        A ZenML service instance.
    """
    filepath = os.path.join(self.uri, SERVICE_CONFIG_FILENAME)
    with self.artifact_store.open(filepath, "r") as f:
        service_id = f.read().strip()

    service = Client().get_service(name_id_or_prefix=uuid.UUID(service_id))
    return BaseDeploymentService.from_model(service)

save(service: BaseService) -> None

Writes a ZenML service.

The configuration and last known status of the input service instance are serialized and saved as an artifact.

Parameters:

Name	Type	Description	Default
service	BaseService	A ZenML service instance.	required

Source code in src/zenml/materializers/service_materializer.py

def save(self, service: BaseService) -> None:
    """Writes a ZenML service.

    The configuration and last known status of the input service instance
    are serialized and saved as an artifact.

    Args:
        service: A ZenML service instance.
    """
    filepath = os.path.join(self.uri, SERVICE_CONFIG_FILENAME)
    with self.artifact_store.open(filepath, "w") as f:
        f.write(str(service.uuid))

StructuredStringMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer for HTML or Markdown strings.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

load(data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS

Loads the data from the HTML or Markdown file.

Parameters:

Name	Type	Description	Default
data_type	Type[STRUCTURED_STRINGS]	The type of the data to read.	required

Returns:

Type	Description
STRUCTURED_STRINGS	The loaded data.

Source code in src/zenml/materializers/structured_string_materializer.py

def load(self, data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS:
    """Loads the data from the HTML or Markdown file.

    Args:
        data_type: The type of the data to read.

    Returns:
        The loaded data.
    """
    with self.artifact_store.open(self._get_filepath(data_type), "r") as f:
        return data_type(f.read())

save(data: STRUCTURED_STRINGS) -> None

Save data as an HTML or Markdown file.

Parameters:

Name	Type	Description	Default
data	STRUCTURED_STRINGS	The data to save as an HTML or Markdown file.	required

Source code in src/zenml/materializers/structured_string_materializer.py

def save(self, data: STRUCTURED_STRINGS) -> None:
    """Save data as an HTML or Markdown file.

    Args:
        data: The data to save as an HTML or Markdown file.
    """
    with self.artifact_store.open(
        self._get_filepath(type(data)), "w"
    ) as f:
        f.write(data)

save_visualizations(data: STRUCTURED_STRINGS) -> Dict[str, VisualizationType]

Save visualizations for the given data.

Parameters:

Name	Type	Description	Default
data	STRUCTURED_STRINGS	The data to save visualizations for.	required

Returns:

Type	Description
Dict[str, VisualizationType]	A dictionary of visualization URIs and their types.

Source code in src/zenml/materializers/structured_string_materializer.py

def save_visualizations(
    self, data: STRUCTURED_STRINGS
) -> Dict[str, VisualizationType]:
    """Save visualizations for the given data.

    Args:
        data: The data to save visualizations for.

    Returns:
        A dictionary of visualization URIs and their types.
    """
    filepath = self._get_filepath(type(data))
    filepath = filepath.replace("\\", "/")
    visualization_type = self._get_visualization_type(type(data))
    return {filepath: visualization_type}

UUIDMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer to handle UUID objects.

Define self.data_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/uuid_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)

Functions

extract_metadata(data: uuid.UUID) -> Dict[str, MetadataType]

Extract metadata from the UUID.

Parameters:

Name	Type	Description	Default
data	UUID	The UUID to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	A dictionary of metadata extracted from the UUID.

Source code in src/zenml/materializers/uuid_materializer.py

def extract_metadata(self, data: uuid.UUID) -> Dict[str, MetadataType]:
    """Extract metadata from the UUID.

    Args:
        data: The UUID to extract metadata from.

    Returns:
        A dictionary of metadata extracted from the UUID.
    """
    return {
        "string_representation": str(data),
    }

load(_: Type[uuid.UUID]) -> uuid.UUID

Read UUID from artifact store.

Parameters:

Name	Type	Description	Default
_	Type[UUID]	The type of the data to be loaded.	required

Returns:

Type	Description
UUID	The loaded UUID.

Source code in src/zenml/materializers/uuid_materializer.py

def load(self, _: Type[uuid.UUID]) -> uuid.UUID:
    """Read UUID from artifact store.

    Args:
        _: The type of the data to be loaded.

    Returns:
        The loaded UUID.
    """
    with self.artifact_store.open(self.data_path, "r") as f:
        uuid_str = f.read().strip()
    return uuid.UUID(uuid_str)

save(data: uuid.UUID) -> None

Write UUID to artifact store.

Parameters:

Name	Type	Description	Default
data	UUID	The UUID to be saved.	required

Source code in src/zenml/materializers/uuid_materializer.py

def save(self, data: uuid.UUID) -> None:
    """Write UUID to artifact store.

    Args:
        data: The UUID to be saved.
    """
    with self.artifact_store.open(self.data_path, "w") as f:
        f.write(str(data))

Modules

base_materializer

Metaclass implementation for registering ZenML BaseMaterializer subclasses.

Classes

BaseMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Base Materializer to realize artifact data.

Initializes a materializer with the given URI.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data will be stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store used to store this artifact.	None

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Attributes

artifact_store: BaseArtifactStore property

Returns the artifact store used to store this artifact.

It either comes from the configuration of the materializer or from the active stack.

Returns:

Type	Description
BaseArtifactStore	The artifact store used to store this artifact.

Functions

can_load_type(data_type: Type[Any]) -> bool classmethod

Whether the materializer can load an artifact as the given type.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type to check.	required

Returns:

Type	Description
bool	Whether the materializer can load an artifact as the given type.

Source code in src/zenml/materializers/base_materializer.py

@classmethod
def can_load_type(cls, data_type: Type[Any]) -> bool:
    """Whether the materializer can load an artifact as the given type.

    Args:
        data_type: The type to check.

    Returns:
        Whether the materializer can load an artifact as the given type.
    """
    return any(
        issubclass(associated_type, data_type)
        # This next condition is not always correct, but better to have a
        # false positive here instead of failing for cases where it would
        # have worked.
        or issubclass(data_type, associated_type)
        for associated_type in cls.ASSOCIATED_TYPES
    )

can_save_type(data_type: Type[Any]) -> bool classmethod

Whether the materializer can save a certain type.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type to check.	required

Returns:

Type	Description
bool	Whether the materializer can save the given type.

Source code in src/zenml/materializers/base_materializer.py

@classmethod
def can_save_type(cls, data_type: Type[Any]) -> bool:
    """Whether the materializer can save a certain type.

    Args:
        data_type: The type to check.

    Returns:
        Whether the materializer can save the given type.
    """
    return any(
        issubclass(data_type, associated_type)
        for associated_type in cls.ASSOCIATED_TYPES
    )

extract_full_metadata(data: Any) -> Dict[str, MetadataType]

Extract both base and custom metadata from the given data.

Parameters:

Name	Type	Description	Default
data	Any	The data to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	A dictionary of metadata.

Source code in src/zenml/materializers/base_materializer.py

def extract_full_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
    """Extract both base and custom metadata from the given data.

    Args:
        data: The data to extract metadata from.

    Returns:
        A dictionary of metadata.
    """
    base_metadata = self._extract_base_metadata(data)
    custom_metadata = self.extract_metadata(data)
    return {**base_metadata, **custom_metadata}

extract_metadata(data: Any) -> Dict[str, MetadataType]

Extract metadata from the given data.

This metadata will be tracked and displayed alongside the artifact.

Example:

return {
    "some_attribute_i_want_to_track": self.some_attribute,
    "pi": 3.14,
}

Parameters:

Name	Type	Description	Default
data	Any	The data to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	A dictionary of metadata.

Source code in src/zenml/materializers/base_materializer.py

def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given data.

    This metadata will be tracked and displayed alongside the artifact.

    Example:
    ```
    return {
        "some_attribute_i_want_to_track": self.some_attribute,
        "pi": 3.14,
    }
    ```

    Args:
        data: The data to extract metadata from.

    Returns:
        A dictionary of metadata.
    """
    # Optionally, extract some metadata from `data` for ZenML to store.
    return {}

get_temporary_directory(delete_at_exit: bool, delete_after_step_execution: bool = True) -> Iterator[str]

Context manager to get a temporary directory.

Parameters:

Name	Type	Description	Default
delete_at_exit	bool	If set to True, the temporary directory will be deleted after the context manager exits.	required
delete_after_step_execution	bool	If delete_at_exit is set to False and this is set to True, the temporary directory will be deleted after the step finished executing. If a materializer is being used outside of the context of a step execution, the temporary directory will not be deleted and the user is responsible for deleting it themselves.	True

Yields:

Type	Description
str	Path to the temporary directory.

Source code in src/zenml/materializers/base_materializer.py

@contextlib.contextmanager
def get_temporary_directory(
    self,
    delete_at_exit: bool,
    delete_after_step_execution: bool = True,
) -> Iterator[str]:
    """Context manager to get a temporary directory.

    Args:
        delete_at_exit: If set to True, the temporary directory will be
            deleted after the context manager exits.
        delete_after_step_execution: If `delete_at_exit` is set to False and
            this is set to True, the temporary directory will be deleted
            after the step finished executing. If a materializer is being
            used outside of the context of a step execution, the temporary
            directory will not be deleted and the user is responsible for
            deleting it themselves.

    Yields:
        Path to the temporary directory.
    """
    temp_dir = tempfile.mkdtemp(prefix="zenml-")

    if delete_after_step_execution and not delete_at_exit:
        # We should not delete the directory when the context manager
        # exits, but cleanup once the step has finished executing.
        self._register_directory_for_deletion_after_step_execution(
            temp_dir
        )

    try:
        yield temp_dir
    finally:
        if delete_at_exit:
            shutil.rmtree(temp_dir)

load(data_type: Type[Any]) -> Any

Write logic here to load the data of an artifact.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	What type the artifact data should be loaded as.	required

Returns:

Type	Description
Any	The data of the artifact.

Source code in src/zenml/materializers/base_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Write logic here to load the data of an artifact.

    Args:
        data_type: What type the artifact data should be loaded as.

    Returns:
        The data of the artifact.
    """
    # read from a location inside self.uri
    return None

save(data: Any) -> None

Write logic here to save the data of an artifact.

Parameters:

Name	Type	Description	Default
data	Any	The data of the artifact to save.	required

Source code in src/zenml/materializers/base_materializer.py

def save(self, data: Any) -> None:
    """Write logic here to save the data of an artifact.

    Args:
        data: The data of the artifact to save.
    """

save_visualizations(data: Any) -> Dict[str, VisualizationType]

Save visualizations of the given data.

If this method is not overridden, no visualizations will be saved.

When overriding this method, make sure to save all visualizations to files within self.uri.

Example:

visualization_uri = os.path.join(self.uri, "visualization.html")
with self.artifact_store.open(visualization_uri, "w") as f:
    f.write("<html><body>data</body></html>")

visualization_uri_2 = os.path.join(self.uri, "visualization.png")
data.save_as_png(visualization_uri_2)

return {
    visualization_uri: ArtifactVisualizationType.HTML,
    visualization_uri_2: ArtifactVisualizationType.IMAGE
}

Parameters:

Name	Type	Description	Default
data	Any	The data of the artifact to visualize.	required

Returns:

Type	Description
Dict[str, VisualizationType]	A dictionary of visualization URIs and their types.

Source code in src/zenml/materializers/base_materializer.py

def save_visualizations(self, data: Any) -> Dict[str, VisualizationType]:
    """Save visualizations of the given data.

    If this method is not overridden, no visualizations will be saved.

    When overriding this method, make sure to save all visualizations to
    files within `self.uri`.

    Example:
    ```
    visualization_uri = os.path.join(self.uri, "visualization.html")
    with self.artifact_store.open(visualization_uri, "w") as f:
        f.write("<html><body>data</body></html>")

    visualization_uri_2 = os.path.join(self.uri, "visualization.png")
    data.save_as_png(visualization_uri_2)

    return {
        visualization_uri: ArtifactVisualizationType.HTML,
        visualization_uri_2: ArtifactVisualizationType.IMAGE
    }
    ```

    Args:
        data: The data of the artifact to visualize.

    Returns:
        A dictionary of visualization URIs and their types.
    """
    # Optionally, save some visualizations of `data` inside `self.uri`.
    return {}

validate_load_type_compatibility(data_type: Type[Any]) -> None

Checks whether the materializer can load the given type.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type to check.	required

Raises:

Type	Description
TypeError	If the materializer cannot load the given type.

Source code in src/zenml/materializers/base_materializer.py

def validate_load_type_compatibility(self, data_type: Type[Any]) -> None:
    """Checks whether the materializer can load the given type.

    Args:
        data_type: The type to check.

    Raises:
        TypeError: If the materializer cannot load the given type.
    """
    if not self.can_load_type(data_type):
        raise TypeError(
            f"Unable to load type {data_type}. {self.__class__.__name__} "
            f"can only load artifacts of the following types: "
            f"{self.ASSOCIATED_TYPES}."
        )

validate_save_type_compatibility(data_type: Type[Any]) -> None

Checks whether the materializer can save the given type.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type to check.	required

Raises:

Type	Description
TypeError	If the materializer cannot save the given type.

Source code in src/zenml/materializers/base_materializer.py

def validate_save_type_compatibility(self, data_type: Type[Any]) -> None:
    """Checks whether the materializer can save the given type.

    Args:
        data_type: The type to check.

    Raises:
        TypeError: If the materializer cannot save the given type.
    """
    if not self.can_save_type(data_type):
        raise TypeError(
            f"Unable to save type {data_type}. {self.__class__.__name__} "
            f"can only save artifacts of the following types: "
            f"{self.ASSOCIATED_TYPES}."
        )

BaseMaterializerMeta

Bases: type

Metaclass responsible for registering different BaseMaterializer subclasses.

Materializers are used for reading/writing artifacts.

Functions

__new__(mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]) -> BaseMaterializerMeta

Creates a Materializer class and registers it at the MaterializerRegistry.

Parameters:

Name	Type	Description	Default
name	str	The name of the class.	required
bases	Tuple[Type[Any], ...]	The base classes of the class.	required
dct	Dict[str, Any]	The dictionary of the class.	required

Returns:

Type	Description
BaseMaterializerMeta	The BaseMaterializerMeta class.

Raises:

Type	Description
MaterializerInterfaceError	If the class was improperly defined.

Source code in src/zenml/materializers/base_materializer.py

def __new__(
    mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
) -> "BaseMaterializerMeta":
    """Creates a Materializer class and registers it at the `MaterializerRegistry`.

    Args:
        name: The name of the class.
        bases: The base classes of the class.
        dct: The dictionary of the class.

    Returns:
        The BaseMaterializerMeta class.

    Raises:
        MaterializerInterfaceError: If the class was improperly defined.
    """
    cls = cast(
        Type["BaseMaterializer"], super().__new__(mcs, name, bases, dct)
    )
    if not cls._DOCS_BUILDING_MODE:
        # Skip the following validation and registration for base classes.
        if cls.SKIP_REGISTRATION:
            # Reset the flag so subclasses don't have it set automatically.
            cls.SKIP_REGISTRATION = False
            return cls

        # Validate that the class is properly defined.
        if not cls.ASSOCIATED_TYPES:
            raise MaterializerInterfaceError(
                f"Invalid materializer class '{name}'. When creating a "
                f"custom materializer, make sure to specify at least one "
                f"type in its ASSOCIATED_TYPES class variable.",
                url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
            )

        # Validate associated artifact type.
        if cls.ASSOCIATED_ARTIFACT_TYPE:
            try:
                cls.ASSOCIATED_ARTIFACT_TYPE = ArtifactType(
                    cls.ASSOCIATED_ARTIFACT_TYPE
                )
            except ValueError:
                raise MaterializerInterfaceError(
                    f"Invalid materializer class '{name}'. When creating a "
                    f"custom materializer, make sure to specify a valid "
                    f"artifact type in its ASSOCIATED_ARTIFACT_TYPE class "
                    f"variable.",
                    url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
                )

        # Validate associated data types.
        for associated_type in cls.ASSOCIATED_TYPES:
            if not inspect.isclass(associated_type):
                raise MaterializerInterfaceError(
                    f"Associated type {associated_type} for materializer "
                    f"{name} is not a class.",
                    url="https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types",
                )

        # Register the materializer.
        for associated_type in cls.ASSOCIATED_TYPES:
            materializer_registry.register_materializer_type(
                associated_type, cls
            )

        from zenml.utils import notebook_utils

        notebook_utils.try_to_save_notebook_cell_code(cls)

    return cls

Functions

Modules

built_in_materializer

Implementation of ZenML's builtin materializer.

Classes

BuiltInContainerMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle built-in container types (dict, list, set, tuple).

Define self.data_path and self.metadata_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/built_in_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path` and `self.metadata_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)
    self.metadata_path = os.path.join(self.uri, DEFAULT_METADATA_FILENAME)

Functions

extract_metadata(data: Any) -> Dict[str, MetadataType]

Extract metadata from the given built-in container object.

Parameters:

Name	Type	Description	Default
data	Any	The built-in container object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/built_in_materializer.py

def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given built-in container object.

    Args:
        data: The built-in container object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    if hasattr(data, "__len__"):
        return {"length": len(data)}
    return {}

load(data_type: Type[Any]) -> Any

Reads a materialized built-in container object.

If the data was serialized to JSON, deserialize it.

Otherwise, reconstruct all elements according to the metadata file: 1. Resolve the data type using find_type_by_str(), 2. Get the materializer via the default_materializer_registry, 3. Initialize the materializer with the desired path, 4. Use load() of that materializer to load the element.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Raises:

Type	Description
RuntimeError	If the data was not found.

Source code in src/zenml/materializers/built_in_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Reads a materialized built-in container object.

    If the data was serialized to JSON, deserialize it.

    Otherwise, reconstruct all elements according to the metadata file:
        1. Resolve the data type using `find_type_by_str()`,
        2. Get the materializer via the `default_materializer_registry`,
        3. Initialize the materializer with the desired path,
        4. Use `load()` of that materializer to load the element.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.

    Raises:
        RuntimeError: If the data was not found.
    """
    # If the data was not serialized, there must be metadata present.
    if not self.artifact_store.exists(
        self.data_path
    ) and not self.artifact_store.exists(self.metadata_path):
        raise RuntimeError(
            f"Materialization of type {data_type} failed. Expected either"
            f"{self.data_path} or {self.metadata_path} to exist."
        )

    # If the data was serialized as JSON, deserialize it.
    if self.artifact_store.exists(self.data_path):
        outputs = yaml_utils.read_json(self.data_path)

    # Otherwise, use the metadata to reconstruct the data as a list.
    else:
        metadata = yaml_utils.read_json(self.metadata_path)
        outputs = []

        # Backwards compatibility for zenml <= 0.37.0
        if isinstance(metadata, dict):
            for path_, type_str in zip(
                metadata["paths"], metadata["types"]
            ):
                type_ = find_type_by_str(type_str)
                materializer_class = materializer_registry[type_]
                materializer = materializer_class(uri=path_)
                element = materializer.load(type_)
                outputs.append(element)

        # New format for zenml > 0.37.0
        elif isinstance(metadata, list):
            for entry in metadata:
                path_ = entry["path"]
                type_ = source_utils.load(entry["type"])
                materializer_class = source_utils.load(
                    entry["materializer"]
                )
                materializer = materializer_class(uri=path_)
                element = materializer.load(type_)
                outputs.append(element)

        else:
            raise RuntimeError(f"Unknown metadata format: {metadata}.")

    # Cast the data to the correct type.
    if issubclass(data_type, dict) and not isinstance(outputs, dict):
        keys, values = outputs
        return data_type(zip(keys, values))
    if issubclass(data_type, tuple) and not isinstance(outputs, tuple):
        return data_type(outputs)
    if issubclass(data_type, set) and not isinstance(outputs, set):
        return data_type(outputs)
    return outputs

save(data: Any) -> None

Materialize a built-in container object.

If the object can be serialized to JSON, serialize it.

Otherwise, use the default_materializer_registry to find the correct materializer for each element and materialize each element into a subdirectory.

Tuples and sets are cast to list before materialization.

For non-serializable dicts, materialize keys/values as separate lists.

Parameters:

Name	Type	Description	Default
data	Any	The built-in container object to materialize.	required

Raises:

Type	Description
Exception	If any exception occurs, it is raised after cleanup.

Source code in src/zenml/materializers/built_in_materializer.py

def save(self, data: Any) -> None:
    """Materialize a built-in container object.

    If the object can be serialized to JSON, serialize it.

    Otherwise, use the `default_materializer_registry` to find the correct
    materializer for each element and materialize each element into a
    subdirectory.

    Tuples and sets are cast to list before materialization.

    For non-serializable dicts, materialize keys/values as separate lists.

    Args:
        data: The built-in container object to materialize.

    Raises:
        Exception: If any exception occurs, it is raised after cleanup.
    """
    # tuple and set: handle as list.
    if isinstance(data, tuple) or isinstance(data, set):
        data = list(data)

    # If the data is serializable, just write it into a single JSON file.
    if _is_serializable(data):
        yaml_utils.write_json(
            self.data_path,
            data,
            ensure_ascii=not ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS,
        )
        return

    # non-serializable dict: Handle as non-serializable list of lists.
    if isinstance(data, dict):
        data = [list(data.keys()), list(data.values())]

    # non-serializable list: Materialize each element into a subfolder.
    # Get path, type, and corresponding materializer for each element.
    metadata: List[Dict[str, str]] = []
    materializers: List[BaseMaterializer] = []
    try:
        for i, element in enumerate(data):
            element_path = os.path.join(self.uri, str(i))
            self.artifact_store.mkdir(element_path)
            type_ = type(element)
            materializer_class = materializer_registry[type_]
            materializer = materializer_class(uri=element_path)
            materializers.append(materializer)
            metadata.append(
                {
                    "path": element_path,
                    "type": source_utils.resolve(type_).import_path,
                    "materializer": source_utils.resolve(
                        materializer_class
                    ).import_path,
                }
            )
        # Write metadata as JSON.
        yaml_utils.write_json(self.metadata_path, metadata)
        # Materialize each element.
        for element, materializer in zip(data, materializers):
            materializer.validate_save_type_compatibility(type(element))
            materializer.save(element)
    # If an error occurs, delete all created files.
    except Exception as e:
        # Delete metadata
        if self.artifact_store.exists(self.metadata_path):
            self.artifact_store.remove(self.metadata_path)
        # Delete all elements that were already saved.
        for entry in metadata:
            self.artifact_store.rmtree(entry["path"])
        raise e

save_visualizations(data: Any) -> Dict[str, VisualizationType]

Save visualizations for the given data.

Parameters:

Name	Type	Description	Default
data	Any	The data to save visualizations for.	required

Returns:

Type	Description
Dict[str, VisualizationType]	A dictionary of visualization URIs and their types.

Source code in src/zenml/materializers/built_in_materializer.py

def save_visualizations(self, data: Any) -> Dict[str, "VisualizationType"]:
    """Save visualizations for the given data.

    Args:
        data: The data to save visualizations for.

    Returns:
        A dictionary of visualization URIs and their types.
    """
    # dict/list type objects are always saved as JSON files
    # doesn't work for non-serializable types as they
    # are saved as list of lists in different files
    if _is_serializable(data):
        return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
    return {}

BuiltInMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle JSON-serializable basic types (bool, float, int, str).

Define self.data_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/built_in_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)

Functions

extract_metadata(data: Union[bool, float, int, str]) -> Dict[str, MetadataType]

Extract metadata from the given built-in container object.

Parameters:

Name	Type	Description	Default
data	Union[bool, float, int, str]	The built-in container object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/built_in_materializer.py

def extract_metadata(
    self, data: Union[bool, float, int, str]
) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given built-in container object.

    Args:
        data: The built-in container object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    # For boolean and numbers, add the string representation as metadata.
    # We don't to this for strings because they can be arbitrarily long.
    if isinstance(data, (bool, float, int)):
        return {"string_representation": str(data)}

    return {}

load(data_type: Union[Type[bool], Type[float], Type[int], Type[str]]) -> Any

Reads basic primitive types from JSON.

Parameters:

Name	Type	Description	Default
data_type	Union[Type[bool], Type[float], Type[int], Type[str]]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Source code in src/zenml/materializers/built_in_materializer.py

def load(
    self, data_type: Union[Type[bool], Type[float], Type[int], Type[str]]
) -> Any:
    """Reads basic primitive types from JSON.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.
    """
    contents = yaml_utils.read_json(self.data_path)
    if type(contents) is not data_type:
        # TODO [ENG-142]: Raise error or try to coerce
        logger.debug(
            f"Contents {contents} was type {type(contents)} but expected "
            f"{data_type}"
        )
    return contents

save(data: Union[bool, float, int, str]) -> None

Serialize a basic type to JSON.

Parameters:

Name	Type	Description	Default
data	Union[bool, float, int, str]	The data to store.	required

Source code in src/zenml/materializers/built_in_materializer.py

def save(self, data: Union[bool, float, int, str]) -> None:
    """Serialize a basic type to JSON.

    Args:
        data: The data to store.
    """
    yaml_utils.write_json(
        self.data_path,
        data,
        ensure_ascii=not ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS,
    )

BytesMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle bytes data type, which is not JSON serializable.

Define self.data_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/built_in_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_BYTES_FILENAME)

Functions

load(data_type: Type[Any]) -> Any

Reads a bytes object from file.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Source code in src/zenml/materializers/built_in_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Reads a bytes object from file.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.
    """
    with self.artifact_store.open(self.data_path, "rb") as file_:
        return file_.read()

save(data: Any) -> None

Save a bytes object to file.

Parameters:

Name	Type	Description	Default
data	Any	The data to store.	required

Source code in src/zenml/materializers/built_in_materializer.py

def save(self, data: Any) -> None:
    """Save a bytes object to file.

    Args:
        data: The data to store.
    """
    with self.artifact_store.open(self.data_path, "wb") as file_:
        file_.write(data)

Functions

find_materializer_registry_type(type_: Type[Any]) -> Type[Any]

For a given type, find the type registered in the registry.

This can be either the type itself, or a superclass of the type.

Parameters:

Name	Type	Description	Default
type_	Type[Any]	The type to find.	required

Returns:

Type	Description
Type[Any]	The type registered in the registry.

Raises:

Type	Description
RuntimeError	If the type could not be resolved.

Source code in src/zenml/materializers/built_in_materializer.py

def find_materializer_registry_type(type_: Type[Any]) -> Type[Any]:
    """For a given type, find the type registered in the registry.

    This can be either the type itself, or a superclass of the type.

    Args:
        type_: The type to find.

    Returns:
        The type registered in the registry.

    Raises:
        RuntimeError: If the type could not be resolved.
    """
    # Check that a unique materializer is registered for this type
    materializer_registry[type_]

    # Check if the type itself is registered
    registered_types = materializer_registry.materializer_types.keys()
    if type_ in registered_types:
        return type_

    # Check if a superclass of the type is registered
    for registered_type in registered_types:
        if issubclass(type_, registered_type):
            return registered_type

    # Raise an error otherwise - this should never happen since
    # `default_materializer_registry[type_]` should have raised an error already
    raise RuntimeError(
        f"Cannot find a materializer for type '{type_}' in the "
        f"materializer registry."
    )

find_type_by_str(type_str: str) -> Type[Any]

Get a Python type, given its string representation.

E.g., "" should resolve to int.

Currently this is implemented by checking all artifact types registered in the default_materializer_registry. This means, only types in the registry can be found. Any other types will cause a RunTimeError.

Parameters:

Name	Type	Description	Default
type_str	str	The string representation of a type.	required

Raises:

Type	Description
RuntimeError	If the type could not be resolved.

Returns:

Type	Description
Type[Any]	The type whose string representation is type_str.

Source code in src/zenml/materializers/built_in_materializer.py

def find_type_by_str(type_str: str) -> Type[Any]:
    """Get a Python type, given its string representation.

    E.g., "<class 'int'>" should resolve to `int`.

    Currently this is implemented by checking all artifact types registered in
    the `default_materializer_registry`. This means, only types in the registry
    can be found. Any other types will cause a `RunTimeError`.

    Args:
        type_str: The string representation of a type.

    Raises:
        RuntimeError: If the type could not be resolved.

    Returns:
        The type whose string representation is `type_str`.
    """
    registered_types = materializer_registry.materializer_types.keys()
    type_str_mapping = {str(type_): type_ for type_ in registered_types}
    if type_str in type_str_mapping:
        return type_str_mapping[type_str]
    raise RuntimeError(f"Cannot resolve type '{type_str}'.")

Modules

cloudpickle_materializer

Implementation of ZenML's cloudpickle materializer.

Classes

CloudpickleMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer using cloudpickle.

This materializer can materialize (almost) any object, but does so in a non-reproducble way since artifacts cannot be loaded from other Python versions. It is recommended to use this materializer only as a last resort.

That is also why it has SKIP_REGISTRATION set to True and is currently only used as a fallback materializer inside the materializer registry.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

load(data_type: Type[Any]) -> Any

Reads an artifact from a cloudpickle file.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The data type of the artifact.	required

Returns:

Type	Description
Any	The loaded artifact data.

Source code in src/zenml/materializers/cloudpickle_materializer.py

def load(self, data_type: Type[Any]) -> Any:
    """Reads an artifact from a cloudpickle file.

    Args:
        data_type: The data type of the artifact.

    Returns:
        The loaded artifact data.
    """
    # validate python version
    source_python_version = self._load_python_version()
    current_python_version = Environment().python_version()
    if source_python_version != current_python_version:
        logger.warning(
            f"Your artifact was materialized under Python version "
            f"'{source_python_version}' but you are currently using "
            f"'{current_python_version}'. This might cause unexpected "
            "behavior since pickle is not reproducible across Python "
            "versions. Attempting to load anyway..."
        )

    # load data
    filepath = os.path.join(self.uri, DEFAULT_FILENAME)
    with self.artifact_store.open(filepath, "rb") as fid:
        data = cloudpickle.load(fid)
    return data

save(data: Any) -> None

Saves an artifact to a cloudpickle file.

Parameters:

Name	Type	Description	Default
data	Any	The data to save.	required

Source code in src/zenml/materializers/cloudpickle_materializer.py

def save(self, data: Any) -> None:
    """Saves an artifact to a cloudpickle file.

    Args:
        data: The data to save.
    """
    # Log a warning if this materializer was not explicitly specified for
    # the given data type.
    if type(self) is CloudpickleMaterializer:
        logger.warning(
            f"No materializer is registered for type `{type(data)}`, so "
            "the default Pickle materializer was used. Pickle is not "
            "production ready and should only be used for prototyping as "
            "the artifacts cannot be loaded when running with a different "
            "Python version. Please consider implementing a custom "
            f"materializer for type `{type(data)}` according to the "
            "instructions at https://docs.zenml.io/how-to/handle-data-artifacts/handle-custom-data-types"
        )

    # save python version for validation on loading
    self._save_python_version()

    # save data
    filepath = os.path.join(self.uri, DEFAULT_FILENAME)
    with self.artifact_store.open(filepath, "wb") as fid:
        cloudpickle.dump(data, fid)

Functions

materializer_registry

Implementation of a default materializer registry.

Classes

MaterializerRegistry()

Matches a Python type to a default materializer.

Initialize the materializer registry.

Source code in src/zenml/materializers/materializer_registry.py

def __init__(self) -> None:
    """Initialize the materializer registry."""
    self.default_materializer: Optional[Type["BaseMaterializer"]] = None
    self.materializer_types: Dict[Type[Any], Type["BaseMaterializer"]] = {}

Functions

get_default_materializer() -> Type[BaseMaterializer]

Get the default materializer that is used if no other is found.

Returns:

Type	Description
Type[BaseMaterializer]	The default materializer.

Source code in src/zenml/materializers/materializer_registry.py

def get_default_materializer(self) -> Type["BaseMaterializer"]:
    """Get the default materializer that is used if no other is found.

    Returns:
        The default materializer.
    """
    from zenml.materializers.cloudpickle_materializer import (
        CloudpickleMaterializer,
    )

    if self.default_materializer:
        return self.default_materializer

    return CloudpickleMaterializer

get_materializer_types() -> Dict[Type[Any], Type[BaseMaterializer]]

Get all registered materializer types.

Returns:

Type	Description
Dict[Type[Any], Type[BaseMaterializer]]	A dictionary of registered materializer types.

Source code in src/zenml/materializers/materializer_registry.py

def get_materializer_types(
    self,
) -> Dict[Type[Any], Type["BaseMaterializer"]]:
    """Get all registered materializer types.

    Returns:
        A dictionary of registered materializer types.
    """
    return self.materializer_types

is_registered(key: Type[Any]) -> bool

Returns if a materializer class is registered for the given type.

Parameters:

Name	Type	Description	Default
key	Type[Any]	Indicates the type of object.	required

Returns:

Type	Description
bool	True if a materializer is registered for the given type, False
bool	otherwise.

Source code in src/zenml/materializers/materializer_registry.py

def is_registered(self, key: Type[Any]) -> bool:
    """Returns if a materializer class is registered for the given type.

    Args:
        key: Indicates the type of object.

    Returns:
        True if a materializer is registered for the given type, False
        otherwise.
    """
    return any(issubclass(key, type_) for type_ in self.materializer_types)

register_and_overwrite_type(key: Type[Any], type_: Type[BaseMaterializer]) -> None

Registers a new materializer and also overwrites a default if set.

Parameters:

Name	Type	Description	Default
key	Type[Any]	Indicates the type of object.	required
type_	Type[BaseMaterializer]	A BaseMaterializer subclass.	required

Source code in src/zenml/materializers/materializer_registry.py

def register_and_overwrite_type(
    self, key: Type[Any], type_: Type["BaseMaterializer"]
) -> None:
    """Registers a new materializer and also overwrites a default if set.

    Args:
        key: Indicates the type of object.
        type_: A BaseMaterializer subclass.
    """
    self.materializer_types[key] = type_
    logger.debug(f"Registered materializer {type_} for {key}")

register_materializer_type(key: Type[Any], type_: Type[BaseMaterializer]) -> None

Registers a new materializer.

Parameters:

Name	Type	Description	Default
key	Type[Any]	Indicates the type of object.	required
type_	Type[BaseMaterializer]	A BaseMaterializer subclass.	required

Source code in src/zenml/materializers/materializer_registry.py

def register_materializer_type(
    self, key: Type[Any], type_: Type["BaseMaterializer"]
) -> None:
    """Registers a new materializer.

    Args:
        key: Indicates the type of object.
        type_: A BaseMaterializer subclass.
    """
    if key not in self.materializer_types:
        self.materializer_types[key] = type_
        logger.debug(f"Registered materializer {type_} for {key}")
    else:
        logger.debug(
            f"Found existing materializer class for {key}: "
            f"{self.materializer_types[key]}. Skipping registration of "
            f"{type_}."
        )

Functions

numpy_materializer

Placeholder for importing the Numpy Materializer from its former path.

Classes

pandas_materializer

Placeholder for importing the Pandas Materializer from its former path.

Classes

pydantic_materializer

Implementation of ZenML's pydantic materializer.

Classes

PydanticMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Handle Pydantic BaseModel objects.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

extract_metadata(data: BaseModel) -> Dict[str, MetadataType]

Extract metadata from the given BaseModel object.

Parameters:

Name	Type	Description	Default
data	BaseModel	The BaseModel object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/pydantic_materializer.py

def extract_metadata(self, data: BaseModel) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given BaseModel object.

    Args:
        data: The BaseModel object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    return {"schema": data.schema()}

load(data_type: Type[BaseModel]) -> Any

Reads BaseModel from JSON.

Parameters:

Name	Type	Description	Default
data_type	Type[BaseModel]	The type of the data to read.	required

Returns:

Type	Description
Any	The data read.

Source code in src/zenml/materializers/pydantic_materializer.py

def load(self, data_type: Type[BaseModel]) -> Any:
    """Reads BaseModel from JSON.

    Args:
        data_type: The type of the data to read.

    Returns:
        The data read.
    """
    data_path = os.path.join(self.uri, DEFAULT_FILENAME)
    contents = yaml_utils.read_json(data_path)
    return data_type.model_validate_json(contents)

save(data: BaseModel) -> None

Serialize a BaseModel to JSON.

Parameters:

Name	Type	Description	Default
data	BaseModel	The data to store.	required

Source code in src/zenml/materializers/pydantic_materializer.py

def save(self, data: BaseModel) -> None:
    """Serialize a BaseModel to JSON.

    Args:
        data: The data to store.
    """
    data_path = os.path.join(self.uri, DEFAULT_FILENAME)
    yaml_utils.write_json(data_path, data.model_dump_json())

Modules

service_materializer

Implementation of a materializer to read and write ZenML service instances.

Classes

ServiceMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer to read/write service instances.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

extract_metadata(service: BaseService) -> Dict[str, MetadataType]

Extract metadata from the given service.

Parameters:

Name	Type	Description	Default
service	BaseService	The service to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

Source code in src/zenml/materializers/service_materializer.py

def extract_metadata(
    self, service: BaseService
) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given service.

    Args:
        service: The service to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    from zenml.metadata.metadata_types import Uri

    if prediction_url := service.get_prediction_url() or None:
        return {"uri": Uri(prediction_url)}
    return {}

load(data_type: Type[Any]) -> BaseService

Creates and returns a service.

This service is instantiated from the serialized service configuration and last known status information saved as artifact.

Parameters:

Name	Type	Description	Default
data_type	Type[Any]	The type of the data to read.	required

Returns:

Type	Description
BaseService	A ZenML service instance.

Source code in src/zenml/materializers/service_materializer.py

def load(self, data_type: Type[Any]) -> BaseService:
    """Creates and returns a service.

    This service is instantiated from the serialized service configuration
    and last known status information saved as artifact.

    Args:
        data_type: The type of the data to read.

    Returns:
        A ZenML service instance.
    """
    filepath = os.path.join(self.uri, SERVICE_CONFIG_FILENAME)
    with self.artifact_store.open(filepath, "r") as f:
        service_id = f.read().strip()

    service = Client().get_service(name_id_or_prefix=uuid.UUID(service_id))
    return BaseDeploymentService.from_model(service)

save(service: BaseService) -> None

Writes a ZenML service.

The configuration and last known status of the input service instance are serialized and saved as an artifact.

Parameters:

Name	Type	Description	Default
service	BaseService	A ZenML service instance.	required

Source code in src/zenml/materializers/service_materializer.py

def save(self, service: BaseService) -> None:
    """Writes a ZenML service.

    The configuration and last known status of the input service instance
    are serialized and saved as an artifact.

    Args:
        service: A ZenML service instance.
    """
    filepath = os.path.join(self.uri, SERVICE_CONFIG_FILENAME)
    with self.artifact_store.open(filepath, "w") as f:
        f.write(str(service.uuid))

structured_string_materializer

Implementation of HTMLString materializer.

Classes

StructuredStringMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer for HTML or Markdown strings.

Source code in src/zenml/materializers/base_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Initializes a materializer with the given URI.

    Args:
        uri: The URI where the artifact data will be stored.
        artifact_store: The artifact store used to store this artifact.
    """
    self.uri = uri
    self._artifact_store = artifact_store

Functions

load(data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS

Loads the data from the HTML or Markdown file.

Parameters:

Name	Type	Description	Default
data_type	Type[STRUCTURED_STRINGS]	The type of the data to read.	required

Returns:

Type	Description
STRUCTURED_STRINGS	The loaded data.

Source code in src/zenml/materializers/structured_string_materializer.py

def load(self, data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS:
    """Loads the data from the HTML or Markdown file.

    Args:
        data_type: The type of the data to read.

    Returns:
        The loaded data.
    """
    with self.artifact_store.open(self._get_filepath(data_type), "r") as f:
        return data_type(f.read())

save(data: STRUCTURED_STRINGS) -> None

Save data as an HTML or Markdown file.

Parameters:

Name	Type	Description	Default
data	STRUCTURED_STRINGS	The data to save as an HTML or Markdown file.	required

Source code in src/zenml/materializers/structured_string_materializer.py

def save(self, data: STRUCTURED_STRINGS) -> None:
    """Save data as an HTML or Markdown file.

    Args:
        data: The data to save as an HTML or Markdown file.
    """
    with self.artifact_store.open(
        self._get_filepath(type(data)), "w"
    ) as f:
        f.write(data)

save_visualizations(data: STRUCTURED_STRINGS) -> Dict[str, VisualizationType]

Save visualizations for the given data.

Parameters:

Name	Type	Description	Default
data	STRUCTURED_STRINGS	The data to save visualizations for.	required

Returns:

Type	Description
Dict[str, VisualizationType]	A dictionary of visualization URIs and their types.

Source code in src/zenml/materializers/structured_string_materializer.py

def save_visualizations(
    self, data: STRUCTURED_STRINGS
) -> Dict[str, VisualizationType]:
    """Save visualizations for the given data.

    Args:
        data: The data to save visualizations for.

    Returns:
        A dictionary of visualization URIs and their types.
    """
    filepath = self._get_filepath(type(data))
    filepath = filepath.replace("\\", "/")
    visualization_type = self._get_visualization_type(type(data))
    return {filepath: visualization_type}

Functions

uuid_materializer

Implementation of ZenML's UUID materializer.

Classes

UUIDMaterializer(uri: str, artifact_store: Optional[BaseArtifactStore] = None)

Bases: BaseMaterializer

Materializer to handle UUID objects.

Define self.data_path.

Parameters:

Name	Type	Description	Default
uri	str	The URI where the artifact data is stored.	required
artifact_store	Optional[BaseArtifactStore]	The artifact store where the artifact data is stored.	None

Source code in src/zenml/materializers/uuid_materializer.py

def __init__(
    self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
):
    """Define `self.data_path`.

    Args:
        uri: The URI where the artifact data is stored.
        artifact_store: The artifact store where the artifact data is stored.
    """
    super().__init__(uri, artifact_store)
    self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)

Functions

extract_metadata(data: uuid.UUID) -> Dict[str, MetadataType]

Extract metadata from the UUID.

Parameters:

Name	Type	Description	Default
data	UUID	The UUID to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	A dictionary of metadata extracted from the UUID.

Source code in src/zenml/materializers/uuid_materializer.py

def extract_metadata(self, data: uuid.UUID) -> Dict[str, MetadataType]:
    """Extract metadata from the UUID.

    Args:
        data: The UUID to extract metadata from.

    Returns:
        A dictionary of metadata extracted from the UUID.
    """
    return {
        "string_representation": str(data),
    }

load(_: Type[uuid.UUID]) -> uuid.UUID

Read UUID from artifact store.

Parameters:

Name	Type	Description	Default
_	Type[UUID]	The type of the data to be loaded.	required

Returns:

Type	Description
UUID	The loaded UUID.

Source code in src/zenml/materializers/uuid_materializer.py

def load(self, _: Type[uuid.UUID]) -> uuid.UUID:
    """Read UUID from artifact store.

    Args:
        _: The type of the data to be loaded.

    Returns:
        The loaded UUID.
    """
    with self.artifact_store.open(self.data_path, "r") as f:
        uuid_str = f.read().strip()
    return uuid.UUID(uuid_str)

save(data: uuid.UUID) -> None

Write UUID to artifact store.

Parameters:

Name	Type	Description	Default
data	UUID	The UUID to be saved.	required

Source code in src/zenml/materializers/uuid_materializer.py

def save(self, data: uuid.UUID) -> None:
    """Write UUID to artifact store.

    Args:
        data: The UUID to be saved.
    """
    with self.artifact_store.open(self.data_path, "w") as f:
        f.write(str(data))