Data Pipeline Orchestration for ML:...

Three Orchestrators Walk Into a Production Cluster

ML data pipelines are more demanding than typical ETL jobs. They need to handle large datasets, coordinate with training infrastructure, manage feature computation across batch and streaming contexts, and produce reproducible outputs. General-purpose workflow orchestrators can do this, but some do it much better than others.

I've deployed production ML pipelines on Airflow, Prefect, and Dagster over the past three years. Each has genuine strengths and genuine pain points that only emerge after you've been running the thing for six months. Here's what I've learned, with less diplomatic hedging than the vendor comparison posts usually offer.

Airflow: The Establishment Pick

Apache Airflow has been the default choice for ML pipeline orchestration since before "MLOps" was a word. It's battle-tested, has a massive ecosystem of operators and providers, and every cloud provider offers a managed version.

Airflow's DAG model maps well to ML workflows. A typical training pipeline looks like: data extraction, validation, feature engineering, training, evaluation, registration. Each step is a task, dependencies are explicit, and the scheduler handles retries and backfills.

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

default_args = {
 "retries": 2,
 "retry_delay": timedelta(minutes=5),
 "execution_timeout": timedelta(hours=4),
}

with DAG(
 "fraud_model_training",
 default_args=default_args,
 schedule_interval="0 4 * * 1",
 start_date=datetime(2024, 1, 1),
 catchup=False,
 tags=["ml", "fraud", "training"],
) as dag:

 extract = PythonOperator(
 task_id="extract_training_data",
 python_callable=extract_training_data,
 op_kwargs={
 "date_range": "{{ ds }}",
 "output_path": "s3://data/fraud/raw/",
 },
 )

 validate = PythonOperator(
 task_id="validate_data",
 python_callable=run_data_validation,
 )

 train = PythonOperator(
 task_id="train_model",
 python_callable=train_fraud_model,
 executor_config={"KubernetesExecutor": {"gpu": "1"}},
 )

 extract >> validate >> train

Where Airflow struggles is with dynamic workflows and local development. DAGs are parsed at import time, which means dynamic task generation based on runtime data requires workarounds. Testing a DAG locally means setting up a local Airflow instance or using the somewhat finicky unit test helpers. And the scheduler's performance degrades noticeably past a few hundred DAGs.

Prefect: The Developer-Friendly Option

Prefect's pitch is that orchestration should feel like writing Python, not configuring XML. Flows are Python functions decorated with @flow, tasks are functions decorated with @task, and the dependency graph is inferred from the function call order. It's genuinely pleasant to write.

from prefect import flow, task
from prefect.tasks import task_input_hash
from datetime import timedelta

@task(retries=2, cache_key_fn=task_input_hash,
 cache_expiration=timedelta(hours=24))
def extract_data(date_range: str, output_path: str) -> str:
 return run_extraction(date_range, output_path)

@task(retries=1)
def validate(data_path: str) -> dict:
 report = run_validation(data_path)
 if report.critical_failures:
 raise ValueError(f"Data validation failed: {report.critical_failures}")
 return report.summary

@task(tags=["gpu"])
def train(features_path: str, config: dict) -> str:
 model_path = run_training(features_path, config)
 return model_path

@flow(name="fraud-model-training", log_prints=True)
def training_pipeline(date: str):
 raw_path = extract_data(date, f"s3://data/fraud/raw/{date}/")
 validation = validate(raw_path)
 features_path = compute_features(raw_path)
 model = train(features_path, {"epochs": 50, "lr": 0.001})
 register_model(model, validation)

The task caching in Prefect is excellent for ML workflows. Feature engineering on a large dataset takes hours; if the training step fails and you rerun the pipeline, the cached feature engineering result means you skip straight to training. Airflow can do this with XComs and custom caching logic, but Prefect makes it native.

The downside is that Prefect's infrastructure layer — the work pools, deployments, and agent model — has changed significantly between major versions. If you started on Prefect 1.x, the migration to 2.x was painful. The API is now stable, but the history of breaking changes makes some teams nervous about long-term commitment.

Dagster: The Data-Aware Orchestrator

Dagster takes a fundamentally different approach. Instead of defining tasks and their dependencies, you define assets — the datasets, models, and artifacts your pipeline produces — and Dagster figures out what needs to recompute when inputs change.

from dagster import asset, DailyPartitionsDefinition
import pandas as pd

partitions = DailyPartitionsDefinition(start_date="2024-01-01")

@asset(partitions_def=partitions)
def raw_transaction_data(context) -> pd.DataFrame:
 date = context.partition_key
 return extract_transactions(date)

@asset(partitions_def=partitions)
def validated_transactions(raw_transaction_data: pd.DataFrame) -> pd.DataFrame:
 report = validate(raw_transaction_data)
 return raw_transaction_data[report.valid_mask]

@asset(partitions_def=partitions)
def fraud_features(validated_transactions: pd.DataFrame) -> pd.DataFrame:
 return compute_fraud_features(validated_transactions)

@asset
def fraud_model(fraud_features: pd.DataFrame) -> dict:
 model = train_model(fraud_features)
 return {"path": model.save(), "metrics": model.evaluate()}

The asset-centric model maps naturally to ML workflows where you care about what data exists and whether it's fresh. The Dagster UI shows a data lineage graph — not just task dependencies, but the actual flow of data through your system. When something goes wrong, you can trace it from the broken model back to the specific data partition that caused the issue.

Dagster's type system and I/O managers are a genuine differentiator. You can define typed inputs and outputs for your assets, and Dagster validates the types at runtime. I/O managers abstract where data is stored (local file, S3, BigQuery), so your pipeline logic doesn't change when you switch storage backends.

We covered a related topic in Document Chunking Strategies for Vector Search Quality.

Which One Should You Pick

If you're already on Airflow and it's working, stay on Airflow. Migration costs are real, and Airflow's ecosystem is unmatched. If you're starting fresh with a small team and value developer experience, Prefect gets you productive fastest. If you're building a data platform where lineage and data quality matter as much as job scheduling, Dagster's asset model is worth the steeper learning curve.

For ML specifically, I'd lean toward Dagster for new projects. The asset model aligns with how ML engineers think about their work — in terms of datasets, features, and models, not tasks and schedules. The type system catches data issues at pipeline boundaries instead of deep in training code. And the partitioned asset support makes backfill and data versioning first-class operations instead of afterthoughts.

One more thing worth mentioning: all three can run on Kubernetes, all three support distributed execution, and all three have commercial managed offerings. The infrastructure requirements are similar enough that they shouldn't be the deciding factor. Pick based on the programming model that fits your team's workflow and the features that solve your specific pain points.

Error Handling and Recovery Patterns

ML pipeline failures have different characteristics than web service failures. A training job that fails after 6 hours of GPU computation is much more expensive to retry than a web request that fails after 200 milliseconds. Smart retry and recovery strategies save both time and money.

Checkpointing is the foundation. Every training job should save checkpoints at regular intervals (we use every 30 minutes or every 1000 steps, whichever comes first). When a job fails and restarts, it resumes from the last checkpoint instead of starting over. The orchestrator needs to support this — Dagster's op retries pass the retry context to the op function, Airflow's retry mechanism restarts the entire task, and Prefect's task retries can be configured to reuse cached intermediate results.

For a related perspective, see Object Detection Model Selection: YOLO vs DETR vs EfficientD.

# Checkpoint-aware training with Prefect retry
from prefect import task
from prefect.tasks import task_input_hash

@task(retries=3, retry_delay_seconds=60)
def train_with_checkpoint(config: dict, checkpoint_dir: str) -> str:
 # Check for existing checkpoint
 latest_ckpt = find_latest_checkpoint(checkpoint_dir)
 if latest_ckpt:
 print(f"Resuming from checkpoint: {latest_ckpt}")
 model = load_checkpoint(latest_ckpt)
 start_epoch = model.current_epoch
 else:
 model = initialize_model(config)
 start_epoch = 0

 for epoch in range(start_epoch, config["total_epochs"]):
 train_one_epoch(model, epoch)
 if epoch % 5 == 0:
 save_checkpoint(model, checkpoint_dir, epoch)

 final_path = save_final_model(model, config["output_dir"])
 return final_path

For data processing stages (feature engineering, data validation), idempotency is more important than checkpointing. If the pipeline reruns a data processing step, the output should be identical to the first run. This means using deterministic operations, avoiding append-mode writes (use overwrite or atomic rename instead), and ensuring that random seeds are fixed.

Pipeline Observability Beyond Logs

Logs tell you what happened. Metrics tell you whether it's normal. For ML pipelines, you need both, plus data-specific observability that tells you whether the pipeline's outputs are correct.

We instrument our pipelines with three categories of metrics. Operational metrics (task duration, memory usage, retry count) tell you how the pipeline is performing. Data metrics (row counts, null rates, value distributions) tell you what the pipeline is processing. Quality metrics (model accuracy on validation set, feature importance rankings) tell you whether the pipeline is producing good outputs.

All three categories flow to the same monitoring system (Prometheus/Grafana in our case), and we've dashboards that show the three views side by side. When a quality metric drops, we look at the data metrics to see if the input changed, then at the operational metrics to see if the pipeline itself had issues. This three-layer view usually narrows the root cause to one or two hypotheses within minutes, instead of the multi-hour investigation that log-only observability requires.

Data Pipeline Orchestration for ML: Airflow, Prefect, and Dagster Compared