hussein-awala opened a new pull request, #76706:
URL: https://github.com/apache/airflow/pull/76706
> **Context.** #76463 fixes, for 3.4.x, the starvation of pending
partition-driven Dag runs without a migration: the scheduler rotates through
every pending run instead of re-checking the 500 oldest. It still evaluates
every pending run again and again, whether or not anything changed for it,
which is acceptable as a stopgap for a patch release but not as the design.
This PR is the 3.5.0 design that #76463 is the stopgap for: the scheduler
evaluates a pending run only when something happened that can change its
result. It is based on `main` without #76463 for now; see "Relation to #76463"
below.
On `main`, each scheduler loop locks the 500 oldest pending partition-driven
Dag runs (`AssetPartitionDagRun`) and evaluates them, whether or not anything
changed for them:
- Runs that keep waiting by design, such as a rollup whose window is not
complete, fill those 500 slots, and newer runs whose partitions are complete
are never evaluated (since 3.3.0, #64571). #76463 fixes that.
- The work is the same on an idle deployment: the claim costs 128 ms per
loop on PostgreSQL with 55,000 pending runs, and each loop evaluates up to 500
runs and decodes the serialized Dags they belong to.
- On MySQL the claim locks every pending row, not 500, so an HA peer finds
nothing to claim.
A natural alternative is for writers to flag the runs whose inputs changed.
A flag on the run row, though, would be an UPDATE inside the producer's
task-success transaction, while the execution API holds the task instance's row
lock and while the scheduler may hold the run's row: the producer would wait on
the scheduler, and a lock timeout would fail the producer task. This PR has the
writer insert a request instead, which never waits.
### How
**A request queue.** The new table `asset_partition_dag_run_wake` (migration
0157) holds requests to evaluate pending runs: `id`, `apdr_id` (NULL for a
request about every pending run of a Dag), `target_dag_id`, `reason` and
`created_at`. It has no foreign key and no unique key, so an insert never waits
on the scheduler's lock of the run and never conflicts with a request the
scheduler already read. Rows are never updated: writers insert them, and the
scheduler deletes the ones it read, by primary key.
**Who writes requests**, each in the transaction of the change it reports:
- **Key logs.** An ORM listener writes one request for every
`PartitionedAssetKeyLog` that a session flush inserts. This covers the
execution API, the REST API, the triggerer and any `[core] asset_manager_class`
that adds key logs to the session, without changing them.
- **Definitions.** `SerializedDagModel.write_dag` writes one request for the
whole Dag when the timetable section of the definition changes, for a
partitioned Dag or a Dag that still has pending runs. The timetable holds the
asset condition and the partition mappers; nothing else in the definition
changes an evaluation.
- **Asset activation.** The scheduler (`_activate_referenced_assets`) and
the Dag processor (`activate_assets_if_possible`) write one request for each
Dag scheduled on an asset they activate, because key logs of an inactive asset
do not count.
**What a scheduler loop evaluates**, within the same budget of 500 runs per
loop:
- **Requested runs**: oldest request first within a Dag, and Dags in turn.
On an idle scheduler this read is the only query of the claim.
- **Fan-outs**: for a request about a whole Dag, every pending run of the
Dag, read once and spread over the next loops.
- **Resync passes**: every pending run, once when the scheduler starts and
then once every `[scheduler] asset_partition_resync_interval` (new, 3600 s by
default). They catch what no request reports: key logs written by components
that still run an older version during an upgrade, key logs written with a bulk
insert, and a new link from an alias to an asset that is already active.
While the fan-outs or a pass have runs left, each keeps a tenth of the loop.
**Claim.** As in #76463:
- Runs are locked by primary key with `FOR UPDATE SKIP LOCKED` and `main`'s
paused, draining and stale filter, with up to 3 refills.
- Satisfied runs fire oldest first within a loop.
**Request lifecycle.**
- **No lost requests.** A loop deletes only requests that it read before the
evaluation that answers them: a run's requests in the loop that evaluates the
run, and a Dag's in the loop that finishes its fan-out. It reads them before
the serialized Dag and the key logs, so a request written during the evaluation
survives and asks again, under READ COMMITTED on PostgreSQL and MySQL; SQLite
has one writer.
- **Retries.** Fan-out and pass progress lives in memory and is kept only
once the loop commits, so a loop that `@retry_db_transaction` retries does it
again.
- **Cleanup.** The claim drops requests whose run fired or was deleted. A
timer on `[scheduler] parsing_cleanup_interval` drops the requests of Dags that
no longer exist. It also merges those of paused, draining or stale Dags into
one request per Dag, in batches committed on their own, so the queue the claim
reads stays small.
**Metrics:**
- `scheduler.asset_partitions.wake_lag`: the age of the oldest waiting
request of a Dag that can run (not paused, draining or stale).
- `scheduler.asset_partitions.evaluated`: the runs evaluated.
- `scheduler.asset_partitions.resync_fired`: runs created that only a pass
found. It is tagged `startup` or `periodic`. A periodic count above 0 means
that a change reached no request: key logs written without one, an alias newly
linked to an asset that was already active, or custom code whose result changes
while its inputs do not.
- `scheduler.asset_partitions.resync_pass_duration`.
### Behaviour changes
- **Evaluation follows requests.** A pending run that nothing asks for is
evaluated by the next pass. A pass starts `asset_partition_resync_interval`
after the previous one started, or when that one ends if it took longer.
- **Contract for custom code.** Evaluation must depend only on the
timetable, the run's key logs, which assets are active, and how asset
references and aliases resolve. Custom partition mappers, windows and
timetables must keep to this: a mapper that reads the clock, or fails now and
then, is evaluated again only by the next pass. Key logs inserted in bulk
rather than through a session wait for the pass too.
- **Errors.** A run whose evaluation fails because of the Dag itself
(deserialization, fingerprint, a mapper that raises) is evaluated again only
when something asks for it (a new key log, a timetable change, an asset
activation) or a pass runs, rather than on every loop. A failing
`create_dagrun` still fails the loop, as on `main`.
- **Upgrade.** No backfill: each scheduler's startup pass evaluates the runs
that were pending at the upgrade. Components that still run 3.4 write no
requests, so upgrade the schedulers after the other components, or restart them
once every component runs 3.5.
- **Order.** Within a loop, runs fire oldest first, as before. Across loops,
requested runs come before the others: oldest request first within a Dag, Dags
in turn.
- **Dags no longer partitioned.** Their pending runs are dropped even when
the fingerprints match. A Dag scheduled only on aliases has an empty
fingerprint either way, and evaluating its runs under a timetable that is not
partitioned raises on `main`.
- **Decoded definitions.** `SerializedDagModel` drops its decoded data when
the instance is expired, as on a rollback (Airflow's sessions do not expire on
commit), so a loop retried after a definition was written in place evaluates
the new one.
### Relation to #76463
- **Purpose.** #76463 (milestone 3.4.1) is the stopgap for 3.4.x: no
migration, and every pending run evaluated in rotation. This PR replaces that
rotation in 3.5.0: evaluation follows requests, and a slow pass is the safety
net.
- **Overlap.** Both PRs change `_create_dagruns_for_partitioned_asset_dags`,
`metrics_template.yaml`, `assets.rst` and a comment in the UI's assets route,
so whichever lands second rebases. If #76463 lands first, the rebase removes
its per-loop rotation. Its metrics (`pass_loops`, `pass_duration`,
`oldest_pending_age`) are then either removed or kept for the resync pass.
#71072, which logs when the per-loop cap is reached, also changes that function.
- **Not in this PR:**
- isolating a failing fire, so that one bad run cannot fail the loop;
- showing the request state in the UI;
- handling key logs that arrive while their run fires (pre-existing).
### Open questions for reviewers
1. The resync interval's default (1 hour), which bounds the delay for
requests that nothing writes.
2. A request table, as here, or a marker column on the key log. A marker
needs no listener and also covers writers outside the ORM, but costs an UPDATE
per key log on the largest partition table and an index over it.
3. Whether this change wants a design discussion on the dev list first.
### Testing
On the final commit:
- **SQLite:** the changed test files and the rest of
`tests/unit/migrations`, 1,623 passed; `test_serialized_dag.py`, 114 passed.
- **PostgreSQL 16:** the partition tests of `test_scheduler_job.py`, 100
passed; the other changed test files except `test_db_cleanup.py`, 278 passed.
- **MySQL 8.0:** the partition tests, 99 passed, with the one
PostgreSQL-only test skipped; the other changed test files except
`test_db_cleanup.py`, 278 passed.
- **`breeze verify`:** the core DB tests (13,009 tests in 6 types, no
failures), the core non-DB tests (5,718 passed), the `standard` provider suites
(DB and non-DB), the `common.compat` suite (non-DB; it has no DB tests), and
`breeze build-docs apache-airflow`. `breeze verify` also lists every provider
suite, because `devel-common` changed: before each provider DB test module, the
`_clear_db` fixture runs `clear_db_apdr()`, which this PR changes. The
`standard` DB suite above ran it; CI runs the other suites, including the runs
against older Airflow versions, which have no `AssetPartitionDagRunWake`.
- **prek:** the pre-commit and manual stages, with the migration round trip.
- **Mutations:** 54 changes to the code, each of which breaks one rule of
this PR, are each caught by the test for that rule.
- **Order:** the partition tests also pass in reverse order and in 18
shuffled orders.
Existing tests that change:
- `test_partitioned_dag_run_resumes_when_asset_reactivates`: an
`AssetActive` row that something other than the scheduler or the Dag processor
inserts now waits for the next pass.
- `test_partition_cap_at_exactly_n_processes_all` and
`test_partition_cap_at_n_minus_one_leaves_one_pending`: the order within a Dag
now comes from the requests, so the `created_at` pin of the second test is
gone, and both docstrings name the per-loop budget.
- `test_pending_apdr_select_locks_rows_for_skip_locked_claim`: the claim no
longer runs a locking query when it has no run to evaluate, so the test creates
a pending run.
- `TestNextRunAssets::test_pending_apdr_uses_fifo_order` (UI), docstring
only: the pending run that waited longest is not necessarily the next to fire.
- `test_no_models_missing`: the request table is not part of `airflow db
clean`; the scheduler deletes its rows.
- `test_pending_apdr_select_breaks_tie_on_id` is removed. It ran its own
`SELECT ... ORDER BY created_at, id`, which stood for the claim query that this
PR removes, and called no scheduler code.
Not from this PR: on `main` too, `test_serialized_dag.py` fails after
`test_scheduler_job.py` in one pytest process unless the database was just
initialized, because an existing scheduler test fills the plugin cache while
`load_examples` is off. The uv runs above keep the two files in separate
processes; the breeze run initializes the database first.
### Cost of the claim
Measured as [the guide for performance-sensitive
components](https://github.com/apache/airflow/blob/main/.agents/skills/airflow-publish-changes/references/performance-sensitive-components.md)
describes: `main`'s claim (one locking `SELECT ... ORDER BY created_at, id
LIMIT 500`) against this PR's claim, in alternating rounds with a server-side
control query. The numbers are medians of 6 rounds of 10 loops, each rolled
back; the control drifted at most 2.5%. Each database holds about 1,050,000
partition-driven Dag runs. Of them, 55,000 (PostgreSQL) and 50,101 (MySQL) are
pending, in about 2,300 and 2,400 Dags, and 50,000 and 45,075 belong to Dags
that are not paused, draining or stale.
| Claim per loop | PostgreSQL 16 | MySQL 8.0 |
|---|---|---|
| `main` | 128–130 ms | 95–100 ms |
| No request (idle) | 4.5 ms, 1 statement | 4.2 ms, 1 statement |
| 50 requests | 9.7 ms | 10.7 ms |
| 500 requests | 18.7 ms | 19.7 ms |
| A request for every pending run of a Dag that can run (50,000 and 45,075)
| 97 ms | 134 ms |
- On MySQL, a large backlog of requests costs more than `main`'s claim: a
third more with a request for every pending run of a Dag that can run, because
the signal read groups every request in a temporary table. The backlog drains
by up to 500 runs per loop, while `main` pays its cost on every loop.
- These numbers are the claim only. On `main`, each loop also evaluates up
to 500 runs; here, only the runs that something asked for, plus every pending
run once in each resync pass.
- A resync pass reads every pending run once: about 0.5 s on PostgreSQL for
55,000 runs and 0.6–0.7 s on MySQL for 50,101 runs, with a peak of about
120–130 bytes of Python memory per pending run.
- Starting a fan-out reads the requests and the pending runs of one Dag: 3–5
ms on PostgreSQL for its largest Dag (528 runs) and 9–18 ms on MySQL for its
largest Dag (4,270 runs).
related: #76463, #64571
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
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)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
