Performance: Eager load DagModel relationship on TaskInstance queries to avoid N+1 query storm

[rajat315315]

Description

Description

During the Airflow scheduling loop, when evaluating task instance scheduling decisions for active DagRuns (inside DagRun.update_state() -> _get_ready_tis()), the scheduler checks various task dependencies by calling are_dependencies_met().

One of the standard checks is DagUnpausedDep which verifies if the parent DAG is paused or not. Previously, DagUnpausedDep performed an ad-hoc query directly on the database:

return session.scalar(select(DagModel.is_paused).where(DagModel.dag_id == dag_id))

This triggers an N+1 queries storm. For every task instance being checked for execution, a separate database round-trip is made to check if its DAG is paused. On instances with hundreds of task instances, this results in hundreds of queries to the dag table in every scheduler tick.

Proposed Solution

1. Relationship Check: Update DagUnpausedDep to check ti.dag_model.is_paused directly using the ORM relationship attribute.
2. Eager Loading: Update DagRun.fetch_task_instances (which is used by DagRun.get_task_instances inside update_state) to eager load the dag_model relationship using joinedload(TI.dag_model).

This keeps the modular dependency check design, but reduces the database round-trips from $O(N)$ queries to $O(1)$ in-memory lookups.

Benchmarking Results

We verified the query savings and performance gains on a mock dataset of 1,000 TaskInstances across 100 distinct DAG Runs:

• Without Eager Loading (Lazy/Ad-Hoc): 101 queries total, taking 610.05 ms.
• With Eager Loading (Joinedload): 1 query total, taking 320.95 ms.
• Query Reduction: 100 fewer database queries (99.01% query reduction).
• Latency Speedup: 1.90x faster on SQLite.

Affected Components

• airflow-core/src/airflow/models/dagrun.py
• airflow-core/src/airflow/ti_deps/deps/dag_unpaused_dep.py

Use case/motivation

No response

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct