Pipeline Orchestration: Airflow vs Prefect vs Dagster

Pipeline orchestration for ML looks like ordinary data orchestration at first: steps, dependencies, a schedule. The differences appear when training gets expensive and multiplies. Feature jobs take an hour and rarely change. One model becomes one model per region or customer segment. A transient storage error should not throw away three hours of training. And six weeks later someone asks which data produced the model in production.

The general comparison of Airflow, Prefect and Dagster on one ETL pipeline is in Airflow vs Dagster vs Prefect. This guide looks only at the ML-specific needs, and runs them. The data is the UCI Bike Sharing daily rentals for 2011 and 2012 (UCI: Bike Sharing), with one model per segment: year by working day or day off.

import os, tempfile, time, warnings
WORK = tempfile.mkdtemp()
os.environ.update(PREFECT_HOME=os.path.join(WORK, "prefect"), PREFECT_LOGGING_LEVEL="CRITICAL")
warnings.filterwarnings("ignore")

import pandas as pd
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.metrics import mean_absolute_error

# UCI Bike Sharing daily rentals (CC BY 4.0), one model per segment: year x working day
SEGMENTS = ["2011_workday", "2011_offday", "2012_workday", "2012_offday"]
FEATURES = ["mnth", "weekday", "weathersit", "temp", "atemp", "hum", "windspeed"]

def make_features(path):
    day = pd.read_csv(path)
    day["segment"] = day.yr.map({0: "2011", 1: "2012"}) + "_" + day.workingday.map({1: "workday", 0: "offday"})
    return day

def fit_segment(features, segment):
    rows = features[features.segment == segment]
    cut = int(len(rows) * 0.8)                                  # train on the first 80% of days, test on the rest
    model = HistGradientBoostingRegressor(random_state=0).fit(rows[FEATURES][:cut], rows.cnt[:cut])
    return model, mean_absolute_error(rows.cnt[cut:], model.predict(rows[FEATURES][cut:]))

# ---- Prefect: cache the expensive step, fan out over segments, retry a flaky one ----
from prefect import flow, task, unmapped
from prefect.cache_policies import INPUTS

attempts = {}

@task(cache_policy=INPUTS, persist_result=True)
def build_features(path):
    time.sleep(1.5)                                             # stands in for an expensive feature job
    return make_features(path)

@task(retries=2, retry_delay_seconds=0)
def train(features, segment):
    attempts[segment] = attempts.get(segment, 0) + 1
    if segment == "2012_workday" and attempts[segment] == 1:
        raise ConnectionError("model store timed out")          # a transient failure on the first attempt
    _, mae = fit_segment(features, segment)
    return segment, round(mae), attempts[segment]

@flow
def train_all(path="day.csv"):
    features = build_features(path, return_state=True)
    results = train.map(unmapped(features.result()), SEGMENTS).result()
    return features.name, results

print("O1 Prefect")
for run in (1, 2):
    attempts.clear()
    feature_state, results = train_all()
    print(f"   run {run}: build_features state = {feature_state}")
    for segment, mae, tries in results:
        print(f"      {segment:<13} MAE {mae:>5} bikes/day, attempts {tries}")

# ---- Dagster: the model as a partitioned asset, one partition per segment ----
import dagster as dg

@dg.asset
def daily_features():
    return make_features("day.csv")

@dg.asset(partitions_def=dg.StaticPartitionsDefinition(SEGMENTS))
def segment_model(context: dg.AssetExecutionContext, daily_features):
    model, mae = fit_segment(daily_features, context.partition_key)
    context.add_output_metadata({"mae": round(mae)})
    return model

instance = dg.DagsterInstance.ephemeral()
io = {"io_manager": dg.FilesystemIOManager(base_dir=os.path.join(WORK, "dagster"))}
dg.materialize([daily_features], instance=instance, resources=io)
result = dg.materialize([daily_features, segment_model], selection=[segment_model], partition_key="2012_workday",
                        instance=instance, resources=io)

print("\nO2 Dagster")
print("   segment_model depends on:", sorted(k.to_user_string() for k in segment_model.dependency_keys))
print("   partitions materialized:", sorted(instance.get_materialized_partitions(dg.AssetKey("segment_model"))))
mae = result.asset_materializations_for_node("segment_model")[0].metadata["mae"].value
print(f"   2012_workday MAE {mae} bikes/day")

O1 Prefect
   run 1: build_features state = Completed
      2011_workday  MAE   555 bikes/day, attempts 1
      2011_offday   MAE   874 bikes/day, attempts 1
      2012_workday  MAE   871 bikes/day, attempts 2
      2012_offday   MAE  1289 bikes/day, attempts 1
   run 2: build_features state = Cached
      2011_workday  MAE   555 bikes/day, attempts 1
      2011_offday   MAE   874 bikes/day, attempts 1
      2012_workday  MAE   871 bikes/day, attempts 2
      2012_offday   MAE  1289 bikes/day, attempts 1

O2 Dagster
   segment_model depends on: ['daily_features']
   partitions materialized: ['2012_workday']
   2012_workday MAE 871 bikes/day

Need 1: skip expensive steps that have not changed

In Prefect, the feature step ran on the first run and came back Cached on the second, so the 1.5-second stand-in for an expensive job was not repeated. Caching needs result persistence, which is off by default, and the cache policy decides what the key is made of. The default key includes the flow run ID, so it only reuses results within one run; the INPUTS policy used here keys on the task's inputs alone, which is what reuse across runs needs (Prefect: caching).

That is also the trap. The input here is a file path. If the file at that path changes, an inputs-only key still returns the old features, and the new model is trained on stale data. For ML, key the cache on something that changes with the data: a content hash, a data version from a tool such as DVC, or a snapshot date. Set an expiry as a safety net.

Dagster approaches the same need from the data side: each asset is materialized and stored, so downstream assets load the stored result instead of recomputing it. In the run, segment_model loaded daily_features from storage rather than rebuilding it.

In Airflow, skipping unchanged work is usually your own logic, for example checking whether the output for the current data version already exists and short-circuiting the rest.

Need 2: fan out over segments, and retry only what failed

Training per segment is a natural fan-out. In Prefect, train.map created one task run per segment. The 2012_workday model hit a simulated timeout on its first attempt; the retry policy ran it again, it succeeded on attempt 2, and the other three segments were unaffected. The second run shows the training tasks running again, as they should, because only the feature step had an inputs-based cache.

Airflow does the same with dynamic task mapping: one task instance per input, created at run time from upstream output, up to 1,024 by default (Airflow: dynamic task mapping). For training DAG structure in Airflow, see Airflow for ML pipelines.

Whatever the tool, set retries per task, and only on errors that can succeed on retry. A timeout is worth retrying; a bad training label is not.

Need 3: models as partitioned assets, with lineage

Dagster models each segment's model as a partition of one asset (Dagster: partitioning assets). The run materialized only the 2012_workday partition, and the instance recorded exactly that: one partition materialized, with its MAE of 871 bikes a day attached as metadata. The asset's dependency on daily_features is part of its definition, so the lineage question, which data produced this model, is answered by the orchestrator rather than by a wiki page.

This shape fits common ML operations: retrain one segment after its data is corrected, backfill a new segment, or see at a glance which segments have a current model. In Prefect or Airflow you can achieve the same by naming outputs per segment and recording metadata yourself; Dagster makes it the default unit of work.

Keep the ML code independent of the orchestrator

Notice that the demo's real work lives in two plain functions, make_features and fit_segment, and both orchestrators call the same ones. The Prefect tasks and Dagster assets are thin wrappers. Keeping it that way pays off three times: the functions can be unit-tested without starting any orchestrator, a data scientist can run them in a notebook, and changing orchestrator later means rewriting wrappers, not models.

Choosing for an ML team

• Mostly Python ML code, quick to adopt: Prefect's decorated functions, caching and retries map closely onto training scripts.
• Many models and datasets to track: Dagster's assets, partitions and lineage suit teams that think in data products and models.
• An existing Airflow platform: Airflow covers fan-out and retries well; add your own data-version checks and metadata, and record model lineage in a registry.

Whichever you choose, the orchestrator schedules and retries; experiment tracking and the model registry still record what was trained. See continuous training for when to retrain, and Kubernetes for ML for running the training jobs at scale.

Pipelines, orchestration and retraining are part of the MLOps course.

Sources

• Caching (Prefect documentation) — Prefect
• Partitioning assets (Dagster documentation) — Dagster Labs
• Dynamic Task Mapping (Apache Airflow documentation) — Apache Software Foundation
• Bike Sharing Dataset (UCI Machine Learning Repository) — UCI Machine Learning Repository