Airflow for ML Pipelines: A Practical Guide

Airflow is usually introduced as a scheduler for data loads, but much of the routine work around a model is also scheduled work: retrain every week, score every day, compare the error with what the model promised, and tell someone when it slips. This guide builds that as two Airflow DAGs and runs them on real bike rental data, including a day when the data simply did not arrive.

If Airflow itself is new to you, start with the Airflow tutorial, which covers DAGs, tasks, XComs and logical dates. This article assumes those and focuses on the ML parts.

The data and the job

The data is the UCI Bike Sharing set: hourly rentals from Capital Bikeshare in Washington, D.C., for 2011 and 2012, with weather and calendar fields, under CC BY 4.0 (UCI: Bike Sharing). Here it stands in for a daily feed. The model estimates rentals per hour from the hour, day type, season and weather, the kind of estimate an operator uses to decide where to move bikes.

Two DAGs, not one

Training and scoring run on different schedules and fail for different reasons, so they are separate DAGs:

• bike_train runs weekly. It fits a model on older data, measures its error on the most recent 13 weeks it did not train on, and publishes the model together with that error. The holdout error is the model's promise about how wrong it will usually be.
• bike_score runs daily. It extracts the day's hours, validates them, scores them, and checks the error against the promise. If the error is more than 1.25 times the holdout error, it raises an alert; otherwise it records an all-clear.

The only connection between them is the published model file. Here it is a file on shared disk; in a larger setup it would be a model registry entry, as in the model registry guide. Either way, a failed training run leaves the last good model in place, and scoring carries on.

The DAGs run in the official Airflow image with the ML libraries added:

FROM apache/airflow:3.3.2-python3.11
RUN pip install --no-cache-dir "apache-airflow==3.3.2" scikit-learn==1.9.1 pandas==3.0.6

Both DAGs, followed by a test harness that runs one training date and five scoring dates:

import os
from datetime import datetime, timedelta, timezone

import joblib
import pandas as pd
from airflow.exceptions import AirflowFailException
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk import dag, task

HOURS = "data/hour.csv"           # UCI Bike Sharing hourly rentals (CC BY 4.0), standing in for the daily feed
MODEL = "models/current.joblib"   # what the training DAG publishes and the scoring DAG reads
SCORES = "scores.csv"             # one row per scored day
FEATURES = ["season", "yr", "mnth", "hr", "holiday", "weekday", "workingday", "weathersit", "temp", "atemp", "hum", "windspeed"]


@dag(schedule="@weekly", start_date=datetime(2012, 1, 1, tzinfo=timezone.utc), catchup=False, tags=["bikes", "ml"])
def bike_train():
    @task
    def train(ds=None):
        from sklearn.ensemble import HistGradientBoostingRegressor
        hours = pd.read_csv(HOURS)
        holdout_start = (pd.Timestamp(ds) - pd.Timedelta(days=91)).strftime("%Y-%m-%d")
        fit, holdout = hours[hours.dteday < holdout_start], hours[(hours.dteday >= holdout_start) & (hours.dteday < ds)]
        model = HistGradientBoostingRegressor(random_state=0).fit(fit[FEATURES], fit.cnt)
        error = float((model.predict(holdout[FEATURES]) - holdout.cnt).abs().mean())
        os.makedirs("models", exist_ok=True)
        joblib.dump({"model": model, "holdout_mae": error, "trained_before": holdout_start}, MODEL)

    train()


@dag(schedule="@daily", start_date=datetime(2012, 1, 1, tzinfo=timezone.utc), catchup=False, tags=["bikes", "ml"],
     default_args={"retries": 2, "retry_delay": timedelta(minutes=10)})
def bike_score():
    @task
    def extract(ds=None):
        hours = pd.read_csv(HOURS)
        hours[hours.dteday == ds].to_csv(f"day_{ds}.csv", index=False)
        return ds

    @task
    def validate(day):
        hours = len(pd.read_csv(f"day_{day}.csv"))
        if hours < 20:                # a retry will not bring missing hours back, so fail without retrying
            raise AirflowFailException(f"{day}: only {hours} of 24 hours arrived")
        return day

    @task
    def score(day):
        saved, hours = joblib.load(MODEL), pd.read_csv(f"day_{day}.csv")
        error = float((saved["model"].predict(hours[FEATURES]) - hours.cnt).abs().mean())
        row = pd.DataFrame([{"day": day, "mae": round(error, 1), "limit": round(1.25 * saved["holdout_mae"], 1)}])
        old = pd.read_csv(SCORES) if os.path.exists(SCORES) else row.iloc[:0]
        pd.concat([old[old.day != day], row]).to_csv(SCORES, index=False)    # a rerun replaces its own row
        return {"mae": error, "limit": 1.25 * saved["holdout_mae"]}

    @task.branch
    def check(result):
        return "raise_alert" if result["mae"] > result["limit"] else "all_clear"

    @task
    def raise_alert(ds=None):
        with open("alerts.txt", "a") as alerts:      # in production: page someone or open a retraining ticket
            alerts.write(f"{ds}\n")

    check(score(validate(extract()))) >> [raise_alert(), EmptyOperator(task_id="all_clear")]


trainer, scorer = bike_train(), bike_score()

if __name__ == "__main__":
    utc = timezone.utc
    print(f"bike_train for 2012-07-01: {trainer.test(logical_date=datetime(2012, 7, 1, tzinfo=utc)).state}")
    saved = joblib.load(MODEL)
    print(f"  trained on hours before {saved['trained_before']}; holdout MAE {saved['holdout_mae']:.1f} bikes/hour, "
          f"alert limit {1.25 * saved['holdout_mae']:.1f}")
    for day in ("2012-10-24", "2012-10-25", "2012-10-28", "2012-10-29", "2012-10-31"):
        run = scorer.test(logical_date=datetime.fromisoformat(day).replace(tzinfo=utc))
        states = {ti.task_id: ti.state for ti in run.get_task_instances()}
        ran = ", ".join(f"{t}={states[t]}" for t in ("validate", "score", "raise_alert", "all_clear"))
        scored = pd.read_csv(SCORES).set_index("day")
        mae = f"MAE {scored.loc[day, 'mae']}" if day in scored.index else f"not scored, {len(pd.read_csv(f'day_{day}.csv'))} of 24 hours arrived"
        print(f"bike_score {day}: {run.state} | {mae} | {ran}")

bike_train for 2012-07-01: success
  trained on hours before 2012-04-01; holdout MAE 47.8 bikes/hour, alert limit 59.8
bike_score 2012-10-24: success | MAE 45.9 | validate=success, score=success, raise_alert=skipped, all_clear=success
bike_score 2012-10-25: success | MAE 62.3 | validate=success, score=success, raise_alert=success, all_clear=skipped
bike_score 2012-10-28: success | MAE 44.9 | validate=success, score=success, raise_alert=skipped, all_clear=success
bike_score 2012-10-29: failed | not scored, 1 of 24 hours arrived | validate=failed, score=upstream_failed, raise_alert=upstream_failed, all_clear=upstream_failed
bike_score 2012-10-31: success | MAE 33.8 | validate=success, score=success, raise_alert=skipped, all_clear=success

What happened on each day

• Training set the bar. The model trained on hours before 1 April 2012 missed by 47.8 bikes an hour on April to June, so the scoring DAG's alert limit became 59.8.
• 24 October was normal. The error was 45.9, under the limit, so the branch task chose the all-clear path and skipped the alert.
• 25 October crossed the line. The error was 62.3, the branch chose the alert path, and the all-clear task was skipped. The alert does not say why the model missed; that is the job of whoever picks it up, using the tools in monitoring ML models in production.
• 28 October looked fine. The storm was already approaching on that Sunday, but weather is one of the model's inputs, and the error stayed at 44.9, under the limit.
• 29 October failed fast. Hurricane Sandy came ashore near Brigantine, New Jersey, late that day (NOAA: Hurricane Sandy report), and the data holds only one hour of rentals for the date. Validation raised an AirflowFailException, the run failed, and scoring, alerting and the all-clear never ran. The day was not scored at all, which is correct: an error computed on one hour would have been noise.
• 31 October recovered. With a full day of data, the error was 33.8 and the all-clear ran.

Three details that matter for ML DAGs

Fail fast when retrying cannot help. The scoring DAG retries failed tasks twice, ten minutes apart, which suits a flaky database or network. Missing data will not appear on a retry, so validation raises AirflowFailException, which marks the task failed and ignores the remaining retries (Airflow docs: tasks). In the first run of this example, a bug made a task raise an ordinary error, and the test run sat waiting out the ten-minute retry delay. With retries configured, the exception type decides whether a failure costs seconds or twenty minutes.

Branch on the result, do not just log it. A task decorated with @task.branch returns the ID of the next task to run; Airflow follows that path and skips the others (Airflow docs: DAGs). Writing the check as a branch makes the decision visible in the run itself: anyone looking at the run can see whether the alert or the all-clear fired, without reading logs.

Make reruns safe. The scoring task replaces its own row for the day instead of appending, so rerunning a day after fixing its data does not create a duplicate score. The same habit is covered for loads in the general Airflow tutorial.

One more trap from the first run: a task parameter called ds fails, because Airflow reserves that name for the logical date it passes in. Task inputs that come from other tasks need their own names, which is why the tasks here take day.

Testing DAGs before the scheduler sees them

Every run above used dag.test(), which executes a real DAG run in one Python process without a scheduler or executor (Airflow docs: debugging). Calling it with specific logical dates turns past days into test cases: a normal day, a bad day and a missing day. Kept as a script in continuous integration, that becomes a regression test for the pipeline logic.

Where Airflow fits for ML

Airflow suits ML work that is mostly scheduled batch steps around a model: retraining, batch scoring, data checks, report generation. It is less suited to running each training step in its own container with tracked artifacts on Kubernetes, which is what Kubeflow Pipelines is built for, or to serving predictions in real time. For how retraining should be triggered and gated, see continuous training.

Common mistakes

• Putting training and scoring in one DAG, so a failed retrain also stops the day's scoring.
• Retrying data problems, which delays the failure without fixing it.
• Scoring incomplete days, which produces error numbers that look real and mean nothing.
• Alerting by log message only, where nobody sees it.
• Appending results, so every rerun adds a duplicate row.

Scheduling training and scoring pipelines is part of the MLOps course.

Sources

• Dags (Apache Airflow documentation) — Apache Airflow
• Tasks (Apache Airflow documentation) — Apache Airflow
• Debugging Airflow Dags (Apache Airflow documentation) — Apache Airflow
• Bike Sharing dataset (UCI Machine Learning Repository) — UCI Machine Learning Repository
• Tropical Cyclone Report: Hurricane Sandy (NOAA National Hurricane Center) — NOAA National Hurricane Center