jason810496 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3610199153
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3676,6 +3682,88 @@ def _remove_unreferenced_triggers(self, *, session:
Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ @provide_session
+ def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) ->
None:
+ """
+ Process pending asset-event registrations enqueued by the execution
API.
+
+ The task-success path enqueues an
:class:`~airflow.models.asset.AssetEventQueue` row when
+ synchronous asset-event registration is skipped or fails under
database lock contention.
+ Here we claim pending rows one at a time (``FOR UPDATE SKIP LOCKED``,
so multiple schedulers
+ do not collide), (re)run ``register_asset_changes_in_db`` to create
the ``AssetEvent`` and
+ ``AssetDagRunQueue`` rows, and delete the queue row in the same
transaction so it is
+ processed exactly once. A row that keeps failing is retried on later
passes until it exceeds
+ ``asset_event_queue_max_attempts``, after which it is left in place
(parked) and surfaced via
+ the ``asset.event_queue.failures`` metric for an operator to inspect.
Review Comment:
I don't feel the current batch implementation is right.
IIUC, we should fetch the current batch in single query and commit in the
same session.
Additionally, we shouldn't call the explicit `commit` or `rollback` within
the function decorated by `provide_session`.
Somehow like:
```python
def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) ->
None:
from airflow.api_fastapi.execution_api.datamodels.asset import
AssetProfile
batch_size = conf.getint("scheduler", "asset_event_queue_batch_size",
fallback=100)
max_attempts = conf.getint("scheduler",
"asset_event_queue_max_attempts", fallback=5)
rows = session.scalars(
with_row_locks(
select(AssetEventQueue)
.where(AssetEventQueue.attempts < max_attempts)
.order_by(AssetEventQueue.created_at)
.limit(batch_size),
of=AssetEventQueue,
session=session,
skip_locked=True,
)
).all()
processed = 0
for row in rows:
try:
with session.begin_nested():
ti = session.get(TaskInstance, row.ti_id)
if ti is not None:
task_outlets = [AssetProfile.model_validate(o) for o in
row.task_outlets]
TaskInstance.register_asset_changes_in_db(
ti, task_outlets, row.outlet_events, session=session
)
session.delete(row)
processed += 1
except Exception:
with session.begin_nested():
row.attempts += 1
if row.attempts >= max_attempts:
self.log.error("Giving up on asset-event registration for TI
%s after %s attempts", row.ti_id, row.attempts)
stats.incr("asset.event_queue.failures")
else:
self.log.warning("Asset-event registration for TI %s failed;
will retry (attempt %s)", row.ti_id, row.attempts)
session.commit()
if processed:
stats.incr("asset.event_queue.processed", processed)
stats.gauge("asset.event_queue.pending",
session.scalar(select(func.count()).select_from(AssetEventQueue)) or 0)
```
--
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]
