Orchestration Options: Bare Pipeline vs Airflow vs Prefect

The Pipeline class in this toolkit is an in-process orchestrator: it runs

your steps in one Python process, in dependency order, right now. That is exactly

what you want for development, CI, batch jobs that fit in memory, and embedding a

pipeline inside a larger service. It is *not* a scheduler, a distributed executor,

or a UI.

This guide compares three ways to run the same Step graph and shows how to move

between them with minimal rewrites.

Table of contents

1. Decision tree

2. Comparison at a glance

3. Option A: the bare Pipeline

4. Option B: Apache Airflow

5. Option C: Prefect

6. Keeping step logic portable

7. Migration checklist

Decision tree

Does the whole job fit in one process and finish in a single run?
|
+-- YES: Do you need scheduling, retries across machines, or a UI?
|        |
|        +-- NO  -> Option A: bare Pipeline (this toolkit). Stop here.
|        +-- YES -> Is your team already on a scheduler?
|                   |
|                   +-- Airflow shop  -> Option B
|                   +-- Greenfield / Pythonic / dynamic -> Option C (Prefect)
|
+-- NO (data too big / steps need different machines / long-running):
         Use a distributed engine (Spark, Ray, Dask) for the heavy step and
         orchestrate the stages with Airflow or Prefect. The bare Pipeline can
         still run *inside* a single task.

The honest default for most tabular problems is Option A. Reach for B or C

only when you actually need scheduling, cross-process retries, backfills, or a

control-plane UI -- not before.

Comparison at a glance

Option A: the bare Pipeline

No extra infrastructure -- import, build, run.

from ml_pipeline import build_standard_pipeline, load_config

pipeline = build_standard_pipeline(load_config("configs/pipeline.yaml"))
run = pipeline.run()                 # executes ingest -> preprocess -> train -> evaluate
if not run.succeeded:
    raise SystemExit(run.report())
print(run["evaluate"].to_markdown())

Run it on a schedule with anything that can launch a process -- cron, a CI job, a

systemd timer, a Kubernetes CronJob:

# Retrain every night at 02:00 and write the evaluation report.
0 2 * * *  cd /opt/ml && /opt/ml/.venv/bin/python examples/run_pipeline.py --config configs/pipeline.yaml

Use when: the job fits in memory, runs start-to-finish, and you do not need a

scheduler/UI. Outgrow it when: you need cross-machine retries, backfills, SLAs,

or a control plane.

Option B: Apache Airflow

Airflow is a mature, battle-tested scheduler. Map each Step to an Airflow task

so Airflow owns scheduling, retries, and observability while your step logic stays

unchanged.

# dags/ml_pipeline_dag.py
from datetime import datetime
from airflow import DAG
from airflow.operators.python import PythonOperator

from ml_pipeline import load_config
from ml_pipeline.steps import (
    IngestionStep, PreprocessingStep, TrainingStep, EvaluationStep,
)

CONFIG = load_config("/opt/airflow/configs/pipeline.yaml")

def run_ingest(**ctx):
    ds = IngestionStep("ingest", **CONFIG["ingestion"]).run({})
    ctx["ti"].xcom_push(key="dataset", value=ds)   # small data only; see note

with DAG(
    "ml_pipeline",
    start_date=datetime(2024, 1, 1),
    schedule="0 2 * * *",
    catchup=False,
    default_args={"retries": 2},
) as dag:
    ingest = PythonOperator(task_id="ingest", python_callable=run_ingest)
    # ... define preprocess / train / evaluate similarly ...
    ingest >> preprocess >> train >> evaluate

XCom note: Airflow passes data between tasks via XCom, which is sized for

small metadata, not DataFrames or model objects. In practice each task should

persist its real output to shared storage (S3/GCS/a database) and push only the

*path* through XCom. The Dataset / PreprocessedData / TrainedModel objects

make natural artifacts to serialize at task boundaries.

Strengths: robust scheduler, backfills, mature UI, huge operator ecosystem.

Costs: you run a scheduler, a metadata DB, and workers; dynamic graphs are

more awkward than in plain Python.

Option C: Prefect

Prefect feels like decorating normal Python, which pairs naturally with this

toolkit's plain-callable steps.

# flows/ml_pipeline_flow.py
from prefect import flow, task

from ml_pipeline import load_config
from ml_pipeline.steps import (
    IngestionStep, PreprocessingStep, TrainingStep, EvaluationStep,
)

CONFIG = load_config("configs/pipeline.yaml")

@task(retries=2, retry_delay_seconds=30)
def ingest():
    return IngestionStep("ingest", **CONFIG["ingestion"]).run({})

@task
def preprocess(dataset):
    return PreprocessingStep("preprocess", dataset_step="ingest",
                             **CONFIG["preprocessing"]).run({"ingest": dataset})

@task
def train(pre):
    return TrainingStep("train", features_step="preprocess").run({"preprocess": pre})

@task
def evaluate(model, pre):
    return EvaluationStep("evaluate", model_step="train",
                          features_step="preprocess").run({"train": model, "preprocess": pre})

@flow(name="ml-pipeline")
def ml_pipeline():
    ds = ingest()
    pre = preprocess(ds)
    model = train(pre)
    return evaluate(model, pre)

if __name__ == "__main__":
    ml_pipeline()

Because the flow is ordinary Python, Prefect infers the dependency graph from how

the task results are passed -- and the data is passed in memory between tasks, so

there is no XCom-style size limit when running in a single environment.

Strengths: Pythonic, first-class dynamic graphs, schedules and event triggers,

gentle local-to-cloud path. Costs: a newer ecosystem; distributed runs still

need workers/agents.

Keeping step logic portable

The single most important rule: **put real work in Step subclasses, keep the

orchestrator thin.** Every option above calls the *same* Step.run(context). If

your logic lives in steps (not in Airflow operators or Prefect tasks), switching

orchestrators is glue code, not a rewrite.

Concretely:

• Steps should read inputs from context and return a plain, serializable-ish

object. They should not call print for results, reach for globals, or assume

a particular scheduler.

• Pass configuration in via constructor arguments (as the bundled steps do) so

the same class works from YAML, Airflow, or Prefect.

• At process boundaries (Airflow tasks, multi-machine runs) persist each step's

output to shared storage and pass references, not the objects themselves.

Migration checklist

When you graduate from the bare Pipeline to Airflow or Prefect:

• [ ] Confirm each step's output can be persisted (CSV/Parquet for data, joblib

or pickle for fitted models/transformers).

• [ ] Replace in-memory context passing with artifact storage + references at any

boundary that crosses a process or machine.

• [ ] Move per-step retry/timeout policy into the scheduler (retries=...) and

stop relying on fail_fast alone.

• [ ] Externalize configuration (env vars / a config store) instead of a local

configs/pipeline.yaml path.

• [ ] Add idempotency: a re-run of ingest/train should overwrite, not append.
• [ ] Keep tests/test_pipeline.py green -- the bare Pipeline remains the

fastest way to unit-test step wiring before it ever reaches a scheduler.

Pipeline Architecture

This guide explains how the templates are structured, why they are split the way

they are, and how to extend them. Read it once and you will be able to add a new

stage, swap a model, or wire the pipeline into a larger system with confidence.

Table of contents

1. The two layers

2. The Step abstraction

3. The Pipeline orchestrator

4. The four ML stages

5. The Dataset interface

6. Data flow through the context

7. Why a DAG instead of a linear script

8. Extending the pipeline

9. Design decisions and trade-offs

The two layers

The package is deliberately split into two layers with different dependency

profiles:

+--------------------------------------------------------------+
|  Orchestration core  (pure standard library)                 |
|    pipeline.py   - Pipeline, PipelineRun, topological_order   |
|    steps.py      - Step, FunctionStep, StepResult, StepStatus |
+--------------------------------------------------------------+
|  ML stages  (pandas / numpy, scikit-learn optional)          |
|    ingestion.py     - Dataset + load_csv / parquet / sql      |
|    preprocessing.py - impute / scale / encode + split        |
|    training.py      - cross-validation + estimators          |
|    evaluation.py    - metrics + EvaluationReport             |
+--------------------------------------------------------------+

The core never imports pandas or scikit-learn. That keeps it importable

anywhere, makes the orchestrator unit-testable in milliseconds, and means you can

reuse Pipeline for non-ML work too. The ML stages are imported lazily (see

__init__.py), so import ml_pipeline stays cheap until you actually touch a

stage.

The Step abstraction

A step is the atomic unit of work. Every step has three things:

• a unique name (its node id in the graph),
• a list of inputs -- the names of the steps it depends on,
• a run(context) method that does the work.

from ml_pipeline import Step

class NormalizeStep(Step):
    def __init__(self, name, source):
        super().__init__(name, inputs=[source])
        self.source = source

    def run(self, context):
        data = self.require(context, self.source)   # safe dependency fetch
        return [x / max(data) for x in data]

require() fetches a dependency's output from the context and raises a clear

StepError if it is missing -- you never index a raw dict and get a cryptic

KeyError.

For quick, functional steps there is FunctionStep, which adapts a plain

callable. By default it passes each dependency's output as a positional argument:

from ml_pipeline import FunctionStep

load  = FunctionStep("load", lambda: [1, 2, 3])           # load()
total = FunctionStep("total", sum, inputs=["load"])        # sum(load_output)

Pass pass_context=True if your function would rather receive the whole context

mapping and read several upstream outputs by name.

The Pipeline orchestrator

Pipeline takes a set of steps and runs them. On run() it:

1. Validates the graph -- duplicate names, dangling dependencies, and cycles

all raise specific PipelineError subclasses *before* anything executes.

2. Topologically sorts the steps with Kahn's algorithm. Among steps that are

simultaneously ready, the lexicographically smallest name runs first, so the

order is deterministic and easy to assert in tests.

3. Executes each step in order, threading a shared context dict (step name

-> output) through them, timing each one, and capturing failures.

The result is a PipelineRun:

run = pipeline.run()
run.succeeded        # bool: did every step succeed?
run.order            # ['ingest', 'preprocess', 'train', 'evaluate']
run["evaluate"]      # the evaluate step's output (an EvaluationReport)
print(run.report())  # a per-step status / timing summary

With fail_fast=True (the default) the first failure stops the run and the

remaining steps are marked SKIPPED; with fail_fast=False independent branches

keep running.

The four ML stages

The bundled steps mirror the canonical supervised-learning lifecycle:

Note that evaluate depends on both train (for the model) and

preprocess (for the held-out test matrix). That makes the graph a small diamond

rather than a straight line:

        ingest
          |
       preprocess
        /      \
     train      \
        \       /
        evaluate

build_standard_pipeline(config) constructs exactly this graph from a config

dict (see configs/pipeline.yaml).

The Dataset interface

Every loader returns a Dataset -- a thin wrapper around a pandas DataFrame plus

the name of the target column. Downstream stages depend on Dataset, never on a

file path, so adding a new source means writing one load_* function:

from ml_pipeline import load_dataset

ds = load_dataset("data/train.csv", target="label")     # csv (suffix inferred)
ds = load_dataset("data/train.parquet", target="label") # parquet
ds = load_dataset("SELECT * FROM t", fmt="sql",
                  connection=conn, target="label")       # sql

ds.features        # DataFrame without the target
ds.y               # the target Series (or None)
ds.feature_names   # column names excluding the target

Data flow through the context

Outputs are passed implicitly through the run context, keyed by step name:

ingest      --> context["ingest"]      = Dataset
preprocess  --> context["preprocess"]  = PreprocessedData(X_train, X_test, y_train, y_test, ...)
train       --> context["train"]       = TrainedModel(estimator, cv, ...)
evaluate    --> context["evaluate"]    = EvaluationReport(task, metrics, ...)

The feature transformer is fit on the training split only and then applied to

the test split, so test-set statistics never leak into training. Unknown

categories that appear only at test time map to an all-zero one-hot segment.

Why a DAG instead of a linear script

A flat script (load(); clean(); train(); score()) works until it doesn't:

• You cannot reorder or parallelize independent work without rewriting glue code.
• A failure halfway through leaves you guessing what ran and what didn't.
• Branching ("train two models, compare them") turns into nested conditionals.

Modelling the pipeline as a DAG gives you explicit dependencies, deterministic

ordering, per-step status/timing, validation before execution, and a structure

that maps one-to-one onto Airflow/Prefect tasks when you outgrow a single

process (see orchestration-options.md).

Extending the pipeline

Add a feature-engineering stage between preprocess and train by subclassing

Step, reading context["preprocess"], returning a modified PreprocessedData,

and pointing train at your new step instead.

Swap the model by changing one config line -- a built-in baseline

(ridge/mean/majority) or any scikit-learn estimator via a dotted path:

training:
  estimator: "sklearn:sklearn.ensemble.GradientBoostingRegressor"
  params: { n_estimators: 200, max_depth: 3 }

Reuse the orchestrator for non-ML work -- ETL, report generation, anything

with dependencies -- since Pipeline/Step know nothing about ML.

Design decisions and trade-offs

• Pure-stdlib core. Keeps orchestration testable and dependency-free; the

cost is that ML helpers live in separate modules rather than on Pipeline.

• Estimator factories, not instances. cross_validate takes a zero-arg

callable so each fold gets a fresh model without needing sklearn.clone; this

is what makes the CV utility backend-agnostic.

• numpy fallbacks for every ML stage. The pipeline runs end-to-end with no

scikit-learn installed (and the test suite proves it). scikit-learn remains the

recommended production backend and is used automatically when present.

• In-process execution. These templates run in one process by design --

simple, debuggable, and ideal for datasets that fit in memory. For

distributed/scheduled execution, wrap the steps in Airflow or Prefect rather

than reinventing a scheduler here.