From 9be01eba4f4126399ffb0771763e317435d1a265 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Wed, 15 Jul 2026 18:32:57 +0530 Subject: [PATCH] Register task-success asset events asynchronously to avoid API-server lock contention On task success, ti_update_state committed the TI state and then registered asset events while still holding the task_instance row lock. Under high fan-out (many tasks emitting assets at once) this held the lock across the long registration work, piling up requests and OOMKilling the API server. Record the asset events as a durable marker in a new asset_event_queue table, committed in the same transaction as the TI state, and let the scheduler drain that queue and run the registration. The task-completion request no longer registers assets itself, so the row lock is held only for the state write plus one insert. The drain resolves the task instance by natural key (dag_id, run_id, task_id, map_index) rather than its surrogate id, so clearing a task (which reassigns the id) does not drop pending events on any backend. A registration that keeps failing is retried up to [scheduler] asset_event_queue_max_attempts times, then parked and surfaced via the asset.event_queue.failures metric. dag.test runs in process without a scheduler, so it drains its own queued events after each task finishes, keeping asset events visible during a test run. Also replace the alias-event ORM append with a direct association insert to avoid an O(n) lazy-load on long-lived aliases. closes: #66853 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../asset-scheduling.rst | 15 ++ airflow-core/docs/migrations-ref.rst | 4 +- .../newsfragments/66854.significant.rst | 34 ++++ .../execution_api/routes/task_instances.py | 91 ++++++++- airflow-core/src/airflow/assets/manager.py | 13 +- .../src/airflow/config_templates/config.yml | 25 +++ .../src/airflow/jobs/scheduler_job_runner.py | 78 +++++++ .../0128_3_4_0_add_asset_event_queue.py | 71 +++++++ airflow-core/src/airflow/models/asset.py | 96 +++++++++ airflow-core/src/airflow/utils/db.py | 2 +- .../versions/head/test_task_instances.py | 192 ++++++++++-------- .../tests/unit/jobs/test_scheduler_job.py | 146 +++++++++++++ .../metrics/metrics_template.yaml | 24 +++ .../tests/task_sdk_tests/conftest.py | 23 ++- task-sdk/src/airflow/sdk/definitions/dag.py | 10 + 15 files changed, 722 insertions(+), 102 deletions(-) create mode 100644 airflow-core/newsfragments/66854.significant.rst create mode 100644 airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py diff --git a/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst b/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst index 1bfe98d715ec2..69eb5070a2d3f 100644 --- a/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst +++ b/airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst @@ -51,6 +51,21 @@ 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 durably records its + asset events and the scheduler registers them shortly afterwards rather than inside the + task-completion request, usually within one scheduler loop and longer under heavy load. + Downstream asset-scheduled Dags are therefore triggered after a brief, eventually consistent + delay. The scheduler retries a queued registration up to ``[scheduler] + asset_event_queue_max_attempts`` times, after which the row is parked and reported via the + ``asset.event_queue.failures`` metric. Use the ``[scheduler] asset_event_queue_drain_interval`` + option to tune how often the scheduler processes queued asset events. + + ``dag.test`` runs without a scheduler, so it registers each task's asset events itself once + the task finishes; the events stay visible during the test, though downstream asset-scheduled + Dags are still only triggered by a running scheduler. + Schedule Dags with assets ------------------------- diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index de64f8123f059..2bb682744be6c 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | +| ``a63ae7baae97`` (head) | ``7a98f1b7dbd3`` | ``3.4.0`` | Add asset_event_queue table. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``7a98f1b7dbd3`` | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``c4e7a1f9b2d0`` | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/newsfragments/66854.significant.rst b/airflow-core/newsfragments/66854.significant.rst new file mode 100644 index 0000000000000..42883e3cd303d --- /dev/null +++ b/airflow-core/newsfragments/66854.significant.rst @@ -0,0 +1,34 @@ +Asset events from successful tasks are registered asynchronously + +When a task succeeds, its asset events (and the resulting scheduling of downstream +asset-triggered Dags) are no longer written synchronously inside the task state-update +request. The execution API now records the events as a durable marker in a new +``asset_event_queue`` table, committed in the same transaction as the task state, and the +scheduler drains that queue and registers the events. This removes the asset registration +work from the task-completion request path, which previously held the ``task_instance`` row +lock while it ran and could pile up requests and exhaust the API server under high fan-out. + +As a result there can be a short delay between a task succeeding and its asset events being +visible (typically one scheduler loop, longer under sustained load). A queued registration is +retried up to ``[scheduler] asset_event_queue_max_attempts`` times; a row that still fails +after that is parked in ``asset_event_queue`` and surfaced via the ``asset.event_queue.failures`` +metric rather than retried forever. ``dag.test`` runs without a scheduler, so it drains the queue +itself after each task, keeping asset events visible during a test run. The drain cadence and +batch size are tunable via ``[scheduler] asset_event_queue_drain_interval``, +``asset_event_queue_batch_size`` and ``asset_event_queue_max_attempts``. + +* Types of change + + * [ ] Dag changes + * [x] Config changes + * [ ] API changes + * [ ] CLI changes + * [x] Behaviour changes + * [ ] Plugin changes + * [ ] Dependency changes + * [ ] Code interface changes + +* Migration rules needed + + * None - a new ``asset_event_queue`` table is created by the schema migration; no user + action is required. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c1bac7960234d..7247fb1aebf87 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -40,6 +40,7 @@ from sqlalchemy.sql import select from structlog.contextvars import bind_contextvars +from airflow._shared.observability.metrics import stats from airflow._shared.observability.traces import override_ids from airflow._shared.state import TaskScope from airflow._shared.timezones import timezone @@ -77,7 +78,7 @@ ) from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound -from airflow.models.asset import AssetActive +from airflow.models.asset import AssetActive, AssetEventQueue from airflow.models.base import ID_LEN from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun as DR @@ -97,6 +98,8 @@ if TYPE_CHECKING: from sqlalchemy.sql.dml import Update + from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile + router = VersionedAPIRouter() ti_id_router = VersionedAPIRouter( @@ -499,6 +502,30 @@ def ti_update_state( extra=json.dumps({"host_name": hostname}) if hostname else None, ) ) + # Durably record the successful task's asset events in the SAME transaction that commits the + # state, so the marker can never be lost in a crash between this commit and their + # registration. The scheduler drain is the single writer that registers them; the request + # never runs asset registration itself, which is what keeps the task_instance row lock short + # under high fan-out. Only a genuine RUNNING->SUCCESS transition reaches here (a duplicate + # SUCCESS->SUCCESS short-circuits earlier), so a completion is never enqueued twice. + if ( + updated_state == TaskInstanceState.SUCCESS + and isinstance(ti_patch_payload, TISuccessStatePayload) + and (ti_patch_payload.task_outlets or ti_patch_payload.outlet_events) + ): + _enqueue_asset_events( + task_instance_id=task_instance_id, + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + map_index=map_index, + task_outlets=ti_patch_payload.task_outlets, + outlet_events=ti_patch_payload.outlet_events, + session=session, + ) + # Commit the state, log entry and durable queue marker together so they land atomically and + # the row lock is released promptly. + session.commit() except DataError: # Let DataErrorHandler return a 422 (not the opaque 500 below). raise @@ -525,7 +552,9 @@ 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, @@ -622,6 +651,58 @@ def _validate_outlet_event_partition_keys(outlet_events: list[dict[str, Any]]) - ) +def _asset_event_payload( + task_outlets: list[AssetProfile], + outlet_events: list[dict[str, Any]], + ti_key: dict[str, Any], +) -> dict[str, Any]: + """Serialize a task's outlets, outlet events, and natural key into the queue row's JSON payload.""" + return { + "task_outlets": [outlet.model_dump(mode="json") for outlet in task_outlets], + "outlet_events": outlet_events, + # The scheduler drain resolves the live task instance by this natural key rather than by the + # surrogate ``ti_id``. Clearing a task reassigns its uuid7 id, and relying on the foreign + # key's ON UPDATE CASCADE to re-point the row only works on Postgres -- SQLite does not + # enforce foreign keys in Airflow's production engine -- so a natural-key lookup is what + # keeps the pending events reachable on every backend. + "ti_key": ti_key, + } + + +def _enqueue_asset_events( + *, + task_instance_id: UUID, + dag_id: str, + run_id: str, + task_id: str, + map_index: int, + task_outlets: list[AssetProfile], + outlet_events: list[dict[str, Any]], + session: SessionDep, +) -> None: + """ + Add a durable queue marker for a completed task's asset events, without committing. + + The caller commits this in the same transaction as the task-state update, so the marker and the + ``SUCCESS`` state land atomically: the events can never be lost in a crash between the state + commit and their registration. The scheduler drain is the single writer that later registers the + events and deletes the row. ``merge`` overwrites any pending row a retried completion left behind + rather than conflicting on the ``ti_id`` primary key. + """ + session.merge( + AssetEventQueue( + ti_id=task_instance_id, + payload=_asset_event_payload( + task_outlets, + outlet_events, + {"dag_id": dag_id, "run_id": run_id, "task_id": task_id, "map_index": map_index}, + ), + attempts=0, + ) + ) + stats.incr("asset.event_queue.enqueued") + + def _create_ti_state_update_query_and_update_state( *, ti_patch_payload: TIStateUpdate, @@ -659,14 +740,6 @@ def _create_ti_state_update_query_and_update_state( # Store retry policy overrides so next_retry_datetime() can read them. # These are cleared when the task enters RUNNING (ti_run). query = query.values(retry_delay_override=retry_delay_override, retry_reason=retry_reason) - 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, - ) try: _emit_task_span(ti, state=updated_state) except Exception: diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index 1dd26379b9d12..2f01daa681fac 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING import structlog -from sqlalchemy import exc, or_, select +from sqlalchemy import exc, insert, or_, select from sqlalchemy.orm import joinedload from airflow._shared.observability.metrics import stats @@ -41,6 +41,7 @@ DagScheduleAssetUriReference, PartitionedAssetKeyLog, TaskOutletAssetReference, + asset_alias_asset_event_association_table, ) from airflow.models.log import Log from airflow.timetables.base import compute_rollup_fingerprint @@ -359,8 +360,14 @@ 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) + # Direct INSERT instead of asset_alias_model.asset_events.append(...), which + # would lazy-load the alias's entire asset_events collection just to add one row. + session.execute( + insert(asset_alias_asset_event_association_table).values( + alias_id=asset_alias_model.id, + event_id=asset_event.id, + ) + ) dags_to_queue_from_asset_alias |= { alias_ref.dag diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index b3dc3f133d229..3520c41310830 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2639,6 +2639,31 @@ scheduler: type: integer example: ~ default: "60" + asset_event_queue_drain_interval: + description: | + How often (in seconds) the scheduler drains the asset event queue, registering asset + events that the execution API deferred to the queue (for example when the synchronous + write on task completion hit database lock contention). + version_added: 3.4.0 + type: float + example: ~ + default: "5.0" + asset_event_queue_batch_size: + description: | + The maximum number of pending asset event registrations the scheduler processes in a + single drain pass. + version_added: 3.4.0 + type: integer + example: ~ + default: "100" + asset_event_queue_max_attempts: + description: | + How many times the scheduler retries a failing asset event registration before parking + the row for inspection and emitting the ``asset.event_queue.failures`` metric. + version_added: 3.4.0 + type: integer + example: ~ + default: "5" pool_metrics_interval: description: | How often (in seconds) should pool usage stats be sent to StatsD (if statsd_on is enabled) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 9b9e85c07e353..85c324a80ce43 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -77,6 +77,7 @@ AssetAliasModel, AssetDagRunQueue, AssetEvent, + AssetEventQueue, AssetModel, AssetPartitionDagRun, AssetWatcherModel, @@ -85,6 +86,7 @@ PartitionedAssetKeyLog, TaskInletAssetReference, TaskOutletAssetReference, + _register_queued_asset_event, ) from airflow.models.asset_state_store import AssetStateStoreModel from airflow.models.backfill import Backfill, BackfillDagRun @@ -335,6 +337,17 @@ def __init__( self._parallelism = conf.getint("core", "parallelism") self._multi_team = conf.getboolean("core", "multi_team") self._dag_tags_in_metrics = conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False) + # Asset-event queue drain tuning. Read here (not per drain tick) so the drain loop stays free + # of conf lookups, which can hit the secret backend and break HA locks (see note above). + self._asset_event_queue_batch_size = conf.getint( + "scheduler", "asset_event_queue_batch_size", fallback=100 + ) + self._asset_event_queue_max_attempts = conf.getint( + "scheduler", "asset_event_queue_max_attempts", fallback=5 + ) + self._asset_event_queue_drain_interval = conf.getfloat( + "scheduler", "asset_event_queue_drain_interval", fallback=5.0 + ) self._max_partition_dag_runs_per_loop = MAX_PARTITION_DAG_RUNS_PER_LOOP self._dag_id_to_team_name: dict[str, str | None] = {} @@ -1788,6 +1801,11 @@ def _run_scheduler_loop(self) -> None: self._remove_unreferenced_triggers, ) + timers.call_regular_interval( + self._asset_event_queue_drain_interval, + self._drain_asset_event_queue, + ) + if any(x.is_local for x in self.executors): bundle_cleanup_mgr = BundleUsageTrackingManager() check_interval = conf.getint( @@ -3698,6 +3716,66 @@ 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 commits an :class:`~airflow.models.asset.AssetEventQueue` marker + atomically with the task state instead of registering the events inline, so the scheduler + is the sole writer that runs ``register_asset_changes_in_db`` in a live deployment. We claim + a batch of pending rows in a single ``FOR UPDATE SKIP LOCKED`` query (so multiple schedulers + do not collide), register each row's events and delete it 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. + """ + max_attempts = self._asset_event_queue_max_attempts + rows = session.scalars( + with_row_locks( + select(AssetEventQueue) + .where(AssetEventQueue.attempts < max_attempts) + .order_by(AssetEventQueue.created_at) + .limit(self._asset_event_queue_batch_size), + of=AssetEventQueue, + session=session, + skip_locked=True, + ) + ).all() + + processed = 0 + for row in rows: + ti_id = row.ti_id + try: + # Each row is registered in its own savepoint so one poison row rolls back alone + # instead of discarding the whole batch; the outer transaction commits once at the end. + with session.begin_nested(): + _register_queued_asset_event(row, session=session) + 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 task instance %s after %s " + "attempts; the row is left in asset_event_queue for inspection.", + ti_id, + row.attempts, + ) + stats.incr("asset.event_queue.failures") + else: + self.log.warning( + "Asset-event registration for task instance %s failed; will retry (attempt %s).", + ti_id, + row.attempts, + ) + + session.commit() + if processed: + stats.incr("asset.event_queue.processed", processed) + pending = session.scalar(select(func.count()).select_from(AssetEventQueue)) + stats.gauge("asset.event_queue.pending", pending or 0) + @provide_session def _update_asset_orphanage(self, *, session: Session = NEW_SESSION) -> None: """ diff --git a/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py new file mode 100644 index 0000000000000..7e0f3fd9aaa9a --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0128_3_4_0_add_asset_event_queue.py @@ -0,0 +1,71 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Add asset_event_queue table. + +Revision ID: a63ae7baae97 +Revises: 7a98f1b7dbd3 +Create Date: 2026-07-15 00:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +from airflow.utils.sqlalchemy import UtcDateTime + +# revision identifiers, used by Alembic. +revision = "a63ae7baae97" +down_revision = "7a98f1b7dbd3" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Apply Add asset_event_queue table.""" + op.create_table( + "asset_event_queue", + sa.Column("ti_id", sa.Uuid(), nullable=False), + # Emitted task outlets and outlet events are stored together in a single JSON payload. The + # queue is a durable buffer only ever read back in full by the scheduler drain, so separate + # columns buy nothing and a single write keeps the enqueue on the task-success path cheap. + sa.Column("payload", sa.JSON(), nullable=False), + sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"), + sa.Column("created_at", UtcDateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("ti_id", name="asset_event_queue_pkey"), + sa.ForeignKeyConstraint( + ["ti_id"], + ["task_instance.id"], + name="aeq_ti_fkey", + ondelete="CASCADE", + # Cascade a TI id change (clear/retry reassigns a new uuid7) so pending asset events + # follow the task instance instead of being orphaned and dropped by the scheduler drain. + onupdate="CASCADE", + ), + ) + with op.batch_alter_table("asset_event_queue", schema=None) as batch_op: + batch_op.create_index("idx_asset_event_queue_created_at", ["created_at"], unique=False) + + +def downgrade(): + """Unapply Add asset_event_queue table.""" + with op.batch_alter_table("asset_event_queue", schema=None) as batch_op: + batch_op.drop_index("idx_asset_event_queue_created_at") + op.drop_table("asset_event_queue") diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index 7aafddb3b9490..f7262df6172d3 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -20,6 +20,7 @@ from datetime import datetime from typing import TYPE_CHECKING from urllib.parse import urlsplit +from uuid import UUID import sqlalchemy as sa from sqlalchemy import ( @@ -788,6 +789,101 @@ def __repr__(self): return f"{self.__class__.__name__}({', '.join(args)})" +class AssetEventQueue(Base): + """ + Durable marker of asset events a successful task emitted, awaiting registration. + + On task success the execution API commits one of these rows atomically with the task + state instead of registering the asset events inline, so the ``ti_update_state`` request + holds the ``task_instance`` row lock only for the state write plus this insert rather than + for the whole ``register_asset_changes_in_db`` call (which was the source of API-server + lock contention under high fan-out). The scheduler drains this table, resolves the live + task instance by natural key, 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. The in-process runner behind ``dag.test`` has no scheduler, so it drains the row + itself via :func:`register_pending_asset_events` right after the task finishes. + + ``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) + # Both the emitted task outlets and the outlet events live in one JSON payload + # (``{"task_outlets": [...], "outlet_events": [...], "ti_key": {...}}``). The queue is a durable + # buffer only ever read back in full when draining, so a single column keeps the enqueue cheap. + payload: Mapped[dict] = mapped_column(sa.JSON(), nullable=False, default=dict) + attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0") + created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) + + __tablename__ = "asset_event_queue" + __table_args__ = ( + PrimaryKeyConstraint(ti_id, name="asset_event_queue_pkey"), + ForeignKeyConstraint( + (ti_id,), + ["task_instance.id"], + name="aeq_ti_fkey", + ondelete="CASCADE", + # Referential cleanup only. Correctness on clear does not rely on these cascades: + # SQLite does not enforce foreign keys in Airflow's production engine, so they silently + # no-op there. The drain re-resolves the task instance by natural key instead, which + # survives the id reassignment a clear performs on every backend. + onupdate="CASCADE", + ), + Index("idx_asset_event_queue_created_at", created_at), + ) + + def __repr__(self): + return f"AssetEventQueue(ti_id={self.ti_id!r}, attempts={self.attempts!r})" + + +def _register_queued_asset_event(row: AssetEventQueue, *, session: Session) -> None: + """ + Register the asset events captured in one :class:`AssetEventQueue` row, then delete it. + + Resolves the live task instance by natural key (``dag_id``/``run_id``/``task_id``/``map_index``) + rather than the surrogate ``ti_id``: clearing a task reassigns its id, so a lookup by the + enqueued id would miss the row on any backend that does not cascade the id change. If the task + instance no longer exists there is nothing to register and the row is simply dropped. The caller + owns the surrounding transaction (the scheduler wraps each row in a savepoint; the in-process + runner commits the session). + """ + from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile + from airflow.models.taskinstance import TaskInstance + + payload = row.payload + ti_key = payload["ti_key"] + ti = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == ti_key["dag_id"], + TaskInstance.run_id == ti_key["run_id"], + TaskInstance.task_id == ti_key["task_id"], + TaskInstance.map_index == ti_key["map_index"], + ) + ) + if ti is not None: + task_outlets = [AssetProfile.model_validate(outlet) for outlet in payload["task_outlets"]] + TaskInstance.register_asset_changes_in_db(ti, task_outlets, payload["outlet_events"], session=session) + session.delete(row) + + +def register_pending_asset_events(*, ti_ids: Iterable[UUID], session: Session) -> None: + """ + Register queued asset events for the given task instances in line and commit. + + The execution API records asset events on task success as durable + :class:`AssetEventQueue` markers that the scheduler drains. The in-process task runner + behind ``dag.test`` has no scheduler, so it calls this to register the events itself right + after a task succeeds, the same way it runs the triggerer in process. Best-effort and + unbatched: production contention handling and retry/park bookkeeping stay in the scheduler + drain (``SchedulerJobRunner._drain_asset_event_queue``). + """ + rows = session.scalars(select(AssetEventQueue).where(AssetEventQueue.ti_id.in_(ti_ids))).all() + for row in rows: + _register_queued_asset_event(row, session=session) + session.commit() + + association_table = Table( "dagrun_asset_event", Base.metadata, diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 1076715ee7180..03713248c6ddd 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "7a98f1b7dbd3", + "3.4.0": "a63ae7baae97", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 542ce7eaaf15b..41994cf59c2a9 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -45,7 +45,7 @@ from airflow.api_fastapi.execution_api.security import require_auth from airflow.exceptions import AirflowSkipException from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger -from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel +from airflow.models.asset import AssetActive, AssetEvent, AssetEventQueue, AssetModel from airflow.models.dag import DagModel from airflow.models.log import Log from airflow.models.task_state_store import TaskStateStoreModel @@ -103,19 +103,6 @@ def _is_select_for_update(statement) -> bool: return getattr(statement, "is_select", False) and "FOR UPDATE" in str(statement.compile()).upper() -def _create_asset_aliases(session, num: int = 2) -> None: - asset_aliases = [ - AssetAliasModel( - id=i, - name=f"simple{i}", - group="alias", - ) - for i in range(1, 1 + num) - ] - session.add_all(asset_aliases) - session.commit() - - @pytest.fixture def _use_real_jwt_bearer(exec_app): """Remove the mock require_auth override so the real JWT validation runs end-to-end.""" @@ -1228,17 +1215,9 @@ def test_ti_update_state_to_terminal_with_rendered_map_index( assert ti.rendered_map_index == rendered_map_index @pytest.mark.parametrize( - "task_outlets", + "outlet_events", [ - pytest.param([{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}], id="asset"), - pytest.param([{"name": "my-task", "type": "AssetNameRef"}], id="name-ref"), - pytest.param([{"uri": "s3://bucket/my-task", "type": "AssetUriRef"}], id="uri-ref"), - ], - ) - @pytest.mark.parametrize( - ("outlet_events", "expected_extra"), - [ - pytest.param([], {}, id="default"), + pytest.param([], id="default"), pytest.param( [ { @@ -1250,23 +1229,18 @@ def test_ti_update_state_to_terminal_with_rendered_map_index( "extra": {"foo": 2}, }, ], - {"foo": 1}, id="extra", ), ], ) - def test_ti_update_state_to_success_with_asset_events( - self, client, session, create_task_instance, task_outlets, outlet_events, expected_extra + def test_ti_update_state_to_success_enqueues_asset_events( + self, client, session, create_task_instance, outlet_events ): - asset = AssetModel( - id=1, - name="my-task", - uri="s3://bucket/my-task", - group="asset", - extra={}, - ) - asset_active = AssetActive.for_asset(asset) - session.add_all([asset, asset_active]) + """A successful task with outlets records a durable ``asset_event_queue`` marker in the same + transaction as the state. Registration is the scheduler drain's job, so no ``AssetEvent`` + exists yet -- the request only enqueues the raw outlets/events plus the TI natural key the + drain resolves the live task instance with.""" + task_outlets = [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}] ti = create_task_instance( task_id="test_ti_update_state_to_success_with_asset_events", @@ -1289,52 +1263,37 @@ def test_ti_update_state_to_success_with_asset_events( assert response.text == "" session.expire_all() - event = session.scalars(select(AssetEvent)).all() - assert len(event) == 1 - assert event[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={}) - assert event[0].extra == expected_extra + # Registration is deferred to the scheduler drain: nothing registered inline. + assert session.scalars(select(AssetEvent)).all() == [] + + rows = session.scalars(select(AssetEventQueue)).all() + assert len(rows) == 1 + assert rows[0].ti_id == ti.id + assert rows[0].attempts == 0 + assert rows[0].payload == { + "task_outlets": task_outlets, + "outlet_events": outlet_events, + "ti_key": { + "dag_id": ti.dag_id, + "run_id": ti.run_id, + "task_id": ti.task_id, + "map_index": ti.map_index, + }, + } - @pytest.mark.parametrize( - ("outlet_events", "expected_extra"), - [ - pytest.param([], None, id="default"), - pytest.param( - [ - { - "dest_asset_key": {"name": "my-task", "uri": "s3://bucket/my-task"}, - "source_alias_name": "simple1", - "extra": {"foo": 1}, - }, - { - "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"}, - "extra": {"foo": 2}, - }, - { - "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"}, - "source_alias_name": "simple2", - "extra": {"foo": 3}, - }, - ], - {"foo": 1}, - id="extra", - ), - ], - ) - def test_ti_update_state_to_success_with_asset_alias_events( - self, client, session, create_task_instance, outlet_events, expected_extra + def test_ti_update_state_to_success_with_asset_alias_events_enqueues( + self, client, session, create_task_instance ): - asset = AssetModel( - id=1, - name="my-task", - uri="s3://bucket/my-task", - group="asset", - extra={}, - ) - asset_active = AssetActive.for_asset(asset) - session.add_all([asset, asset_active]) - - _create_asset_aliases(session, num=2) - + """Alias outlets are enqueued verbatim; the scheduler drain resolves the alias to concrete + AssetEvents when it registers (see the manager/drain tests), so the request itself creates + only the queue marker and registers nothing inline.""" + outlet_events = [ + { + "dest_asset_key": {"name": "my-task", "uri": "s3://bucket/my-task"}, + "source_alias_name": "simple1", + "extra": {"foo": 1}, + }, + ] ti = create_task_instance( task_id="test_ti_update_state_to_success_with_asset_events", start_date=DEFAULT_START_DATE, @@ -1356,13 +1315,72 @@ def test_ti_update_state_to_success_with_asset_alias_events( assert response.text == "" session.expire_all() - events = session.scalars(select(AssetEvent)).all() - if expected_extra is None: - assert events == [] - else: - assert len(events) == 1 - assert events[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={}) - assert events[0].extra == expected_extra + assert session.scalars(select(AssetEvent)).all() == [] + rows = session.scalars(select(AssetEventQueue)).all() + assert len(rows) == 1 + assert rows[0].payload["task_outlets"] == [{"name": "simple1", "uri": None, "type": "AssetAlias"}] + assert rows[0].payload["outlet_events"] == outlet_events + + def test_ti_update_state_forced_failure_skips_asset_queue(self, client, session, create_task_instance): + """When an unexpected error while building the state update forces the TI to FAILED, the + success-only asset path must not run: nothing is registered and nothing is enqueued.""" + ti = create_task_instance( + task_id="test_ti_update_state_forced_failure_skips_asset_queue", + start_date=DEFAULT_START_DATE, + state=State.RUNNING, + ) + session.commit() + + with mock.patch( + "airflow.api_fastapi.execution_api.routes.task_instances." + "_create_ti_state_update_query_and_update_state", + side_effect=Exception("kaboom"), + ): + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": "success", + "end_date": DEFAULT_END_DATE.isoformat(), + "task_outlets": [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}], + "outlet_events": [], + }, + ) + + assert response.status_code == 204 + assert response.text == "" + session.expire_all() + + assert session.get(TaskInstance, ti.id).state == State.FAILED + assert session.scalars(select(AssetEventQueue)).all() == [] + assert session.scalars(select(AssetEvent)).all() == [] + + def test_ti_update_state_success_no_outlets_no_queue(self, client, session, create_task_instance): + """A successful TI with no outlets and no outlet events short-circuits: no registration, no + queue row, no AssetEvent.""" + ti = create_task_instance( + task_id="test_ti_update_state_success_no_outlets_no_queue", + start_date=DEFAULT_START_DATE, + state=State.RUNNING, + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": "success", + "end_date": DEFAULT_END_DATE.isoformat(), + "task_outlets": [], + "outlet_events": [], + }, + ) + + assert response.status_code == 204 + assert response.text == "" + session.expire_all() + + assert session.get(TaskInstance, ti.id).state == State.SUCCESS + assert session.scalars(select(AssetEventQueue)).all() == [] + assert session.scalars(select(AssetEvent)).all() == [] @pytest.mark.parametrize( ("partition_key", "expected_status"), diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index dff3adea48b21..a10fd26b79de4 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -63,6 +63,7 @@ AssetAliasModel, AssetDagRunQueue, AssetEvent, + AssetEventQueue, AssetModel, AssetPartitionDagRun, DagScheduleAssetReference, @@ -5712,6 +5713,151 @@ def dict_from_obj(obj): assert created_run.creating_job_id == scheduler_job.id + def _seed_asset_producer_and_queue_row(self, session, dag_maker, *, outlet_events=None): + """Create a producer DAG whose task has an asset outlet, a downstream DAG scheduled on that + asset, a successful producer TI, and a matching ``AssetEventQueue`` row. Returns the TI.""" + asset = Asset(uri="test://drain-asset", name="drain_asset", group="test_group") + + with dag_maker(dag_id="drain-producer", start_date=timezone.utcnow(), session=session) as dag: + BashOperator(task_id="task", bash_command="echo 1", outlets=[asset]) + dr = dag_maker.create_dagrun(run_id="producer-run") + [ti] = dr.get_task_instances(session=session) + ti.state = State.SUCCESS + # Capture the outlet profiles before creating other DAGs staled the serialized reference. + task_outlets = [o.asprofile().model_dump(mode="json") for o in dag.get_task("task").outlets] + session.flush() + + # Downstream DAG scheduled on the asset -> registration should enqueue it (ADRQ row). + with dag_maker(dag_id="drain-consumer", schedule=[asset], session=session): + EmptyOperator(task_id="downstream") + session.add( + AssetEventQueue( + ti_id=ti.id, + payload={ + "task_outlets": task_outlets, + "outlet_events": outlet_events or [], + "ti_key": { + "dag_id": ti.dag_id, + "run_id": ti.run_id, + "task_id": ti.task_id, + "map_index": ti.map_index, + }, + }, + attempts=0, + ) + ) + session.commit() + return ti + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_registers_and_deletes(self, session, dag_maker): + """A pending queue row for a successful producer TI is registered (AssetEvent + downstream + AssetDagRunQueue), the queue row is deleted, and the processed metric is emitted.""" + ti = self._seed_asset_producer_and_queue_row(session, dag_maker) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + with mock.patch("airflow.jobs.scheduler_job_runner.stats") as mock_stats: + self.job_runner._drain_asset_event_queue(session=session) + + # The queue row was consumed. + assert session.scalars(select(AssetEventQueue)).all() == [] + # An AssetEvent was produced by the (now drained) TI. + event = session.scalar(select(AssetEvent).where(AssetEvent.source_run_id == ti.run_id)) + assert event is not None + # The downstream asset-scheduled DAG was queued. + adrq = session.scalars( + select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id == "drain-consumer") + ).all() + assert len(adrq) == 1 + mock_stats.incr.assert_any_call("asset.event_queue.processed", 1) + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_retries_then_parks(self, session, dag_maker): + """A row whose registration keeps failing is retried across drain passes, one attempt per + pass, until it reaches max_attempts, at which point it is parked (no longer claimed) and the + failures metric is emitted exactly once. Further passes leave attempts capped.""" + self._seed_asset_producer_and_queue_row(session, dag_maker) + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + max_attempts = self.job_runner._asset_event_queue_max_attempts + + with mock.patch( + "airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db", + side_effect=RuntimeError("boom"), + ): + # Run one extra pass beyond max_attempts to prove the counter caps and does not overflow. + for cur_attempt in range(max_attempts + 1): + with mock.patch("airflow.jobs.scheduler_job_runner.stats") as mock_stats: + self.job_runner._drain_asset_event_queue(session=session) + session.expire_all() + row = session.scalars(select(AssetEventQueue)).one() + # attempts increases by one per pass until it reaches max_attempts, then stays there + # because the drain only claims rows WHERE attempts < max_attempts. + assert row.attempts == min(cur_attempt + 1, max_attempts) + if cur_attempt + 1 == max_attempts: + # The pass that reaches the cap parks the row and reports it once. + mock_stats.incr.assert_any_call("asset.event_queue.failures") + else: + assert mock.call("asset.event_queue.failures") not in mock_stats.incr.mock_calls + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_missing_ti_purged(self, session, dag_maker): + """A queue row whose task instance can no longer be resolved by natural key is dropped + without error or registration. We point the row's ``ti_key`` at a run that does not exist so + the drain's lookup misses, exercising the purge branch.""" + self._seed_asset_producer_and_queue_row(session, dag_maker) + row = session.scalars(select(AssetEventQueue)).one() + row.payload = {**row.payload, "ti_key": {**row.payload["ti_key"], "run_id": "does-not-exist"}} + session.commit() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + self.job_runner._drain_asset_event_queue(session=session) + + assert session.scalars(select(AssetEventQueue)).all() == [] + # Nothing was registered because the task instance could not be resolved. + assert session.scalars(select(AssetEvent)).all() == [] + + @pytest.mark.need_serialized_dag + def test_drain_asset_event_queue_empty_noop(self, session, dag_maker): + """With no pending rows the drain is a no-op and creates nothing.""" + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + self.job_runner._drain_asset_event_queue(session=session) + + assert session.scalars(select(AssetEventQueue)).all() == [] + assert session.scalars(select(AssetEvent)).all() == [] + + @pytest.mark.need_serialized_dag + def test_clearing_ti_still_registers_via_natural_key(self, session, dag_maker): + """Clearing/retrying a TI reassigns its uuid7 id. The drain resolves the live task instance + by natural key (dag_id, run_id, task_id, map_index), not by the enqueued surrogate id, so the + pending asset events are still registered after a clear. This runs on the default sqlite + engine (no ``PRAGMA foreign_keys=ON``), which does not honor the FK's ON UPDATE CASCADE -- + exactly the production condition under which a cascade-dependent re-point would drop them.""" + ti = self._seed_asset_producer_and_queue_row(session, dag_maker) + old_id = ti.id + + # Clearing reassigns the id; the same run_id keeps the DagRun partition context intact. + ti.prepare_db_for_next_try(session=session) + session.commit() + assert ti.id != old_id + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + self.job_runner._drain_asset_event_queue(session=session) + + # The events were registered despite the id change, and the queue row was consumed. + assert session.scalars(select(AssetEventQueue)).all() == [] + event = session.scalar(select(AssetEvent).where(AssetEvent.source_run_id == ti.run_id)) + assert event is not None + @pytest.mark.need_serialized_dag @pytest.mark.parametrize( ("catchup", "expects_old_event"), diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index c0d43641559d6..012e588d6b5e6 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -318,6 +318,30 @@ metrics: legacy_name: "-" name_variables: [] + - name: "asset.event_queue.enqueued" + description: "Number of task-completion asset registrations written as a durable queue marker" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "asset.event_queue.processed" + description: "Number of asset event registrations drained from the queue by the scheduler" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "asset.event_queue.failures" + description: "Number of asset event registrations parked after exceeding the retry limit" + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "asset.event_queue.pending" + description: "Number of asset event registrations currently pending in the queue" + type: "gauge" + legacy_name: "-" + name_variables: [] + - name: "deadline_alerts.deadline_created" description: "Number of deadline alerts created for a Dag run" type: "counter" diff --git a/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py b/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py index b2c75db65ef52..f290411a094b7 100644 --- a/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py +++ b/task-sdk-integration-tests/tests/task_sdk_tests/conftest.py @@ -725,10 +725,12 @@ def sdk_client_for_assets(asset_test_setup): @pytest.fixture(scope="session") def asset_test_setup(docker_compose_setup, airflow_ready): """Setup assets for testing by triggering asset_producer_dag.""" + import time + headers = airflow_ready["headers"] auth_token = airflow_ready["auth_token"] - return setup_dag_and_get_client( + setup = setup_dag_and_get_client( dag_id="asset_producer_dag", headers=headers, auth_token=auth_token, @@ -741,3 +743,22 @@ def asset_test_setup(docker_compose_setup, airflow_ready): "alias_name": "test_asset_alias", }, ) + + # A successful task only enqueues its asset events; the scheduler registers them when it + # drains the queue, so there is a short delay between the producer Dag reaching "success" + # and the events being queryable. Wait for the first event to appear before handing the + # setup to the tests, otherwise they race the scheduler drain and see zero events. + sdk_client = setup["sdk_client"] + for attempt in range(30): + response = sdk_client.asset_events.get(name=setup["name"]) + if response.asset_events: + console.print(f"[green]✅ Asset event registered after {attempt + 1} attempt(s)") + break + console.print( + f"[blue]Waiting for asset event registration (attempt {attempt + 1}/30, drained by scheduler)" + ) + time.sleep(2) + else: + raise TimeoutError("Asset events for asset_producer_dag were not registered within timeout period") + + return setup diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index cadfe6978300e..f46b3d859d120 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -1571,6 +1571,16 @@ def _run_task( break raise + # On success the execution API only enqueues the task's asset events as a durable marker; + # in a live deployment the scheduler drains that queue. dag.test runs in process with no + # scheduler, so register the events in line here (the same way the triggerer is run in + # process above) rather than leaving them stuck in the queue. + if ti.state == TaskInstanceState.SUCCESS: + from airflow.models.asset import register_pending_asset_events + + with create_session() as session: + register_pending_asset_events(ti_ids=[ti.id], session=session) + log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, ti.map_index) return taskrun_result