This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 345ddb8dfb5 Fix max_active_tis_per_dag for deferred task instances
(#61769)
345ddb8dfb5 is described below
commit 345ddb8dfb554aa3a31f1f761d5ba1dbaafc831a
Author: Saksham Singhal <[email protected]>
AuthorDate: Mon May 4 12:29:25 2026 +0530
Fix max_active_tis_per_dag for deferred task instances (#61769)
Co-authored-by: Tzu-ping Chung <[email protected]>
---
.../src/airflow/jobs/scheduler_job_runner.py | 21 +-
airflow-core/src/airflow/models/taskinstance.py | 56 ++++-
.../src/airflow/ti_deps/dependencies_states.py | 10 +
.../airflow/ti_deps/deps/task_concurrency_dep.py | 4 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 268 +++++++++++++++++++++
.../tests/unit/models/test_taskinstance.py | 70 ++++--
.../unit/ti_deps/deps/test_task_concurrency.py | 13 +-
7 files changed, 396 insertions(+), 46 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 5076ab1608a..640060158e5 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -87,7 +87,7 @@ from airflow.models.trigger import TRIGGER_FAIL_REPR,
Trigger, TriggerFailureRea
from airflow.observability.metrics import stats_utils
from airflow.serialization.definitions.assets import SerializedAssetUniqueKey
from airflow.serialization.definitions.notset import NOTSET
-from airflow.ti_deps.dependencies_states import EXECUTION_STATES
+from airflow.ti_deps.dependencies_states import ACTIVE_STATES, EXECUTION_STATES
from airflow.timetables.simple import AssetTriggeredTimetable
from airflow.utils.event_scheduler import EventScheduler
from airflow.utils.log.logging_mixin import LoggingMixin
@@ -206,14 +206,19 @@ class ConcurrencyMap:
self.task_concurrency_map.clear()
self.task_dagrun_concurrency_map.clear()
query = session.execute(
- select(TI.dag_id, TI.task_id, TI.run_id, func.count("*"))
- .where(TI.state.in_(EXECUTION_STATES))
- .group_by(TI.task_id, TI.run_id, TI.dag_id)
+ select(TI.dag_id, TI.task_id, TI.run_id, TI.state, func.count("*"))
+ .where(TI.state.in_(ACTIVE_STATES))
+ .group_by(TI.dag_id, TI.task_id, TI.run_id, TI.state)
)
- for dag_id, task_id, run_id, c in query:
- self.dag_run_active_tasks_map[dag_id, run_id] += c
- self.task_concurrency_map[(dag_id, task_id)] += c
- self.task_dagrun_concurrency_map[(dag_id, run_id, task_id)] += c
+ for dag_id, task_id, run_id, state, count in query:
+ # Always count towards task-level concurrency
(max_active_tis_per_dag /
+ # max_active_tis_per_dagrun), including DEFERRED.
+ self.task_concurrency_map[(dag_id, task_id)] += count
+ self.task_dagrun_concurrency_map[(dag_id, run_id, task_id)] +=
count
+ # Only count non-deferred states towards DAG-run active tasks
+ # (max_active_tasks / worker slot accounting).
+ if state != TaskInstanceState.DEFERRED:
+ self.dag_run_active_tasks_map[dag_id, run_id] += count
def _is_parent_process() -> bool:
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 4e15a0def06..e5b19f2768d 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -22,6 +22,7 @@ import itertools
import json
import logging
import math
+import warnings
from collections import defaultdict
from collections.abc import Collection, Iterable
from datetime import datetime, timedelta
@@ -71,6 +72,7 @@ from airflow._shared.observability.traces import
new_dagrun_trace_carrier, new_t
from airflow._shared.timezones import timezone
from airflow.assets.manager import asset_manager
from airflow.configuration import conf
+from airflow.exceptions import RemovedInAirflow4Warning
from airflow.executors.workloads import BaseWorkload
from airflow.listeners.listener import get_listener_manager
from airflow.models.asset import AssetModel
@@ -1954,22 +1956,52 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
@provide_session
def get_num_running_task_instances(self, session: Session, same_dagrun:
bool = False) -> int:
- """Return Number of running TIs from the DB."""
- # .count() is inefficient
- num_running_task_instances_query = (
+ """Count running TIs from the DB."""
+ warnings.warn(
+ "This function is deprecated and will be removed in Airflow.",
+ RemovedInAirflow4Warning,
+ stacklevel=2,
+ )
+ return self._get_num_task_instances_of_state(
+ [TaskInstanceState.RUNNING],
+ same_dagrun=same_dagrun,
+ session=session,
+ )
+
+ def get_num_active_task_instances(self, *, same_dagrun: bool = False,
session: Session) -> int:
+ """
+ Count active (running or deferred) TIs for this task from the DB.
+
+ Deferred TIs are included because they are still logically in-flight
+ and must count against max_active_tis_per_dag /
max_active_tis_per_dagrun.
+
+ :meta private:
+ """
+ return self._get_num_task_instances_of_state(
+ [TaskInstanceState.RUNNING, TaskInstanceState.DEFERRED],
+ same_dagrun=same_dagrun,
+ session=session,
+ )
+
+ def _get_num_task_instances_of_state(
+ self,
+ states: Collection[TaskInstanceState],
+ *,
+ same_dagrun: bool = False,
+ session: Session,
+ ) -> int:
+ stmt = (
select(func.count())
.select_from(TaskInstance)
- .where(
- TaskInstance.dag_id == self.dag_id,
- TaskInstance.task_id == self.task_id,
- TaskInstance.state == TaskInstanceState.RUNNING,
- )
+ .where(TaskInstance.dag_id == self.dag_id, TaskInstance.task_id ==
self.task_id)
)
+ if states:
+ stmt = stmt.where(or_(*(TaskInstance.state == s for s in states)))
+ else:
+ return 0
if same_dagrun:
- num_running_task_instances_query =
num_running_task_instances_query.where(
- TaskInstance.run_id == self.run_id
- )
- return session.scalar(num_running_task_instances_query) or 0
+ stmt = stmt.where(TaskInstance.run_id == self.run_id)
+ return session.scalar(stmt) or 0
@staticmethod
def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) ->
ColumnElement[bool] | None:
diff --git a/airflow-core/src/airflow/ti_deps/dependencies_states.py
b/airflow-core/src/airflow/ti_deps/dependencies_states.py
index ebf581ab48e..7d5bd5091e0 100644
--- a/airflow-core/src/airflow/ti_deps/dependencies_states.py
+++ b/airflow-core/src/airflow/ti_deps/dependencies_states.py
@@ -23,6 +23,16 @@ EXECUTION_STATES = {
TaskInstanceState.QUEUED,
}
+# States counted for task-level concurrency limits (max_active_tis_per_dag /
+# max_active_tis_per_dagrun). DEFERRED is included because a deferred task
+# instance is still logically in-flight and must block additional instances
+# from being scheduled. This is intentionally separate from EXECUTION_STATES
+# so that DAG-level max_active_tasks and pool slot calculations are unaffected.
+ACTIVE_STATES = {
+ *EXECUTION_STATES,
+ TaskInstanceState.DEFERRED,
+}
+
# In order to be able to get queued a task must have one of these states
SCHEDULEABLE_STATES = {
None,
diff --git a/airflow-core/src/airflow/ti_deps/deps/task_concurrency_dep.py
b/airflow-core/src/airflow/ti_deps/deps/task_concurrency_dep.py
index 1f1416214c7..9b2ecaba888 100644
--- a/airflow-core/src/airflow/ti_deps/deps/task_concurrency_dep.py
+++ b/airflow-core/src/airflow/ti_deps/deps/task_concurrency_dep.py
@@ -36,13 +36,13 @@ class TaskConcurrencyDep(BaseTIDep):
if (
ti.task.max_active_tis_per_dag is not None
- and ti.get_num_running_task_instances(session) >=
ti.task.max_active_tis_per_dag
+ and ti.get_num_active_task_instances(session=session) >=
ti.task.max_active_tis_per_dag
):
yield self._failing_status(reason="The max task concurrency has
been reached.")
return
if (
ti.task.max_active_tis_per_dagrun is not None
- and ti.get_num_running_task_instances(session, same_dagrun=True)
+ and ti.get_num_active_task_instances(same_dagrun=True,
session=session)
>= ti.task.max_active_tis_per_dagrun
):
yield self._failing_status(reason="The max task concurrency per
run has been reached.")
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 53c64e0ea2d..22272694880 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -2147,6 +2147,274 @@ class TestSchedulerJob:
assert len(res) == 1
session.rollback()
+ def
test_find_executable_task_instances_max_active_tis_per_dag_deferred_blocks(self,
dag_maker, session):
+ """
+ A DEFERRED TI should count against max_active_tis_per_dag.
+
+ When one TI is deferred, no additional TI of the same task should be
+ scheduled if the limit is already reached.
+ Regression test for https://github.com/apache/airflow/issues/61700
+ """
+ dag_id = "SchedulerJobTest.test_max_active_tis_per_dag_deferred_blocks"
+ with dag_maker(dag_id=dag_id, max_active_tasks=16, session=session):
+ task1 = EmptyOperator(task_id="deferrable_task",
max_active_tis_per_dag=1)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
run_id="run_1", session=session)
+ dr2 = dag_maker.create_dagrun_after(
+ dr1, run_type=DagRunType.SCHEDULED, run_id="run_2", session=session
+ )
+
+ # DR1's TI is deferred (waiting on a trigger)
+ ti1 = dr1.get_task_instance(task1.task_id, session)
+ ti1.state = TaskInstanceState.DEFERRED
+ session.merge(ti1)
+
+ # DR2's TI is scheduled and wants to run
+ ti2 = dr2.get_task_instance(task1.task_id, session)
+ ti2.state = State.SCHEDULED
+ session.merge(ti2)
+ session.flush()
+
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ # ti2 should be blocked because ti1 is deferred and counts as active
+ assert len(res) == 0
+ session.rollback()
+
+ def
test_find_executable_task_instances_max_active_tis_per_dag_deferred_plus_running(
+ self, dag_maker, session
+ ):
+ """
+ Deferred + running TIs together fill the max_active_tis_per_dag limit.
+
+ With max_active_tis_per_dag=2 and one RUNNING + one DEFERRED, a third
+ SCHEDULED TI should be blocked.
+ """
+ dag_id =
"SchedulerJobTest.test_max_active_tis_per_dag_deferred_plus_running"
+ with dag_maker(dag_id=dag_id, max_active_tasks=16, session=session):
+ task1 = EmptyOperator(task_id="task", max_active_tis_per_dag=2)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
run_id="run_1", session=session)
+ dr2 = dag_maker.create_dagrun_after(
+ dr1, run_type=DagRunType.SCHEDULED, run_id="run_2", session=session
+ )
+ dr3 = dag_maker.create_dagrun_after(
+ dr2, run_type=DagRunType.SCHEDULED, run_id="run_3", session=session
+ )
+
+ ti1 = dr1.get_task_instance(task1.task_id, session)
+ ti1.state = TaskInstanceState.RUNNING
+ session.merge(ti1)
+
+ ti2 = dr2.get_task_instance(task1.task_id, session)
+ ti2.state = TaskInstanceState.DEFERRED
+ session.merge(ti2)
+
+ ti3 = dr3.get_task_instance(task1.task_id, session)
+ ti3.state = State.SCHEDULED
+ session.merge(ti3)
+ session.flush()
+
+ # 1 running + 1 deferred = 2, which equals the limit
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ assert len(res) == 0
+ session.rollback()
+
+ def
test_find_executable_task_instances_max_active_tis_per_dag_deferred_with_room(
+ self, dag_maker, session
+ ):
+ """
+ With max_active_tis_per_dag=2 and only 1 deferred, one more TI
+ should be allowed to schedule.
+ """
+ dag_id =
"SchedulerJobTest.test_max_active_tis_per_dag_deferred_with_room"
+ with dag_maker(dag_id=dag_id, max_active_tasks=16, session=session):
+ task1 = EmptyOperator(task_id="task", max_active_tis_per_dag=2)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
run_id="run_1", session=session)
+ dr2 = dag_maker.create_dagrun_after(
+ dr1, run_type=DagRunType.SCHEDULED, run_id="run_2", session=session
+ )
+ dr3 = dag_maker.create_dagrun_after(
+ dr2, run_type=DagRunType.SCHEDULED, run_id="run_3", session=session
+ )
+
+ ti1 = dr1.get_task_instance(task1.task_id, session)
+ ti1.state = TaskInstanceState.DEFERRED
+ session.merge(ti1)
+
+ ti2 = dr2.get_task_instance(task1.task_id, session)
+ ti2.state = State.SCHEDULED
+ session.merge(ti2)
+
+ ti3 = dr3.get_task_instance(task1.task_id, session)
+ ti3.state = State.SCHEDULED
+ session.merge(ti3)
+ session.flush()
+
+ # 1 deferred -> room for 1 more (limit is 2)
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ assert len(res) == 1
+ session.rollback()
+
+ def
test_find_executable_task_instances_deferred_does_not_block_different_task(self,
dag_maker, session):
+ """
+ A DEFERRED TI of task A should NOT block task B from scheduling.
+
+ max_active_tis_per_dag is per-task, not per-DAG.
+ """
+ dag_id = "SchedulerJobTest.test_deferred_does_not_block_different_task"
+ with dag_maker(dag_id=dag_id, max_active_tasks=16, session=session):
+ task_a = EmptyOperator(task_id="task_a", max_active_tis_per_dag=1)
+ task_b = EmptyOperator(task_id="task_b")
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
run_id="run_1", session=session)
+ dr2 = dag_maker.create_dagrun_after(
+ dr1, run_type=DagRunType.SCHEDULED, run_id="run_2", session=session
+ )
+
+ # task_a in DR1 is deferred
+ ti_a1 = dr1.get_task_instance(task_a.task_id, session)
+ ti_a1.state = TaskInstanceState.DEFERRED
+ session.merge(ti_a1)
+
+ # task_a in DR2 is scheduled (should be blocked by deferred ti_a1)
+ ti_a2 = dr2.get_task_instance(task_a.task_id, session)
+ ti_a2.state = State.SCHEDULED
+ session.merge(ti_a2)
+
+ # task_b in DR1 is scheduled (should NOT be blocked)
+ ti_b1 = dr1.get_task_instance(task_b.task_id, session)
+ ti_b1.state = State.SCHEDULED
+ session.merge(ti_b1)
+ session.flush()
+
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ queued_task_ids = [ti.task_id for ti in res]
+ # task_b should be queued, task_a should be blocked
+ assert "task_b" in queued_task_ids
+ assert "task_a" not in queued_task_ids
+ session.rollback()
+
+ def test_find_executable_task_instances_deferred_to_success_unblocks(self,
dag_maker, session):
+ """
+ When a deferred TI completes (SUCCESS), the next TI should be
unblocked.
+ """
+ dag_id = "SchedulerJobTest.test_deferred_to_success_unblocks"
+ with dag_maker(dag_id=dag_id, max_active_tasks=16, session=session):
+ task1 = EmptyOperator(task_id="task", max_active_tis_per_dag=1)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
run_id="run_1", session=session)
+ dr2 = dag_maker.create_dagrun_after(
+ dr1, run_type=DagRunType.SCHEDULED, run_id="run_2", session=session
+ )
+
+ ti1 = dr1.get_task_instance(task1.task_id, session)
+ ti2 = dr2.get_task_instance(task1.task_id, session)
+
+ # Step 1: ti1 is deferred, ti2 scheduled -> ti2 blocked
+ ti1.state = TaskInstanceState.DEFERRED
+ ti2.state = State.SCHEDULED
+ session.merge(ti1)
+ session.merge(ti2)
+ session.flush()
+
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ assert len(res) == 0
+
+ # Step 2: ti1 completes -> ti2 should be unblocked
+ ti1.state = TaskInstanceState.SUCCESS
+ session.merge(ti1)
+ session.flush()
+
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ assert len(res) == 1
+ assert res[0].key == ti2.key
+ session.rollback()
+
+ def
test_find_executable_task_instances_max_active_tis_per_dagrun_deferred(self,
dag_maker, session):
+ """
+ DEFERRED TIs should also count against max_active_tis_per_dagrun.
+
+ With max_active_tis_per_dagrun=1 and 2 mapped instances in the same
+ dagrun, if one is deferred, the other should be blocked.
+ """
+ dag_id = "SchedulerJobTest.test_max_active_tis_per_dagrun_deferred"
+ with dag_maker(dag_id=dag_id, max_active_tasks=16, session=session):
+ task_a = EmptyOperator.partial(task_id="task_a",
max_active_tis_per_dagrun=1).expand_kwargs(
+ [{"inputs": 1}, {"inputs": 2}]
+ )
+ EmptyOperator(task_id="task_b")
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
session=session)
+
+ ti_a0 = dr.get_task_instance(task_a.task_id, session, map_index=0)
+ ti_a1 = dr.get_task_instance(task_a.task_id, session, map_index=1)
+
+ ti_a0.state = TaskInstanceState.DEFERRED
+ ti_a1.state = State.SCHEDULED
+ session.merge(ti_a0)
+ session.merge(ti_a1)
+ session.flush()
+
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ queued_task_ids = [(ti.task_id, ti.map_index) for ti in res]
+ # ti_a1 should be blocked, task_b may be queued
+ assert ("task_a", 1) not in queued_task_ids
+ session.rollback()
+
+ def
test_find_executable_task_instances_deferred_does_not_affect_max_active_tasks(
+ self, dag_maker, session
+ ):
+ """
+ Deferred TIs should NOT count toward max_active_tasks.
+
+ max_active_tasks is about worker-level parallelism, while deferred
tasks
+ don't consume worker slots. With max_active_tasks=2 and 1 deferred TI,
+ 2 more SCHEDULED TIs should be allowed.
+ """
+ dag_id =
"SchedulerJobTest.test_deferred_does_not_affect_max_active_tasks"
+ with dag_maker(dag_id=dag_id, max_active_tasks=2, session=session):
+ EmptyOperator(task_id="task_1")
+ EmptyOperator(task_id="task_2")
+ EmptyOperator(task_id="task_3")
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED,
session=session)
+ t1, t2, t3 = sorted(dr.get_task_instances(session=session), key=lambda
t: t.task_id)
+
+ t1.state = TaskInstanceState.DEFERRED
+ t2.state = State.SCHEDULED
+ t3.state = State.SCHEDULED
+ session.merge(t1)
+ session.merge(t2)
+ session.merge(t3)
+ session.flush()
+
+ res = self.job_runner._executable_task_instances_to_queued(max_tis=32,
session=session)
+ # Deferred doesn't count toward max_active_tasks=2, so both scheduled
can run
+ assert len(res) == 2
+ session.rollback()
+
def
test_change_state_for_executable_task_instances_no_tis_with_state(self,
dag_maker):
dag_id = "SchedulerJobTest.test_change_state_for__no_tis_with_state"
task_id_1 = "dummy"
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index 8dfbf997e66..b9e5452855f 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -35,6 +35,7 @@ from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace.propagation.tracecontext import
TraceContextTextMapPropagator
from sqlalchemy import delete, func, select
from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm.attributes import set_committed_value
from airflow import settings
from airflow._shared.observability.metrics.base_stats_logger import StatsLogger
@@ -1575,11 +1576,11 @@ class TestTaskInstance:
ti.state = State.SUCCESS
assert ti.try_number == 2 # unaffected by state
- def test_get_num_running_task_instances(self, dag_maker,
create_task_instance):
+ def test_get_num_active_task_instances(self, dag_maker,
create_task_instance):
session = settings.Session()
ti1 = create_task_instance(
- dag_id="test_get_num_running_task_instances", task_id="task1",
session=session
+ dag_id="test_get_num_active_task_instances", task_id="task1",
session=session
)
logical_date = DEFAULT_DATE + datetime.timedelta(days=1)
@@ -1598,7 +1599,7 @@ class TestTaskInstance:
ti2.task = ti1.task
ti3 = create_task_instance(
- dag_id="test_get_num_running_task_instances_dummy",
task_id="task2", session=session
+ dag_id="test_get_num_active_task_instances_dummy",
task_id="task2", session=session
)
assert ti3 in session
assert ti1 in session
@@ -1609,11 +1610,11 @@ class TestTaskInstance:
assert ti3 in session
session.commit()
- assert ti1.get_num_running_task_instances(session=session) == 1
- assert ti2.get_num_running_task_instances(session=session) == 1
- assert ti3.get_num_running_task_instances(session=session) == 1
+ assert ti1.get_num_active_task_instances(session=session) == 1
+ assert ti2.get_num_active_task_instances(session=session) == 1
+ assert ti3.get_num_active_task_instances(session=session) == 1
- def test_get_num_running_task_instances_per_dagrun(self,
create_task_instance, dag_maker):
+ def test_get_num_active_task_instances_per_dagrun(self,
create_task_instance, dag_maker):
session = settings.Session()
with dag_maker(dag_id="test_dag"):
@@ -1652,15 +1653,49 @@ class TestTaskInstance:
session.commit()
- assert tis1[("task_1",
0)].get_num_running_task_instances(session=session, same_dagrun=True) == 1
- assert tis1[("task_1",
1)].get_num_running_task_instances(session=session, same_dagrun=True) == 1
- assert tis1[("task_2",
0)].get_num_running_task_instances(session=session) == 2
- assert tis1[("task_3",
0)].get_num_running_task_instances(session=session, same_dagrun=True) == 1
+ assert tis1[("task_1",
0)].get_num_active_task_instances(session=session, same_dagrun=True) == 1
+ assert tis1[("task_1",
1)].get_num_active_task_instances(session=session, same_dagrun=True) == 1
+ assert tis1[("task_2",
0)].get_num_active_task_instances(session=session) == 2
+ assert tis1[("task_3",
0)].get_num_active_task_instances(session=session, same_dagrun=True) == 1
- assert tis2[("task_1",
0)].get_num_running_task_instances(session=session, same_dagrun=True) == 1
- assert tis2[("task_1",
1)].get_num_running_task_instances(session=session, same_dagrun=True) == 1
- assert tis2[("task_2",
0)].get_num_running_task_instances(session=session) == 2
- assert tis2[("task_3",
0)].get_num_running_task_instances(session=session, same_dagrun=True) == 1
+ assert tis2[("task_1",
0)].get_num_active_task_instances(session=session, same_dagrun=True) == 1
+ assert tis2[("task_1",
1)].get_num_active_task_instances(session=session, same_dagrun=True) == 1
+ assert tis2[("task_2",
0)].get_num_active_task_instances(session=session) == 2
+ assert tis2[("task_3",
0)].get_num_active_task_instances(session=session, same_dagrun=True) == 1
+
+ def test_get_num_active_task_instances_includes_deferred(self, dag_maker,
create_task_instance):
+ """
+ get_num_active_task_instances should count DEFERRED TIs.
+
+ Regression test for https://github.com/apache/airflow/issues/61700
+ """
+ session = settings.Session()
+
+ ti1 = create_task_instance(
+ dag_id="test_get_num_active_task_instances_deferred",
task_id="task1", session=session
+ )
+
+ logical_date = DEFAULT_DATE + datetime.timedelta(days=1)
+ dr = dag_maker.create_dagrun(
+ logical_date=logical_date,
+ run_type=DagRunType.MANUAL,
+ state=None,
+ run_id="2",
+ session=session,
+ data_interval=(logical_date, logical_date),
+ run_after=logical_date,
+ triggered_by=DagRunTriggeredByType.TEST,
+ )
+ ti2 = dr.task_instances[0]
+ ti2.task = ti1.task
+
+ ti1.state = TaskInstanceState.RUNNING
+ ti2.state = TaskInstanceState.DEFERRED
+ session.commit()
+
+ # Both RUNNING and DEFERRED should be counted
+ assert ti1.get_num_active_task_instances(session=session) == 2
+ assert ti2.get_num_active_task_instances(session=session) == 2
def test_log_url(self, create_task_instance):
ti = create_task_instance(dag_id="my_dag", task_id="op",
logical_date=timezone.datetime(2018, 1, 1))
@@ -3570,11 +3605,6 @@ def
test_get_dagrun_loaded_but_none_returns_dagrun(dag_maker, session):
Test that `get_dagrun()` fetches `DagRun` from DB when the `dag_run`
relationship is marked as loaded but unset (`None`).
"""
- from sqlalchemy.orm.attributes import set_committed_value
-
- from airflow.operators.empty import EmptyOperator
- from airflow.utils.state import State
-
with dag_maker(dag_id="test_get_dagrun_loaded_none"):
EmptyOperator(task_id="test_task")
diff --git a/airflow-core/tests/unit/ti_deps/deps/test_task_concurrency.py
b/airflow-core/tests/unit/ti_deps/deps/test_task_concurrency.py
index 4b7150f2cf7..564e6246521 100644
--- a/airflow-core/tests/unit/ti_deps/deps/test_task_concurrency.py
+++ b/airflow-core/tests/unit/ti_deps/deps/test_task_concurrency.py
@@ -35,7 +35,7 @@ class TestTaskConcurrencyDep:
return BaseOperator(task_id="test_task", dag=DAG("test_dag",
schedule=None), **kwargs)
@pytest.mark.parametrize(
- ("kwargs", "num_running_tis", "is_task_concurrency_dep_met"),
+ ("kwargs", "num_active_tis", "is_task_concurrency_dep_met"),
[
({}, None, True),
({"max_active_tis_per_dag": 1}, 0, True),
@@ -47,12 +47,17 @@ class TestTaskConcurrencyDep:
({"max_active_tis_per_dag": 1, "max_active_tis_per_dagrun": 2}, 1,
False),
({"max_active_tis_per_dag": 2, "max_active_tis_per_dagrun": 1}, 1,
False),
({"max_active_tis_per_dag": 1, "max_active_tis_per_dagrun": 1}, 1,
False),
+ # Deferred-specific scenarios: the count returned by
+ # 1 deferred TI fills a limit of 1 -> blocked
+ ({"max_active_tis_per_dag": 1}, 1, False),
+ # 1 deferred + 1 running = 2, limit 3 -> allowed
+ ({"max_active_tis_per_dag": 3}, 2, True),
],
)
- def test_concurrency(self, kwargs, num_running_tis,
is_task_concurrency_dep_met):
+ def test_concurrency(self, kwargs, num_active_tis,
is_task_concurrency_dep_met):
task = self._get_task(start_date=datetime(2016, 1, 1), **kwargs)
dep_context = DepContext()
ti = Mock(task=task, logical_date=datetime(2016, 1, 1))
- if num_running_tis is not None:
- ti.get_num_running_task_instances.return_value = num_running_tis
+ if num_active_tis is not None:
+ ti.get_num_active_task_instances.return_value = num_active_tis
assert TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) ==
is_task_concurrency_dep_met
