matthieuauger opened a new pull request, #72185: URL: https://github.com/apache/airflow/pull/72185
related: #71488 (its follow-up) Marking a dag run failed or success visits every pending task instance with `TaskInstance.set_state`, which merges and flushes each of them: one UPDATE round-trip plus the unit-of-work overhead per task instance — and `set_state()` does the same again for the running ones. #71488 removed the per-task-instance *session*; the per-task-instance *flush* remained. Measured on a run with 3,000 mapped task instances (Airflow 3.3.1 + #71488, PostgreSQL 17 behind pgbouncer, 2.5 ms RTT): `PATCH state=failed` took **84 s** server-side, well past the proxy timeout the UI request dies at. This change adds `set_task_instances_state()`, the bulk counterpart of `TaskInstance.set_state` in the shape of `clear_task_instances`: the same state and dates on each loaded task instance — the date logic is shared with `set_state` through `_set_state_and_dates`, which reuses `set_duration` — then one flush, so the unit of work emits the UPDATEs in batches (`executemany`, which the psycopg2 dialect already batches with `values_plus_batch`). Both loops of `mark_tasks` use it. Same run, same machine: **84 s → 2.6 s** (3.7 s when the task instances carry a previous try), the remainder being the response and the run itself (measured with py-spy on the API server). It also addresses the running task instances by `(task, map_index)`. Passing the bare task to `set_state` selected *every* map index of a mapped task, so the pending siblings of one running mapped task instance ended up FAILED instead of SKIPPED — contrary to the documented behaviour (“non-finished ones to SKIPPED”) — and were updated twice. With 2 of 3,000 mapped task instances running, the 2,998 others are now SKIPPED; the listener hooks still only see the 2 killed ones (`killed_tis` is unchanged). Tests: the two new tests cover the dates/duration set on never-started and already-started pending task instances, and the mapped case (only the running map index fails). --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (please specify the tool below) Generated-by: Claude Code 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]
