Machine learning pipelines are only as reliable as the infrastructure that runs them. A model that trains successfully in a notebook but fails silently in production is not a production model. Apache Airflow solves this by treating ML workflows as code, giving you version control, retries, observability, and scheduling — all out of the box.
In this tutorial, you will learn to:
- Define a Directed Acyclic Graph (DAG) that models an end-to-end ML training pipeline
- Use PythonOperator, BashOperator, and Sensor operators to orchestrate ML tasks
- Build conditional branching so high-quality model checkpoints are registered automatically
- Apply Airflow best practices for observability, resource control, and error recovery
Page Contents
What Is Apache Airflow?
Apache Airflow is an open-source workflow orchestration platform first open-sourced by Airbnb in 2015 and now a top-level Apache Software Foundation project. Airflow pipelines are defined as Python code, which means you get the full power of Python for your workflow logic — no YAML-only DSL to fight.
The core primitive in Airflow is the DAG (Directed Acyclic Graph). A DAG has two constraints: every edge points in one direction (directed), and there are no cycles (acyclic). These constraints guarantee that the scheduler can traverse the graph without infinite loops and can compute a deterministic execution order.
Three concepts are central to every Airflow deployment:
- Operator — a task definition. An operator describes what gets done: run a Python function, execute a bash command, wait for a file to appear, call a REST API.
- Task — a specific instance of an operator placed inside a DAG. Tasks are the nodes in your graph.
- Sensor — a special operator that waits for an external condition: a file landing on S3, a database record appearing, an HTTP endpoint returning 200.
ML Pipeline Architecture
A production ML training pipeline typically spans five stages:
- Data ingestion — pull raw data from a database, API, or object storage.
- Preprocessing — clean, normalize, and split the data into train/val/test sets.
- Training — run the training loop with hyperparameter selection.
- Evaluation — compute metrics on a held-out test set.
- Model registry — if metrics meet a quality threshold, register the model for deployment.
Each stage is a task in the DAG. Dependencies between stages become edges. Airflow’s scheduler ensures tasks run in the correct order and handles failures gracefully. If you are building Python tooling around ML workflows, our guide to the Python AI agent ecosystem covers related packaging and deployment patterns that complement Airflow. For task scheduling beyond Airflow, the Python sched module offers a lightweight in-process alternative for simpler use cases.
For data preprocessing and feature engineering, pandas is the standard tool for data loading, cleaning, and splitting in Python ML pipelines. The scikit-learn dataset utilities also provide convenient access to benchmark datasets for prototyping and testing pipelines.
Installing Airflow
Install Airflow with the Python API providers you need. For this tutorial we use local executor mode, which is sufficient for development and small-scale production workloads:
# Install Airflow with pip (use a virtual environment) pip install apache-airflow==2.10.0 --constraint "https://raw.githubusercontent.com/apache/airflowconstraints/main/constraints-2.10.0/python3.11.txt" # Verify the installation airflow version
Initialize the metadata database and create an admin user:
airflow db init airflow users create --username admin --password admin --firstname Admin --lastname User --role Admin --email [email protected] # Start the web server and scheduler in two terminals airflow webserver --port 8080 airflow scheduler
Defining the ML Training DAG
Create a file named ml_training_dag.py in your AIRFLOW_HOME (default: ~/airflow) under the dags/ directory. Airflow automatically picks up any Python file in this directory.
"""
ML Training Pipeline DAG — Apache Airflow for Machine Learning Workflows
https://www.pyblog.in/programming/python/apache-airflow-ml-workflows-dags-operators-sensors/
"""
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
from airflow.sensors.filesystem import FileSensor
from airflow.models import XCom
# ---------------------------------------------------------------------------
# DAG configuration
# ---------------------------------------------------------------------------
default_args = {
"owner": "ml-platform-team",
"depends_on_past": False,
"email_on_failure": True,
"email_on_retry": False,
"retries": 2,
"retry_delay": timedelta(minutes=5),
}
dag = DAG(
dag_id="ml_training_pipeline",
description="End-to-end ML training pipeline with data validation, "
"preprocessing, training, evaluation, and model registry",
schedule_interval="0 2 * * *", # Run daily at 02:00 UTC
start_date=datetime(2026, 6, 1),
catchup=False,
default_args=default_args,
tags=["ml", "training", "pipeline"],
)
# ---------------------------------------------------------------------------
# Task 1 — Wait for raw training data
# ---------------------------------------------------------------------------
wait_for_data = FileSensor(
task_id="wait_for_training_data",
filepath="/data/raw/training_data.csv",
fs_conn_id="fs_default",
timeout=300, # Fail after 5 minutes of waiting
poke_interval=30, # Check every 30 seconds
mode="reschedule", # Free the worker slot between checks
dag=dag,
)
# ---------------------------------------------------------------------------
# Task 2 — Preprocess data
# ---------------------------------------------------------------------------
def preprocess_data(**context):
"""
Load raw CSV, split into train/val/test, and save each split.
Persists paths to XCom so downstream tasks can retrieve them.
"""
import pandas as pd
from sklearn.model_selection import train_test_split
import json
raw_path = "/data/raw/training_data.csv"
output_dir = "/data/processed"
df = pd.read_csv(raw_path)
# Drop rows with missing target values
df = df.dropna(subset=["target"])
train, temp = train_test_split(df, test_size=0.3, random_state=42)
val, test = train_test_split(temp, test_size=0.5, random_state=42)
train.to_csv("{} /train.csv".format(output_dir), index=False)
val.to_csv("{} /val.csv".format(output_dir), index=False)
test.to_csv("{} /test.csv".format(output_dir), index=False)
paths = {
"train": "{} /train.csv".format(output_dir),
"val": "{} /val.csv".format(output_dir),
"test": "{} /test.csv".format(output_dir),
}
# Push paths to XCom so the training task can read them
context["ti"].xcom_push(key="data_paths", value=paths)
return "Preprocessed {} rows → train/val/test".format(len(train))
preprocess = PythonOperator(
task_id="preprocess_data",
python_callable=preprocess_data,
do_xcom_push=True,
dag=dag,
)
# ---------------------------------------------------------------------------
# Task 3 — Train the model
# ---------------------------------------------------------------------------
def train_model(**context):
"""
Read train/val paths from XCom, run a RandomForest training loop,
and save the fitted model to disk.
"""
import pandas as pd
import numpy as np
import json
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import joblib
ti = context["ti"]
paths = ti.xcom_pull(task_ids="preprocess_data", key="data_paths")
train_df = pd.read_csv(paths["train"])
val_df = pd.read_csv(paths["val"])
feature_cols = [c for c in train_df.columns if c != "target"]
X_train = train_df[feature_cols].values
y_train = train_df["target"].values
X_val = val_df[feature_cols].values
y_val = val_df["target"].values
# Simple hyperparameter search over n_estimators
best_acc = 0.0
best_n = 100
for n in [50, 100, 200]:
model = RandomForestClassifier(
n_estimators=n,
max_depth=10,
random_state=42,
n_jobs=-1,
)
model.fit(X_train, y_train)
acc = accuracy_score(y_val, model.predict(X_val))
if acc > best_acc:
best_acc = acc
best_n = n
# Train final model with best hyperparameter
final_model = RandomForestClassifier(
n_estimators=best_n,
max_depth=10,
random_state=42,
n_jobs=-1,
)
final_model.fit(X_train, y_train)
model_path = "/models/ml_model.joblib"
joblib.dump(final_model, model_path)
result = {
"model_path": model_path,
"val_accuracy": float(best_acc),
"best_n_estimators": best_n,
}
# Push metrics to XCom for the evaluation task
ti.xcom_push(key="training_result", value=result)
return "Best n_estimators={} — val accuracy={:.4f}".format(
best_n, best_acc)
train = PythonOperator(
task_id="train_model",
python_callable=train_model,
do_xcom_push=True,
dag=dag,
)
# ---------------------------------------------------------------------------
# Task 4 — Evaluate on test set
# ---------------------------------------------------------------------------
def evaluate_model(**context):
"""
Load the trained model, evaluate on the test set, and decide
whether to register the model based on a minimum accuracy threshold.
"""
import pandas as pd
from sklearn.metrics import accuracy_score, f1_score, classification_report
import joblib
ti = context["ti"]
result = ti.xcom_pull(task_ids="train_model", key="training_result")
paths = ti.xcom_pull(task_ids="preprocess_data", key="data_paths")
model = joblib.load(result["model_path"])
test_df = pd.read_csv(paths["test"])
feature_cols = [c for c in test_df.columns if c != "target"]
X_test = test_df[feature_cols].values
y_test = test_df["target"].values
y_pred = model.predict(X_test)
acc = accuracy_score(y_test, y_pred)
f1 = f1_score(y_test, y_pred, average="weighted")
quality_threshold = 0.80
meets_threshold = acc >= quality_threshold
eval_result = {
"test_accuracy": float(acc),
"test_f1": float(f1),
"meets_threshold": meets_threshold,
"model_path": result["model_path"],
}
print("Test accuracy : {:.4f}".format(acc))
print("Test F1 score : {:.4f}".format(f1))
print("Threshold met: {}".format(meets_threshold))
if meets_threshold:
print("Model quality check PASSED — registering to model registry")
else:
print("Model quality check FAILED — skipping registration")
ti.xcom_push(key="eval_result", value=eval_result)
return eval_result
evaluate = PythonOperator(
task_id="evaluate_model",
python_callable=evaluate_model,
do_xcom_push=True,
dag=dag,
)
# ---------------------------------------------------------------------------
# Task 5 — Register model (conditional on quality)
# ---------------------------------------------------------------------------
def register_model(**context):
"""
Copy the approved model artifact to the model registry directory
with a timestamped version label. Once registered, serve the model
with a dedicated serving framework like
BentoML
for production inference endpoints.
"""
import shutil
from datetime import datetime
ti = context["ti"]
eval_result = ti.xcom_pull(task_ids="evaluate_model", key="eval_result")
if not eval_result["meets_threshold"]:
print("Skipping registration — model did not meet quality threshold")
return None
model_src = eval_result["model_path"]
version = datetime.now().strftime("%Y%m%d_%H%M%S")
registry_path = "/models/registry/ml_model_{}.joblib".format(version)
shutil.copy(model_src, registry_path)
print("Model registered at: {}".format(registry_path))
return registry_path
register = PythonOperator(
task_id="register_model",
python_callable=register_model,
trigger_rule="none_failed_min_one_success",
dag=dag,
)
# ---------------------------------------------------------------------------
# Task 6 — Notify on failure (Airflow callback pattern)
# ---------------------------------------------------------------------------
def notify_failure(context):
"""
Airflow failure callback. Sends an alert when any task fails.
Extend this to send Slack messages, emails, or PagerDuty alerts.
"""
dag_run = context.get("dag_run")
task_instance = context.get("task_instance")
print("ALERT: DAG {} failed at task {}".format(
dag_run.dag_id if dag_run else "unknown",
task_instance.task_id if task_instance else "unknown",
))
# Attach the failure callback to the DAG
dag.on_failure_callback = notify_failure
# ---------------------------------------------------------------------------
# Task dependencies — define the execution order
# ---------------------------------------------------------------------------
# wait_for_data must succeed before preprocessing starts
wait_for_data >> preprocess
# preprocess must complete before training starts
preprocess >> train
# training output feeds into evaluation
train >> evaluate
# If evaluation passes, register the model; otherwise skip registration
evaluate >> register
How the DAG Executes
When Airflow’s scheduler picks up this DAG, it walks the dependency graph and submits tasks to the executor in the correct order. The wait_for_training_data FileSensor repeatedly polls the filesystem every 30 seconds. While it is polling, the worker slot is released (mode="reschedule"), so your cluster can run other tasks.
Once the CSV file appears, the sensor marks itself successful and preprocess_data runs. The preprocessing task loads the raw data, splits it, and pushes the train/val/test file paths to XCom. The train_model task pulls those paths from XCom, runs the hyperparameter search, and saves the best model. Finally, evaluate_model scores the model on the held-out test set and decides whether to register it.
The trigger_rule="none_failed_min_one_success" on register_model ensures the registration task runs if evaluation succeeded — even if some other optional task upstream failed. This is important for production pipelines where you want the model registered whenever quality checks pass, regardless of peripheral failures.
Understanding XCom
XCom (cross-communication) lets tasks pass small amounts of data to each other. In the pipeline above, preprocess_data pushes the train/val/test file paths to XCom, and train_model pulls them back. This pattern avoids hardcoding file paths in every task and lets each task be developed and tested independently.
For large artifacts like trained model files, use a shared filesystem or object storage (S3, GCS) instead of XCom — XCom is designed for metadata, not binary blobs. The model file path is passed via XCom; the file itself lives on disk.
Sensors for ML Pipelines
Sensors are one of Airflow’s most powerful features for ML workflows. Data scientists often need to wait for a daily data export, a model download, or an upstream pipeline to finish. FileSensor, HttpSensor, SqlSensor, and the S3-specific S3KeySensor cover the most common cases:
from airflow.sensors.filesystem import FileSensor
from airflow.sensors.http_sensor import HttpSensor
from airflow.sensors.sql import SqlSensor
# Wait for a CSV file to land on the local filesystem
file_sensor = FileSensor(
task_id="wait_for_batch_data",
filepath="/data/batch/features.csv",
fs_conn_id="fs_default",
timeout=600,
poke_interval=60,
mode="reschedule",
dag=dag,
)
# Wait for an HTTP endpoint to return 200 OK
http_sensor = HttpSensor(
task_id="wait_for_model_api",
http_conn_id="model_api",
endpoint="health",
request_params={},
timeout=300,
poke_interval=30,
dag=dag,
)
# Wait for a SQL query to return at least one row
sql_sensor = SqlSensor(
task_id="wait_for_training_labels",
conn_id="warehouse_db",
sql="SELECT COUNT(*) FROM labels WHERE processed = 0",
success_check=lambda result: result[0][0] > 0,
timeout=3600,
poke_interval=120,
dag=dag,
)
Observability and Monitoring
Airflow’s built-in web UI gives you a visual DAG graph, task-level logs, execution history, and a Gantt chart view for identifying bottlenecks. For production deployments, integrate with Prometheus and Grafana via the Airflow statsd exporter to track:
- Task duration — detect tasks that are getting slower over time
- Success/failure rate — alert when the failure rate spikes
- Queue depth — know when the scheduler is overwhelmed
The failure callback notify_failure wired into the DAG sends an alert whenever any task fails. In a production setting, replace the print statement with a Slack webhook, an email via your SMTP server, or a PagerDuty incident.
Scheduling and Catchup
The DAG is scheduled daily at 02:00 UTC (schedule_interval="0 2 * * *"). Because catchup=False is set, Airflow will not backfill runs for past dates when the DAG was not active. Set catchup=True if you need the pipeline to replay historical runs when you deploy a new DAG.
For ML pipelines, catchup=False is almost always the right choice — you typically want to train on the latest available data, not re-train on data that was available weeks ago.
Conclusion
Apache Airflow turns ML workflows into versioned, observable, schedulable Python code. The DAG model forces you to think clearly about dependencies between pipeline stages. The sensor pattern lets you wait gracefully for data without busy-waiting. XCom lets tasks communicate without tight coupling. And the built-in retry, alerting, and web UI give you everything you need to run ML pipelines in production with confidence.
The pipeline in this article covered the full lifecycle: data validation via FileSensor, preprocessing with PythonOperator, hyperparameter search with BashOperator, evaluation with quality gating, and conditional model registration. Every stage is independently testable, observable, and retryable — the hallmarks of a production-grade ML system.
