In this project, you will build and orchestrate a machine learning (ML) workflow using Apache Airflow. The pipeline will:

  1. Extract and preprocess training data.

  2. Train a machine learning model.

  3. Evaluate the model’s performance.

  4. Deploy the trained model (store in file system or cloud).

This project demonstrates how Airflow can orchestrate not just ETL pipelines, but also end-to-end ML workflows.

Learning Objectives

By the end of this project, you will be able to:

  • Orchestrate ML pipelines in Airflow.

  • Automate data preprocessing, training, and evaluation tasks.

  • Use XComs to pass model artifacts between tasks.

  • Deploy trained models for downstream usage.

Project Requirements

  • Apache Airflow installed (Docker or local).

  • Python libraries: pandas, scikit-learn, joblib.

  • Example dataset: Iris dataset (from scikit-learn).

  • Storage for trained models (local or cloud).

Project Architecture

Pipeline Steps (DAG):

  1. Extract & Preprocess → Load dataset, clean, and split into train/test.

  2. Train Model → Train a classifier (e.g., Logistic Regression).

  3. Evaluate Model → Compute accuracy, precision, recall.

  4. Deploy Model → Save the trained model for production use.

Visually:

 
Extract & Preprocess → Train Model → Evaluate Model → Deploy Model

DAG Implementation

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
import pandas as pd
import joblib
import os

# Step 1: Extract & Preprocess
def extract_preprocess(**context):
    iris = load_iris(as_frame=True)
    df = iris.frame
    X = df.drop("target", axis=1)
    y = df["target"]

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    context['ti'].xcom_push(key='X_train', value=X_train.to_json())
    context['ti'].xcom_push(key='X_test', value=X_test.to_json())
    context['ti'].xcom_push(key='y_train', value=y_train.to_json())
    context['ti'].xcom_push(key='y_test', value=y_test.to_json())

# Step 2: Train Model
def train_model(**context):
    X_train = pd.read_json(context['ti'].xcom_pull(key='X_train', task_ids='extract_task'))
    y_train = pd.read_json(context['ti'].xcom_pull(key='y_train', task_ids='extract_task'))

    model = LogisticRegression(max_iter=200)
    model.fit(X_train, y_train)

    model_path = "/tmp/iris_model.pkl"
    joblib.dump(model, model_path)
    context['ti'].xcom_push(key='model_path', value=model_path)

# Step 3: Evaluate Model
def evaluate_model(**context):
    model_path = context['ti'].xcom_pull(key='model_path', task_ids='train_task')
    model = joblib.load(model_path)

    X_test = pd.read_json(context['ti'].xcom_pull(key='X_test', task_ids='extract_task'))
    y_test = pd.read_json(context['ti'].xcom_pull(key='y_test', task_ids='extract_task'))

    predictions = model.predict(X_test)
    accuracy = accuracy_score(y_test, predictions)
    print(f"Model Accuracy: {accuracy:.2f}")

# Step 4: Deploy Model
def deploy_model(**context):
    model_path = context['ti'].xcom_pull(key='model_path', task_ids='train_task')
    final_path = "/opt/airflow/models/iris_model.pkl"
    os.makedirs(os.path.dirname(final_path), exist_ok=True)
    os.replace(model_path, final_path)
    print(f"Model deployed to {final_path}")

# DAG definition
with DAG(
    dag_id="ml_workflow_dag",
    description="End-to-end ML pipeline with Airflow",
    start_date=datetime(2025, 1, 1),
    schedule_interval="@daily",
    catchup=False
) as dag:

    extract_task = PythonOperator(
        task_id="extract_task",
        python_callable=extract_preprocess,
        provide_context=True
    )

    train_task = PythonOperator(
        task_id="train_task",
        python_callable=train_model,
        provide_context=True
    )

    evaluate_task = PythonOperator(
        task_id="evaluate_task",
        python_callable=evaluate_model,
        provide_context=True
    )

    deploy_task = PythonOperator(
        task_id="deploy_task",
        python_callable=deploy_model,
        provide_context=True
    )

    extract_task >> train_task >> evaluate_task >> deploy_task

Key Notes

  • Extract → Loads Iris dataset, splits into training/testing sets.

  • Train → Fits a Logistic Regression model.

  • Evaluate → Prints accuracy in Airflow logs.

  • Deploy → Saves model artifact (.pkl) to a models folder.

Expected Output

  • Logs showing model accuracy (e.g., Model Accuracy: 0.97).

  • A deployed model file at /opt/airflow/models/iris_model.pkl.

  • A DAG graph: Extract → Train → Evaluate → Deploy.

Project Extensions (Optional)

  • Add hyperparameter tuning (GridSearchCV).

  • Log metrics to a database or MLflow.

  • Deploy the model to S3, GCS, or a model registry.

  • Add a BranchPythonOperator: retrain model only if accuracy < threshold.

Deliverables

  1. ml_workflow_dag.py DAG file in Airflow dags/ folder.

  2. A deployed model file in your environment.

  3. Screenshots of Graph View, Logs, and Deployed Model.

✨ Congratulations! You’ve completed all three real-world projects in this course. You now know how to build ETL pipelines, warehouse loaders, and ML workflows with Apache Airflow.

Previous