jason810496 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3610071679
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -622,6 +646,95 @@ def _validate_outlet_event_partition_keys(outlet_events:
list[dict[str, Any]]) -
)
+# Lock timeout (seconds) for the opportunistic synchronous asset-event write.
Kept short so a
+# task completion sheds to the durable queue quickly under asset-row
contention rather than
+# holding an API-server worker. This only bounds lock waits; an uncontended
write does not wait
+# on locks and so is unaffected by the timeout.
+_ASSET_SYNC_LOCK_TIMEOUT = 2
+
+
+def _register_asset_changes_or_enqueue(
Review Comment:
Why not make the `_register_asset_changes_or_enqueue` itself only append the
`AssetEventQueue` record for further processing and skip the
`TI.register_asset_changes_in_db` at all?
##########
airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst:
##########
@@ -51,6 +51,17 @@ In addition to scheduling Dags based on time, you can also
schedule Dags to run
:ref:`asset_definitions` for how to declare assets.
+.. note::
+
+ Asset events are registered asynchronously. When a task succeeds, Airflow
registers its
+ asset events shortly afterwards rather than strictly inside the
task-completion request,
+ usually within one scheduler loop and longer under heavy load or database
contention.
+ Downstream asset-scheduled Dags are therefore triggered after a brief,
eventually
+ consistent delay. The events are never dropped: if the registration cannot
complete
+ inline it is handed to the scheduler, which retries until it succeeds. Use
the
+ ``[scheduler] asset_event_queue_drain_interval`` option to tune how often
the scheduler
+ processes deferred asset events.
+
Review Comment:
It would be nice to point out the new limitation for the local `DAG.test`
that the Asset-triggered downstream runs can silently never be scheduled. (Not
sure do we even list out all the limitation of `DAG.test` or not).
##########
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.
+ """
+ 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)
Review Comment:
The config should be set at class attribute level so that we won't retrieve
the conf every time when the `_drain_asset_event_queue` tick (get `conf` might
access the secret backend).
##########
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 (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)
```
##########
airflow-core/src/airflow/models/asset.py:
##########
@@ -788,6 +789,45 @@ def __repr__(self):
return f"{self.__class__.__name__}({', '.join(args)})"
+class AssetEventQueue(Base):
+ """
+ Pending asset-event registrations awaiting processing by the scheduler.
+
+ A row is enqueued on the task-success path in the execution API when
synchronous
+ asset-event registration is skipped or fails (for example under database
lock
+ contention). The scheduler drains this table, (re)runs
+
:meth:`~airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db`
to
+ create the ``AssetEvent`` and ``AssetDagRunQueue`` rows, and deletes the
queue row
+ once that write commits. This keeps the long-running asset-registration
work off the
+ synchronous ``ti_update_state`` request and guarantees the events are
eventually
+ delivered instead of being dropped when the write cannot complete inline.
+
+ ``ti_id`` is the primary key: at most one pending registration exists per
task
+ instance, and the row is cascade-deleted if the task instance is removed.
+ """
+
+ ti_id: Mapped[UUID] = mapped_column(sa.Uuid(), primary_key=True,
nullable=False)
+ task_outlets: Mapped[list] = mapped_column(sa.JSON(), nullable=False,
default=list)
+ outlet_events: Mapped[list] = mapped_column(sa.JSON(), nullable=False,
default=list)
Review Comment:
Additionally, the `AssetEventQueue` table will only be used for the further
async Asset Event records ingestion purpose. I don't need we need the
`task_outlets` and `outlet_events` columns at all.
I'd like to collapse the both `task_outlets` and `outlet_events` into a
single `JSON` type `payload` or `assets_profiles` column (or other better
naming). Or even just `text` column? Need double check the write path for the
append-only purpose among the DB backend we supported.
The time we need to access the `AssetEventQueue` records, we will need all
the `task_outlets` and `outlet_events` anyway. The scheduler side also
validation the integrity again, so IMO it's fine to collapse the columns.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -525,14 +535,28 @@ def ti_update_state(
task_id=task_id,
map_index=map_index,
)
+ session.commit()
except Exception:
+ session.rollback()
log.warning(
Review Comment:
> I think the transactional outbox (durable marker) approach is the right
direction, but IIUC it is not yet implemented atomically here.
Exactly, the `_enqueue_asset_events` should be in the same transaction to
ensure the consistency.
The new `AssetEventQueue` append-only table we introduced should shorten the
lock time comparing to executing the `TI.register_asset_changes_in_db` in the
API request.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -622,6 +646,95 @@ def _validate_outlet_event_partition_keys(outlet_events:
list[dict[str, Any]]) -
)
+# Lock timeout (seconds) for the opportunistic synchronous asset-event write.
Kept short so a
+# task completion sheds to the durable queue quickly under asset-row
contention rather than
+# holding an API-server worker. This only bounds lock waits; an uncontended
write does not wait
+# on locks and so is unaffected by the timeout.
+_ASSET_SYNC_LOCK_TIMEOUT = 2
+
+
+def _register_asset_changes_or_enqueue(
+ *,
+ task_instance_id: UUID,
+ task_outlets: list[AssetProfile],
+ outlet_events: list[dict[str, Any]],
+ session: SessionDep,
+) -> None:
+ """
+ Register a successful task's asset events, falling back to the durable
queue on contention.
+
+ Runs after the task-instance state has been committed and its row lock
released. The write is
+ attempted synchronously under a short lock timeout so the common,
uncontended case stays fast.
+ If the write cannot acquire its locks in time (or fails for any other
reason) the events are
+ written to ``asset_event_queue`` and the scheduler drains them, retrying
until they land. The
+ events are therefore never dropped, and the request never holds the
task_instance lock for the
+ registration work.
+ """
+ if not task_outlets and not outlet_events:
+ return
+
+ ti = session.get(TI, task_instance_id)
+ if ti is None:
+ return
+
+ try:
+ with with_db_lock_timeout(session,
lock_timeout=_ASSET_SYNC_LOCK_TIMEOUT):
+ TI.register_asset_changes_in_db(ti, task_outlets, outlet_events,
session=session)
+ # Supersede any row a previous failed attempt for this TI left in
the queue so the
+ # scheduler drain does not re-register the same events after a
synchronous success.
+
session.execute(delete(AssetEventQueue).where(AssetEventQueue.ti_id ==
task_instance_id))
+ session.commit()
+ return
+ except OperationalError as e:
+ session.rollback()
+ if not is_lock_not_available_error(e):
+ log.warning(
+ "Asset registration failed with a database error; deferring to
the asset event "
+ "queue for the scheduler to retry.",
+ task_instance_id=str(task_instance_id),
+ exc_info=True,
+ )
+ except Exception:
+ session.rollback()
+ log.warning(
+ "Asset registration failed; deferring to the asset event queue for
the scheduler to retry.",
+ task_instance_id=str(task_instance_id),
+ exc_info=True,
+ )
+
+ _enqueue_asset_events(
+ task_instance_id=task_instance_id,
+ task_outlets=task_outlets,
+ outlet_events=outlet_events,
+ session=session,
+ )
+
+
+def _enqueue_asset_events(
+ *,
+ task_instance_id: UUID,
+ task_outlets: list[AssetProfile],
+ outlet_events: list[dict[str, Any]],
+ session: SessionDep,
+) -> None:
+ """
+ Durably record a completed task's asset events for the scheduler to
register later.
+
+ Uses ``merge`` so a retried completion overwrites the pending row rather
than conflicting on
+ the ``ti_id`` primary key.
+ """
+ session.merge(
+ AssetEventQueue(
+ ti_id=task_instance_id,
+ task_outlets=[outlet.model_dump(mode="json") for outlet in
task_outlets],
+ outlet_events=outlet_events,
+ attempts=0,
+ )
Review Comment:
Then we could further prevent the serialization overhead here for storing
the `AssetEventQueue` outbox records and to shorten the total lock time as well.
##########
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.
+ """
+ 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)
+
+ processed = 0
+ # Track rows attempted this pass so a freshly-failed row (whose
attempts were just bumped)
+ # is not re-claimed and hammered within the same drain; it waits for
the next pass instead.
+ seen: set[UUID] = set()
+ for _ in range(batch_size):
+ conditions = [AssetEventQueue.attempts < max_attempts]
+ if seen:
+ conditions.append(AssetEventQueue.ti_id.notin_(seen))
+ row = session.scalars(
+ with_row_locks(
+
select(AssetEventQueue).where(*conditions).order_by(AssetEventQueue.created_at).limit(1),
+ of=AssetEventQueue,
+ session=session,
+ skip_locked=True,
+ )
+ ).first()
+ if row is None:
+ break
+
+ ti_id = row.ti_id
+ seen.add(ti_id)
+ ti = session.get(TaskInstance, ti_id)
+ try:
Review Comment:
We should fix the case that the TI being clear as we set the FK constraint
between `AssetEventQueue` and `TaskInstance`.
When every "TaskInstance clear", we will set the new `ti_id` UUID. If the
`AssetEventQueue` records are not yet consumed up, we will face the enter the
`if ti is not None:` path and drop the AssetEvent ingestion.
https://github.com/apache/airflow/blob/ffde96203e76c0c27ccf1ae6eeaad270653248bf/airflow-core/src/airflow/models/taskinstance.py#L1034-L1040
---
To fix the above scenario, we should update the `AssetEventQueue.ti_id` to
new ti_id.
```python
def prepare_db_for_next_try(self, session: Session):
from airflow.models.taskinstancehistory import TaskInstanceHistory
TaskInstanceHistory.record_ti(self, session=session)
session.execute(delete(TaskReschedule).filter_by(ti_id=self.id))
session.execute(
update(AssetEventQueue).where(AssetEventQueue.ti_id ==
self.id).values(ti_id=new_id)
) # new
self.id = uuid7()
```
--
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]
