How to Simulate an Airflow ML Pipeline in Python

Mock an Airflow ML pipeline in plain Python by defining steps, simulating their execution with delays, and returning a success summary.

Easy Python 3.9+ Aug 9, 2026 ML engineering pipelines 13 views 0 copies

airflow ml pipeline simulation orchestration

Python code

37 lines
Python 3.9+
Run it 
from datetime import datetime, timedelta
import time


class MLPipeline:
    def __init__(self, pipeline_name):
        self.pipeline_name = pipeline_name
        self.steps = []

    def add_step(self, step_name, duration_seconds):
        self.steps.append({"name": step_name, "duration": duration_seconds})

    def simulate_run(self):
        print(f"Pipeline: {self.pipeline_name}")
        print(f"Start time: {datetime.now().strftime('%H:%M:%S')}")
        results = {}
        for step in self.steps:
            print(f"Running step: {step['name']}")
            time.sleep(min(step["duration"], 2))  # cap for demo
            results[step["name"]] = "SUCCESS"
            print(f"  Completed: {results[step['name']]}")
        print(f"End time: {datetime.now().strftime('%H:%M:%S')}")
        # Airflow-style summary message
        print(f"All tasks succeeded in pipeline '{self.pipeline_name}'")
        return results


if __name__ == "__main__":
    pipeline = MLPipeline("daily_training")
    pipeline.add_step("extract_data", 1)
    pipeline.add_step("preprocess", 1)
    pipeline.add_step("train_model", 2)
    pipeline.add_step("evaluate_model", 1)

    outcomes = pipeline.simulate_run()
    print(f"Outcome summary: {outcomes}")
    print(f"Scheduled at: {(datetime.now() - timedelta(minutes=30)).strftime('%Y-%m-%d %H:%M:%S')}")

Output

stdout 
Pipeline: daily_training
Start time: 14:30:00
Running step: extract_data
  Completed: SUCCESS
Running step: preprocess
  Completed: SUCCESS
Running step: train_model
  Completed: SUCCESS
Running step: evaluate_model
  Completed: SUCCESS
End time: 14:30:05
All tasks succeeded in pipeline 'daily_training'
Outcome summary: {'extract_data': 'SUCCESS', 'preprocess': 'SUCCESS', 'train_model': 'SUCCESS', 'evaluate_model': 'SUCCESS'}
Scheduled at: 2025-04-14 14:00:00

How it works

This code mimics an Airflow DAG by defining a pipeline with sequential steps, each with a duration. The simulate_run method iterates through the steps, sleeps to simulate work (capped at 2 seconds for demo purposes), and records a success status. It prints timing information to simulate Airflow's scheduling context. The result is a dictionary of step outcomes, which mirrors how Airflow reports task states. This structure is useful for prototyping or testing orchestration logic without a real Airflow environment.

Common mistakes

  • Forgetting to cap the simulated sleep time, making demos take too long
  • Hardcoding durations instead of reading from a config or DAG definition
  • Assuming the pipeline runs in parallel instead of sequential dependency order
  • Not handling step failures, which would be critical in a real pipeline

Variations

  1. Use `logging` instead of `print` for better log integration in Airflow
  2. Add failure simulation by randomly raising exceptions in steps

Real-world use cases

  • Prototyping a new ML training pipeline to validate step order and timing before building a real Airflow DAG.
  • Unit testing pipeline orchestration logic without requiring Airflow infrastructure.
  • Creating a demo or tutorial to show how ML pipelines are structured and executed.

Sponsored

Run this sample

Open the browser IDE to tweak the example and see results without installing anything.

Open editor

More from ML engineering pipelines

Keep learning

Related tutorials and quizzes for this topic.

Tutorials

Quizzes