apache / apache/airflow

Scheduler resolves dag versions per-item in the critical section and via listener lazy loads

Open
#69,275 1 comment 0 reactions 0 assignees View on GitHub
area:core area:scheduler kind:feature
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 9h
Merged PRs (30d)
472

Description

### Description

Since scheduler dag resolution became version-aware, resolving a dag run's dag goes through `DagVersion.get_latest_version` — an uncached point `SELECT` issued once per call. Two scheduler code paths still resolve per item, so their query count scales with the number of items processed per tick:

1. **Critical section, per-task-instance resolution** — `SchedulerJobRunner._executable_task_instances_to_queued` calls `scheduler_dag_bag.get_dag_for_run(...)` per candidate task instance when the dag has task-concurrency limits (the `has_task_concurrency_limits` branch in `airflow-core/src/airflow/jobs/scheduler_job_runner.py`). Each call issues one `DagVersion.get_latest_version` `SELECT`; with `max_tis_per_query` (default 512) candidates, deployments using `max_active_tis_per_dag` / `max_active_tis_per_dagrun` can pay hundreds of point queries per tick inside the scheduler's most latency-sensitive section.

2. **`DagRun.version_number` lazy loads in dag-run state listeners** — `on_dag_run_running` implementations that read `dag_run.version_number` (the example-plugin listener does; the OpenLineage provider listener shares the hook) trigger per-run lazy loads of the `dag_versions` relationship — roughly 1.7 extra queries per started run when measured. Eager-loading the relationship in `DagRun.get_queued_dag_runs_to_set_running`, or batching at the listener call sites, would remove this.

Both are measurable with `tests_common.test_utils.asserts.CountQueries` around the respective loops: query counts grow linearly with the number of items instead of staying constant per loop.

### Use case/motivation

The scheduler runs these loops on every tick, so redundant per-item version lookups add database round-trips and latency directly on the scheduling hot path, and the cost grows with the number of task instances and dag runs a deployment examines per tick. The goal is for version resolution within a single scheduler loop to cost a constant number of queries regardless of how many items the loop processes — resolved once per batch (or at minimum once per dag) — mirroring the batching the main scheduling loops already use.

Acceptance criteria:

- Version lookups in the critical section do not scale with the number of candidate task instances.
- Starting queued runs does not trigger per-run `dag_versions` lazy loads from core-shipped listeners.
- Query-count regression coverage using `CountQueries` asserting per-loop constants.

### Related issues

- #30704 added per-loop caching of dag resolution in the scheduler to avoid repeated lookups.
- #49097 made scheduler dag resolution version-aware (routing it through `DagVersion.get_latest_version`).

### 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

Start with the has_task_concurrency_limits branch in airflow-core/src/airflow/jobs/scheduler_job_runner.py and inspect DagRun.get_queued_dag_runs_to_set_running plus the on_dag_run_running listeners. Use tests_common.test_utils.asserts.CountQueries around the respective loops to establish baseline counts. Done means version lookups stay constant per batch and queued-run startup avoids per-run dag_versions lazy loads, with regression coverage.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, databases, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.