Airflow does not train a model, and setting it up for ML pipeline orchestration is a different job than pointing it at a script and adding a schedule. The pipeline has to survive four things a plain ETL DAG never faces: retraining triggers that depend on data changing, not time passing; a training step running on different compute than the scheduler; an evaluation gate that can refuse a bad model; and backfills that must never duplicate a feature table. This guide covers the setup stage by stage, building on Airflow orchestration for AI, using the TaskFlow API, dynamic task mapping, and asset-aware scheduling.
Prerequisites
Permalink to “Prerequisites”A defined pipeline scope. Decide which stages become separate DAGs and which stay as tasks inside one. Most teams split ingestion and feature engineering into one DAG that produces a governed table, and training into a second DAG the first triggers.
Airflow running somewhere real. A local instance through the Astro CLI works for development. Production needs a managed or self-hosted deployment with an executor matched to your concurrency.
Compute outside the scheduler for training. Airflow workers should trigger training, not run it. Have a Kubernetes cluster, a Databricks workspace, or a Spark cluster ready first.
An experiment tracker. MLflow or a comparable tool needs to be reachable from wherever training runs, so every run logs its parameters, metrics, and a registered model version.
A governed feature store or table. Feast, or a warehouse table on a data lakehouse for AI with a named owner. Ad hoc files are the most common reason feature engineering for AI breaks between environments; if the source has sensitive fields, decide how to handle PII in AI pipelines up front.
A stated evaluation threshold. Decide the metric and the bar a new model must clear before you write the gate.
Why does orchestrating an ML pipeline take more than a scheduled script?
Permalink to “Why does orchestrating an ML pipeline take more than a scheduled script?”A cron job runs a script on a timer and calls it done. Orchestrating an ML pipeline means managing dependencies across stages that behave nothing alike: ingestion is idempotent and cheap to rerun, training is stateful and expensive, and deployment is a decision, not a transformation. According to Astronomer’s State of Airflow 2026 report, based on responses from more than 5,800 data professionals across 122 countries, 32% of Airflow users now have GenAI or MLOps use cases in production, doubling among the platform’s own customers. Worth naming early: LLMOps vs MLOps, since most of this applies to both, but LLM pipelines add prompt versioning steps a classic training DAG does not need. Tools built for training, such as Kubeflow Pipelines or SageMaker Pipelines, handle that step well inside their own platform. Airflow earns its place when the pipeline crosses systems those tools do not own: a data infrastructure for AI ingestion job and a feature store materialization sharing one dependency graph. A broader view of the alternatives is in ML pipeline orchestration patterns for production; this guide assumes that call is made.
Step 1: Map your ML pipeline stages onto DAG boundaries
Permalink to “Step 1: Map your ML pipeline stages onto DAG boundaries”Decide what becomes a DAG and what becomes a task inside a DAG before writing anything. The pattern that holds up: one DAG owns ingestion through feature engineering and writes to a governed system of record for data and knowledge; a second DAG owns training, evaluation, and the deployment gate, starting when the first DAG’s output changes rather than on a fixed schedule.
This split matters because ingestion and training fail differently. Ingestion failures are usually a source system being unavailable, cheap to retry. Training failures are expensive: a bad run can burn a GPU allocation for hours, and keeping them separate means a training failure does not block tomorrow’s ingestion. Write down which DAG produces which data asset for AI, with a named owner, or debugging “the pipeline failed” turns into guessing which of six DAGs actually broke.
☐ Ingestion and feature engineering are scoped to one DAG, training and deployment to another
☐ Every DAG has a named owner and a documented output asset
☐ The boundary between DAGs is a governed table, not a shared in-memory object
Step 2: Install and configure Airflow for ML workloads
Permalink to “Step 2: Install and configure Airflow for ML workloads”Get Airflow running before writing ML-specific code. For local development, the Astro CLI or the official Docker Compose file gets a working instance up in minutes with the LocalExecutor. Production deployments need the CeleryExecutor or KubernetesExecutor, depending on whether you want workers as a fixed pool or scaled per task.
Install only the provider packages the pipeline calls: the Kubernetes provider if training runs there, the Databricks provider if it runs there, and the mlflow package wherever training executes. Airflow 3, released in April 2025, is worth planning for even if you are not on it yet. According to Astronomer, 26% of Airflow users have already upgraded, and the asset-based scheduling in Step 7 is native to that version; Airflow 2.x has the same concept under the older name, Datasets.
Keep DAG files light. A DAG definition should describe structure and dependencies, not do the actual data work inline; heavy imports at the top level slow down every scheduler parse cycle.
Step 3: Write ingestion and feature engineering tasks with the TaskFlow API
Permalink to “Step 3: Write ingestion and feature engineering tasks with the TaskFlow API”The TaskFlow API turns a plain Python function into an Airflow task with the @task decorator, removing most of the manual XCom code that ingestion pipelines otherwise accumulate. A function’s return value passes to the next task automatically, and Airflow infers the dependency graph from which functions call which.
from airflow.decorators import dag, task
from datetime import datetime
@dag(schedule="@daily", start_date=datetime(2026, 1, 1), catchup=False)
def ingest_and_engineer():
@task
def ingest(logical_date=None) -> str:
# pull raw records for logical_date
return f"s3://ml-staging/raw/{logical_date}.parquet"
@task
def clean(raw_path: str) -> str:
# enforce schema, drop malformed rows
return raw_path.replace("raw", "clean")
@task
def engineer_features(clean_path: str) -> str:
# write model-ready features to the governed table
return clean_path.replace("clean", "features")
engineer_features(clean(ingest()))
ingest_and_engineer()
Notice logical_date, not datetime.now(), drives the ingestion window, the single choice that makes this task safe to rerun or backfill, a point Step 8 returns to. Enrich every output row with source, timestamp, and an owner reference, and enforce access rules before a training job reads the result; see how to give AI agents access to enterprise data for the same pattern applied to inference.
Step 4: Parallelize feature engineering with dynamic task mapping
Permalink to “Step 4: Parallelize feature engineering with dynamic task mapping”Real feature engineering rarely runs as one function over one blob of data. It runs per customer segment, per region, or per date partition, and the count is not known until the DAG executes. Dynamic task mapping solves this: .expand() creates one parallel task instance per item in a list, and .partial() holds arguments constant across all of them.
@task
def engineer_features_for_partition(partition: str, model_version: str) -> str:
# compute features for one partition, return its output path
return f"s3://ml-features/{partition}/{model_version}.parquet"
partitions = get_partitions() # a task returning a list, resolved at runtime
engineer_features_for_partition.partial(model_version="v12").expand(
partition=partitions
)
This scales feature engineering without a hardcoded loop, and Airflow tracks each mapped instance’s state independently, so one failed partition retries on its own instead of forcing a full rerun. The same pattern parallelizes work across a data mesh of domain-owned tables as well as it does across date ranges, and it applies equally when generating synthetic training data across scenario configurations.
Step 5: Trigger training and log runs to MLflow
Permalink to “Step 5: Trigger training and log runs to MLflow”Training is where Airflow’s job changes from doing work to delegating it. Route training to a compute layer built for it, using KubernetesPodOperator to launch a training container, or the Databricks provider’s DatabricksSubmitRunOperator to submit a job to an existing cluster. Airflow’s own workers should never hold a GPU allocation for a training run.
Inside the training job, log every run to MLflow: hyperparameters, the metric the evaluation gate will check in Step 6, and the artifact itself, registered as a new model version rather than overwriting the last one. Registering versions, not files, is what makes Step 6’s comparison against a baseline possible. If the model consumes a knowledge graph for AI agents or a semantic layer for AI agents as a feature source, log which version of that upstream definition it trained against, not just the model’s own hyperparameters.
Step 6: Gate deployment with a branching evaluation check
Permalink to “Step 6: Gate deployment with a branching evaluation check”A pipeline that trains a model and deploys it unconditionally has no stop button. Add a branch task, using @task.branch or BranchPythonOperator, that reads the new model’s evaluation metric from MLflow, compares it against the currently deployed model’s metric, and returns the deployment task’s ID only if the new model clears the bar. Otherwise it returns a “hold” task that logs the result and stops there.
@task.branch
def evaluate_and_gate(new_metric: float, baseline_metric: float) -> str:
if new_metric >= baseline_metric:
return "deploy_model"
return "hold_model"
This is the change that separates a demo pipeline from one a team trusts. Without it, a training run that regresses on a data drift issue ships anyway, silently. Real-time data for AI agents raises the stakes further, since a model serving live traffic has no grace period to catch a bad deployment the way a batch report does.
Step 7: Schedule retraining with asset-aware triggers
Permalink to “Step 7: Schedule retraining with asset-aware triggers”A cron-scheduled retraining DAG runs whether or not anything changed, wasting compute some days and missing a real change that lands between scheduled runs on others. Airflow assets fix this by making the schedule data-aware: define the feature table your training DAG depends on as an asset, and trigger on updates to it.
from airflow.sdk import Asset
feature_table = Asset("s3://ml-features/customer_churn/")
@dag(schedule=[feature_table], catchup=False)
def retrain_on_feature_update():
...
Every time the ingestion DAG from Step 1 writes a new version of that table, this DAG queues automatically. Airflow 3 elevated this from the “Datasets” concept in Airflow 2.x to a first-class Asset, with the same outlet and inlet pattern plus a new @asset decorator; on Airflow 2.x, the equivalent mechanism is available under the Dataset name. Either version beats a blind schedule for a feature whose freshness matters, the same way an internal knowledge assistant depends on how current its source is, not how long the calendar says it has been.
Step 8: Build in idempotency before you run a backfill
Permalink to “Step 8: Build in idempotency before you run a backfill”Airflow reruns tasks constantly: on retry after a failure, on a manual rerun, and on a backfill across a historical date range. A non-idempotent task risks duplicate rows, double-counted features, or a training set that silently teaches a model the wrong distribution.
Three habits fix this. Use the logical date from the DAG context, never the current timestamp, so a run for a historical date always processes that date’s data no matter when it executes. Delete existing rows for that date before inserting, or use a merge, instead of a bare INSERT. Use CREATE TABLE IF NOT EXISTS so a rerun does not fail on an object that already exists.
Backfills carry a second risk: scheduler load. Running a year of a daily DAG at once, with no concurrency cap, can overwhelm the scheduler and any source system the ingestion task queries. Set max_active_runs deliberately.
☐ Every task uses logical_date, not datetime.now(), for its data window
☐ Writes use delete-then-insert or merge, never a bare append
☐ max_active_runs is set before any historical backfill, not after the first one stalls the scheduler
Common pitfalls when orchestrating ML pipelines in Airflow
Permalink to “Common pitfalls when orchestrating ML pipelines in Airflow”Most Airflow ML pipelines that break in production fail for a small set of repeated reasons, invisible without real data observability for AI pipelines watching the tables in between.
Running training on the scheduler’s own workers. Works in a demo, breaks the moment a job needs a GPU or runs past the worker’s timeout. Route it to dedicated compute, as in Step 5.
Non-idempotent writes. A bare INSERT instead of delete-then-insert turns every retry and backfill into a data quality incident.
No evaluation gate. A pipeline that deploys whatever it just trained, unconditionally, will eventually deploy a regression. The fix is the branch task from Step 6.
Retraining on a blind schedule. Time-based retraining wastes compute or misses an update landing off-schedule. Asset triggers, from Step 7, fix both.
No owner on the feature table a DAG produces. When accuracy drops, the first question is “who owns this,” and an undocumented owner turns that into an investigation.
Uncapped backfills. Reprocessing a year of history with no concurrency limit can take down the scheduler for every other DAG running.
How a context layer keeps an orchestrated ML pipeline trustworthy
Permalink to “How a context layer keeps an orchestrated ML pipeline trustworthy”Airflow answers when each stage runs and in what order. It does not answer what a feature means, who owns the table it wrote to, or whether the model serving traffic now was trained against a definition that has since changed; those are business context questions, and they sit above the DAG.
A context layer connects a DAG run back to the metadata for AI it produced and consumed: which source tables fed a run, which feature definitions it used, and which model version is live now, extending Airflow OpenLineage for AI into column-level detail. That turns “the model regressed” into “the model regressed because a source table’s schema changed three tasks upstream.” Different types of metadata for AI agents matter at different points: data lineage for tracing a bad output, ownership for knowing who to page, freshness for judging a retraining trigger. The same pattern used to build a knowledge base for AI agents applies here, against feature tables instead of documents. None of this replaces Airflow; it turns a DAG run history into an answer instead of a log file.
Orchestration schedules the pipeline; it does not explain what the pipeline produced
Permalink to “Orchestration schedules the pipeline; it does not explain what the pipeline produced”Airflow, configured well, will run the DAGs in this guide reliably for years: ingestion on schedule, feature engineering in parallel, training on dedicated compute, retraining triggered by real data changes, and a gate that refuses a regression. That reliability is necessary and it is not the whole job. A DAG history tells you a task ran and succeeded. It does not tell you whether the feature table fifteen tasks upstream was redefined last week, or whether the model serving production traffic was trained against a table since deprecated.
That is the gap between orchestration and trust. The strongest Airflow setups pair a well-built DAG with a context engineering framework that tracks what the outputs mean, not just whether they ran. Getting the DAGs right, as this guide covers, is the infrastructure. Knowing what they produced is what makes it worth trusting.
FAQs about Airflow ML pipeline orchestration
Permalink to “FAQs about Airflow ML pipeline orchestration”1. What is Airflow used for in an ML pipeline?
Permalink to “1. What is Airflow used for in an ML pipeline?”Airflow schedules and sequences the pipeline: ingestion, feature engineering, training, evaluation, and the deployment or retraining trigger. It rarely trains the model itself; it calls out to Kubernetes, Databricks, or a cloud ML service.
2. Do I need Airflow if my ML platform already has a built-in pipeline tool?
Permalink to “2. Do I need Airflow if my ML platform already has a built-in pipeline tool?”Not always. Kubeflow Pipelines and SageMaker Pipelines handle training orchestration well inside their own platform. Teams add Airflow when the pipeline spans systems the ML platform does not own.
3. What is the TaskFlow API and do I have to use it?
Permalink to “3. What is the TaskFlow API and do I have to use it?”A set of decorators, led by @task, that turns a Python function into an Airflow task and passes its return value forward automatically. It is not mandatory, but it removes most boilerplate XCom code.
4. How do I trigger retraining automatically in Airflow?
Permalink to “4. How do I trigger retraining automatically in Airflow?”Define the feature table your training DAG depends on as an Airflow asset, and schedule the DAG on updates to that asset instead of a fixed cron interval.
5. Can Airflow run model training itself, or does it just call out to other systems?
Permalink to “5. Can Airflow run model training itself, or does it just call out to other systems?”It can run small jobs on its own workers, but production pipelines route anything resource-heavy to dedicated compute through KubernetesPodOperator or a Databricks provider.
6. What is the difference between Airflow orchestration and a feature store?
Permalink to “6. What is the difference between Airflow orchestration and a feature store?”Airflow orchestrates when feature engineering runs. A feature store, such as Feast, is where the computed features live, split into offline and online stores.
7. How do I avoid duplicate or corrupted data when backfilling an ML pipeline?
Permalink to “7. How do I avoid duplicate or corrupted data when backfilling an ML pipeline?”Use the logical date instead of the current timestamp, delete existing rows before inserting or use a merge, and cap concurrent backfill runs.
Sources
Permalink to “Sources”- Best practices for orchestrating MLOps pipelines with Airflow, Astronomer
- State of Airflow 2026, Astronomer
- MLOps, Apache Airflow
- Asset Definitions, Apache Airflow Documentation
- Dynamic Task Mapping, Apache Airflow Documentation
- Running Feast in production, Feast Documentation
- 10 Airflow Best Practices, Astronomer
- Backfilling Historical Data With Idempotent Data Pipelines, ml4devs
