Apache Airflow for ML orchestration in MLOps - Commands & Configuration

Introduction

Apache Airflow helps you automate and schedule tasks in machine learning projects. It solves the problem of running complex workflows step-by-step without manual work.

When you want to run data preprocessing, model training, and evaluation in order automatically.

When you need to retry failed ML tasks without starting everything over.

When you want to track the order and timing of ML pipeline steps.

When you want to run ML workflows on a schedule, like daily model retraining.

When you want to visualize the progress and status of your ML tasks.

Config File - my_dag.py

my_dag.py

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

def preprocess():
    print('Data preprocessing step')

def train_model():
    print('Model training step')

def evaluate():
    print('Model evaluation step')

def notify():
    print('Notify completion')

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

dag = DAG(
    'ml_pipeline',
    default_args=default_args,
    description='Simple ML pipeline',
    schedule_interval=timedelta(days=1),
    catchup=False,
)

preprocess_task = PythonOperator(
    task_id='preprocess',
    python_callable=preprocess,
    dag=dag,
)

train_task = PythonOperator(
    task_id='train_model',
    python_callable=train_model,
    dag=dag,
)

evaluate_task = PythonOperator(
    task_id='evaluate',
    python_callable=evaluate,
    dag=dag,
)

notify_task = PythonOperator(
    task_id='notify',
    python_callable=notify,
    dag=dag,
)

preprocess_task >> train_task >> evaluate_task >> notify_task

This file defines a Directed Acyclic Graph (DAG) named 'ml_pipeline' that runs daily starting from January 1, 2024.

Each PythonOperator runs a Python function representing a step: preprocess, train_model, evaluate, and notify.

The tasks are linked in order so they run one after another.

Retries are set to 1 with a 5-minute delay if a task fails.

Commands

This command lists all the DAGs Airflow knows about, to check if your ML pipeline DAG is recognized.

Terminal

airflow dags list

Expected Output

dag_id    | owner   | paused

ml_pipeline | airflow | False

This command manually starts the ML pipeline DAG to run the workflow immediately.

Terminal

airflow dags trigger ml_pipeline

Expected Output

Created <DagRun ml_pipeline @ 2024-06-01T12:00:00+00:00: manual__2024-06-01T12:00:00+00:00, externally triggered: True>

This command shows all the tasks defined in the ml_pipeline DAG to verify the steps.

Terminal

airflow tasks list ml_pipeline

Expected Output

preprocess

train_model

evaluate

notify

This command runs the 'preprocess' task for the given date without scheduling the whole DAG, useful for testing a single step.

Terminal

airflow tasks test ml_pipeline preprocess 2024-06-01

Expected Output

[2024-06-01 12:00:00,000] {taskinstance.py:1234} INFO - Running preprocess

Data preprocessing step

[2024-06-01 12:00:01,000] {taskinstance.py:1234} INFO - Task preprocess succeeded

Key Concept

If you remember nothing else from this pattern, remember: Airflow runs your ML steps in order automatically and lets you retry and schedule them easily.

Common Mistakes

Not setting start_date or setting it to a future date

Airflow will not run the DAG because it thinks the start date is not reached yet.

Set start_date to a past date or today to allow the DAG to run immediately.

Not linking tasks with >> or set_upstream/set_downstream

Tasks run independently and not in the intended order, breaking the workflow.

Use >> to chain tasks in the correct sequence.

Running airflow tasks test without specifying the correct execution date

The task may not run or run with wrong context, causing confusion.

Always provide a valid execution date matching your DAG schedule.

Summary

Define your ML workflow steps as Python functions and link them in a DAG file.

Use airflow commands to list, trigger, and test your ML pipeline tasks.

Set start_date and retries in default_args to control scheduling and failure handling.

Practice

1. What is the main purpose of Apache Airflow in ML orchestration?

easy

A. To store large datasets for ML training

B. To write ML model code in Python

C. To visualize ML model performance metrics

D. To automate and schedule ML workflows as directed tasks

Solution

1. Step 1: Understand Airflow's role

   Apache Airflow is designed to automate workflows by scheduling and running tasks in order.

2. Step 2: Differentiate from other ML tools

   It does not store data, visualize metrics, or write model code but manages task execution.

3. Final Answer:

   To automate and schedule ML workflows as directed tasks -> Option D

4. Quick Check:

   Airflow = workflow automation [OK]

Hint: Airflow schedules tasks, not data or model code [OK]

Common Mistakes:

• Confusing Airflow with data storage tools
• Thinking Airflow writes ML model code
• Assuming Airflow visualizes model metrics

2. Which of the following is the correct way to define a DAG in Apache Airflow using Python?

easy

A. dag = DAG('my_dag', run_every='daily')

B. dag = DAG('my_dag', schedule_interval='@daily')

C. dag = DAG('my_dag', interval='daily')

D. dag = DAG('my_dag', schedule='daily')

Solution

1. Step 1: Recall DAG initialization syntax

   The correct parameter to set schedule is schedule_interval, not run_every, interval, or schedule.

2. Step 2: Verify the example

   dag = DAG('my_dag', schedule_interval='@daily') is the standard syntax to schedule daily runs.

3. Final Answer:

   dag = DAG('my_dag', schedule_interval='@daily') -> Option B

4. Quick Check:

   Use schedule_interval to set DAG timing [OK]

Hint: Use schedule_interval to set DAG timing [OK]

Common Mistakes:

• Using incorrect parameter names like run_every
• Confusing schedule_interval with schedule
• Forgetting to use quotes around '@daily'

3. Given the following Airflow DAG snippet, what will be the order of task execution?

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

def task_a():
    print('Task A')

def task_b():
    print('Task B')

def task_c():
    print('Task C')

dag = DAG('example_dag', start_date=datetime(2024, 1, 1), schedule_interval='@once')

t1 = PythonOperator(task_id='a', python_callable=task_a, dag=dag)
t2 = PythonOperator(task_id='b', python_callable=task_b, dag=dag)
t3 = PythonOperator(task_id='c', python_callable=task_c, dag=dag)

t1 >> t2 >> t3

medium

A. Task A, then Task B, then Task C

B. Task C, then Task B, then Task A

C. Task A, Task B, and Task C run in parallel

D. Task B, then Task A, then Task C

Solution

1. Step 1: Understand task dependencies

   The operator chaining t1 >> t2 >> t3 means t1 runs first, then t2, then t3.

2. Step 2: Confirm execution order

   Tasks print in order: Task A, Task B, Task C.

3. Final Answer:

   Task A, then Task B, then Task C -> Option A

4. Quick Check:

   Operator chaining sets task order [OK]

Hint: >> means run left task before right task [OK]

Common Mistakes:

• Assuming tasks run in parallel without dependencies
• Misreading the >> operator order
• Confusing task IDs with execution order

4. You wrote this Airflow DAG code but get an error: TypeError: DAG.__init__() got an unexpected keyword argument 'start'
What is the likely cause?

dag = DAG('my_dag', start='2024-01-01', schedule_interval='@daily')

medium

A. The parameter should be start_date, not start

B. The schedule_interval value '@daily' is invalid

C. DAG name cannot be 'my_dag'

D. Missing import for datetime module

Solution

1. Step 1: Identify incorrect parameter

   The error says start is unexpected; Airflow expects start_date.

2. Step 2: Confirm correct parameter usage

   Replacing start with start_date fixes the error.

3. Final Answer:

   The parameter should be start_date, not start -> Option A

4. Quick Check:

   Use start_date, not start [OK]

Hint: Use start_date, not start, for DAG start time [OK]

Common Mistakes:

• Using 'start' instead of 'start_date'
• Assuming '@daily' is invalid schedule
• Ignoring error message details

5. You want to create an Airflow DAG that runs an ML training task only if data preprocessing succeeded. Which Airflow feature should you use to enforce this dependency?

hard

A. Schedule both tasks to run at the same time

B. Use Airflow Variables to store task status

C. Set task dependencies using >> operator between preprocessing and training tasks

D. Write a single Python function combining both tasks

Solution

1. Step 1: Understand task dependency in Airflow

   Airflow uses task dependencies to control execution order, ensuring one task runs after another succeeds.

2. Step 2: Apply dependency operator

   Using the >> operator sets the training task to run only after preprocessing completes successfully.

3. Final Answer:

   Set task dependencies using >> operator between preprocessing and training tasks -> Option C

4. Quick Check:

   Use >> to enforce task order [OK]

Hint: Use >> to link tasks in order [OK]

Common Mistakes:

• Thinking Variables control task order
• Scheduling tasks simultaneously without dependencies
• Combining tasks loses modularity and control