Building an Automated ML Retraining Pipeline with Apache Airflow

Welcome to the fourth lesson of our "Automating Retraining with Apache Airflow" course! You've made fantastic progress — after learning the basics of Airflow, designing pipeline structures, and testing workflows, you're now ready to bring everything together. In this lesson, we'll focus on building a complete, automated ML retraining pipeline using Airflow and the modular pipeline components we implemented in Course 1. This is where all your previous work pays off: you'll see how to orchestrate the entire machine learning process, from data loading to model evaluation and saving, in a single, scheduled workflow.

By the end of this lesson, you'll be able to connect all the pieces into a robust pipeline that retrains and evaluates your model automatically. Let's dive in and see how it all fits together!

To build our Airflow pipeline, we'll use the functions from our data, model, and evaluation modules. Here's a quick refresher on the interfaces we built previously:

• load_diamonds_data(file_path): Loads the diamonds dataset from a CSV file and returns a DataFrame.
• preprocess_diamonds_data(df): Preprocesses the data, splits it into training and test sets, and returns the processed features, targets, and preprocessor.
• train_model(X_train, y_train, model_type, **params): Trains a machine learning model (like a Random Forest) and returns the trained model.
• evaluate_model(model, X_test, y_test): Evaluates the model's performance and returns a dictionary of metrics, as well as the model's predictions in the test set.
• save_model(model, preprocessor, model_dir, model_name, metadata): Saves the model, preprocessor, and metadata to disk.

These functions are the building blocks of our pipeline. By chaining them together as Airflow tasks, we can automate the entire retraining process.

Let's start by setting up the structure of our Airflow DAG. The DAG defines the workflow and the order in which tasks are executed. Here's how we set up the basic configuration and constants:

from datetime import datetime, timedelta
import os
from airflow.decorators import dag, task

# Import our custom ML modules
from data import load_diamonds_data, preprocess_diamonds_data
from model import train_model, save_model
from evaluation import evaluate_model

# Set default arguments for the DAG
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

# Define paths for data and model storage
DATA_PATH = "diamonds.csv"
MODEL_DIR = "./saved_models"
os.makedirs(MODEL_DIR, exist_ok=True)  # Ensure the model directory exists

Here, we import the necessary modules and set up the DAG's default arguments, such as retry behavior and notification preferences. We also define constants for the data file and model storage directory. This setup ensures that our pipeline has a consistent environment every time it runs.

Now, let's define the pipeline steps as tasks and chain them together.

A key design choice here is to pass complex objects (like DataFrames and models) between tasks. Airflow uses XComs for this, and by default, it serializes objects using pickle. While this is convenient for prototyping and small pipelines, using pickle in production can be risky: it may introduce security vulnerabilities and compatibility issues, especially if your code or dependencies change over time. For production pipelines, consider using more robust serialization formats (like Parquet for DataFrames, or storing models in object storage) and/or passing only references (e.g., file paths) between tasks.

With that in mind, here's how we define the first two tasks: loading and processing the data. We'll stick to pickle-based XCom for simplicity:

@dag(
    dag_id='mlops_pipeline',
    description='ML model training pipeline for diamond price prediction',
    default_args=default_args,
    schedule='@daily',
    start_date=datetime(2025, 1, 1),
    catchup=False,
    tags=['ml', 'training', 'diamonds'],
)
def diamond_training_pipeline():
    @task(task_id="load_data")
    def task_load_data():
        # Load the diamonds dataset from CSV
        print(f"Loading diamonds data from {DATA_PATH}")
        df = load_diamonds_data(DATA_PATH)
        print(f"Dataset loaded with shape: {df.shape}")
        return df

    @task(task_id="process_data")
    def task_process_data(df):
        # Preprocess the data and split into train/test sets
        print("Preprocessing data")
        X_train, X_test, y_train, y_test, preprocessor = preprocess_diamonds_data(df)
        print(f"Data preprocessing complete. Training set size: {X_train.shape[0]}")
        return {
            "X_train": X_train,
            "X_test": X_test,
            "y_train": y_train,
            "y_test": y_test,
            "preprocessor": preprocessor
        }

• Each task is focused on a single responsibility, making the pipeline modular and easy to maintain.
• The output of each task is structured as a dictionary, which makes it clear what data is being passed downstream.
• As noted above, passing large objects via XComs is fine for learning and small-scale use, but be cautious in production.

With the data prepared, the next step is to train the model. This task consumes the processed data and outputs the trained model along with the data needed for evaluation.

    @task(task_id="train_model")
    def task_train_model(processed_data):
        # Train a Random Forest model using the training set
        print("Training model")
        model = train_model(
            processed_data["X_train"],
            processed_data["y_train"],
            model_type="random_forest",
            n_estimators=10,
            max_depth=5,
            random_state=42
        )
        print("Model training complete.")
        return {"model": model, "processed_data": processed_data}

• The model type and hyperparameters are hardcoded for demonstration, but could be parameterized for flexibility.
• By returning both the model and the processed data, we ensure the next task has everything it needs for evaluation.

After training, we need to evaluate the model's performance on the test set. This step is crucial for monitoring model quality and detecting issues early.

    @task(task_id="evaluate_model")
    def task_evaluate_model(training_result):
        # Evaluate the trained model on the test set
        model = training_result["model"]
        processed_data = training_result["processed_data"]
        print("Evaluating model")
        metrics, _ = evaluate_model(model, processed_data["X_test"], processed_data["y_test"])
        print(f"Model evaluation metrics: {metrics}")
        return {
            "model": model,
            "preprocessor": processed_data["preprocessor"],
            "metrics": metrics
        }

• Evaluation metrics are returned as part of the output, making it easy to log or trigger alerts if performance drops.
• The preprocessor is included for saving alongside the model, ensuring reproducibility.

The final step is to persist the trained model, preprocessor, and metadata. This ensures that each retrained model is versioned and traceable.

    @task(task_id="save_model")
    def task_save_model(evaluation_result):
        # Save the trained model, preprocessor, and metadata
        model = evaluation_result["model"]
        preprocessor = evaluation_result["preprocessor"]
        metrics = evaluation_result["metrics"]
        print(f"Saving model to {MODEL_DIR}")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        model_name = f"diamond_price_model_{timestamp}"
        metadata = {
            "metrics": metrics,
            "model_type": "RandomForestRegressor",
            "timestamp": timestamp
        }
        model_path = save_model(
            model, 
            preprocessor,
            MODEL_DIR, 
            model_name, 
            metadata
        )
        print(f"Model saved to {model_path}")
        return model_path

• Each model is saved with a unique timestamp, supporting model versioning and traceability.
• Metadata (including metrics and model type) is stored alongside the model, which is essential for auditing and future analysis.

To connect all the tasks, we simply pass the output of one function as the input to the next. This chaining defines the dependencies and execution order for the entire pipeline.

    df = task_load_data()
    processed_data = task_process_data(df)
    training_result = task_train_model(processed_data)
    evaluation_result = task_evaluate_model(training_result)
    final_model_path = task_save_model(evaluation_result)
    
diamond_training_pipeline()

• The pipeline is linear and easy to follow, with each step depending on the output of the previous one.
• This design makes it straightforward to add new steps (e.g., model monitoring, notification) or swap out components as needed.

You've now seen how to bring together all the essential components of an automated ML retraining pipeline using Airflow. By defining each step as a task and connecting them in a clear sequence, you've built a workflow that is both powerful and maintainable.

As you move forward, you'll have the opportunity to put these concepts into practice and deepen your understanding. Keep exploring, and you'll soon be able to design and automate even more advanced machine learning workflows!