Skip to content

Reduce scheduler memory for asset-triggered Dag runs #69849

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
viiccwen wants to merge 2 commits into
base: main
Choose a base branch
from
+112 −30
Open

Reduce scheduler memory for asset-triggered Dag runs #69849

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f1963ec
Reduce scheduler memory for asset-triggered Dag runs
viiccwen Jul 14, 2026
bb68de6
Preserve asset event ordering in bulk Dag run associations
viiccwen Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 59 additions & 26 deletions airflow-core/src/airflow/jobs/scheduler_job_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,9 @@
delete,
exists,
func,
insert,
inspect,
literal,
or_,
select,
text,
Expand Down Expand Up @@ -85,6 +87,7 @@
PartitionedAssetKeyLog,
TaskInletAssetReference,
TaskOutletAssetReference,
association_table,
)
from airflow.models.asset_state_store import AssetStateStoreModel
from airflow.models.backfill import Backfill, BackfillDagRun
Expand Down Expand Up @@ -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
Expand All @@ -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: Select[tuple[int]],
session: Session,
) -> None:
"""Associate selected asset events without materializing ORM objects."""
selected_event_ids = asset_event_ids_select.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.
Expand Down Expand Up @@ -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 = 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_select=asset_event_ids_select,
session=session,
)
dag_run.consumed_asset_events.extend(asset_events)
session.flush()
apdr.created_dag_run_id = dag_run.id
session.flush()
Expand Down Expand Up @@ -2619,27 +2642,25 @@ 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 = (
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),
)
.order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc())
)

dag_run = dag.create_dagrun(
Expand All @@ -2661,12 +2682,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_select=asset_event_ids_select,
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
Expand Down
57 changes: 53 additions & 4 deletions airflow-core/tests/unit/jobs/test_scheduler_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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=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):
"""
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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
Expand Down