ashb commented on code in PR #71737:
URL: https://github.com/apache/airflow/pull/71737#discussion_r4359011003
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1479,8 +1561,39 @@ def recalculate(self) -> _UnfinishedStates:
@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. Finished ones are never mutated
during
+ # scheduling, so hydrating them is pure overhead.
Review Comment:
"Finished ones are never mutated during scheduling" doesn't hold when a
mapped task's expansion shrinks. `_revise_map_indexes_if_mapped` (reached from
`_get_ready_tis` for a schedulable mapped TI) bulk-UPDATEs the indexes past the
new length to REMOVED. On main those finished TIs were identity-mapped ORM
objects, so the UPDATE's default `synchronize_session` refreshed them, and
downstream tasks evaluated later in the same `_get_ready_tis` pass saw REMOVED.
A `FinishedTI` keeps the old state.
Repro: `work.expand(arg=producer())` over 3 items, with `zreport`
(ALL_SUCCESS) downstream. `work[1]` is SUCCESS and `work[2]` FAILED; the user
clears `work[0]` with upstream, so `producer` reruns and now returns 2 items.
On the next pass the revise for `work[0]` marks `work[2]` REMOVED, but
`zreport` still counts it as FAILED and goes `upstream_failed`, so the run
fails. On main `zreport` stays unset and waits for `work[0]`. With ONE_FAILED,
`zreport` gets scheduled here and not on main.
Main already gets this wrong when the downstream's task_id sorts before
`work`; this change makes it happen regardless of task order.
One fix: have `_revise_map_indexes_if_mapped` return the `(task_id,
region_id, indexes)` it removed, and swap the matching entries in
`finished_tis` (the same list object as `dep_context.finished_tis`) for
`dataclasses.replace(fti, state=TaskInstanceState.REMOVED)`, with a test for
the scenario above. At minimum this comment shouldn't say they're never mutated.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1479,8 +1561,39 @@ def recalculate(self) -> _UnfinishedStates:
@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. Finished ones are never mutated
during
+ # scheduling, so hydrating them is pure overhead.
+ #
+ # 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()
Review Comment:
On the scheduler path the only pending change at this point is
`last_scheduling_decision`, which `update_state` sets just before calling this
(line 1404). So this flush turns what used to be one batched `UPDATE dag_run`
at commit into one UPDATE per dag run per loop. That is the cost the "We do not
flush here for performance reasons" comment at the end of `update_state` is
avoiding.
Actual statement counts for the `@once` 10x10 rows in
`TestSchedulerJobQueriesCount` (SQLite, recorded rather than checked against
the margin):
- main: 208/67/60/60 (no_structure), 218/74/74/74 (linear)
- this PR: 227/86/79/79, 237/93/93/93
- this PR without the flush: 218/77/70/70, 228/84/84/84
So about half the extra statements per loop come from this flush; the rest
is the second SELECT, which is the expected cost of the split.
Moving `self.last_scheduling_decision = start_dttm` to after the
`task_instance_scheduling_decisions` call keeps this flush for the case it's
there for (pending in-memory TI state changes) but makes it a no-op on the
scheduler path. Alternatively, only flush when the session holds a new or dirty
`TaskInstance`. Either way the expected counts in `test_scheduler_job.py` can
then come back down; even as it stands, the linear row would pass on SQLite
without being raised, because of `margin=15`.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -5041,6 +5056,35 @@ def my_teardown():
}
+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."""
+ with dag_maker(session=session):
+ EmptyOperator(task_id="stays")
+ EmptyOperator(task_id="vanishes")
+ dr = dag_maker.create_dagrun()
+ for ti in dr.get_task_instances(session=session):
+ ti.state = TaskInstanceState.SUCCESS
+ session.flush()
+
+ del dr.dag.task_dict["vanishes"]
+ decision = dr.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 dr.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 = dr.get_task_instance("stays", session=session)
+ assert (finished_ti.id, finished_ti.dag_id, finished_ti.run_id) ==
(stays.id, dr.dag_id, dr.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 dr.get_task_instance("vanishes", session=session).state ==
TaskInstanceState.REMOVED
Review Comment:
This pins the orphan path but not the UPDATE's scope: dropping the
`TI.dag_id`/`TI.run_id` predicates or the `TI.state != REMOVED` guard from that
UPDATE leaves the suite green. A second dag run where `vanishes` is SUCCESS
(asserting it stays SUCCESS), and making `vanishes` mapped with 2-3 indexes
(asserting all become REMOVED), would cover it.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1495,10 +1608,38 @@ def _filter_tis_and_exclude_removed(dag: SerializedDAG,
tis: list[TI]) -> Iterab
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._mapping, task=task))
Review Comment:
Nit: `FINISHED_TI_COLUMNS` is built from `fields(FinishedTI)` in order, so
`FinishedTI(*row, task=task)` is equivalent and skips building a dict per row.
In a quick benchmark (1000 rows, SQLite) that halved the construction cost.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -163,14 +164,95 @@ 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 = None
+
+ @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.
+FINISHED_TI_COLUMNS = tuple(getattr(TI, f.name) for f in fields(FinishedTI) if
f.name != "task")
+
+
+class SchedulingTI(Protocol):
Review Comment:
Non-blocking question: `SchedulingTI` repeats the six properties of
`TaskCoordinate`, which this PR made read-only so `FinishedTI` could satisfy
it. Would `class SchedulingTI(TaskCoordinate, Protocol)` adding only `id`,
`map_index`, `state`, `start_date`, `end_date` and `task` be better? It would
have to live in `task_coordinates.py` (that module imports `dagrun` at runtime)
and be imported under `TYPE_CHECKING` here.
Related: `NotPreviouslySkippedDep._mapped_group_skip_decisions`
(`not_previously_skipped_dep.py:214`) also receives `ensure_finished_tis()` but
is still annotated `list[TaskInstance]`. mypy doesn't flag it because its
caller is unannotated.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1479,8 +1561,39 @@ def recalculate(self) -> _UnfinishedStates:
@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. Finished ones are never mutated
during
+ # scheduling, so hydrating them is pure overhead.
+ #
+ # 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_stmt = (
+ select(*FINISHED_TI_COLUMNS)
+ .where(
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.state.in_(State.finished),
+ )
+ .order_by(TI.task_id, TI.region_id, TI.region_index)
+ )
+ if task_ids is not None:
+ finished_stmt = finished_stmt.where(TI.task_id.in_(task_ids))
Review Comment:
This repeats `fetch_task_instances`' WHERE/ORDER BY by hand, and the orphan
marking below relies on the two staying identical (as the comment at 1572
says). They have already drifted once in this PR: the missing partial-Dag
filter in an earlier revision would have marked rows outside the subset
REMOVED. Could `fetch_task_instances` take the columns to select (or a shared
helper build the statement) so the filter exists in one place?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
