Runs with many consumed asset events lead to Task execution API timeouts #76558

Under which category would you file this issue?

Airflow Core

Apache Airflow version

3.3.2

What happened and how to reproduce it?

Issue Description

We've run into an issue relating to runs triggered by a large number of asset events. PATCH requests to the /execution/task-instances/{task_instance_id}/run (ti_run) takes time proportional to the number of asset events the DagRun consumes, and it does that work while holding a FOR UPDATE lock on the task_instance row.

When an asset-triggered run consumes thousands of events, every task start in that run is slow enough to hit statement timeouts, while the Task SDK's retries queue up behind the lock held by the previous attempt.

Code path on main (task_instances.py @ bca1e14):

1. L182-L212: SELECT ... FROM task_instance JOIN dag_run JOIN dag WHERE task_instance.id = :id FOR UPDATE OF task_instance.
2. The TI is updated to RUNNING.
3. L322-L334: the DagRun is loaded with .options(joinedload(DR.consumed_asset_events), ...). Only the collection itself is eager-loaded.
4. L366-L367: TIRunContext(dag_run=dr, ...) is built. All of this runs in one transaction, so the row lock is held until commit.

This endpoint runs DagRun datamodel validations with asset lazy-loading (datamodels/taskinstance.py L362-L396) which causes numerous reads for asset and source_aliases on every AssetEventDagRunReference. Because of this the validation issues about one source_aliases SELECT per consumed event, all while the TI row lock is still held.

The scheduler does not run into the same issue because it avoids the pitfall (see _eager_load_dag_run_for_validation() )

Semi-related past issues and PRs: #65422, #59693, #66190.

How To Reproduce

1. Create a DAG scheduled on a single asset, then pause it.
2. Emit several thousand events for that asset, e.g. with repeated POST /api/v2/assets/events.
3. Unpause the DAG. The scheduler creates one asset-triggered run that consumes all of the queued events.
4. Watch the ti_run latency for the first task of that run. With SQLAlchemy echo or Postgres statement logging on, you see one asset_alias_asset_event / source_aliases SELECT per consumed event inside the transaction that holds FOR UPDATE OF task_instance. Latency grows with the number of consumed events.

What you think should happen instead?

The ti_run execution API's endpoint latency, and particularly how long it holds the TI row lock, should not grow with the number of consumed asset events.

Possible fixes:

1. Eager-load the nested relations in ti_run in a similar way that the scheduler does via _eager_load_dag_run_for_validation() (see link in issue description above). That turns the N+1 queries into a fixed number of queries.
2. Load and serialize the consumed events after the TI state update is committed, so that work does not run under FOR UPDATE. Alternatively, stop sending every consumed event in TIRunContext and have the Task SDK fetch them lazily when triggering_asset_events is accessed.

Fix (1) is small and self-contained. Fix (2) also covers the cost of serializing large extra payloads.

Operating System

Linux (Debian-based container image)

Anything else?

This happens every time an asset-triggered run consumes a large number of events, which is common after a long pause of an asset-scheduled DAG whose upstream keeps producing events. Runs that consume only a few events are not noticeably affected.

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct