Stop failing scheduled tasks when a serialized Dag is briefly missing

A serialized Dag that cannot be read is almost always transient -- the Dag processor
is mid-write, or the row was briefly unreadable -- but the scheduler treated it as
final and failed every SCHEDULED task instance of the Dag, across all of its Dag
runs, bypassing retries and failure callbacks. The Dag run was not finished either,
so those tasks could only be recovered by clearing them by hand.

The bulk update also doubled as a starvation filter: flipping the rows to FAILED is
what kept them out of the next iteration of the critical-section query. Starving the
run preserves that effect, so an unresolvable Dag cannot hold up the Dags queued
behind it. The error is reported once per Dag rather than once per task instance, so
a Dag with many stuck runs cannot flood the scheduler log now that its task
instances survive the round.

closes: #62050

Author: rjgoyln

airflow-core/newsfragments/72814.bugfix.rst

@@ -0,0 +1 @@
+Stop failing every ``SCHEDULED`` task instance of a Dag when its serialized Dag cannot be read while evaluating task-level concurrency limits; such a lookup is usually transient, so the task instances are now left ``SCHEDULED`` and reconsidered on the next scheduling loop.

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -587,6 +587,7 @@ def _task_concurrency_allows_execution(
         task_instance: TI,
         concurrency_map: ConcurrencyMap,
         session: Session,
+        starved_dag_runs: dict[str, set[str]],
         starved_tasks: set[tuple[str, str]],
         starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]],
     ) -> bool:
@@ -595,25 +596,30 @@ def _task_concurrency_allows_execution(
         task_id = task_instance.task_id
         run_id = task_instance.run_id
 
+        if run_id in starved_dag_runs.get(dag_id, ()):
+            return False
+
         serialized_dag = self.scheduler_dag_bag.get_dag_for_run(
             dag_run=task_instance.dag_run,
             session=session,
         )
 
-        # If the DAG is missing, fail all scheduled TIs for this DAG.
+        # Leave the task instances SCHEDULED: the miss is usually transient, and when it is not,
+        # _schedule_dag_run cannot progress the run either, so failing its task instances costs
+        # them their retries and callbacks without unsticking the run. Starving the run is what
+        # keeps this loop live: unstarved, these task instances refill every query iteration and
+        # no other Dag is ever examined.
         if not serialized_dag:
-            self.log.error(
-                "DAG '%s' for task instance %s not found in serialized_dag table",
-                dag_id,
-                task_instance,
-            )
+            if dag_id not in starved_dag_runs:
+                self.log.error(
+                    "Dag '%s' (run %s) for task instance %s not found in serialized_dag table; "
+                    "skipping its task instances for this scheduling round",
+                    dag_id,
+                    run_id,
+                    task_instance,
+                )
 
-            session.execute(
-                update(TI)
-                .where(TI.dag_id == dag_id, TI.state == TaskInstanceState.SCHEDULED)
-                .values(state=TaskInstanceState.FAILED)
-                .execution_options(synchronize_session="fetch")
-            )
+            starved_dag_runs.setdefault(dag_id, set()).add(run_id)
 
             return False
 
@@ -723,12 +729,17 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
         starved_dags: set[str] = set()
         starved_tasks: set[tuple[str, str]] = set()
         starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
+        # Run ids, per Dag, whose serialized Dag could not be resolved. Grouped that way because
+        # resolution is per run -- a run pinned to a deleted version misses while another run of
+        # the same Dag still resolves -- so the runs starve one by one but the Dag is logged once.
+        starved_dag_runs: dict[str, set[str]] = {}
 
         pool_num_starving_tasks: dict[str, int] = Counter()
 
         for loop_count in itertools.count(start=1):
             num_starved_pools = len(starved_pools)
             num_starved_dags = len(starved_dags)
+            num_starved_dag_runs = sum(len(run_ids) for run_ids in starved_dag_runs.values())
             num_starved_tasks = len(starved_tasks)
             num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
 
@@ -771,6 +782,17 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
             if starved_dags:
                 query = query.where(TI.dag_id.not_in(starved_dags))
 
+            if starved_dag_runs:
+                query = query.where(
+                    tuple_(TI.dag_id, TI.run_id).not_in(
+                        [
+                            (dag_id, run_id)
+                            for dag_id, run_ids in starved_dag_runs.items()
+                            for run_id in run_ids
+                        ]
+                    )
+                )
+
             if starved_tasks:
                 query = query.where(tuple_(TI.dag_id, TI.task_id).not_in(starved_tasks))
 
@@ -987,6 +1009,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
                         task_instance=task_instance,
                         concurrency_map=concurrency_map,
                         session=session,
+                        starved_dag_runs=starved_dag_runs,
                         starved_tasks=starved_tasks,
                         starved_tasks_task_dagrun_concurrency=(starved_tasks_task_dagrun_concurrency),
                     )
@@ -1035,6 +1058,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
             found_new_filters = (
                 len(starved_pools) > num_starved_pools
                 or len(starved_dags) > num_starved_dags
+                or sum(len(run_ids) for run_ids in starved_dag_runs.values()) > num_starved_dag_runs
                 or len(starved_tasks) > num_starved_tasks
                 or len(starved_tasks_task_dagrun_concurrency) > num_starved_tasks_task_dagrun_concurrency
             )

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -2188,8 +2188,8 @@ def test_find_executable_task_instances_in_default_pool(self, dag_maker, mock_ex
         session.rollback()
         session.close()
 
-    def test_queued_task_instances_fails_with_missing_dag(self, dag_maker, session):
-        """Check that task instances of missing DAGs are failed"""
+    def test_queued_task_instances_skipped_when_serialized_dag_missing(self, dag_maker, session):
+        """An unresolvable serialized Dag leaves every run's task instances SCHEDULED for a retry."""
         dag_id = "SchedulerJobTest.test_find_executable_task_instances_not_in_dagbag"
         task_id_1 = "dummy"
         task_id_2 = "dummydummy"
@@ -2201,22 +2201,135 @@ def test_queued_task_instances_fails_with_missing_dag(self, dag_maker, session):
         scheduler_job = Job()
         self.job_runner = SchedulerJobRunner(job=scheduler_job)
 
-        self.job_runner.scheduler_dag_bag = mock.MagicMock()
+        self.job_runner.scheduler_dag_bag = mock.MagicMock(spec=CachedDBDagBag)
         self.job_runner.scheduler_dag_bag.get_dag_for_run.return_value = None
 
+        dr1 = dag_maker.create_dagrun(run_id="run_1", state=DagRunState.RUNNING)
+        dr2 = dag_maker.create_dagrun_after(dr1, run_id="run_2", state=DagRunState.RUNNING)
+
+        for dr in (dr1, dr2):
+            for ti in dr.task_instances:
+                ti.state = State.SCHEDULED
+                session.merge(ti)
+        session.flush()
+        res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session)
+        session.flush()
+        assert len(res) == 0
+        for dr in (dr1, dr2):
+            tis = dr.get_task_instances(session=session)
+            assert len(tis) == 2
+            assert all(ti.state == State.SCHEDULED for ti in tis)
+
+        # The lookup is made once per run, not once for each of the four task instances.
+        assert self.job_runner.scheduler_dag_bag.get_dag_for_run.call_count == 2
+
+    def test_missing_serialized_dag_for_one_run_does_not_block_another_run(self, dag_maker, session):
+        """Resolution is per run, so one run's miss must not hold back the Dag's other runs."""
+        dag_id = "SchedulerJobTest.test_missing_serialized_dag_one_run"
+
+        with dag_maker(dag_id=dag_id, session=session, default_args={"max_active_tis_per_dag": 2}):
+            EmptyOperator(task_id="dummy")
+
+        dr1 = dag_maker.create_dagrun(run_id="run_1", state=DagRunState.RUNNING)
+        dr2 = dag_maker.create_dagrun_after(dr1, run_id="run_2", state=DagRunState.RUNNING)
+        serialized_dag = dag_maker.serialized_dag
+
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(job=scheduler_job)
+        self.job_runner.scheduler_dag_bag = mock.MagicMock(spec=CachedDBDagBag)
+        self.job_runner.scheduler_dag_bag.get_dag_for_run.side_effect = lambda dag_run, session: (
+            None if dag_run.run_id == "run_1" else serialized_dag
+        )
+
+        for dr in (dr1, dr2):
+            for ti in dr.task_instances:
+                ti.state = State.SCHEDULED
+                session.merge(ti)
+        session.flush()
+
+        res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session)
+        session.flush()
+
+        assert [(ti.dag_id, ti.run_id) for ti in res] == [(dag_id, "run_2")]
+        assert all(ti.state == State.SCHEDULED for ti in dr1.get_task_instances(session=session))
+
+    def test_missing_serialized_dag_queues_task_instances_once_it_resolves(self, dag_maker, session):
+        """Task instances left SCHEDULED are queued on a later round once the serialized Dag resolves."""
+        dag_id = "SchedulerJobTest.test_missing_serialized_dag_resolves"
+
+        with dag_maker(dag_id=dag_id, session=session, default_args={"max_active_tis_per_dag": 2}):
+            EmptyOperator(task_id="dummy")
+
         dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+        serialized_dag = dag_maker.serialized_dag
 
-        tis = dr.task_instances
-        for ti in tis:
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(job=scheduler_job)
+        self.job_runner.scheduler_dag_bag = mock.MagicMock(spec=CachedDBDagBag)
+        self.job_runner.scheduler_dag_bag.get_dag_for_run.side_effect = [None, serialized_dag]
+
+        for ti in dr.task_instances:
             ti.state = State.SCHEDULED
             session.merge(ti)
         session.flush()
+
+        assert self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session) == []
+        session.flush()
+        assert all(ti.state == State.SCHEDULED for ti in dr.get_task_instances(session=session))
+
         res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session)
+        assert [ti.task_id for ti in res] == ["dummy"]
+
+    def test_missing_serialized_dag_does_not_starve_other_dags(self, dag_maker, session, caplog):
+        """A Dag that cannot be resolved must not hold up the Dags queued behind it."""
+        broken_dag_id = "SchedulerJobTest.test_missing_serialized_dag_starvation_broken"
+        healthy_dag_id = "SchedulerJobTest.test_missing_serialized_dag_starvation_healthy"
+
+        with dag_maker(
+            dag_id=broken_dag_id,
+            session=session,
+            default_args={"max_active_tis_per_dag": 1, "priority_weight": 100},
+        ):
+            EmptyOperator(task_id="broken")
+        broken_runs = [dag_maker.create_dagrun(run_id="broken_run_0", state=DagRunState.RUNNING)]
+        for run_number in range(1, 4):
+            broken_runs.append(
+                dag_maker.create_dagrun_after(
+                    broken_runs[-1], run_id=f"broken_run_{run_number}", state=DagRunState.RUNNING
+                )
+            )
+
+        with dag_maker(dag_id=healthy_dag_id, session=session, default_args={"priority_weight": 1}):
+            EmptyOperator(task_id="healthy")
+        healthy_run = dag_maker.create_dagrun(run_id="healthy_run", state=DagRunState.RUNNING)
+
+        for dr in (*broken_runs, healthy_run):
+            for ti in dr.task_instances:
+                ti.state = State.SCHEDULED
+                session.merge(ti)
+        session.commit()
+
+        session.execute(delete(SerializedDagModel).where(SerializedDagModel.dag_id == broken_dag_id))
+        session.commit()
+        session.expire_all()
+
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+        # The broken Dag has more runs than one query can return, and its task instances outrank the
+        # healthy Dag's, so the healthy Dag is only reachable if every broken run is starved out over
+        # successive query iterations.
+        res = self.job_runner._executable_task_instances_to_queued(max_tis=2, session=session)
         session.flush()
-        assert len(res) == 0
-        tis = dr.get_task_instances(session=session)
-        assert len(tis) == 2
-        assert all(ti.state == State.FAILED for ti in tis)
+
+        assert [(ti.dag_id, ti.task_id) for ti in res] == [(healthy_dag_id, "healthy")]
+        for dr in broken_runs:
+            assert all(ti.state == State.SCHEDULED for ti in dr.get_task_instances(session=session))
+
+        # One error for the Dag, not one per missing run.
+        errors = [entry for entry in caplog.entries if entry["log_level"] == "error"]
+        assert len(errors) == 1
+        assert f"Dag '{broken_dag_id}'" in errors[0]["event"]
 
     def test_nonexistent_pool(self, dag_maker):
         dag_id = "SchedulerJobTest.test_nonexistent_pool"