Nataneljpwd commented on code in PR #71737:
URL: https://github.com/apache/airflow/pull/71737#discussion_r3856794824
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -140,14 +141,41 @@
tracer = trace.get_tracer(__name__)
+@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 the same identity-based equality and hashing a
TaskInstance has.
+
+ Anything needing a full TaskInstance must re-fetch it (see
``_ensure_type_task_instance``).
+ """
+
+ task_id: str
+ map_index: int
+ state: str | None
+ start_date: datetime | None
+ end_date: datetime | None
+ task: Operator | None = None
+
+
+# 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.
Review Comment:
This comment is a little interesting I don't think that a reader would
understand the implications, at first it looks like it just described what's
done in code but then it states te SSOT part, I think It should be higher up
with the field definitions and remove the part before it that just explains
what's done in the line bellow
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -140,14 +141,41 @@
tracer = trace.get_tracer(__name__)
+@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 the same identity-based equality and hashing a
TaskInstance has.
+
+ Anything needing a full TaskInstance must re-fetch it (see
``_ensure_type_task_instance``).
+ """
+
+ task_id: str
+ map_index: int
+ state: str | None
+ start_date: datetime | None
+ end_date: datetime | None
+ task: Operator | None = None
+
+
+# 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")
Review Comment:
Here if the field name does not exist for the TI, this will silently fail,
maybe rather add it in the init to throw an error in that case
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1465,10 +1512,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: list[str] = []
+ for row in rows:
+ try:
+ task = dag.get_task(row.task_id)
Review Comment:
Idk how lightweight this is, but if we order by task Id and map index, we
may optimize it by just saving the last task and task id, if it is the same we
don't need to query the database again as it's the same task
Just an optional optimization that might speed it up even more
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1355,15 +1383,15 @@ def recalculate(self) -> _UnfinishedStates:
self.notify_dagrun_state_changed(msg="success")
if dag.has_on_success_callback:
- last_succeeded_ti: TI | None = max(
+ last_succeeded_ti: TI | FinishedTI | 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,
)
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),
Review Comment:
Same as above, always finished ti
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1737,10 +1812,28 @@ def _expand_mapped_task_if_needed(ti: TI) ->
Iterable[TI] | None:
return ready_tis, changed_tis, expansion_happened
+ def _ensure_type_task_instance(self, ti: TI | FinishedTI | None, *,
session: Session) -> TI | None:
+ """
+ Return ``ti`` as a full TaskInstance, re-fetching it when it is a
FinishedTI 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.
+ """
+ if not isinstance(ti, FinishedTI):
+ return ti
+ return DagRun.fetch_task_instance(
+ dag_id=self.dag_id,
+ dag_run_id=self.run_id,
+ task_id=ti.task_id,
+ map_index=ti.map_index,
+ session=session,
+ )
Review Comment:
Is this used anywhere? As the 3 places I saw used it for no reason as it
always was a finished ti
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1756,7 +1849,9 @@ def _are_premature_tis(
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[TI | FinishedTI]
Review Comment:
Same as above
##########
airflow-core/src/airflow/ti_deps/dep_context.py:
##########
@@ -79,7 +79,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[TaskInstance | FinishedTI] | None = None
Review Comment:
Same as above
##########
airflow-core/src/airflow/ti_deps/dep_context.py:
##########
@@ -107,20 +107,22 @@ class DepContext:
fresh empty dict, so they would neither read the memo nor warm it for
anything else.
"""
- def ensure_finished_tis(self, dag_run: DagRun, session: Session) ->
list[TaskInstance]:
+ def ensure_finished_tis(self, dag_run: DagRun, session: Session) ->
list[TaskInstance | FinishedTI]:
"""
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[TaskInstance | FinishedTI]
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:
Review Comment:
Why was the variable renamed? It had a better name before
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -58,7 +59,7 @@ class _UpstreamTIStates(NamedTuple):
skipped_setup: int
@classmethod
- def calculate(cls, finished_upstreams: Iterator[TaskInstance]) ->
_UpstreamTIStates:
+ def calculate(cls, finished_upstreams: Iterator[TaskInstance |
FinishedTI]) -> _UpstreamTIStates:
Review Comment:
Isn't it always finished tis? As above?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -140,14 +141,41 @@
tracer = trace.get_tracer(__name__)
+@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 the same identity-based equality and hashing a
TaskInstance has.
+
+ Anything needing a full TaskInstance must re-fetch it (see
``_ensure_type_task_instance``).
+ """
+
+ task_id: str
+ map_index: int
+ state: str | None
+ start_date: datetime | None
+ end_date: datetime | None
+ task: Operator | None = None
+
+
+# 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.
Review Comment:
Also, makw it clear that the columns are taken from the TI
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1465,10 +1512,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: list[str] = []
+ for row in rows:
+ try:
+ task = dag.get_task(row.task_id)
+ except TaskNotFound:
+ if row.state != TaskInstanceState.REMOVED:
+ orphaned.append(row.task_id)
+ continue
+ finished.append(FinishedTI(**row._mapping, 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))
Review Comment:
Can't we just not append the removed here? Or pop them on runtime?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1355,15 +1383,15 @@ def recalculate(self) -> _UnfinishedStates:
self.notify_dagrun_state_changed(msg="success")
if dag.has_on_success_callback:
- last_succeeded_ti: TI | None = max(
+ last_succeeded_ti: TI | FinishedTI | None = max(
Review Comment:
Isn't it always a finished ti? If it's the last succeeded ti?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1333,7 +1361,7 @@ def recalculate(self) -> _UnfinishedStates:
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),
Review Comment:
This ti is always failed, so always finished
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1737,10 +1812,28 @@ def _expand_mapped_task_if_needed(ti: TI) ->
Iterable[TI] | None:
return ready_tis, changed_tis, expansion_happened
+ def _ensure_type_task_instance(self, ti: TI | FinishedTI | None, *,
session: Session) -> TI | None:
+ """
+ Return ``ti`` as a full TaskInstance, re-fetching it when it is a
FinishedTI 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.
+ """
+ if not isinstance(ti, FinishedTI):
+ return ti
+ return DagRun.fetch_task_instance(
+ dag_id=self.dag_id,
+ dag_run_id=self.run_id,
+ task_id=ti.task_id,
+ map_index=ti.map_index,
+ session=session,
+ )
+
def _are_premature_tis(
self,
unfinished_tis: Sequence[TI],
- finished_tis: list[TI],
+ finished_tis: list[TI | FinishedTI],
Review Comment:
Can they not be finished tis? If so the type hint is wrong
##########
airflow-core/src/airflow/ti_deps/dep_context.py:
##########
@@ -107,20 +107,22 @@ class DepContext:
fresh empty dict, so they would neither read the memo nor warm it for
anything else.
"""
- def ensure_finished_tis(self, dag_run: DagRun, session: Session) ->
list[TaskInstance]:
+ def ensure_finished_tis(self, dag_run: DagRun, session: Session) ->
list[TaskInstance | FinishedTI]:
"""
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[TaskInstance | FinishedTI]
Review Comment:
Same as above
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1629,7 +1704,7 @@ def execute_dag_callbacks(
def _get_ready_tis(
self,
schedulable_tis: list[TI],
- finished_tis: list[TI],
+ finished_tis: list[TI | FinishedTI],
Review Comment:
How can a finished ti list contain non finished tis?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1465,10 +1512,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: list[str] = []
+ for row in rows:
+ try:
+ task = dag.get_task(row.task_id)
Review Comment:
We can also bulk update it on case a lot of mapped tasks are removed but I
don't think it's worth it as it does not happen a lot
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -34,6 +34,7 @@
from sqlalchemy.orm import Session
from sqlalchemy.sql import ColumnElement
+ from airflow.models.dagrun import FinishedTI
Review Comment:
Why is it in models.dagrun? Maybe it is worth moving to the task instance
model file?
--
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]
