Resolve dag versions once per scheduler loop

[steveahnahn]

The scheduler resolves each dag run's dag through DBDagBag.get_dag_for_run, which issues one uncached DagVersion.get_latest_version SELECT per call. Since #49097 every scheduling loop pays that cost once per dag run per tick — and the main scheduling loop pays it twice per run, because _verify_integrity_if_dag_changed repeats the same lookup.

#30704 ("Save scheduler execution time by caching dags") deliberately added a per-loop lru_cache for exactly this, keyed by dag_id. #49097 changed the cached call from cached_get_dag(dag_run.dag_id) to cached_get_dag(dag_run) to make resolution version-aware — which silently disabled the cache: the key became the DagRun ORM instance (identity hashing, no custom __eq__/__hash__), and each instance appears exactly once per loop. Rebuilding the old construct and passing 20 runs of one dag, cache_info() reports hits=0, misses=20. The "cache saves time during scheduling of many dag_runs for same dag" comment has not been true since.

Measured with CountQueries (which counts real after_cursor_execute events), 20 runs of one dag, verified on both SQLite and PostgreSQL:

Loop	version SELECTs before	after
_start_queued_dagruns	20	1
_schedule_all_dag_runs	40	1

At default per-loop caps this removes tens of point queries per scheduler tick per scheduler in every deployment; the waste scales with max_dagruns_per_loop_to_schedule.

What changed

• DagVersion.get_latest_versions(dag_ids): batch latest-version lookup (one GROUP BY max(version_number) query, backed by the (dag_id, version_number) unique constraint), semantics identical to get_latest_version per dag.
• DBDagBag.get_dags_for_runs(dag_runs, ...): batch counterpart of get_dag_for_run, mirroring _version_from_dag_run per run (bundle-pinned runs still resolve via created_dag_version_id); results keyed by (dag_id, run_id).
• _schedule_all_dag_runs resolves versions and dags once for the batch, threads them into _schedule_dag_run -> _verify_integrity_if_dag_changed as optional parameters (per-call behavior unchanged for all other callers), and returns the resolved dag so the callback-dispatch loop no longer re-resolves; _start_queued_dagruns uses the batch resolver; both dead lru_caches are removed.

Notes

• Version resolution is now consistent within one loop. On MySQL (InnoDB REPEATABLE READ) the old per-run SELECTs already read a single snapshot, so this only aligns PostgreSQL (READ COMMITTED) with the behavior MySQL already had; a version landing mid-loop is picked up on the next tick by the existing verify-integrity flow, as before.
• If a dag is deleted mid-tick, its dag-level callback is now dispatched using the dag resolved during scheduling instead of being dropped at dispatch time with an error log.

Tests

• Query-count regression test asserting one version lookup per loop (fails on the previous code with 3 and 6).
• Batch-method unit tests including a single/batch parity assertion and get_dags_for_runs resolution tests.
• Full test_scheduler_job.py suite passes.

related: #69275 (remaining per-item resolution paths in the critical section and dag-run listeners)

---

Was generative AI tooling used to co-author this PR?

• Yes - Claude Code (Opus 4.8)

Generated-by: Claude Code (Opus 4.8) following the guidelines

[Vamsi-klu]

Nice scheduler cleanup. One thing I think we should align before merge: get_latest_version() currently defines latest by created_at DESC, but get_latest_versions() uses max(version_number). If those can diverge, this changes which serialized Dag the scheduler resolves in the batched path. Could you either make the batch query match _latest_version_select() semantics, or add a test proving version_number is the intended source of truth here?

---

Drafted-by: Codex (GPT-5); reviewed by @Vamsi-klu before posting

[steveahnahn]

_latest_version_select() (used by get_latest_version()) already orders by version_number.desc() and not created_at, so both paths use version_number. max(version_number) joined on the (dag_id, version_number) unique constraint resolves the same single row as order_by(version_number.desc()).limit(1) therefore unique + non-null means no ties, so no divergence. The batch query references created_at nowhere.

test_get_latest_versions_resolves_each_dag_in_one_query already asserts get_latest_versions()[dag].id == get_latest_version(dag).id (same version resolved by both paths), and test_latest_version_uses_version_number_not_created_at pins version_number as the source of truth.