vatsrahul1001 commented on code in PR #71737:
URL: https://github.com/apache/airflow/pull/71737#discussion_r4306257554
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1479,8 +1517,27 @@ 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()
+ unfinished_tis = self.get_task_instances(session=session,
state=State.unfinished)
+ finished_rows = session.execute(
Review Comment:
Question on small runs: the query-count expectations in
`test_scheduler_job.py` go from `[218, 69, 69, 69]` to `[218, 97, 89, 89]`
(`@once`, `no_structure`) and from `[228, 84, 84, 84]` to `[228, 104, 104,
104]` (`linear`), about two extra queries per Dag run per loop. The benchmark
covers wide mapped runs where hydration dominates. For typical small runs,
especially with a remote database, the extra round trips may cost more than the
hydration saves. Is there a measurement for small and medium runs, or a way to
skip the second query when a run has few task instances?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1479,8 +1517,27 @@ 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()
+ unfinished_tis = self.get_task_instances(session=session,
state=State.unfinished)
+ finished_rows = session.execute(
+ 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.map_index)
Review Comment:
Nit: `fetch_task_instances` orders by `(task_id, region_id, region_index)`,
but this orders by `(task_id, map_index)`. Rows of the same task that share a
`region_index` across loop regions now come back in no defined order, which can
matter for the `max(..., key=end_date)` tie-breaks used for the Dag callbacks.
Using the same ordering would keep the output stable.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1479,8 +1517,27 @@ 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()
+ unfinished_tis = self.get_task_instances(session=session,
state=State.unfinished)
+ finished_rows = session.execute(
+ select(*FINISHED_TI_COLUMNS)
+ .where(
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.state.in_(State.finished),
Review Comment:
Optional: the old path went through `get_task_instances`, which applies
`DagRun._get_partial_task_ids(self.dag)` for a partial Dag. This query doesn't,
so for a partial Dag, finished rows outside the subset would hit `TaskNotFound`
in `_build_finished_tis` and be updated to `REMOVED`, which the old code never
did. The scheduler and `dag.test` use full Dags today, so I don't think it is
reachable, but adding the same `task_ids` filter here would keep the invariant.
--
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]
