apache / apache/airflow

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

Open
#70,050 0 comments 0 reactions 0 assignees View on GitHub
area:core area:performance area:scheduler kind:feature needs-triage
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 9h
Merged PRs (30d)
472

Description

### Description

### Description

During the Airflow scheduling loop, when evaluating task instance scheduling decisions for active `DagRun`s (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:
```python
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?

- [x] Yes I am willing to submit a PR!

### Code of Conduct

- [x] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Contributor guide

Open the contributing guide

Research direction

Read airflow-core/src/airflow/ti_deps/deps/dag_unpaused_dep.py and airflow-core/src/airflow/models/dagrun.py, starting at DagUnpausedDep and DagRun.fetch_task_instances as used by update_state. Confirm the relationship is available during _get_ready_tis() and that eager loading removes the per-task DAG query. Done means the dependency reads the relationship and the query count matches the proposed reduction.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, databases
Issue type
Refactor
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.