Post

Orchestration Tool: Airflow (Section One)

Posted Updated
By
6 min read
Orchestration Tool: Airflow (Section One)

Guide

Table of Contents

Overview

Apache Airflow is a powerful open-source orchestration tool designed to automate, schedule, and monitor complex data and AI workflows. It enables users to define workflows as code, manage dependencies between tasks, and ensure reliable execution with robust error handling and observability. Airflow is widely used for productionizing pipelines that require repeatability, scalability, and integration with diverse data and compute systems.

Why Do We Need an Orchestration Tool?

When developing data or AI workflows, tools like Jupyter notebooks are excellent for experimentation, prototyping, and debugging. However, notebooks are not ideal for production systems because they do not provide reliable automation, scheduling, monitoring, or robust failure handling. This is where orchestration tools become essential.

An orchestration tool such as Apache Airflow converts notebook logic into structured, automated pipelines. Instead of manually executing each step, every logical unit of work (such as reading data, generating embeddings, or storing outputs) is defined as a task. These tasks run with dependency rules, can be scheduled, and can recover from failures.

A simple AI pipeline often includes these steps:

  • Reading text data
  • Generating embeddings
  • Storing vectors in a database

Airflow ensures these steps run in the correct order, retries failed tasks, and gives visibility into each run. It also helps answer production questions such as what to do when tasks fail, whether retries are needed, and how to track outcomes.

Core Concepts in Airflow

Airflow introduces several core concepts that structure pipelines.

1- DAG (Directed Acyclic Graph)

A DAG is the main unit in Airflow. It defines:

  • What tasks to run
  • In what order they run
  • On what schedule they run

Each DAG is written as a Python script. A DAG defines how a pipeline runs over time. It includes a schedule (for example, hourly), a set of tasks, and dependency rules between tasks. Each task contains Python logic, often adapted from notebook prototypes. Airflow executes tasks in the correct order and adds reliability through retries, logging, and run history.

2- Task

A task is one unit of work, often implemented as a Python function. Examples include:

  • Extracting data
  • Transforming text
  • Querying a database

3- Dependencies

Dependencies connect tasks so one task runs only after another has successfully completed.

4- Airflow UI

The Airflow UI lets you:

  • Visualize DAGs
  • Trigger workflows manually
  • Monitor task status
  • Debug failures

Orchestration in GenAI Pipelines

Airflow is especially useful for Generative AI pipelines, including:

  • Retrieval-Augmented Generation (RAG)
  • Model inference pipelines
  • Batch inference
  • Model training and retraining
  • Fine-tuning workflows
  • Streaming or asynchronous inference

For a RAG system, the typical flow is data ingestion, embedding generation, vector storage, and retrieval during inference. Airflow orchestrates these steps with repeatability and operational visibility.

Best Practices for Pipeline Design

1- Atomicity

Each task should perform one clear responsibility. This improves observability, retry precision, and efficiency because you only rerun failed parts.

2- Idempotency

If a task runs multiple times with the same input, it should ideally produce the same output. Full idempotency can be harder in GenAI systems, but it is still important for ingestion and transformation stages.

3- Software Engineering Practices

Airflow code should be treated as production code:

  • Use version control with Git
  • Keep code modular and readable
  • Add tests where practical
  • Use clear naming and comments

Airflow Architecture

Airflow includes multiple components:

  • DAG Processor: Parses DAG files
  • Scheduler: Decides when tasks should run
  • Workers: Execute tasks
  • Metadata Database: Stores DAG run state, task state, and metadata
  • Web Server (UI): Displays pipeline status and logs

Execution flow:

  1. A DAG file is parsed.
  2. The Scheduler checks whether it should run.
  3. Runnable tasks are queued.
  4. Workers execute tasks.
  5. Results and status are stored and displayed in the UI.

Setting Up Airflow Locally

There are two common approaches.

Astro CLI simplifies local setup.

Steps:

  1. Clone the project repository.
  2. Create a .env file with Airflow configs and required API keys.
  3. Run:
1

astro dev start

This starts local services such as Scheduler, Webserver, Metadata DB, and Triggerer. Airflow UI is available at http://localhost:8080.

Method 2: Docker Setup

You can run Airflow manually using Docker containers and docker-compose.

Adding a Vector Database (Weaviate)

To support RAG, use a vector database such as Weaviate.

Run Weaviate:

1

docker run -d -p 8081:8080 semitechnologies/weaviate:latest

Connect in Python:

1
2

import weaviate
client = weaviate.Client("http://localhost:8081")

Project Example: Book Recommendation System (RAG)

Goal: build a system that reads book descriptions, converts them into embeddings, stores them in Weaviate, and returns recommendations from user queries.

Step 1: Read Data

Read book descriptions from text files.

1
2

import os
files = os.listdir("include/data")

Step 2: Extract Book Data

Parse each line into fields like title, author, and description, then store each row as a dictionary.

Step 3: Generate Embeddings

1
2
3
4

from fastembed import TextEmbedding

model = TextEmbedding("BAAI/bge-small-en-v1.5")
embeddings = list(model.embed(["book description"]))

Step 4: Store in Weaviate

Each book object stores both vector and metadata.

1
2
3
4
5
6
7
8

collection.data.insert({
    "vector": embedding,
    "properties": {
        "title": title,
        "author": author,
        "description": description
    }
})

Step 5: Query for Recommendations

1
2
3
4

results = collection.query.near_vector(
    near_vector=query_embedding,
    limit=1
)

Converting Prototype to an Airflow Pipeline

After validating prototype logic, automate it with DAGs.

DAG 1: fetch_data

Purpose:

  • Read new book data
  • Generate embeddings
  • Store vectors in Weaviate

Time-based schedule example:

1
2
3

    start_date=datetime(2025, 4, 1),
    schedule="@hourly"
)

DAG 2: query_data

Purpose:

  • Accept user query
  • Retrieve best matches
  • Return recommendation payload

Data-Aware Scheduling (Advanced)

Instead of running only on a clock, use data-aware scheduling with Airflow assets.

In fetch_data:

1

@task(outlets=[Asset("my_book_vector_data")])

In query_data:

1

@dag(schedule=[Asset("my_book_vector_data")])

Now the query DAG runs when upstream vector data is updated.

Dynamic Task Mapping (Parallel Processing)

Dynamic task mapping allows each file to be processed in parallel.

1
2
3

_transform = transform.expand(
    book_description_file=list_files
)

Benefits:

  • Faster execution
  • Better fault isolation
  • Improved scalability

Handling Failures

Use retries and retry delays in DAG configuration:

1
2
3
4
5
6

@dag(
    default_args={
        "retries": 1,
        "retry_delay": timedelta(seconds=10)
    }
)

If a task fails, Airflow retries automatically according to policy, reducing manual intervention.

Custom Docker Image for Airflow

To include project-specific dependencies, define them in requirements.txt and build a custom image.

requirements.txt:

weaviate-client
pandas
numpy

Dockerfile:

1
2
3
4

FROM apache/airflow:2.8.3

COPY requirements.txt /requirements.txt
RUN pip install -r /requirements.txt

Final Architecture Summary

A complete RAG recommendation system with Airflow can work as follows:

fetch_data DAG (hourly)

  • Reads book files
  • Generates embeddings
  • Stores vectors in Weaviate

query_data DAG (event-based)

  • Triggered when vector data updates
  • Runs semantic search
  • Returns recommendation results

This architecture combines:

  • Automation
  • Scalability
  • Fault tolerance
  • Near real-time responsiveness

Quick Notes:

Note 1: Delete dags and restart them in airflow using powershell

1
2
3

docker compose exec airflow-worker airflow dags delete fetch_data_new --yes

docker compose exec airflow-worker airflow dags reserialize
Cloud AI and MLOPs, Data Workflow Orchestration & Automation
ML Python Airflow RAG Weaviate FastAPI Recommendation Vector Search Docker Data Engineering Orchestration
This post is licensed under CC BY 4.0 by the author.