Repository navigation
Conversation
In ti_run (PATCH /execution/task-instances/{id}/run), when a DAG Run is triggered by many asset events, serializing consumed_asset_events triggers lazy N+1 queries for AssetEvent.asset and AssetEvent.source_aliases for each event inside a transaction holding a FOR UPDATE OF task_instance row lock. This causes task execution API statement timeouts.
This change uses selectinload to eager-load AssetEvent.asset and AssetEvent.source_aliases along with DR.consumed_asset_events, mirroring the pattern used in _eager_load_dag_run_for_validation() in scheduler_job_runner.py.
2 tasks done
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
closes: #76558
Description
In
ti_run(PATCH /execution/task-instances/{id}/run), when a DAG Run is triggered by many asset events, serializingconsumed_asset_eventsonDagRuntriggers lazy N+1 queries forAssetEvent.assetandAssetEvent.source_aliasesfor each event inside a transaction holding aFOR UPDATE OF task_instancerow lock. When runs have thousands of asset events, this leads to task execution API statement timeouts.This PR implements Option 1 outlined in #76558 by using
selectinloadto eager-loadAssetEvent.assetandAssetEvent.source_aliasesalong withDR.consumed_asset_events, mirroring the pattern already used in_eager_load_dag_run_for_validation()inscheduler_job_runner.py. This serves as a lightweight, non-breaking fix/PoC that resolves the database row lock bottleneck without requiring changes to the Task SDK wire protocol.Testing
test_ti_run_consumed_asset_events_eager_loadedinairflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.pyverifying that bothAssetEvent.assetandAssetEvent.source_aliasesare preloaded in the run context response across multiple events.Empirical Query Profiling on PostgreSQL 16 (N+1 vs selectinload)
Benchmarked against PostgreSQL 16 with 1 DagRun consuming 5 AssetEvents (each linked to an Asset and an AssetAlias):
FOR UPDATElock)SELECT dag_run ... WHERE dag_id = ... AND run_id = ...SELECT ... FROM dagrun_asset_event JOIN asset_event ... WHERE dag_run_id IN (...)SELECT ... FROM asset_alias_event JOIN asset_alias ... WHERE event_id IN ($1, $2, $3, $4, $5)SELECT ... FROM asset WHERE asset.id IN ($1, $2, $3, $4, $5)Was generative AI tooling used to co-author this PR?
Generated-by: Antigravity following the guidelines