from zenml import step

@step
def train_model(X_train: pd.DataFrame, lr: float = 0.01) -> sklearn.base.BaseEstimator:
    """lr is a parameter (literal value); X_train is an artifact (from upstream step)."""
    model = LogisticRegression(C=1/lr).fit(X_train.drop("target", axis=1), X_train["target"])
    return model

from typing import Annotated, Tuple
from zenml import step
import pandas as pd

@step
def split_data(df: pd.DataFrame, ratio: float = 0.8) -> Tuple[
    Annotated[pd.DataFrame, "train"],
    Annotated[pd.DataFrame, "test"],
]:
    idx = int(len(df) * ratio)
    return df.iloc[:idx], df.iloc[idx:]

from zenml import pipeline

@pipeline
def training_pipeline(dataset_path: str = "data.csv", lr: float = 0.01) -> None:
    df = load_data(path=dataset_path)
    train, test = split_data(df=df)
    model = train_model(X_train=train, lr=lr)
    evaluate(model=model, X_test=test)

if __name__ == "__main__":
    training_pipeline()

my_pipeline_project/
├── steps/                    # One file per step
│   ├── load_data.py
│   ├── preprocess.py
│   ├── train_model.py
│   └── evaluate.py
├── pipelines/
│   └── training.py           # Pipeline definition(s)
├── materializers/            # Custom materializers (if any)
│   └── my_data_materializer.py
├── visualizations/           # HTML/CSS templates for dashboard visualizations
│   └── metrics_report.html
├── configs/                  # One YAML config per environment
│   ├── dev.yaml
│   ├── staging.yaml
│   └── prod.yaml
├── run.py                    # CLI entry point (argparse, not click)
├── README.md                 # How to run, what stacks to use, etc.
└── pyproject.toml            # Dependencies — always pyproject.toml, not requirements.txt

[project]
name = "my-pipeline-project"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
    "zenml>=0.93",
    "pandas>=2.0",
    # Add pipeline-specific dependencies here
]

[project.optional-dependencies]
dev = [
    "pytest",
    "ruff",
    "mypy",
]

import argparse
from pipelines.training import training_pipeline

def main():
    parser = argparse.ArgumentParser(description="Run the training pipeline")
    parser.add_argument("--config", default="configs/dev.yaml", help="Path to YAML config")
    parser.add_argument("--no-cache", action="store_true", help="Disable caching")
    parser.add_argument("--sample-size", type=int, default=None,
                        help="Use only N rows (for quick local iteration)")
    args = parser.parse_args()

    pipeline_instance = training_pipeline.with_options(
        config_path=args.config,
        enable_cache=not args.no_cache,
    )
    pipeline_instance(sample_size=args.sample_size)

if __name__ == "__main__":
    main()

# ANTI-PATTERN: works locally, fails on cloud
@step
def preprocess(input_path: str) -> str:
    df = pd.read_csv(input_path)           # Reads from local disk
    output_path = "/tmp/processed.csv"
    df.to_csv(output_path)
    return output_path                      # Next step can't access /tmp on a different pod

@step
def train(data_path: str) -> None:
    df = pd.read_csv(data_path)             # FileNotFoundError on cloud!

# CORRECT: data flows as artifacts
@step
def preprocess(input_path: str) -> pd.DataFrame:
    return pd.read_csv(input_path)          # ZenML serializes the DataFrame to the artifact store

@step
def train(data: pd.DataFrame) -> None:
    ...                                     # ZenML loads it from the artifact store — works everywhere

from zenml import pipeline, step

@step
def get_count() -> int:
    return 3

@step
def process(index: int) -> None:
    print(f"Processing {index}")

@pipeline(dynamic=True)
def my_dynamic_pipeline() -> None:
    count = get_count()
    count_data = count.load()       # .load() gets actual Python value
    for idx in range(count_data):
        process(index=idx)

items = produce_list()
for i, val in enumerate(items.load()):   # load to iterate
    if val > threshold:
        chunk = items.chunk(index=i)      # chunk to wire
        process(chunk)

from zenml import ExternalArtifact, pipeline, step
import pandas as pd

@step
def train(data: pd.DataFrame) -> None:
    ...

@pipeline
def my_pipeline() -> None:
    df = pd.read_csv("local_data.csv")
    train(data=ExternalArtifact(value=df))

# run.py (client-side, runs on your machine)
from zenml.artifacts.utils import save_artifact
import pandas as pd

df = pd.read_csv("local_data.csv")
art = save_artifact(data=df, name="my_dataset")
print(art.id)  # Pass this UUID to the pipeline

# pipeline.py (runs inside the orchestrator pod)
from zenml.client import Client
from uuid import UUID

@pipeline
def my_pipeline(dataset_id: str) -> None:
    artifact = Client().get_artifact_version(UUID(dataset_id))
    train(data=artifact)

# configs/dev.yaml
enable_cache: false
parameters:
  dataset_path: "data/small.csv"
  lr: 0.05
steps:
  train_model:
    settings:
      resources:
        cpu_count: 2

training_pipeline.with_options(config_path="configs/dev.yaml")()

from zenml.config import DockerSettings

# Install pip packages
docker = DockerSettings(
    requirements=["scikit-learn>=1.0", "pandas>=2.0"],
    apt_packages=["libgomp1"],
    environment={"PYTHONUNBUFFERED": "1"},
)

@pipeline(settings={"docker": docker})
def my_pipeline() -> None:
    ...

from zenml import step
from zenml.utils.metadata_utils import log_metadata

@step
def train_model(X_train: pd.DataFrame, y_train: pd.Series) -> sklearn.base.BaseEstimator:
    model = LogisticRegression().fit(X_train, y_train)
    log_metadata({"accuracy": model.score(X_train, y_train), "n_samples": len(X_train)})
    return model

# One-time setup (CLI)
zenml secret create db_credentials --host=db.example.com --username=admin --password=secret123

from zenml import step
from zenml.client import Client

@step
def load_from_database(query: str) -> pd.DataFrame:
    secret = Client().get_secret("db_credentials")
    host = secret.secret_values["host"]
    username = secret.secret_values["username"]
    password = secret.secret_values["password"]
    # Use credentials to connect...

from zenml.types import HTMLString, MarkdownString, CSVString

@step
def report(metrics: dict) -> HTMLString:
    # For polished HTML, load from a separate template file rather than inline strings
    html_path = Path(__file__).parent.parent / "visualizations" / "metrics_report.html"
    template = html_path.read_text()
    return HTMLString(template.format(**metrics))

import json, os
from typing import Any, Type
from zenml.enums import ArtifactType
from zenml.materializers.base_materializer import BaseMaterializer

class MyDataMaterializer(BaseMaterializer):
    ASSOCIATED_TYPES = (MyData,)
    ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA

    def save(self, data: MyData) -> None:
        path = os.path.join(self.uri, "data.json")
        with self.artifact_store.open(path, "w") as f:   # Always use self.artifact_store.open()
            json.dump(data.to_dict(), f)

    def load(self, data_type: Type[Any]) -> MyData:
        path = os.path.join(self.uri, "data.json")
        with self.artifact_store.open(path, "r") as f:
            return MyData(**json.load(f))

import logging
from zenml import step

logger = logging.getLogger(__name__)

@step
def train_model(data: pd.DataFrame) -> Model:
    logger.info(f"Training on {len(data)} samples")  # Captured by ZenML
    print("Starting training...")                      # Also captured
    ...