Scheduler resolves dag versions per-item in the critical section and via listener lazy loads
- 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
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