From 144a105071dffbef72306ca390679b844f33961d Mon Sep 17 00:00:00 2001 From: Akshat Anand <40275336+cipheraxat@users.noreply.github.com> Date: Sun, 12 Jul 2026 04:58:23 +0530 Subject: [PATCH] Count post-deferral scheduled tasks toward include_deferred pool slots When include_deferred is enabled, tasks that return to scheduled after a trigger fires (next_method set) kept dropping out of pool occupancy, undercounting slots until they were queued again. Count those TIs as occupied without treating first-lifetime scheduled tasks the same way, which would deadlock scheduling. --- airflow-core/newsfragments/69187.bugfix.rst | 1 + airflow-core/src/airflow/models/pool.py | 59 ++++++++++++- .../ti_deps/deps/pool_slots_available_dep.py | 2 +- airflow-core/tests/unit/models/test_pool.py | 88 +++++++++++++++++++ .../deps/test_pool_slots_available_dep.py | 27 ++++++ 5 files changed, 174 insertions(+), 3 deletions(-) create mode 100644 airflow-core/newsfragments/69187.bugfix.rst diff --git a/airflow-core/newsfragments/69187.bugfix.rst b/airflow-core/newsfragments/69187.bugfix.rst new file mode 100644 index 0000000000000..b062feb743196 --- /dev/null +++ b/airflow-core/newsfragments/69187.bugfix.rst @@ -0,0 +1 @@ +Fix pool slot undercounting for tasks that return to ``scheduled`` after deferral when ``include_deferred`` is enabled. Those task instances keep occupying a pool slot (identified via ``next_method``) until they are queued or running again, so pools continue to protect external resources across the deferral boundary. diff --git a/airflow-core/src/airflow/models/pool.py b/airflow-core/src/airflow/models/pool.py index d6a4915ea2bb2..3f04d29e74ab3 100644 --- a/airflow-core/src/airflow/models/pool.py +++ b/airflow-core/src/airflow/models/pool.py @@ -21,7 +21,7 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any, TypedDict -from sqlalchemy import Boolean, ForeignKey, Integer, String, Text, func, select +from sqlalchemy import Boolean, ForeignKey, Integer, String, Text, and_, func, or_, select from sqlalchemy.orm import Mapped, mapped_column from airflow._shared.observability.metrics.stats import normalize_name_for_stats @@ -36,6 +36,8 @@ from sqlalchemy.orm.session import Session from sqlalchemy.sql import Select + from airflow.models.taskinstance import TaskInstance as TaskInstanceModel + logger = logging.getLogger(__name__) @@ -242,11 +244,32 @@ def slots_stats( else: raise AirflowException(f"Unexpected state. Expected values: {allowed_execution_states}.") + # Scheduled TIs that are resuming after deferral still hold a pool slot when + # include_deferred is enabled (see get_occupied_states / occupied_slots). + # Distinguish them via next_method so first-lifetime scheduled TIs are not counted + # (that would deadlock scheduling when open = total - scheduled). + post_deferral_scheduled_by_pool: dict[str, int] = {} + if any(pool_includes_deferred.values()): + post_deferral_rows = session.execute( + select(TaskInstance.pool, func.sum(TaskInstance.pool_slots)) + .where( + TaskInstance.state == TaskInstanceState.SCHEDULED, + TaskInstance.next_method.is_not(None), + ) + .group_by(TaskInstance.pool) + ) + post_deferral_scheduled_by_pool = { + pool_name: int(decimal_count) + for pool_name, decimal_count in post_deferral_rows + if pool_name is not None + } + # calculate open metric for pool_name, stats_dict in pools.items(): stats_dict["open"] = stats_dict["total"] - stats_dict["running"] - stats_dict["queued"] if pool_includes_deferred[pool_name]: stats_dict["open"] -= stats_dict["deferred"] + stats_dict["open"] -= post_deferral_scheduled_by_pool.get(pool_name, 0) return pools @@ -270,29 +293,61 @@ def occupied_slots(self, *, session: Session = NEW_SESSION) -> int: """ Get the number of slots used by running/queued tasks at the moment. + When ``include_deferred`` is enabled, deferred tasks and scheduled tasks that are + resuming after deferral (``next_method`` is set) also count as occupied. + :param session: SQLAlchemy ORM Session :return: the used number of slots """ from airflow.models.taskinstance import TaskInstance # Avoid circular import occupied_states = self.get_occupied_states() + occupancy_filter = TaskInstance.state.in_(occupied_states) + if self.include_deferred: + # Post-deferral scheduled TIs keep holding the slot until they are queued/running. + occupancy_filter = or_( + occupancy_filter, + and_( + TaskInstance.state == TaskInstanceState.SCHEDULED, + TaskInstance.next_method.is_not(None), + ), + ) return int( session.scalar( select(func.sum(TaskInstance.pool_slots)) .filter(TaskInstance.pool == self.pool) - .filter(TaskInstance.state.in_(occupied_states)) + .filter(occupancy_filter) ) or 0 ) def get_occupied_states(self): + """ + Return states that always occupy pool slots for this pool. + + Post-deferral ``scheduled`` occupancy (``next_method`` set) is handled separately in + :meth:`occupied_slots` and :meth:`task_instance_occupies_slot` because it cannot be + expressed as a state-only set without also counting first-lifetime scheduled tasks. + """ if self.include_deferred: return EXECUTION_STATES | { TaskInstanceState.DEFERRED, } return EXECUTION_STATES + def task_instance_occupies_slot(self, ti: TaskInstanceModel) -> bool: + """ + Return whether ``ti`` currently occupies a slot in this pool. + + Used by pool availability deps so a TI that already holds a slot does not block itself. + """ + if ti.state in self.get_occupied_states(): + return True + return bool( + self.include_deferred and ti.state == TaskInstanceState.SCHEDULED and ti.next_method is not None + ) + @provide_session def running_slots(self, *, session: Session = NEW_SESSION) -> int: """ diff --git a/airflow-core/src/airflow/ti_deps/deps/pool_slots_available_dep.py b/airflow-core/src/airflow/ti_deps/deps/pool_slots_available_dep.py index b70443cf2c14c..67d4d387c6a64 100644 --- a/airflow-core/src/airflow/ti_deps/deps/pool_slots_available_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/pool_slots_available_dep.py @@ -71,7 +71,7 @@ def _get_dep_statuses(self, ti, dep_context=None, *, session): return open_slots = pool.open_slots(session=session) - if ti.state in pool.get_occupied_states(): + if pool.task_instance_occupies_slot(ti): open_slots += ti.pool_slots if open_slots <= (ti.pool_slots - 1): diff --git a/airflow-core/tests/unit/models/test_pool.py b/airflow-core/tests/unit/models/test_pool.py index 2c7db9d9800e4..47e20dd3675d9 100644 --- a/airflow-core/tests/unit/models/test_pool.py +++ b/airflow-core/tests/unit/models/test_pool.py @@ -178,6 +178,94 @@ def test_open_slots_including_deferred(self, dag_maker): }, } + def test_open_slots_including_deferred_counts_post_deferral_scheduled(self, dag_maker): + """Scheduled TIs resuming after deferral (next_method set) occupy slots when include_deferred.""" + pool = Pool(pool="test_pool", slots=5, include_deferred=True) + with dag_maker( + dag_id="test_open_slots_post_deferral_scheduled", + start_date=DEFAULT_DATE, + ): + op1 = EmptyOperator(task_id="dummy1", pool="test_pool") + op2 = EmptyOperator(task_id="dummy2", pool="test_pool") + op3 = EmptyOperator(task_id="dummy3", pool="test_pool") + + dr = dag_maker.create_dagrun() + + ti1 = dr.get_task_instance(task_id=op1.task_id) + ti2 = dr.get_task_instance(task_id=op2.task_id) + ti3 = dr.get_task_instance(task_id=op3.task_id) + ti1.state = State.RUNNING + # Post-deferral resume: scheduled with next_method set. + ti2.state = State.SCHEDULED + ti2.next_method = "execute_complete" + # First-lifetime scheduled: must not occupy a slot. + ti3.state = State.SCHEDULED + ti3.next_method = None + + session = settings.Session() + session.add(pool) + session.merge(ti1) + session.merge(ti2) + session.merge(ti3) + session.commit() + session.close() + + assert pool.open_slots() == 3 + assert pool.running_slots() == 1 + assert pool.scheduled_slots() == 2 + assert pool.occupied_slots() == 2 + assert pool.task_instance_occupies_slot(ti2) is True + assert pool.task_instance_occupies_slot(ti3) is False + assert pool.slots_stats() == { + "default_pool": { + "open": 128, + "queued": 0, + "total": 128, + "running": 0, + "scheduled": 0, + "deferred": 0, + }, + "test_pool": { + "open": 3, + "queued": 0, + "running": 1, + "deferred": 0, + "scheduled": 2, + "total": 5, + }, + } + + def test_open_slots_excluding_deferred_ignores_post_deferral_scheduled(self, dag_maker): + """Without include_deferred, post-deferral scheduled TIs do not occupy slots.""" + pool = Pool(pool="test_pool", slots=5, include_deferred=False) + with dag_maker( + dag_id="test_open_slots_exclude_post_deferral_scheduled", + start_date=DEFAULT_DATE, + ): + op1 = EmptyOperator(task_id="dummy1", pool="test_pool") + op2 = EmptyOperator(task_id="dummy2", pool="test_pool") + + dr = dag_maker.create_dagrun() + + ti1 = dr.get_task_instance(task_id=op1.task_id) + ti2 = dr.get_task_instance(task_id=op2.task_id) + ti1.state = State.RUNNING + ti2.state = State.SCHEDULED + ti2.next_method = "execute_complete" + + session = settings.Session() + session.add(pool) + session.merge(ti1) + session.merge(ti2) + session.commit() + session.close() + + assert pool.open_slots() == 4 + assert pool.occupied_slots() == 1 + assert pool.task_instance_occupies_slot(ti2) is False + assert pool.slots_stats()["test_pool"]["open"] == 4 + assert pool.slots_stats()["test_pool"]["scheduled"] == 1 + def test_infinite_slots(self, dag_maker): pool = Pool(pool="test_pool", slots=-1, include_deferred=False) with dag_maker( diff --git a/airflow-core/tests/unit/ti_deps/deps/test_pool_slots_available_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_pool_slots_available_dep.py index c65f96589f6c6..9812da7765106 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_pool_slots_available_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_pool_slots_available_dep.py @@ -70,6 +70,33 @@ def test_deferred_pooled_task_pass(self, mock_open_slots): ti_to_fail = Mock(pool="test_pool", state=TaskInstanceState.DEFERRED, pool_slots=1) assert not PoolSlotsAvailableDep().is_met(ti=ti_to_fail) + @patch("airflow.models.Pool.open_slots", return_value=0) + def test_post_deferral_scheduled_pooled_task_pass(self, mock_open_slots): + """Scheduled TIs resuming after deferral already hold a slot when include_deferred is on.""" + ti = Mock( + pool="test_includes_deferred_pool", + state=TaskInstanceState.SCHEDULED, + pool_slots=1, + next_method="execute_complete", + ) + assert PoolSlotsAvailableDep().is_met(ti=ti) + # First-lifetime scheduled (no next_method) must still require an open slot. + ti_first_schedule = Mock( + pool="test_includes_deferred_pool", + state=TaskInstanceState.SCHEDULED, + pool_slots=1, + next_method=None, + ) + assert not PoolSlotsAvailableDep().is_met(ti=ti_first_schedule) + # include_deferred=False never treats scheduled as occupied. + ti_excluded = Mock( + pool="test_pool", + state=TaskInstanceState.SCHEDULED, + pool_slots=1, + next_method="execute_complete", + ) + assert not PoolSlotsAvailableDep().is_met(ti=ti_excluded) + def test_task_with_nonexistent_pool(self): ti = Mock(pool="nonexistent_pool", pool_slots=1) assert not PoolSlotsAvailableDep().is_met(ti=ti)