This is an automated email from the ASF dual-hosted git repository.
ashb 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 b5ce03962cc Convert finished TIs into a lightweight read-only view
during scheduling decisions (#71737)
b5ce03962cc is described below
commit b5ce03962cc92b38d5047999cacfed42847bec33
Author: Christos Bisias <[email protected]>
AuthorDate: Wed Oct 28 00:54:00 2026 +0200
Convert finished TIs into a lightweight read-only view during scheduling
decisions (#71737)
---
airflow-core/src/airflow/models/dagrun.py | 233 +++++++++++++++++----
.../src/airflow/models/task_coordinates.py | 59 +++++-
airflow-core/src/airflow/ti_deps/dep_context.py | 12 +-
.../ti_deps/deps/not_previously_skipped_dep.py | 3 +-
.../src/airflow/ti_deps/deps/trigger_rule_dep.py | 5 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 4 +-
airflow-core/tests/unit/models/test_dagrun.py | 179 +++++++++++++++-
.../tests/unit/models/test_dynamic_region.py | 2 +-
.../tests/unit/models/test_mappedoperator.py | 9 +-
.../tests/unit/models/test_taskinstance.py | 8 +-
10 files changed, 447 insertions(+), 67 deletions(-)
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index ac0a4f03487..7be6c87be56 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -23,6 +23,7 @@ import os
import re
from collections import defaultdict
from collections.abc import Callable, Iterable, Iterator, Sequence
+from dataclasses import dataclass, field, fields, replace
from datetime import datetime
from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar, cast, overload
from uuid import UUID
@@ -120,8 +121,8 @@ if TYPE_CHECKING:
from typing import Literal, TypeAlias
from pydantic import NonNegativeInt
- from sqlalchemy.engine import ScalarResult
- from sqlalchemy.orm import Session
+ from sqlalchemy.engine import Row, ScalarResult
+ from sqlalchemy.orm import InstrumentedAttribute, Session
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.selectable import Select
@@ -130,6 +131,7 @@ if TYPE_CHECKING:
TaskInstance as TIDataModel,
)
from airflow.models.dag_version import DagVersion
+ from airflow.models.task_coordinates import SchedulingTI
from airflow.sdk import DAG as SDKDAG
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.serialization.definitions.mappedoperator import Operator
@@ -163,14 +165,57 @@ def validate_loop_decision(stop: bool | None,
region_index: int, group: Serializ
raise InvalidLoopDecision("Fixed-count loop decision does not match
its iteration limit")
+@dataclass(frozen=True, slots=True, eq=False)
+class FinishedTI:
+ """
+ Immutable view of a task instance in a terminal state.
+
+ Keeps the dag_run rescan lightweight and fast: scheduling only reads these
columns
+ off finished task instances, so a fully instrumented TaskInstance per row
is
+ unnecessary. ``task`` is attached from the serialized dag rather than the
row.
+
+ ``eq=False`` keeps identity equality and ``__hash__`` mirrors the key
hash, the same as
+ a TaskInstance.
+
+ Anything needing a full TaskInstance must re-fetch it (see
``_ensure_type_task_instance``).
+ """
+
+ id: UUID
+ dag_id: str
+ run_id: str
+ task_id: str
+ region_id: UUID
+ region_index: int
+ dag_version_id: UUID | None
+ state: str | None
+ start_date: datetime | None
+ end_date: datetime | None
+ task: Operator | None = field(default=None, kw_only=True)
+
+ @property
+ def map_index(self) -> int:
+ """Mirror the ``TaskInstance.map_index`` synonym of
``region_index``."""
+ return self.region_index
+
+ def __hash__(self) -> int:
+ return hash((self.task_id, self.dag_id, self.run_id,
self.region_index))
+
+
+# Get the necessary columns from the FinishedTI fields so that the class stays
+# the single source of truth. For any needed change, add or remove a field
there.
+# Rows are unpacked into FinishedTI positionally, so the column fields must
stay
+# first and in declaration order, with ``task`` last.
+FINISHED_TI_COLUMNS = tuple(getattr(TI, f.name) for f in fields(FinishedTI) if
f.name != "task")
+
+
class TISchedulingDecision(NamedTuple):
"""Type of return for DagRun.task_instance_scheduling_decisions."""
- tis: list[TI]
+ tis: list[SchedulingTI]
schedulable_tis: list[TI]
changed_tis: bool
unfinished_tis: list[TI]
- finished_tis: list[TI]
+ finished_tis: list[SchedulingTI]
def _default_run_after(ctx):
@@ -978,6 +1023,30 @@ class DagRun(Base, LoggingMixin):
return
DagRunType(run_type).generate_run_id(suffix=run_after.isoformat())
return
DagRunType(run_type).generate_run_id(suffix=f"{run_after.isoformat()}_{get_random_string()}")
+ @overload
+ @staticmethod
+ def fetch_task_instances(
+ dag_id: str | None = ...,
+ run_id: str | None = ...,
+ task_ids: list[str] | None = ...,
+ state: TaskInstanceState | Iterable[TaskInstanceState | None] | None =
...,
+ *,
+ columns: None = None,
+ session: Session = ...,
+ ) -> list[TI]: ...
+
+ @overload
+ @staticmethod
+ def fetch_task_instances(
+ dag_id: str | None = ...,
+ run_id: str | None = ...,
+ task_ids: list[str] | None = ...,
+ state: TaskInstanceState | Iterable[TaskInstanceState | None] | None =
...,
+ *,
+ columns: Sequence[InstrumentedAttribute[Any]],
+ session: Session = ...,
+ ) -> Sequence[Row[Any]]: ...
+
@staticmethod
@provide_session
def fetch_task_instances(
@@ -986,36 +1055,40 @@ class DagRun(Base, LoggingMixin):
task_ids: list[str] | None = None,
state: TaskInstanceState | Iterable[TaskInstanceState | None] | None =
None,
*,
+ columns: Sequence[InstrumentedAttribute[Any]] | None = None,
session: Session = NEW_SESSION,
- ) -> list[TI]:
- """Return the task instances for this dag run."""
- tis = (
- select(TI)
- .options(joinedload(TI.dag_run))
- .where(
- TI.dag_id == dag_id,
- TI.run_id == run_id,
- )
+ ) -> list[TI] | Sequence[Row[Any]]:
+ """
+ Return the task instances for this dag run.
+
+ If ``columns`` is provided then the result is plain rows with the
+ values of these columns, otherwise the result is ORM objects.
+ """
+ stmt = (
+ (select(*columns) if columns else
select(TI).options(joinedload(TI.dag_run)))
+ .where(TI.dag_id == dag_id, TI.run_id == run_id)
.order_by(TI.task_id, TI.region_id, TI.region_index)
)
if state:
if isinstance(state, str):
- tis = tis.where(TI.state == state)
+ stmt = stmt.where(TI.state == state)
else:
# this is required to deal with NULL values
if None in state:
if all(x is None for x in state):
- tis = tis.where(TI.state.is_(None))
+ stmt = stmt.where(TI.state.is_(None))
else:
not_none_state = (s for s in state if s)
- tis = tis.where(or_(TI.state.in_(not_none_state),
TI.state.is_(None)))
+ stmt = stmt.where(or_(TI.state.in_(not_none_state),
TI.state.is_(None)))
else:
- tis = tis.where(TI.state.in_(state))
+ stmt = stmt.where(TI.state.in_(state))
if task_ids is not None:
- tis = tis.where(TI.task_id.in_(task_ids))
- return list(session.scalars(tis).all())
+ stmt = stmt.where(TI.task_id.in_(task_ids))
+ if columns:
+ return session.execute(stmt).all()
+ return list(session.scalars(stmt).all())
def _check_last_n_dagruns_failed(self, dag_id,
max_consecutive_failed_dag_runs, session):
"""Check if last N dags failed."""
@@ -1361,7 +1434,7 @@ class DagRun(Base, LoggingMixin):
callback = self.produce_dag_callback(
dag=dag,
success=False,
- relevant_ti=ti_causing_failure,
+
relevant_ti=self._ensure_type_task_instance(ti_causing_failure,
session=session),
reason="task_failure",
execute=execute_callbacks,
session=session,
@@ -1384,7 +1457,7 @@ class DagRun(Base, LoggingMixin):
self.notify_dagrun_state_changed(msg="success")
if dag.has_on_success_callback:
- last_succeeded_ti: TI | None = max(
+ last_succeeded_ti: SchedulingTI | None = max(
(ti for ti in tis if ti.state ==
TaskInstanceState.SUCCESS),
key=lambda ti: ti.end_date or
timezone.make_aware(datetime.min),
default=None,
@@ -1392,7 +1465,7 @@ class DagRun(Base, LoggingMixin):
callback = self.produce_dag_callback(
dag=dag,
success=True,
- relevant_ti=last_succeeded_ti,
+
relevant_ti=self._ensure_type_task_instance(last_succeeded_ti, session=session),
reason="success",
execute=execute_callbacks,
session=session,
@@ -1480,8 +1553,41 @@ class DagRun(Base, LoggingMixin):
@provide_session
def task_instance_scheduling_decisions(self, *, session: Session =
NEW_SESSION) -> TISchedulingDecision:
- tis = self.get_task_instances(session=session, state=State.task_states)
- self.log.debug("number of tis tasks for %s: %s task(s)", self,
len(tis))
+ # Fetch finished task instances as lightweight read-only rows and only
the
+ # unfinished ones as full ORM objects. Scheduling only reads finished
ones, so
+ # hydrating them is pure overhead.
+ #
+ # Check if there are any pending TI changes.
+ has_pending_ti_changes = any(
+ isinstance(obj, TI) and obj.dag_id == self.dag_id and obj.run_id
== self.run_id
+ for obj in itertools.chain(session.new, session.dirty,
session.deleted)
+ )
+ if has_pending_ti_changes:
+ # Flush here to avoid task instances ending up in the wrong list
or in both,
+ # in case of a pending in-session state change.
+ # The airflow session maker sets autoflush=False (ref
settings.configure_orm).
+ session.flush()
+ # Both queries restrict a partial Dag to the same subset, so no row
outside it is
+ # mistaken for an orphan and marked REMOVED.
+ task_ids = DagRun._get_partial_task_ids(self.dag)
+ unfinished_tis = DagRun.fetch_task_instances(
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ task_ids=task_ids,
+ state=State.unfinished,
+ session=session,
+ )
+ finished_rows = DagRun.fetch_task_instances(
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ task_ids=task_ids,
+ state=State.finished,
+ columns=FINISHED_TI_COLUMNS,
+ session=session,
+ )
+ self.log.debug(
+ "number of tis tasks for %s: %s task(s)", self,
len(unfinished_tis) + len(finished_rows)
+ )
def _filter_tis_and_exclude_removed(dag: SerializedDAG, tis: list[TI])
-> Iterable[TI]:
"""Populate ``ti.task`` while excluding those missing one, marking
them as REMOVED."""
@@ -1496,10 +1602,37 @@ class DagRun(Base, LoggingMixin):
else:
yield ti
- tis = list(_filter_tis_and_exclude_removed(self.get_dag(), tis))
+ def _build_finished_tis(dag: SerializedDAG, rows) -> list[FinishedTI]:
+ """Attach the serialized task, dropping (and marking REMOVED) rows
whose task is gone."""
+ finished: list[FinishedTI] = []
+ orphaned: set[str] = set()
+ for row in rows:
+ try:
+ task = dag.get_task(row.task_id)
+ except TaskNotFound:
+ if row.state != TaskInstanceState.REMOVED:
+ orphaned.add(row.task_id)
+ continue
+ finished.append(FinishedTI(*row, task=task))
+ if orphaned:
+ self.log.error("Failed to get task for finished tis %s.
Marking them as removed.", orphaned)
+ session.execute(
+ update(TI)
+ .where(
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.task_id.in_(orphaned),
+ TI.state != TaskInstanceState.REMOVED,
+ )
+ .values(state=TaskInstanceState.REMOVED)
+ .execution_options(synchronize_session=False)
+ )
+ return finished
+
+ dag = self.get_dag()
+ unfinished_tis = list(_filter_tis_and_exclude_removed(dag,
unfinished_tis))
+ finished_tis: list[SchedulingTI] = list(_build_finished_tis(dag,
finished_rows))
- unfinished_tis = [t for t in tis if t.state in State.unfinished]
- finished_tis = [t for t in tis if t.state in State.finished]
if unfinished_tis:
schedulable_tis = [ut for ut in unfinished_tis if ut.state in
SCHEDULEABLE_STATES]
self.log.debug("number of scheduleable tasks for %s: %s task(s)",
self, len(schedulable_tis))
@@ -1520,6 +1653,7 @@ class DagRun(Base, LoggingMixin):
schedulable_tis = []
changed_tis = False
+ tis: list[SchedulingTI] = [*unfinished_tis, *finished_tis]
return TISchedulingDecision(
tis=tis,
schedulable_tis=schedulable_tis,
@@ -1721,7 +1855,7 @@ class DagRun(Base, LoggingMixin):
def _get_ready_tis(
self,
schedulable_tis: list[TI],
- finished_tis: list[TI],
+ finished_tis: list[SchedulingTI],
session: Session,
) -> tuple[list[TI], bool, bool]:
old_states: dict[UUID, Any] = {}
@@ -1802,10 +1936,20 @@ class DagRun(Base, LoggingMixin):
# checking the map index for each mapped task significantly
slows down scheduling
expansion_key = (schedulable.task.task_id,
schedulable.region_id)
if expansion_key not in revised_expansion_keys:
- revised_tis = self._revise_map_indexes_if_mapped(
+ revised_tis, removed_indexes =
self._revise_map_indexes_if_mapped(
schedulable,
session=session,
)
+ if removed_indexes:
+ # Mirror the bulk UPDATE onto the frozen views, as the
session does for ORM rows.
+ for i, fti in enumerate(finished_tis):
+ if (
+ isinstance(fti, FinishedTI)
+ and fti.task_id == schedulable.task_id
+ and fti.region_id == schedulable.region_id
+ and fti.region_index in removed_indexes
+ ):
+ finished_tis[i] = replace(fti,
state=TaskInstanceState.REMOVED)
ready_tis.extend(revised_tis)
revised_expansion_keys.add(expansion_key)
if revised_tis:
@@ -1827,10 +1971,25 @@ class DagRun(Base, LoggingMixin):
return ready_tis, changed_tis, expansion_happened
+ def _ensure_type_task_instance(self, ti: SchedulingTI | None, *, session:
Session) -> TI | None:
+ """
+ Return ``ti`` as a full TaskInstance, re-fetching it by primary key
when it is a view.
+
+ Full instances (and ``None``) pass through unchanged. Only the
dag-callback path
+ needs the full row, and only when a dag_run reaches a terminal state,
so the
+ re-fetch costs one query per finished dag_run rather than one per loop.
+
+ The lookup is by primary key: a public ``map_index`` lookup would miss
a finished
+ instance that lives in an earlier loop iteration's region.
+ """
+ if ti is None or isinstance(ti, TI):
+ return ti
+ return session.get(TI, ti.id)
+
def _are_premature_tis(
self,
unfinished_tis: Sequence[TI],
- finished_tis: list[TI],
+ finished_tis: list[SchedulingTI],
session: Session,
) -> tuple[bool, bool]:
dep_context = DepContext(
@@ -1846,7 +2005,7 @@ class DagRun(Base, LoggingMixin):
dep_context.have_changed_ti_states,
)
- def _emit_true_scheduling_delay_stats_for_finished_state(self,
finished_tis: list[TI]) -> None:
+ def _emit_true_scheduling_delay_stats_for_finished_state(self,
finished_tis: list[SchedulingTI]) -> None:
"""
Emit the true scheduling delay stats.
@@ -2543,7 +2702,7 @@ class DagRun(Base, LoggingMixin):
ti: TI,
*,
session: Session,
- ) -> list[TI]:
+ ) -> tuple[list[TI], set[int]]:
"""
Check if task increased or reduced in length and handle appropriately.
@@ -2560,7 +2719,7 @@ class DagRun(Base, LoggingMixin):
if task is None:
raise ValueError(f"Task definition is required to revise {ti}")
if not task.get_needs_expansion():
- return []
+ return [], set()
region_id = ti.region_id
dag_version_id = ti.dag_version_id
try:
@@ -2571,9 +2730,9 @@ class DagRun(Base, LoggingMixin):
producer_contexts=TaskCoordinateResolver.for_dag(self.dag,
session).producer_contexts(ti),
)
except NotMapped:
- return [] # Not a mapped task, don't need to do anything.
+ return [], set() # Not a mapped task, don't need to do anything.
except NotFullyPopulated:
- return [] # Upstreams not ready, don't need to revise this yet.
+ return [], set() # Upstreams not ready, don't need to revise this
yet.
query = session.scalars(
select(TI.region_index).where(
@@ -2602,7 +2761,7 @@ class DagRun(Base, LoggingMixin):
missing_indexes = [index for index in range(total_length) if index not
in existing_indexes]
if not missing_indexes:
- return []
+ return [], removed_indexes
last_tries = TI.get_last_try_numbers(
dag_id=self.dag_id,
task_id=task.task_id,
@@ -2629,7 +2788,7 @@ class DagRun(Base, LoggingMixin):
new_tis.append(ti)
if new_tis:
session.flush()
- return new_tis
+ return new_tis, removed_indexes
@classmethod
@provide_session
diff --git a/airflow-core/src/airflow/models/task_coordinates.py
b/airflow-core/src/airflow/models/task_coordinates.py
index fd85cbddc60..83a0b842c10 100644
--- a/airflow-core/src/airflow/models/task_coordinates.py
+++ b/airflow-core/src/airflow/models/task_coordinates.py
@@ -46,12 +46,14 @@ from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
from collections.abc import Collection, Iterable, Sequence
+ from datetime import datetime
from sqlalchemy import Select
from sqlalchemy.orm import Session
from sqlalchemy.sql.elements import ColumnElement
from airflow.serialization.definitions.dag import SerializedDAG,
SerializedOperator
+ from airflow.serialization.definitions.mappedoperator import Operator
from airflow.utils.sqlalchemy import UUIDSelect
@@ -73,14 +75,57 @@ class _ProducerRequest:
class TaskCoordinate(Protocol):
- """Stored task identity shared by live and retained task data."""
+ """
+ Stored task identity shared by live and retained task data.
- dag_id: str
- run_id: str
- task_id: str
- region_id: UUID
- region_index: int
- dag_version_id: UUID | None
+ Read-only on purpose: resolvers only ever read coordinates, and declaring
them as
+ properties lets immutable views (such as a frozen dataclass) satisfy the
protocol.
+ """
+
+ @property
+ def dag_id(self) -> str: ...
+
+ @property
+ def run_id(self) -> str: ...
+
+ @property
+ def task_id(self) -> str: ...
+
+ @property
+ def region_id(self) -> UUID: ...
+
+ @property
+ def region_index(self) -> int: ...
+
+ @property
+ def dag_version_id(self) -> UUID | None: ...
+
+
+class SchedulingTI(TaskCoordinate, Protocol):
+ """
+ Shape of a task instance during scheduling, satisfied by ``TaskInstance``
and the frozen ``FinishedTI``.
+
+ Extends the stored coordinates with the identity, state and timing that
scheduling
+ decisions and trigger-rule evaluation read.
+ """
+
+ @property
+ def id(self) -> UUID: ...
+
+ @property
+ def map_index(self) -> int: ...
+
+ @property
+ def state(self) -> str | None: ...
+
+ @property
+ def start_date(self) -> datetime | None: ...
+
+ @property
+ def end_date(self) -> datetime | None: ...
+
+ @property
+ def task(self) -> Operator | None: ...
def enclosing_loop(task: SerializedOperator) -> SerializedLoopTaskGroup | None:
diff --git a/airflow-core/src/airflow/ti_deps/dep_context.py
b/airflow-core/src/airflow/ti_deps/dep_context.py
index 4083793c067..790d47f2c07 100644
--- a/airflow-core/src/airflow/ti_deps/dep_context.py
+++ b/airflow-core/src/airflow/ti_deps/dep_context.py
@@ -31,7 +31,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm.session import Session
from airflow.models.dagrun import DagRun
- from airflow.models.task_coordinates import TaskCoordinateResolver
+ from airflow.models.task_coordinates import SchedulingTI,
TaskCoordinateResolver
from airflow.models.taskinstance import TaskInstance
@@ -82,7 +82,7 @@ class DepContext:
ignore_task_deps: bool = False
ignore_ti_state: bool = False
ignore_unmapped_tasks: bool = False
- finished_tis: list[TaskInstance] | None = None
+ finished_tis: list[SchedulingTI] | None = None
description: str | None = None
have_changed_ti_states: bool = False
@@ -182,20 +182,22 @@ class DepContext:
task mid-pass adds task instances but not skip decisions, so it never
needs invalidating.
"""
- def ensure_finished_tis(self, dag_run: DagRun, session: Session) ->
list[TaskInstance]:
+ def ensure_finished_tis(self, dag_run: DagRun, session: Session) ->
list[SchedulingTI]:
"""
Ensure finished_tis is populated if it's currently None, which allows
running tasks without dag_run.
:param dag_run: The DagRun for which to find finished tasks
:return: A list of all the finished tasks of this DAG and logical_date
"""
+ finished_tis: list[SchedulingTI]
if self.finished_tis is None:
- finished_tis = dag_run.get_task_instances(state=State.finished,
session=session)
- for ti in finished_tis:
+ fetched = dag_run.get_task_instances(state=State.finished,
session=session)
+ for ti in fetched:
if getattr(ti, "task", None) is not None or (dag :=
dag_run.dag) is None:
continue
with contextlib.suppress(TaskNotFound):
ti.task = dag.get_task(ti.task_id)
+ finished_tis = list(fetched)
self.finished_tis = finished_tis
else:
finished_tis = self.finished_tis
diff --git
a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
index 80c3380f877..07c438c195a 100644
--- a/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
+++ b/airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py
@@ -30,6 +30,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from airflow.models.task_coordinates import SchedulingTI
from airflow.models.taskinstance import TaskInstance
from airflow.serialization.definitions.taskgroup import
SerializedMappedTaskGroup
from airflow.ti_deps.dep_context import DepContext
@@ -211,7 +212,7 @@ def _shares_map_indexes(task, mapped_group:
SerializedMappedTaskGroup) -> bool:
def _mapped_group_skip_decisions(
ti: TaskInstance,
mapped_group: SerializedMappedTaskGroup,
- finished_tis: list[TaskInstance],
+ finished_tis: list[SchedulingTI],
dep_context: DepContext,
session: Session,
) -> dict[int, list[tuple[str, bool, dict]]]:
diff --git a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
index 8ed5d25e425..ec99c2aa692 100644
--- a/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -35,6 +35,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
from sqlalchemy.sql import ColumnElement
+ from airflow.models.task_coordinates import SchedulingTI
from airflow.models.taskinstance import TaskInstance
from airflow.serialization.definitions.mappedoperator import Operator
from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
@@ -59,7 +60,7 @@ class _UpstreamTIStates(NamedTuple):
skipped_setup: int
@classmethod
- def calculate(cls, finished_upstreams: Iterator[TaskInstance]) ->
_UpstreamTIStates:
+ def calculate(cls, finished_upstreams: Iterator[SchedulingTI]) ->
_UpstreamTIStates:
"""
Calculate states for a task instance.
@@ -228,7 +229,7 @@ class TriggerRuleDep(BaseTIDep):
),
)
- def _is_relevant_upstream(upstream: TaskInstance, relevant_ids:
set[str] | KeysView[str]) -> bool:
+ def _is_relevant_upstream(upstream: SchedulingTI, relevant_ids:
set[str] | KeysView[str]) -> bool:
"""
Whether a task instance is a "relevant upstream" of the current
task.
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 5ebbb8be9dc..0689ee02d80 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -13542,8 +13542,8 @@ class TestSchedulerJobQueriesCount:
# 10 DAGs with 10 tasks per DAG file.
([10, 10, 10, 10], 10, 10, "1d", "None", "no_structure"),
([10, 10, 10, 10], 10, 10, "1d", "None", "linear"),
- ([218, 69, 69, 69], 10, 10, "1d", "@once", "no_structure"),
- ([228, 84, 84, 84], 10, 10, "1d", "@once", "linear"),
+ ([229, 80, 79, 79], 10, 10, "1d", "@once", "no_structure"),
+ ([239, 95, 95, 95], 10, 10, "1d", "@once", "linear"),
([217, 119, 119, 119], 10, 10, "1d", "30m", "no_structure"),
([2227, 145, 145, 145], 10, 10, "1d", "30m", "linear"),
([227, 139, 139, 139], 10, 10, "1d", "30m", "binary_tree"),
diff --git a/airflow-core/tests/unit/models/test_dagrun.py
b/airflow-core/tests/unit/models/test_dagrun.py
index cb39a14303b..8d583b4c0da 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -30,6 +30,7 @@ from uuid import UUID
import pendulum
import pytest
+import time_machine
from opentelemetry import trace as otel_trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
@@ -62,6 +63,7 @@ from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import (
DagRun,
DagRunNote,
+ FinishedTI,
InvalidLoopDecision,
clear_partition_runs,
get_or_create_dagrun,
@@ -789,6 +791,10 @@ class TestDagRun:
# Make sure the correct TI is passed on success
call_args = execute_dag_callbacks.call_args
ti_passed = call_args.kwargs["relevant_ti"]
+ # During scheduling, finished tis are accessed as a minimal FinishedTI
view,
+ # and only the ti handed to a dag callback is fetched as a full
TaskInstance.
+ # Check that the callback receives a TaskInstance and not a FinishedTI.
+ assert isinstance(ti_passed, TaskInstance)
assert ti_passed.task_id == "test_state_succeeded2"
assert dag_run.state == DagRunState.SUCCESS
@@ -823,6 +829,10 @@ class TestDagRun:
# Make sure the correct TI is passed on failure
call_args = execute_dag_callbacks.call_args
ti_passed = call_args.kwargs["relevant_ti"]
+ # During scheduling, finished tis are accessed as a minimal FinishedTI
view,
+ # and only the ti handed to a dag callback is fetched as a full
TaskInstance.
+ # Check that the callback receives a TaskInstance and not a FinishedTI.
+ assert isinstance(ti_passed, TaskInstance)
assert ti_passed.task_id == "test_state_failed2"
assert dag_run.state == DagRunState.FAILED
@@ -3409,9 +3419,15 @@ def test_ti_scheduling_mapped_zero_length(dag_maker,
session):
decision = dr.task_instance_scheduling_decisions(session=session)
- # ti1 finished execution. ti2 goes directly to finished state because it's
- # expanded against a zero-length XCom.
- assert decision.finished_tis == [ti1, ti2]
+ # The finished_tis list contains both FinishedTI and TaskInstance types,
+ # so compare by content rather than identity.
+ #
+ # ti1 finished execution (before the decision - FinishedTI). ti2 goes
directly to finished state
+ # (during the decision - TaskInstance) because it's expanded against a
zero-length XCom.
+ assert [(ti.task_id, ti.state) for ti in decision.finished_tis] == [
+ (ti1.task_id, TaskInstanceState.SUCCESS),
+ (ti2.task_id, TaskInstanceState.SKIPPED),
+ ]
indices = session.execute(
select(TI.map_index, TI.state)
@@ -4431,6 +4447,81 @@ def
test_mapped_task_length_reduction_rerun_downstream_not_deadlocked(session, d
assert finish_ti.state == TaskInstanceState.SUCCESS
[email protected]("trigger_rule", [TriggerRule.ALL_SUCCESS,
TriggerRule.ONE_FAILED])
+def
test_mapped_task_length_reduction_marks_removed_for_downstream_in_same_pass(
+ session, dag_maker, trigger_rule
+):
+ """
+ Shrinking a mapped task mid-pass is visible to a downstream evaluated
later in the same pass.
+
+ The revise marks the surplus index REMOVED with a bulk UPDATE; the
finished-row views taken
+ at the start of the pass must reflect it, or the downstream counts a stale
FAILED upstream.
+ """
+
+ @task
+ def producer():
+ if get_current_context()["ti"].try_number == 1:
+ return [0, 1, 2]
+ return [0, 1]
+
+ @task
+ def work(arg):
+ return arg
+
+ @task(trigger_rule=trigger_rule)
+ def zreport():
+ return None
+
+ with dag_maker(session=session):
+ work.expand(arg=producer()) >> zreport()
+
+ dr = dag_maker.create_dagrun()
+ dag_maker.run_ti("producer", dr)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert [(ti.task_id, ti.map_index) for ti in decision.schedulable_tis] == [
+ ("work", 0),
+ ("work", 1),
+ ("work", 2),
+ ]
+ end_states = {0: TaskInstanceState.SUCCESS, 1: TaskInstanceState.SUCCESS,
2: TaskInstanceState.FAILED}
+ for ti in decision.schedulable_tis:
+ ti.set_state(end_states[ti.map_index], session=session)
+ session.flush()
+
+ # Clear work[0] together with its upstream producer, which now yields two
items.
+ clear_task_instances(
+ [
+ dr.get_task_instance("producer", session=session),
+ dr.get_task_instance("work", map_index=0, session=session),
+ ],
+ session=session,
+ )
+ session.flush()
+ dag_maker.run_ti("producer", dr)
+
+ decision = dr.task_instance_scheduling_decisions(session=session)
+
+ work_states = session.execute(
+ select(TI.map_index, TI.state)
+ .where(TI.task_id == "work", TI.dag_id == dr.dag_id, TI.run_id ==
dr.run_id)
+ .order_by(TI.map_index)
+ ).all()
+ assert work_states == [
+ (0, None),
+ (1, TaskInstanceState.SUCCESS),
+ (2, TaskInstanceState.REMOVED),
+ ]
+ # The downstream waits for work[0] instead of acting on the stale FAILED
state of work[2].
+ assert ("zreport", -1) not in [(ti.task_id, ti.map_index) for ti in
decision.schedulable_tis]
+ assert dr.get_task_instance("zreport", session=session).state is None
+
+ def _get_state_of_ti(tis, task_id, map_index):
+ return next(ti.state for ti in tis if ti.task_id == task_id and
ti.map_index == map_index)
+
+ assert _get_state_of_ti(decision.finished_tis, "work", 2) ==
TaskInstanceState.REMOVED
+ assert _get_state_of_ti(decision.tis, "work", 2) ==
TaskInstanceState.REMOVED
+
+
def test_rerun_with_upstream_task_removed(session, dag_maker):
def _task_ids(tis):
return [(ti.task_id, ti.map_index) for ti in tis]
@@ -5078,6 +5169,88 @@ def test_teardown_and_fail_fast(dag_maker):
}
+def test_scheduling_decisions_mark_orphaned_finished_tis_removed(dag_maker,
session):
+ """
+ A finished ti whose task no longer exists in the dag is marked REMOVED and
excluded.
+
+ Only the scheduled run is touched, neither another run of the Dag nor
another Dag sharing
+ the run_id and task_id loses its row.
+ """
+ with dag_maker("test_dag_1", session=session):
+ EmptyOperator(task_id="stays")
+ EmptyOperator(task_id="vanishes")
+ dag_run_1 = dag_maker.create_dagrun()
+ dag_run_1_later = dag_maker.create_dagrun_after(dag_run_1,
run_id="later_run")
+ with dag_maker("test_dag_2", session=session):
+ EmptyOperator(task_id="vanishes")
+ dag_run_2 = dag_maker.create_dagrun()
+ assert dag_run_2.run_id == dag_run_1.run_id
+ for run in (dag_run_1, dag_run_1_later, dag_run_2):
+ for ti in run.get_task_instances(session=session):
+ ti.state = TaskInstanceState.SUCCESS
+ session.flush()
+
+ del dag_run_1.dag.task_dict["vanishes"]
+ decision = dag_run_1.task_instance_scheduling_decisions(session=session)
+
+ assert len(decision.finished_tis) == 1
+ finished_ti = decision.finished_tis[0]
+ assert isinstance(finished_ti, FinishedTI)
+ assert finished_ti.task_id == "stays"
+ assert finished_ti.state == TaskInstanceState.SUCCESS
+ assert finished_ti.task is dag_run_1.dag.task_dict["stays"]
+ # The view carries the row's identity and coordinates so that trigger-rule
evaluation
+ # and the dag-callback re-fetch never need the full TaskInstance.
+ stays = dag_run_1.get_task_instance("stays", session=session)
+ assert (finished_ti.id, finished_ti.dag_id, finished_ti.run_id) == (
+ stays.id,
+ dag_run_1.dag_id,
+ dag_run_1.run_id,
+ )
+ assert (finished_ti.region_id, finished_ti.region_index) ==
(stays.region_id, stays.region_index)
+ assert finished_ti.dag_version_id == stays.dag_version_id
+ session.expire_all()
+ assert dag_run_1.get_task_instance("vanishes", session=session).state ==
TaskInstanceState.REMOVED
+ assert dag_run_1_later.get_task_instance("vanishes",
session=session).state == TaskInstanceState.SUCCESS
+ assert dag_run_2.get_task_instance("vanishes", session=session).state ==
TaskInstanceState.SUCCESS
+
+
+def
test_scheduling_decisions_orphan_removal_covers_every_map_index_but_skips_removed_rows(
+ dag_maker, session
+):
+ """Every index of an orphaned mapped task is marked REMOVED, rows already
REMOVED are not rewritten."""
+
+ @task
+ def vanishes(arg): ...
+
+ with dag_maker(session=session):
+ EmptyOperator(task_id="stays")
+ vanishes.expand(arg=[1, 2, 3])
+ dr = dag_maker.create_dagrun()
+ earlier = pendulum.datetime(2026, 1, 1, tz="UTC")
+ with time_machine.travel(earlier, tick=False):
+ for ti in dr.get_task_instances(session=session):
+ ti.state = TaskInstanceState.REMOVED if ti.map_index == 2 else
TaskInstanceState.SUCCESS
+ session.flush()
+
+ del dr.dag.task_dict["vanishes"]
+ later = earlier.add(hours=1)
+ with time_machine.travel(later, tick=False):
+ dr.task_instance_scheduling_decisions(session=session)
+
+ session.expire_all()
+ rows = session.execute(
+ select(TI.map_index, TI.state, TI.updated_at)
+ .where(TI.dag_id == dr.dag_id, TI.run_id == dr.run_id, TI.task_id ==
"vanishes")
+ .order_by(TI.map_index)
+ ).all()
+ assert rows == [
+ (0, TaskInstanceState.REMOVED, later),
+ (1, TaskInstanceState.REMOVED, later),
+ (2, TaskInstanceState.REMOVED, earlier),
+ ]
+
+
class TestDagRunHandleDagCallback:
"""Test the execute_dag_callbacks method (only uses in dag.test)."""
diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py
b/airflow-core/tests/unit/models/test_dynamic_region.py
index 591cbbadd24..2255448ecd2 100644
--- a/airflow-core/tests/unit/models/test_dynamic_region.py
+++ b/airflow-core/tests/unit/models/test_dynamic_region.py
@@ -431,7 +431,7 @@ def
test_mapping_revision_only_changes_selected_expansion(dag_maker, session):
regional = TaskInstance(task, version, run_id=dr.run_id, map_index=3,
region_id=uuid4())
session.add_all([ordinary, regional])
session.flush()
- added = dr._revise_map_indexes_if_mapped(regional, session=session)
+ added, _ = dr._revise_map_indexes_if_mapped(regional, session=session)
assert [(ti.region_id, ti.region_index) for ti in added] == [
(regional.region_id, 0),
(regional.region_id, 1),
diff --git a/airflow-core/tests/unit/models/test_mappedoperator.py
b/airflow-core/tests/unit/models/test_mappedoperator.py
index 807b0aa2587..032d9699b6a 100644
--- a/airflow-core/tests/unit/models/test_mappedoperator.py
+++ b/airflow-core/tests/unit/models/test_mappedoperator.py
@@ -416,7 +416,7 @@ def
test_stable_mapped_indexes_do_not_query_historical_max_try(dag_maker, sessio
live = run.get_task_instance(mapped.task_id, map_index=0, session=session)
live.task = dag.task_dict[mapped.task_id]
try:
- new_tis = run._revise_map_indexes_if_mapped(live, session=session)
+ new_tis, _ = run._revise_map_indexes_if_mapped(live, session=session)
finally:
event.remove(session.bind, "before_cursor_execute", capture_sql)
@@ -437,7 +437,7 @@ def
test_missing_mapped_index_uses_retained_max_try(dag_maker, session):
live = run.get_task_instance(mapped.task_id, map_index=0, session=session)
live.task = dag.task_dict[mapped.task_id]
- new_tis = run._revise_map_indexes_if_mapped(live, session=session)
+ new_tis, _ = run._revise_map_indexes_if_mapped(live, session=session)
assert len(new_tis) == 1
assert new_tis[0].map_index == 1
@@ -2280,10 +2280,11 @@ def
test_expansion_grows_and_shrinks_within_one_region(mapped_dag, session):
push_mapped_length(dr.get_task_instance(upstream_id, session=session),
list(range(5)), session=session)
first[0].task = serialized
- assert [ti.region_index for ti in
dr._revise_map_indexes_if_mapped(first[0], session=session)] == [3, 4]
+ new_tis, removed = dr._revise_map_indexes_if_mapped(first[0],
session=session)
+ assert ([ti.region_index for ti in new_tis], removed) == ([3, 4], set())
push_mapped_length(dr.get_task_instance(upstream_id, session=session),
list(range(2)), session=session)
- assert dr._revise_map_indexes_if_mapped(first[0], session=session) == []
+ assert dr._revise_map_indexes_if_mapped(first[0], session=session) == ([],
{2, 3, 4})
session.expire_all()
states = {ti.region_index: ti.state for ti in _live_tis(session, dr,
serialized.task_id)}
assert states[0] is None
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index 286a04574db..f5b8116038c 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -4105,11 +4105,9 @@ class TestMappedTaskInstanceReceiveValue:
session.flush()
with assert_queries_count(5, margin=1):
- new_tis = list(
- dag_run._revise_map_indexes_if_mapped(
- mapped_tis[0],
- session=session,
- )
+ new_tis, _ = dag_run._revise_map_indexes_if_mapped(
+ mapped_tis[0],
+ session=session,
)
# Only the two brand-new indexes (3, 4) are created, contiguous and
usable.
assert sorted(ti.map_index for ti in new_tis) == [3, 4]
