Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 1 addition & 3 deletions airflow-core/src/airflow/jobs/scheduler_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -2616,6 +2616,7 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) -
session=session,
active_non_backfill_runs=active_runs_of_dags[dag_model.dag_id],
)
session.flush()

# Exceptions like ValueError, ParamValidationError, etc. are raised by
# DagModel.create_dagrun() when dag is misconfigured. The scheduler should not
Expand All @@ -2628,9 +2629,6 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) -
# commit after every dag run or use savepoints.
# https://github.com/apache/airflow/issues/59120

# TODO[HA]: Should we do a session.flush() so we don't have to keep lots of state/object in
# memory for larger dags? or expunge_all()

def _create_dag_runs_asset_triggered(
self,
*,
Expand Down
38 changes: 38 additions & 0 deletions airflow-core/tests/unit/jobs/test_scheduler_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -4116,6 +4116,44 @@ def test_queued_dagruns_stops_creating_when_max_active_is_reached(self, dag_make
assert session.scalar(select(func.count(DagRun.state)).where(DagRun.state == State.QUEUED)) == 0
assert orm_dag.next_dagrun_create_after is not None

def test_create_dag_runs_flushes_without_retaining_per_task_objects(self, dag_maker, session):
dag_count = 5
task_count = 20
dag_ids = []
for dag_index in range(dag_count):
dag_id = f"test_create_dag_runs_flushes_without_retaining_per_task_objects_{dag_index}"
with dag_maker(dag_id=dag_id, session=session):
for task_index in range(task_count):
EmptyOperator(task_id=f"task_{task_index}")
dag_ids.append(dag_id)

scheduler_job = Job()
self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec])

orm_dags = list(session.scalars(select(DagModel).where(DagModel.dag_id.in_(dag_ids))))
assert len(orm_dags) == dag_count
session.flush()
assert not session.new
assert not session.dirty

self.job_runner._create_dag_runs(orm_dags, session)

identity_map_counts = Counter(type(obj) for obj in session.identity_map.values())
assert not session.new
assert not session.dirty
assert identity_map_counts[DagRun] == 0
assert identity_map_counts[TaskInstance] == 0
assert (
session.scalar(select(func.count()).select_from(DagRun).where(DagRun.dag_id.in_(dag_ids)))
== dag_count
)
assert (
session.scalar(
select(func.count()).select_from(TaskInstance).where(TaskInstance.dag_id.in_(dag_ids))
)
== dag_count * task_count
)

def test_runs_are_created_after_max_active_runs_was_reached(self, dag_maker, session):
"""
Test that when creating runs once max_active_runs is reached the runs does not stick
Expand Down