Materializers

A ZenML pipeline is built in a data-centric way. The outputs and inputs of steps define how steps are connected and the order in which they are executed. Each step should be considered as its very own process that reads and writes its inputs and outputs from and to the artifact store. This is where materializers come into play.
Materializers: Serializing and deserializing your artifacts
A materializer dictates how a given artifact can be written to and retrieved from the artifact store. It contains all serialization and deserialization logic.
What is a materializer?
1
from typing import Type, Any
2
from zenml.materializers.base_materializer import BaseMaterializerMeta
3
​
4
​
5
class BaseMaterializer(metaclass=BaseMaterializerMeta):
6
    """Base Materializer to realize artifact data."""
7
​
8
    ASSOCIATED_ARTIFACT_TYPES = ()
9
    ASSOCIATED_TYPES = ()
10
​
11
    def __init__(self, artifact: "BaseArtifact"):
12
        """Initializes a materializer with the given artifact."""
13
        self.artifact = artifact
14
​
15
    def handle_input(self, data_type: Type[Any]) -> Any:
16
        """Write logic here to handle input of the step function.
17
​
18
        Args:
19
            data_type: What type the input should be materialized as.
20
        Returns:
21
            Any object that is to be passed into the relevant artifact in the
22
            step.
23
        """
24
        # read from self.artifact.uri
25
        ...
26
​
27
    def handle_return(self, data: Any) -> None:
28
        """Write logic here to handle return of the step function.
29
​
30
        Args:
31
            data: Any object that is specified as an input artifact of the step.
32
        """
33
        # write `data` to self.artifact.uri
34
        ...
Copied!
Above you can see the basic definition of the BaseMaterializer. All other materializers inherit from this class, and this class defines the interface of all materializers.
Each materializer has an artifact object. The most important property of an artifact object is the uri. The uri is created by ZenML at pipeline run time and points to the directory of a file system (the artifact store).
The handle_input and handle_return functions are important.
handle_input is responsible for reading the artifact from the artifact store.
handle_return is responsible for writing the artifact to the artifact store.
Each materializer has ASSOCIATED_TYPES and ASSOCIATED_ARTIFACT_TYPES.
ASSOCIATED_TYPES is the data type that is being stored. ZenML uses this information to call the right materializer at the right time. i.e. If a ZenML step returns a pd.DataFrame, ZenML will try to find any materializer that has pd.DataFrame (or its subclasses) in its ASSOCIATED_TYPES.
ASSOCIATED_ARTIFACT_TYPES simply define what type of artifacts are being stored. This can be DataArtifact, StatisticsArtifact, DriftArtifact, etc. This is simply a tag to query certain artifact types in the post-execution workflow.
Writing a custom materializer
Let's say you have a custom class called MyObject that flows between two steps in a pipeline:
1
import logging
2
from zenml.steps import step
3
from zenml.pipelines import pipeline
4
​
5
​
6
class MyObj:
7
    def __init__(self, name: str):
8
        self.name = name
9
​
10
​
11
@step
12
def my_first_step() -> MyObj:
13
    """Step that returns an object of type MyObj"""
14
    return MyObj("my_object")
15
​
16
​
17
@step
18
def my_second_step(my_obj: MyObj) -> None:
19
    """Step that logs the input object and returns nothing."""
20
    logging.info(
21
        f"The following object was passed to this step: `{my_obj.name}`")
22
​
23
​
24
@pipeline
25
def first_pipeline(
26
        step_1,
27
        step_2
28
):
29
    output_1 = step_1()
30
    step_2(output_1)
31
​
32
​
33
first_pipeline(step_1=my_first_step(), step_2=my_second_step()).run()
Copied!
Running the above without a custom materializer will result in the following error:
1
zenml.exceptions.StepInterfaceError: Unable to find materializer for output 'output' of 
2
type `<class '__main__.MyObj'>` in step 'step1'. Please make sure to either explicitly set a materializer for step 
3
outputs using `step.with_return_materializers(...)` or registering a default materializer for specific types by 
4
subclassing `BaseMaterializer` and setting its `ASSOCIATED_TYPES` class variable. 
5
For more information, visit https://docs.zenml.io/guides/common-usecases/custom-materializer.
Copied!
The above basically means that ZenML does not know how to persist the object of type MyObj between steps (how could it? We just created this!). Therefore, we have to create our own materializer. To do this you can simply extend the BaseMaterializer by sub-classing it.
1
import os
2
from typing import Type
3
​
4
from zenml.artifacts import DataArtifact
5
from zenml.io import fileio
6
from zenml.materializers.base_materializer import BaseMaterializer
7
​
8
​
9
class MyMaterializer(BaseMaterializer):
10
    ASSOCIATED_TYPES = (MyObj,)
11
    ASSOCIATED_ARTIFACT_TYPES = (DataArtifact,)
12
​
13
    def handle_input(self, data_type: Type[MyObj]) -> MyObj:
14
        """Read from artifact store"""
15
        super().handle_input(data_type)
16
        with fileio.open(os.path.join(self.artifact.uri, 'data.txt'), 'r') as f:
17
            name = f.read()
18
        return MyObj(name=name)
19
​
20
    def handle_return(self, my_obj: MyObj) -> None:
21
        """Write to artifact store"""
22
        super().handle_return(my_obj)
23
        with fileio.open(os.path.join(self.artifact.uri, 'data.txt'), 'w') as f:
24
            f.write(my_obj.name)
Copied!
Pro-tip: Use the ZenML fileio module to ensure your materialization logic works across artifact stores (local and remote like S3 buckets).
Now ZenML can use this materializer to handle outputs and inputs of your customs object. Edit the pipeline as follows to see this in action:
1
first_pipeline(
2
    step_1=my_first_step().with_return_materializers(MyMaterializer),
3
    step_2=my_second_step()).run()
Copied!
Due to the typing of the in- and outputs and the ASSOCIATED_TYPES attribute of the materializer you won't necessarily have to add `.with_return_materializers(MyMaterializer)` to the step. It should automatically be detected. It doesn't hurt to be explicit though.
For multiple outputs a dictionary can be supplied of type {OUTPUT_NAME: MATERIALIZER_CLASS} to the with_return_materializers function.
Also, notice that with_return_materializers is only called on step1, all downstream steps will use the same materializer by default.
This will yield the proper response as follows:
1
Creating run for pipeline: `first_pipeline`
2
Cache enabled for pipeline `first_pipeline`
3
Using stack `default` to run pipeline `first_pipeline`...
4
Step `my_first_step` has started.
5
Step `my_first_step` has finished in 0.081s.
6
Step `my_second_step` has started.
7
The following object was passed to this step: `my_object`
8
Step `my_second_step` has finished in 0.048s.
9
Pipeline run `first_pipeline-22_Apr_22-10_58_51_135729` has finished in 0.153s.
Copied!
Summary in Code
Code Example of this Section
Skip materialization
While in most cases, materializers should be used to control how artifacts are consumed and output from steps in a pipeline, you will sometimes need to have a completely non-materialized artifact in a step.
A non-materialized artifact is a BaseArtifact (or any of its subclasses) and has a property uri that points to the unique path in the artifact store where the artifact is stored. One can use a non-materialized artifact by simply specifying it as the type in the step:
1
from zenml.artifacts import DataArtifact
2
from zenml.steps import step
3
​
4
​
5
@step
6
def my_step(my_artifact: DataArtifact)  # rather than pd.DataFrame
7
    pass
Copied!
The list of raw artifact types can be found in zenml.artifacts.* and include ModelArtifact, DataArtifact etc. Materializers link pythonic types to these artifact types implicitly, e.g., a keras.model or torch.nn.Module are pythonic types that are both linked to ModelArtifact implicitly via their materializers. When using artifacts directly, one must be aware of which type they are by looking at the previous step's materializer: if the previous step produces a ModelArtifact then you should specify ModelArtifact in a non-materialized step.
Be careful: Using artifacts directly like this might have unintended consequences for downstream tasks that rely on materialized artifacts.
A simple example
A simple example can suffice to showcase how to use non-materialized artifacts:
1
from typing import Dict, List
2
from zenml.artifacts import DataArtifact, ModelArtifact
3
from zenml.pipelines import pipeline
4
from zenml.steps import Output, step
5
​
6
​
7
@step
8
def step_1() -> Output(dict_=Dict, list_=List):
9
    return {"some": "data"}, []
10
​
11
​
12
@step
13
def step_2() -> Output(dict_=Dict, list_=List):
14
    return {"some": "data"}, []
15
​
16
​
17
@step
18
def step_3(dict_: Dict, list_: List) -> None:
19
    assert type(dict_) is dict
20
    assert type(list_) is list
21
​
22
​
23
@step
24
def step_4(dict_: DataArtifact, list_: ModelArtifact) -> None:
25
    assert hasattr(dict_, "uri")
26
    assert hasattr(list_, "uri")
27
​
28
​
29
@pipeline
30
def p(s1, s2, s3, s4):
31
    s3(*s1())
32
    s4(*s2())
33
​
34
​
35
p(step_1(), step_2(), step_3(), step_4()).run()
Copied!
In the above the pipeline looks as follows:
1
s1 -> s3 
2
s2 -> s4
Copied!
s1 and s2 produce identical artifacts, however s3 consumes materialized artifacts while s4 consumes non-materialized artifacts. s4 can now use the dict_.uri and list_.uri paths directly rather than their materialized counterparts.

Developer Guide - Previous

Post-Execution Workflow

Next - Developer Guide

Caching

Last modified 10d ago

Export as PDF

Copy link

Edit on GitHub

Contents

Materializers: Serializing and deserializing your artifacts

What is a materializer?

Writing a custom materializer

Skip materialization