Put it together: a DAG that ingests data, validates it, trains a model, evaluates against a baseline, and registers the winner — the pipeline the MLOps checklist describes.
Why: nearly every ML training pipeline is the same shape — ingest → validate → preprocess → train → evaluate → register — with a quality gate that stops a bad model from being registered. When: model this as a DAG so it runs on a schedule or on new data, with retries and visibility. Where: this is exactly the pipeline the ML Training Pipeline checklist on the roadmap describes.
ingest ─► validate ─► preprocess ─► train ─► evaluate ─► register
│ │
(fail if (gate: only register
schema/nulls if it beats the
are wrong) baseline metric)
Each step is an idempotent task; the gate makes promotion conditional.Why: writing the pipeline as decorated functions makes the data flow — and thus the dependencies — obvious, and each task can log to MLflow so the run is tracked. When: return references (paths, run IDs) between tasks, not big objects. Where: combine this with the MLflow course — the train task logs the run, the register task promotes it.
from airflow.decorators import dag, task
import pendulum
@dag(start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
schedule="@daily", catchup=False)
def ml_training():
@task
def ingest():
return "s3://data/raw/2024-01-01.parquet"
@task
def validate(path):
# fail loudly on schema/null violations -> stops the pipeline
return path
@task
def train(path):
# ... train, log params/metrics/model to MLflow ...
return {"run_id": "abc123", "accuracy": 0.91}
@task
def evaluate_and_register(result):
if result["accuracy"] > 0.88: # the quality gate
print("register + promote", result["run_id"])
else:
raise ValueError("below baseline — not registered")
evaluate_and_register(train(validate(ingest())))
ml_training()Why: an ML pipeline should be reachable three ways — on a schedule, when new data arrives, and on demand — and all three should run the SAME DAG so behavior never diverges. When: schedule for regular retraining, a sensor/dataset trigger for data-driven runs, and manual trigger for ad-hoc. Where: Airflow Datasets let one DAG trigger when another produces new data, wiring pipelines together.
# Scheduled: set schedule="@daily" (or a cron expression) on the DAG.
# Manual, on demand:
airflow dags trigger ml_training
# Data-driven: Airflow Datasets trigger a DAG when upstream data updates —
# @dag(schedule=[Dataset("s3://data/raw/")])
# so the training DAG fires when the ingestion DAG publishes new data.