Return multiple outputs from a step

Use Annotated to return multiple outputs from a step and name them for easy retrieval and dashboard display.

You can use the Annotated type to return multiple outputs from a step and give each output a name. Naming your step outputs will help you retrieve the specific artifact later and also improves the readability of your pipeline's dashboard.

from typing import Annotated, Tuple

import pandas as pd

from zenml import step

@step

def clean_data(

data: pd.DataFrame,

) -> Tuple[

Annotated[pd.DataFrame, "x_train"],

Annotated[pd.DataFrame, "x_test"],

Annotated[pd.Series, "y_train"],

Annotated[pd.Series, "y_test"],

]:

from sklearn.model_selection import train_test_split

x = data.drop("target", axis=1)

y = data["target"]

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)

return x_train, x_test, y_train, y_test

Inside the step, we split the input data into features (x) and target (y), and then use train_test_split from scikit-learn to split the data into training and testing sets. The resulting DataFrames and Series are returned as a tuple, with each element annotated with its respective name.

By using Annotated, we can easily identify and retrieve specific artifacts later in the pipeline. Additionally, the names will be displayed on the pipeline's dashboard, making it more readable and understandable.

PreviousHow ZenML stores data

NextDelete an artifact

Last updated 15 days ago