The Task SDK Boundary

Concepts covered: Scheduler and executors

A line through the middle of Airflow

Airflow 3 draws a line between the parts that own the metadata database and the parts that run your code. The scheduler and the DAG processor hold database credentials, and so does the API server; they sit on one side. Workers sit on the other with none. A worker's only route to Airflow's state is HTTP to the API server, through a set of endpoints called the Execution API, and every request carries a token issued for 1 task instance and 1 try. The Task SDK, the airflow.sdk package, is the code that lives on the worker side: the names a DAG file imports, plus the machinery that runs a task and talks to the API on its behalf.

What a task process is handed

When a worker takes a task, it starts a supervisor process, and the supervisor starts the task itself in a child process. The supervisor asks the API server for the task's details, sends the heartbeat that tells the scheduler the task is alive, and relays any request the task makes. Each line the task prints passes to the supervisor, which writes it to the task's log file on the worker and, when the deployment has a remote log store, uploads that file there as the try ends; the UI's log view asks the API server, which fetches the lines from the worker or the remote store. The task sees a context dict whose entries are small objects built from the API's reply rather than database rows. ti.task_id and ti.map_index are attributes read from that reply; there's no session behind ti, and a call that assumes one fails.

from airflow.sdk import get_current_context, task


@task
def describe():
    ctx = get_current_context()
    ti = ctx["ti"]
    print(ctx["run_id"], ti.task_id, ti.map_index)

The next file is what a task written for an older deployment often looks like. It opens a session and counts runs. On Airflow 3 the call to create_session raises RuntimeError with a message that direct database access from a task is not allowed, before any query is sent, because the worker has no connection string to send it on.

from airflow.models import DagRun
from airflow.sdk import task
from airflow.utils.session import create_session


@task
def count_runs():
    with create_session() as session:
        return session.query(DagRun).count()

Losing the session does not leave a task blind to the rest of its run. The Execution API answers a short list of read requests, and the supervisor relays them: the value an upstream task returned, the number of task instances that match a filter through ti.get_ti_count, and their states through ti.get_task_states. The API server runs each lookup against the database and sends back only the result. A task can ask these questions; it can't write its own query, and it cannot change the state of any task instance but its own.

Reachable from a task

• The context dict and its objects
• Its own return value, which flows downstream
• Values upstream tasks returned, relayed by the supervisor
• Counts and states of other task instances, asked through the API
• The task log
• Whatever the worker host has on disk or can reach over its network

Kept away

• The metadata database and every ORM model
• A database session, even one a helper opens for you
• A query the Execution API doesn't offer
• The state of any task instance but its own, to change
• A token for any task instance but its own

Credentials, versioning and network reach

A worker runs whatever Python the DAG authors wrote, and a process that runs arbitrary code should not hold the credentials to a database that stores every DAG's state; the token scoped to 1 task instance is the smallest grant that still lets the task report back. airflow.sdk is also released as its own package, apache-airflow-task-sdk, with its own version number, so an operator written against it keeps working across core upgrades, and the database schema can change without any task noticing. And because a worker needs only HTTP to the API server, it can run on a host that cannot reach the database at all.

• dag and DAG, task and task_group, TaskGroup: the shape of a file
• BaseOperator: the class a custom operator extends
• get_current_context: the context dict from inside a function
• Everything else a task needs arrives through the supervisor

Do

• Take dag and task from airflow.sdk, along with DAG and BaseOperator
• Read run_id from the context dict; task and ti.map_index are there too
• Pass what a downstream task needs as a return value
• Keep libraries a task imports installed on every worker

Avoid

• Import from airflow.models or airflow.utils.session inside a task
• Open a session to inspect a sibling task's state
• Assume ti is backed by a database row
• Reuse one task's token or context in another

About this lesson

This section is part of the DAG Anatomy: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.