From f1963ec777bf9137e367c65f018ded0d36b06f7e Mon Sep 17 00:00:00 2001 From: viiccwen Date: Tue, 14 Jul 2026 12:33:25 +0800 Subject: [PATCH 1/2] Reduce scheduler memory for asset-triggered Dag runs Large asset event windows can exhaust scheduler memory because every event, including JSON metadata, is materialized as an ORM object before its Dag run association is persisted. Signed-off-by: viiccwen --- .../src/airflow/jobs/scheduler_job_runner.py | 85 +++++++++++++------ .../tests/unit/jobs/test_scheduler_job.py | 57 ++++++++++++- 2 files changed, 110 insertions(+), 32 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 7d5f7f58020c7..18d7d24e9c611 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -43,7 +43,9 @@ delete, exists, func, + insert, inspect, + literal, or_, select, text, @@ -85,6 +87,7 @@ PartitionedAssetKeyLog, TaskInletAssetReference, TaskOutletAssetReference, + association_table, ) from airflow.models.asset_state_store import AssetStateStoreModel from airflow.models.backfill import Backfill, BackfillDagRun @@ -138,6 +141,7 @@ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session from sqlalchemy.orm.interfaces import LoaderOption + from sqlalchemy.sql import Select from sqlalchemy.sql.selectable import Subquery from airflow._shared.logging.types import Logger @@ -162,6 +166,23 @@ MAX_PARTITION_DAG_RUNS_PER_LOOP = 500 +def _associate_asset_events_with_dag_run( + *, + dag_run: DagRun, + asset_event_ids: Select[tuple[int]], + session: Session, +) -> None: + """Associate selected asset events without materializing ORM objects.""" + selected_event_ids = asset_event_ids.subquery() + session.execute( + insert(association_table).from_select( + ["dag_run_id", "event_id"], + select(literal(dag_run.id), selected_event_ids.c.event_id), + ) + ) + session.expire(dag_run, ["consumed_asset_events"]) + + def _eager_load_dag_run_for_validation() -> tuple[LoaderOption, LoaderOption]: """ Eager-load DagRun relations required for execution API datamodel validation. @@ -2341,13 +2362,15 @@ def _create_dagruns_for_partitioned_asset_dags(self, session: Session) -> set[st creating_job_id=self.job.id, session=session, ) - asset_events = session.scalars( - select(AssetEvent).where( - PartitionedAssetKeyLog.asset_partition_dag_run_id == apdr.id, - PartitionedAssetKeyLog.asset_event_id == AssetEvent.id, - ) + asset_event_ids = select(AssetEvent.id.label("event_id")).where( + PartitionedAssetKeyLog.asset_partition_dag_run_id == apdr.id, + PartitionedAssetKeyLog.asset_event_id == AssetEvent.id, + ) + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids=asset_event_ids, + session=session, ) - dag_run.consumed_asset_events.extend(asset_events) session.flush() apdr.created_dag_run_id = dag_run.id session.flush() @@ -2619,27 +2642,21 @@ def _create_dag_runs_asset_triggered( ) event_window_floor.append(date.min) - asset_events = list( - session.scalars( - select(AssetEvent) - .where( - or_( - AssetEvent.asset_id.in_( - select(DagScheduleAssetReference.asset_id).where( - DagScheduleAssetReference.dag_id == dag.dag_id - ) - ), - AssetEvent.source_aliases.any( - AssetAliasModel.scheduled_dags.any( - DagScheduleAssetAliasReference.dag_id == dag.dag_id - ) - ), + asset_event_ids = select(AssetEvent.id.label("event_id")).where( + or_( + AssetEvent.asset_id.in_( + select(DagScheduleAssetReference.asset_id).where( + DagScheduleAssetReference.dag_id == dag.dag_id ), - AssetEvent.timestamp <= triggered_date, - AssetEvent.timestamp > func.coalesce(*event_window_floor), - ) - .order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc()) - ) + ), + AssetEvent.source_aliases.any( + AssetAliasModel.scheduled_dags.any( + DagScheduleAssetAliasReference.dag_id == dag.dag_id + ) + ), + ), + AssetEvent.timestamp <= triggered_date, + AssetEvent.timestamp > func.coalesce(*event_window_floor), ) dag_run = dag.create_dagrun( @@ -2661,12 +2678,24 @@ def _create_dag_runs_asset_triggered( else None ) stats.incr("asset.triggered_dagruns", tags=prune_dict({"team_name": team_name})) - dag_run.consumed_asset_events.extend(asset_events) + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids=asset_event_ids, + session=session, + ) + consumed_asset_event_count = ( + session.scalar( + select(func.count()) + .select_from(association_table) + .where(association_table.c.dag_run_id == dag_run.id) + ) + or 0 + ) self.log.info( "Created asset-triggered DagRun for '%s': run_id=%s, consumed %d asset events", dag.dag_id, dag_run.run_id, - len(asset_events), + consumed_asset_event_count, ) # Delete only consumed ADRQ rows to avoid dropping newly queued events diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 1f3b6992c3e1c..5d47fb6e92c71 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -36,7 +36,7 @@ import psutil import pytest import time_machine -from sqlalchemy import delete, func, inspect, select, update +from sqlalchemy import delete, event as sqlalchemy_event, func, inspect, select, update from sqlalchemy.dialects import mysql from sqlalchemy.orm import joinedload @@ -57,7 +57,7 @@ from airflow.executors.executor_utils import ExecutorName from airflow.executors.local_executor import LocalExecutor from airflow.jobs.job import Job, run_job -from airflow.jobs.scheduler_job_runner import SchedulerJobRunner +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner, _associate_asset_events_with_dag_run from airflow.models.asset import ( AssetActive, AssetAliasModel, @@ -189,6 +189,29 @@ / "standard" / "example_dags" ) + + +@contextmanager +def assert_no_asset_event_materialization(session: Session) -> Generator[None, None, None]: + """Assert a scheduler operation does not select full AssetEvent entities.""" + materializing_queries = [] + + def track_orm_execute(orm_execute_state): + if orm_execute_state.is_select and any( + description.get("expr") is AssetEvent + for description in orm_execute_state.statement.column_descriptions + ): + materializing_queries.append(orm_execute_state.statement) + + sqlalchemy_event.listen(session, "do_orm_execute", track_orm_execute) + try: + yield + finally: + sqlalchemy_event.remove(session, "do_orm_execute", track_orm_execute) + + assert materializing_queries == [] + + DEFAULT_DATE = timezone.datetime(2016, 1, 1) DEFAULT_LOGICAL_DATE = timezone.coerce_datetime(DEFAULT_DATE) TRY_NUMBER = 1 @@ -5557,6 +5580,30 @@ def test_create_dag_runs(self, dag_maker): assert dr.start_date is None assert dr.creating_job_id == scheduler_job.id + def test_associate_asset_events_with_dag_run_keeps_relationship_coherent(self, session, dag_maker): + with dag_maker(dag_id="asset-event-association"): + pass + dag_run = dag_maker.create_dagrun() + asset = AssetModel(name="association-asset", uri="association-asset", group="asset") + session.add(asset) + session.flush() + asset_event = AssetEvent(asset_id=asset.id) + session.add(asset_event) + session.flush() + asset_event_id = asset_event.id + session.expunge(asset_event) + + with assert_queries_count(1, session=session): + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids=select(AssetEvent.id.label("event_id")).where( + AssetEvent.id == asset_event_id + ), + session=session, + ) + + assert [event.id for event in dag_run.consumed_asset_events] == [asset_event_id] + @pytest.mark.need_serialized_dag def test_create_dag_runs_assets(self, session, dag_maker): """ @@ -5632,7 +5679,8 @@ def test_create_dag_runs_assets(self, session, dag_maker): self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) with create_session() as session: - self.job_runner._create_dagruns_for_dags(session, session) + with assert_no_asset_event_materialization(session): + self.job_runner._create_dagruns_for_dags(session, session) def dict_from_obj(obj): """Get dict of column attrs from SqlAlchemy object.""" @@ -10727,7 +10775,8 @@ def test_partitioned_dag_run_with_customized_mapper( dag_maker=dag_maker, expected_partition_key="key-1", ) - partition_dags = runner._create_dagruns_for_partitioned_asset_dags(session=session) + with assert_no_asset_event_materialization(session): + partition_dags = runner._create_dagruns_for_partitioned_asset_dags(session=session) session.refresh(apdr) # Since asset event for Asset(name="asset-2") with key "key-1" has not yet been created, # no Dag run will be created From bb68de6546639d51d135e088689077cb22a1ffea Mon Sep 17 00:00:00 2001 From: viiccwen Date: Tue, 14 Jul 2026 22:48:49 +0800 Subject: [PATCH 2/2] Preserve asset event ordering in bulk Dag run associations Asset-triggered Dag runs should retain the established event selection order while avoiding ORM materialization. Scoping the ordering to scheduler association creation avoids adding sort costs to every consumed_asset_events relationship load. Signed-off-by: viiccwen --- .../src/airflow/jobs/scheduler_job_runner.py | 40 ++++++++++--------- .../tests/unit/jobs/test_scheduler_job.py | 2 +- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 18d7d24e9c611..d15ab59b0dec2 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -169,11 +169,11 @@ def _associate_asset_events_with_dag_run( *, dag_run: DagRun, - asset_event_ids: Select[tuple[int]], + asset_event_ids_select: Select[tuple[int]], session: Session, ) -> None: """Associate selected asset events without materializing ORM objects.""" - selected_event_ids = asset_event_ids.subquery() + selected_event_ids = asset_event_ids_select.subquery() session.execute( insert(association_table).from_select( ["dag_run_id", "event_id"], @@ -2362,13 +2362,13 @@ def _create_dagruns_for_partitioned_asset_dags(self, session: Session) -> set[st creating_job_id=self.job.id, session=session, ) - asset_event_ids = select(AssetEvent.id.label("event_id")).where( + asset_event_ids_select = select(AssetEvent.id.label("event_id")).where( PartitionedAssetKeyLog.asset_partition_dag_run_id == apdr.id, PartitionedAssetKeyLog.asset_event_id == AssetEvent.id, ) _associate_asset_events_with_dag_run( dag_run=dag_run, - asset_event_ids=asset_event_ids, + asset_event_ids_select=asset_event_ids_select, session=session, ) session.flush() @@ -2642,21 +2642,25 @@ def _create_dag_runs_asset_triggered( ) event_window_floor.append(date.min) - asset_event_ids = select(AssetEvent.id.label("event_id")).where( - or_( - AssetEvent.asset_id.in_( - select(DagScheduleAssetReference.asset_id).where( - DagScheduleAssetReference.dag_id == dag.dag_id + asset_event_ids_select = ( + select(AssetEvent.id.label("event_id")) + .where( + or_( + AssetEvent.asset_id.in_( + select(DagScheduleAssetReference.asset_id).where( + DagScheduleAssetReference.dag_id == dag.dag_id + ), + ), + AssetEvent.source_aliases.any( + AssetAliasModel.scheduled_dags.any( + DagScheduleAssetAliasReference.dag_id == dag.dag_id + ) ), ), - AssetEvent.source_aliases.any( - AssetAliasModel.scheduled_dags.any( - DagScheduleAssetAliasReference.dag_id == dag.dag_id - ) - ), - ), - AssetEvent.timestamp <= triggered_date, - AssetEvent.timestamp > func.coalesce(*event_window_floor), + AssetEvent.timestamp <= triggered_date, + AssetEvent.timestamp > func.coalesce(*event_window_floor), + ) + .order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc()) ) dag_run = dag.create_dagrun( @@ -2680,7 +2684,7 @@ def _create_dag_runs_asset_triggered( stats.incr("asset.triggered_dagruns", tags=prune_dict({"team_name": team_name})) _associate_asset_events_with_dag_run( dag_run=dag_run, - asset_event_ids=asset_event_ids, + asset_event_ids_select=asset_event_ids_select, session=session, ) consumed_asset_event_count = ( diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 5d47fb6e92c71..79e67b98d8128 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -5596,7 +5596,7 @@ def test_associate_asset_events_with_dag_run_keeps_relationship_coherent(self, s with assert_queries_count(1, session=session): _associate_asset_events_with_dag_run( dag_run=dag_run, - asset_event_ids=select(AssetEvent.id.label("event_id")).where( + asset_event_ids_select=select(AssetEvent.id.label("event_id")).where( AssetEvent.id == asset_event_id ), session=session,