hkc-8010 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3510932198
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -634,13 +673,7 @@ def _create_ti_state_update_query_and_update_state(
retry_reason=(ti_patch_payload.retry_reason[:500] if
ti_patch_payload.retry_reason else None),
)
elif isinstance(ti_patch_payload, TISuccessStatePayload):
- if ti is not None:
- TI.register_asset_changes_in_db(
- ti,
- ti_patch_payload.task_outlets,
- ti_patch_payload.outlet_events,
- session=session,
- )
+ pass # Asset registration happens after the TI state is
committed; see ti_update_state.
Review Comment:
Done in 4c168a78. Removed the branch entirely rather than leaving a bare
pass; it did nothing so there was no reason to keep it or explain it.
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -356,8 +357,17 @@ def register_asset_change(
).unique()
for asset_alias_model in asset_alias_models:
- asset_alias_model.asset_events.append(asset_event)
- session.add(asset_alias_model)
+ # Use a direct INSERT rather than ORM .append() to avoid
lazy-loading the
+ # entire asset_events collection. On long-running deployments
that collection
+ # can contain thousands of rows; loading it on the
task-success hot path can
+ # leave DB connections idle-in-transaction for minutes,
blocking other workers.
+ # This intentionally leaves asset_alias_model.asset_events
unsynced in-session.
Review Comment:
Done in 4c168a78. Trimmed to two lines: what the fix is (direct INSERT
instead of .append()) and why it's safe (nothing reads
asset_alias_model.asset_events again before commit). Dropped the restated
OOMKill narrative, that's already covered in the PR description.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -517,14 +526,44 @@ def ti_update_state(
task_id=task_id,
map_index=map_index,
)
+ session.commit()
except Exception:
+ session.rollback()
log.warning(
"Failed to clear task state on success",
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
)
+ # Asset registration runs outside the TI row lock. Failures are logged and
counted;
+ # raising HTTP 500 here would be misleading because the task already
succeeded and
+ # would make the worker retry a state update that has already completed.
Durable
+ # retry/reconciliation for dropped asset events is out of scope for this
hot-path fix.
Review Comment:
Done in 4c168a78. Removed the comment, the guard condition reads fine on its
own.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -491,6 +492,14 @@ def ti_update_state(
extra=json.dumps({"host_name": hostname}) if hostname else
None,
)
)
+ # Commit the TI state update now to release the task_instance row lock
before
+ # running asset-event queries. The direct-INSERT fix in AssetManager
removes
+ # the O(n) lazy-load on the alias-event table, but
register_asset_changes_in_db
+ # also queries scheduled dags and inserts AssetDagRunQueue rows - all
of which
+ # would otherwise hold the row lock and cause idle-in-transaction
pile-up that
+ # exhausts API server memory and triggers OOMKill under high
concurrency.
+ # The task outcome is durable from this point on.
Review Comment:
Done in 4c168a78. Trimmed to one line covering the non-obvious part, that
committing here is what releases the row lock before asset registration runs.
Point taken on reviewing AI-drafted comments before pushing, I should have
caught this myself the first time.
--
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]
