Attach metadata to steps

You might want to log metadata and have that be attached to a specific step during the course of your work. This is possible by using the log_step_metadata method. This method allows you to attach a dictionary of key-value pairs as metadata to a step. The metadata can be any JSON-serializable value, including custom classes such as Uri, Path, DType, and StorageSize.

You can call this method from within a step or from outside. If you call it from within it will attach the metadata to the step and run that is currently being executed.

from zenml import step, log_step_metadata, ArtifactConfig, get_step_context

from typing import Annotated

import pandas as pd

from sklearn.ensemble import RandomForestClassifier

from sklearn.base import ClassifierMixin

@step

def train_model(dataset: pd.DataFrame) -> Annotated[ClassifierMixin, ArtifactConfig(name="sklearn_classifier", is_model_artifact=True)]:

"""Train a model"""

# Fit the model and compute metrics

classifier = RandomForestClassifier().fit(dataset)

accuracy, precision, recall = ...

# Log metadata at the step level

# This associates the metadata with the ZenML step run

log_step_metadata(

metadata={

"evaluation_metrics": {

"accuracy": accuracy,

"precision": precision,

"recall": recall

},

return classifier

If you call it from outside you can attach the metadata to a specific step run from any pipeline and step. This is useful if you want to attach the metadata after you've run the step.

from zenml import log_step_metadata

# run some step

# subsequently log the metadata for the step

log_step_metadata(

metadata={

"some_metadata": {"a_number": 3}

},

pipeline_name_id_or_prefix="my_pipeline",

step_name="my_step",

run_id="my_step_run_id"

Fetching logged metadata

Once metadata has been logged in an artifact, model, we can easily fetch the metadata with the ZenML Client:

from zenml.client import Client

client = Client()