buhuipao opened a new pull request, #76017: URL: https://github.com/apache/airflow/pull/76017
Airflow 3 removed the mini scheduler (#43741), so when a task finishes, the tasks after it wait until a scheduler loop examines their Dag run. Each loop examines only the `max_dagruns_per_loop_to_schedule` running Dag runs that were examined least recently (20 by default). With thousands of running Dag runs, every step from one task to the next waits for a full round, which can take minutes. This adds `DagRun.get_running_dag_runs_with_finished_tis`, modelled on `get_running_dag_runs_to_examine`. Right after the usual query, `_do_scheduling` also locks up to `max_dagruns_per_loop_to_schedule` more running Dag runs that have a task instance in a finished state with `updated_at` later than the run's `last_scheduling_decision`. These are the runs that have new work. `_update_dag_run_state_for_paused_dags` already uses the same comparison. Notes for review: - Runs picked by the first query are excluded by id. A session's own row locks do not make `SKIP LOCKED` skip them, so without this the same run would be examined twice in one loop. - The new method has no `@retry_db_transaction`. Its rollback would release the row locks that the first query took in the same transaction. - `OFFSET 0` in the `EXISTS` subquery is a Postgres planner fence. Without it, Postgres can turn `EXISTS` into a semi-join, such as a hash join over a sequential scan of every finished task instance, instead of checking each running Dag run through the `ti_dag_run` index. - One loop can now examine up to twice `max_dagruns_per_loop_to_schedule` Dag runs. The config description and the scheduler docs say so. #### Query cost Postgres in Breeze (amd64 under emulation on Apple silicon, so absolute numbers are slow), 2,000 running Dag runs with 15 task instances each, 20,000 finished Dag runs with 300,000 finished task instances, 20 running Dag runs with a newly finished task instance. 9 alternating runs; medians: | | median | Dag runs returned | |---|---|---| | base: `get_running_dag_runs_to_examine` | 2.96 ms | 20 | | this PR: both queries | 15.65 ms | 36 | | new query alone | 10.64 ms | 20 | | new query without `OFFSET 0` | 11.41 ms | 20 | | control: `get_queued_dag_runs_to_set_running` (unchanged) | 2.36 ms | 0 | The new query checks every running Dag run through the `ti_dag_run` index, so its cost grows with the number of running Dag runs. With this synthetic data Postgres chose the index plan even without `OFFSET 0`; the fence keeps that plan when table statistics would make Postgres pick the hash join. #### Tests run - New tests in `test_dagrun.py` and `test_scheduler_job.py`, plus the existing `get_running_dag_runs_to_examine` tests: pass on SQLite, Postgres and MySQL 8 (Breeze). - `airflow-core/tests/unit/jobs` and `airflow-core/tests/unit/models` on SQLite: 2876 passed, 26 skipped, 1 xfailed. - `breeze build-docs apache-airflow`: pass. - `prek` pre-commit hooks for the changed files: pass. --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes — Claude Code (Opus 5.5) Generated-by: Claude Code (Opus 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
