Getting Started with ML Pipelines

Build and run your first ML pipeline on Michelangelo AI in minutes. This guide walks you through a complete example -- from defining tasks and workflows to running locally and deploying remotely.

What you'll build

By the end of this guide you'll have a working pipeline that looks like this:

dataset_cols
    └─▶ feature_prep()    # Ray task: downloads data, splits into train/validation
            └─▶ train()   # Ray task: trains a PyTorch Lightning model on the splits

Each step runs as an isolated, containerized task. Michelangelo AI handles data passing between them, caches intermediate results, and retries on transient failures — your code stays plain Python.

What you'll learn

How to define tasks with the @task decorator
How to compose tasks into a workflow with @workflow
How to run pipelines locally and remotely
How to register and manage pipelines with the ma CLI

Prerequisites

Python 3.9+
Poetry installed
Java 17 with JAVA_HOME set — required for the Spark preprocessing step. Java 21 is not compatible with PySpark 3.5 + Hadoop 3.3 (getSubject is not supported error). On macOS: brew install openjdk@17 then export JAVA_HOME=$(brew --prefix openjdk@17)/libexec/openjdk.jdk/Contents/Home
For remote runs: Docker and access to a Kubernetes cluster (or use the local sandbox)
Create a project

Environment setup

Clone the repository and install dependencies:

git clone https://github.com/michelangelo-ai/michelangelo.git
cd michelangelo/python
poetry install -E example

This creates a .venv directory with all dependencies installed. You can activate it directly or run commands via poetry run.

Core concepts

Before writing code, understand the three building blocks of an ML pipeline:

Concept	What It Does	Defined With
Task	A discrete unit of work (data prep, training, evaluation) that runs in a container	`@uniflow.task()`
Workflow	Orchestrates tasks in sequence, with branching and loops	`@uniflow.workflow()`
Pipeline	A deployable instance of a workflow with specific configuration	`pipeline.yaml`

Step 1: Define your tasks

Tasks are Python functions decorated with @uniflow.task(). Each task runs independently in its own container with configurable compute resources.

import michelangelo.uniflow.core as uniflow
from michelangelo.uniflow.plugins.ray import RayTask
from michelangelo.workflow.variables import DatasetVariable

@uniflow.task(
    config=RayTask(
        head_cpu=1,
        head_gpu=0,
        head_memory="4Gi",
        worker_cpu=1,
        worker_gpu=0,
        worker_memory="4Gi",
        worker_instances=0,
    ),
    cache_enabled=True,
)
def feature_prep(
    columns: list[str],
    test_size: float = 0.25,
    seed: int = 1,
) -> tuple[DatasetVariable, DatasetVariable]:
    """Prepare features from the California Housing dataset.

    Loads the California Housing dataset via scikit-learn, performs a
    train/test split, and converts to Ray Datasets for distributed processing.

    Args:
        columns: List of column names to select (features + ``"target"``).
        test_size: Fraction of data to use for validation. Defaults to 0.25.
        seed: Random seed for reproducibility. Defaults to 1.

    Returns:
        Tuple of (train_dataset, validation_dataset) as DatasetVariables.
    """
    import ray.data
    from sklearn.datasets import fetch_california_housing

    housing = fetch_california_housing(as_frame=True)
    df = housing.frame.rename(columns={"MedHouseVal": "target"})

    data = ray.data.from_pandas(df).select_columns(columns)

    train_data, validation_data = data.train_test_split(
        test_size=test_size, shuffle=True, seed=seed
    )

    train_dv = DatasetVariable.create(train_data)
    train_dv.save_ray_dataset()

    validation_dv = DatasetVariable.create(validation_data)
    validation_dv.save_ray_dataset()

    return train_dv, validation_dv

Choosing a task type

Michelangelo AI supports two compute backends for tasks:

Task Type	Best For	Example Config
RayTask	Distributed training, GPU workloads, general-purpose compute	`RayTask(head_cpu=2, worker_instances=4)`
SparkTask	Large-scale data preprocessing, ETL, SQL-based transformations	`SparkTask(driver_cpu=2, executor_instances=4)`

You can mix both types in a single workflow. Data types automatically convert between frameworks.

Step 2: Compose a workflow

A workflow orchestrates your tasks. It defines the execution order and passes data between tasks.

@uniflow.workflow()
def train_workflow(dataset_cols: str):
    """Orchestrate the full training pipeline."""
    columns = dataset_cols.split(",")

    # Step 1: Prepare features
    train_dv, validation_dv = feature_prep(columns=columns)

    # Step 2: Train the model
    result = train(
        train_dv=train_dv,
        validation_dv=validation_dv,
        params={
            "max_epochs": 5,
            "batch_size": 64,
            "precision": "32",
        },
    )
    return result

Important: Workflow code has restrictions because it compiles to Starlark for production execution. Inside a workflow function, you can only call task functions, other workflows, and built-in functions. See Workflow constraints for details.

Step 3: Run locally

Add a main block to create a context and run the workflow:

if __name__ == "__main__":
    ctx = uniflow.create_context()

    # Optionally set different parameters for local vs remote runs
    if ctx.is_local_run():
        ctx.environ["DATASET_SIZE"] = "100"  # small dataset for local testing
    else:
        ctx.environ["DATASET_SIZE"] = "1000000"  # full dataset for remote

    ctx.run(
        train_workflow,
        dataset_cols="MedInc,HouseAge,AveRooms,AveBedrms,Population,AveOccup,Latitude,Longitude,target",
    )

The context provides three methods:

Method	Purpose
`ctx.run(fn, **params)`	Execute the workflow function
`ctx.environ`	Dict of environment variables to set before execution
`ctx.is_local_run()`	Returns `True` if running locally (useful for conditional logic)

Then run it:

The feature_prep/train workflow above is exactly what the pytorch_train pipeline in michelangelo-examples implements — this tutorial's example now lives there, in its own pip-installable package, rather than in this repo. Install and run it:

pip install "michelangelo-examples[california-housing]"
python -m michelangelo_examples.california_housing.pipelines.pytorch_train.pipeline

Local runs execute everything in your Python interpreter with zero additional infrastructure setup (beyond the one-time pip install above) — the fastest way to iterate on your workflow logic.

Why a separate repo? Heavier example pipelines with their own dependency footprint (here, Spark + PyTorch Lightning) live in michelangelo-examples to keep this core repo's own dependencies lean. Lighter in-repo examples like bert_cola (python/examples/bert_cola/) are unrelated to this tutorial's pipeline but still run directly from this checkout:
cd michelangelo/python
PYTHONPATH=. poetry run python examples/bert_cola/bert_cola.py
pytorch_train also ships a separate, lighter-weight local-only trainer (python -m michelangelo_examples.california_housing.pipelines.pytorch_train, no .pipeline suffix) that trains directly with plain PyTorch Lightning — no Ray, Spark, or Uniflow context at all. That script is a standalone quick-start, not the ctx.run(train_workflow) pattern this tutorial teaches; use the .pipeline-suffixed command above to run the actual workflow you just wrote.

Step 4: Run remotely

When you need more compute power or want to validate against production infrastructure, switch to a remote run.

Build and push a Docker image

# From the michelangelo-examples repo checkout:
docker build -t michelangelo-examples:local .
k3d image import michelangelo-examples:local -c michelangelo-sandbox

Run with remote execution

python -m michelangelo_examples.california_housing.pipelines.pytorch_train.pipeline \
  remote-run \
  --image docker.io/library/michelangelo-examples:local \
  --storage-url s3://michelangelo/workflows \
  --environ AWS_ENDPOINT_URL=http://minio:9091 \
  --environ AWS_ACCESS_KEY_ID=michelangeloadmin \
  --environ AWS_SECRET_ACCESS_KEY=michelangeloadmin \
  --environ REGISTRY_ENDPOINT=michelangelo-apiserver:15566 \
  --yes

Sandbox storage URL: the michelangelo bucket is created automatically by ma sandbox create. The --environ flags inject credentials into the Cadence/Temporal workflow so they reach remote Ray/Spark workers. See the pytorch_train README for the full environment variable reference.

Remote runs execute workflow code in a Cadence/Temporal worker and task code in Kubernetes containers with full resource isolation. For detailed remote setup instructions including sandbox configuration, see Running Uniflow pipelines.

Step 5: Register as a pipeline

To manage your workflow through MA Studio and the ma CLI, register it as a pipeline.

Create a project

Pipelines belong to a project. Create one first using the example project config:

# From the michelangelo-examples repo:
ma project apply -f src/michelangelo_examples/california_housing/config/project.yaml

This creates the california-housing project with its namespace. The project config is specific to michelangelo-examples' California Housing pipelines and is separate from the core repo's python/examples/config/project.yaml (which registers the ma-examples project for examples that remain in this monorepo, such as workflow_patterns). For details on project configuration, see Project Management for ML Pipelines.

Create pipeline.yaml

Create a pipeline manifest that references your workflow. The michelangelo/uniflow-image annotation specifies the Docker image that runs your task code in Kubernetes:

apiVersion: michelangelo.api/v2
kind: Pipeline
metadata:
  namespace: "california-housing"
  name: "pytorch-train"
  annotations:
    michelangelo/uniflow-image: ghcr.io/michelangelo-ai/michelangelo-examples:california-housing
spec:
  type: "PIPELINE_TYPE_TRAIN"
  manifest:
    filePath: michelangelo_examples.california_housing.pipelines.pytorch_train.pipeline

The California Housing PyTorch Lightning example includes this manifest at src/michelangelo_examples/california_housing/pipelines/pytorch_train/pipeline.yaml in the michelangelo-examples repo.

Register the pipeline

ma pipeline apply -f src/michelangelo_examples/california_housing/pipelines/pytorch_train/pipeline.yaml

Run the registered pipeline

ma pipeline run -n california-housing --name pytorch-train

Workflow constraints

Workflow functions compile to Starlark for production execution. This means some Python features are not available inside @workflow functions:

Not Supported	Use Instead
`import` statements	Call task functions (imports go inside tasks)
`try-except` blocks	Use task-level retry
`is` / `is not`	`==` / `!=`
f-strings	`"{}".format(value)`
Chained comparisons (`1 < x < 5`)	`1 < x and x < 5`
Standard library calls	Built-in functions (e.g., `uniflow.time()`)

These constraints only apply to workflow code. Task code runs in containers and can use any Python code.

Task features

Caching

Cache task results to skip re-execution on subsequent runs:

@uniflow.task(
    config=RayTask(head_cpu=1, head_memory="4Gi"),
    cache_enabled=True,
    cache_version="v1",
)
def feature_prep(columns: list[str]):
    ...

Cached results are available for approximately 28 days (platform-managed). Change cache_version to force re-execution.

Retry

Automatically retry failed tasks:

@uniflow.task(
    config=RayTask(head_cpu=1, head_memory="4Gi"),
    retry_attempts=3,
)
def train(params: dict):
    ...

Each retry creates a fresh cluster for better isolation.

Task overrides

Create task variants with different resource configurations using with_overrides():

@uniflow.workflow()
def train_workflow(dataset_cols: str):
    columns = dataset_cols.split(",")

    # Run feature_prep with more resources
    large_feature_prep = feature_prep.with_overrides(
        alias="large_feature_prep",
        config=RayTask(head_cpu=2, worker_instances=2),
    )
    train_dv, validation_dv = large_feature_prep(columns=columns)
    ...

Important: The override config replaces the original config entirely -- it does not merge fields. Specify all required fields in the override.

Complete example

See the full California Housing PyTorch Lightning example at pytorch_train/ in the michelangelo-examples repo. This example demonstrates:

Heterogeneous workflow: Ray tasks for data prep and training, Spark task for preprocessing
Task caching: Reuse feature preparation results across runs
Task overrides: Customize resource allocation per workflow
DatasetVariable: Pass datasets between tasks across different compute frameworks

Next steps

Pipeline Running Modes -- Understand when to use Local, Remote, Dev, and Pipeline runs
Pipeline Management -- Learn about standard vs custom workflows
Caching and Resume -- Resume failed pipeline runs from a specific step
Data Preparation -- Deep dive into data preprocessing patterns
Model Training -- Advanced distributed training with Lightning Trainer SDK

Troubleshooting

Out of memory during training? Increase head_memory or worker_memory in your task config, or reduce your dataset size for local runs.
Remote run fails to start? Verify your Docker image exists and is accessible. Check that --storage-url points to a valid S3-compatible bucket.
Workflow code errors with "not supported in Starlark"? Move the unsupported syntax (imports, try-except, f-strings) into a task function. See Workflow constraints.
Spark fails with getSubject is not supported? Java 21 is incompatible with PySpark 3.5 + Hadoop 3.3. Switch to Java 17: brew install openjdk@17 then export JAVA_HOME=$(brew --prefix openjdk@17)/libexec/openjdk.jdk/Contents/Home.

What you'll build​

What you'll learn​

Prerequisites​

Environment setup​

Core concepts​

Step 1: Define your tasks​

Choosing a task type​

Step 2: Compose a workflow​

Step 3: Run locally​

Step 4: Run remotely​

Build and push a Docker image​

Run with remote execution​

Step 5: Register as a pipeline​

Create a project​

Create pipeline.yaml​

Register the pipeline​

Run the registered pipeline​

Workflow constraints​

Task features​

Caching​

Retry​

Task overrides​

Complete example​

Next steps​

Troubleshooting​