ZenML: ML Pipeline Orchestration & MLOps Framework

Defining Steps & Pipelines

Install with pip install zenml and zenml init. Define steps with type annotations — ZenML uses them for artifact serialization:

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from zenml import step, pipeline
from zenml.integrations.mlflow.experiment_trackers import MLFlowExperimentTracker

# Each @step is an isolated, cacheable unit with typed inputs/outputs
@step
def load_data() -> pd.DataFrame:
    """Load and return raw training data."""
    return pd.read_csv("data/train.csv")

@step
def preprocess(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series]:
    """Clean and split features/labels."""
    X = df.drop("target", axis=1).fillna(0)
    y = df["target"]
    return X, y

@step(experiment_tracker="mlflow_tracker")  # attach experiment tracker
def train_model(
    X: pd.DataFrame, y: pd.Series
) -> RandomForestClassifier:
    """Train and return a RandomForest model."""
    from zenml.integrations.mlflow.flavors.mlflow_experiment_tracker_flavor import MLFlowExperimentTrackerSettings
    import mlflow

    clf = RandomForestClassifier(n_estimators=100, max_depth=8)
    clf.fit(X, y)
    mlflow.log_metric("train_accuracy", clf.score(X, y))
    return clf

@step
def evaluate(model: RandomForestClassifier, X: pd.DataFrame, y: pd.Series) -> float:
    """Evaluate model and return accuracy."""
    acc = accuracy_score(y, model.predict(X))
    print(f"Accuracy: {acc:.4f}")
    return acc

# Pipeline wires steps together
@pipeline
def training_pipeline():
    df = load_data()
    X, y = preprocess(df)
    model = train_model(X, y)
    evaluate(model, X, y)

# Run locally
training_pipeline()

Stacks: Switch Infrastructure Without Code Changes

Register Stack components once via CLI, then activate the Stack you need:

# Register individual components
zenml artifact-store register s3_store \
  --flavor=s3 \
  --path=s3://my-ml-bucket/artifacts

zenml orchestrator register kubeflow_orch \
  --flavor=kubeflow \
  --kubeflow_pipelines_ui_port=8080

zenml experiment-tracker register mlflow_tracker \
  --flavor=mlflow \
  --tracking_uri=http://mlflow.internal:5000

zenml container-registry register ecr \
  --flavor=aws \
  --uri=123456789.dkr.ecr.us-east-1.amazonaws.com

# Build a production Stack from components
zenml stack register prod_stack \
  --artifact-store=s3_store \
  --orchestrator=kubeflow_orch \
  --experiment-tracker=mlflow_tracker \
  --container-registry=ecr

# Activate the Stack (same pipeline code, different infrastructure)
zenml stack set prod_stack
python run_pipeline.py    # now runs on Kubeflow, artifacts in S3

# Switch back to local for debugging
zenml stack set default
python run_pipeline.py    # runs locally

Artifact Versioning & Pipeline Lineage

Every step output is versioned — inspect and reuse artifacts programmatically:

from zenml.client import Client

client = Client()

# List all pipeline runs
runs = client.list_pipeline_runs(pipeline_name="training_pipeline")
for run in runs:
    print(f"Run: {run.name} | Status: {run.status} | Created: {run.created}")

# Get a specific artifact from a run
run = client.get_pipeline_run("training_pipeline-2024_04_01")
model_artifact = run.steps["train_model"].output

# Load the artifact (ZenML deserializes it using the appropriate Materializer)
model = model_artifact.load()
print(type(model))  # RandomForestClassifier

# Get artifact metadata and lineage
print(model_artifact.version)
print(model_artifact.producer_step_run_id)

# Compare artifacts across runs (for model promotion decisions)
for run in runs[:5]:
    acc_artifact = run.steps["evaluate"].output
    print(f"{run.name}: accuracy = {acc_artifact.load():.4f}")

ZenML vs Metaflow vs Kubeflow Pipelines vs Prefect