Re: [PR] Scheduler: prevent stale upstream_failed after task recovery [airflow]

["Copilot (via GitHub)" <gi...@apache.org>]

Copilot commented on code in PR #64025:
URL: https://github.com/apache/airflow/pull/64025#discussion_r3025366367


##########
airflow-core/src/airflow/ti_deps/dep_context.py:
##########
@@ -89,11 +89,15 @@ def ensure_finished_tis(self, dag_run: DagRun, session: Session) -> list[TaskIns
         """
         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
+        :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

Review Comment:
   Docstring indentation is inconsistent: the `:return:` line is misaligned (extra leading space) compared to `:param dag_run:`. Please align the docstring fields so they render correctly in generated docs.
   ```suggestion
           :return: A list of all the finished tasks of this DAG and logical_date
   ```



##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -800,6 +800,8 @@ def fetch_task_instances(
         task_ids: list[str] | None = None,
         state: Iterable[TaskInstanceState | None] | None = None,
         session: Session = NEW_SESSION,
+        *,
+        refresh_from_db: bool = False,
     ) -> list[TI]:
         """Return the task instances for this dag run."""

Review Comment:
   `DagRun.fetch_task_instances()` / `get_task_instances()` now add a `refresh_from_db` keyword-only parameter. Since these methods are widely used (and callable by plugins/providers), this is a public-ish API change. Please ensure the docstring explicitly documents what “refresh” means (e.g., uses `populate_existing` to overwrite identity-map values) and when callers should use it—otherwise it’s hard to reason about correctness/perf tradeoffs.
   ```suggestion
           """
           Return the task instances for the given DAG run.
   
           When ``refresh_from_db`` is ``False`` (the default), SQLAlchemy will return any
           existing ``TaskInstance`` objects already present in the session's identity map
           without overwriting their in-memory state. This is usually what you want for
           short‑lived sessions created per request/operation, and it avoids the extra work
           of re-populating already-loaded ORM instances.
   
           When ``refresh_from_db`` is ``True``, the underlying query is executed with
           ``execution_options(populate_existing=True)``. This instructs SQLAlchemy to
           overwrite the attributes of any ``TaskInstance`` objects that are already present
           in the identity map with the current values from the database. Use this when:
   
           * You are reusing a long‑lived session and may already have ``TaskInstance``
             rows for this DAG run loaded, and you need to ensure they reflect the latest
             committed database state (e.g. after other code has modified TIs in the same
             session or in a different session/process).
           * Correctness (reading the most up‑to‑date state) is more important than the
             additional overhead of refreshing ORM instances.
   
           Leave ``refresh_from_db`` at its default of ``False`` when you either know the
           session does not already contain the relevant ``TaskInstance`` objects, or you
           are willing to tolerate potentially stale in‑memory values in exchange for
           slightly better performance.
           """
   ```



##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -2254,6 +2254,68 @@ def test_schedule_tis_only_one_scheduler_update_succeeds_when_competing(dag_make
     assert refreshed_ti.try_number == 1
 
 
+def test_task_instance_scheduling_decisions_refresh_finished_tis_before_setting_upstream_failed(
+    dag_maker, session
+):
+    with dag_maker(session=session, dag_id="refresh_finished_tis_before_upstream_failed", serialized=True):
+        fail_task = EmptyOperator(task_id="fail_task")
+        t0 = EmptyOperator(task_id="t0")
+        fail_task >> t0
+
+    dag = dag_maker.serialized_dag
+    dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED, state=DagRunState.RUNNING)
+    dr.dag = dag
+
+    tis = {ti.task_id: ti for ti in dr.get_task_instances(session=session, state=State.task_states)}
+    tis["fail_task"].state = TaskInstanceState.FAILED
+    session.commit()
+
+    with create_session() as other_session:
+        dag.set_task_instance_state(
+            task_id="fail_task",
+            run_id=dr.run_id,
+            state=TaskInstanceState.SUCCESS,
+            session=other_session,
+        )
+
+    decision = dr.task_instance_scheduling_decisions(session=session)
+    assert [ti.task_id for ti in decision.schedulable_tis] == ["t0"]
+    assert tis["fail_task"].state == TaskInstanceState.SUCCESS
+    assert tis["t0"].state is None
+
+
+def test_task_instance_scheduling_decisions_refresh_finished_tis_after_api_clear(dag_maker, session):
+    with dag_maker(session=session, dag_id="refresh_finished_tis_after_api_clear", serialized=True):
+        fail_task = EmptyOperator(task_id="fail_task")
+        t0 = EmptyOperator(task_id="t0")
+        t1 = EmptyOperator(task_id="t1")
+        t2 = EmptyOperator(task_id="t2")
+        fail_task >> t0 >> t1 >> t2
+
+    dag = dag_maker.serialized_dag
+    dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED, state=DagRunState.RUNNING)
+    dr.dag = dag
+
+    tis = {ti.task_id: ti for ti in dr.get_task_instances(session=session, state=State.task_states)}
+    tis["fail_task"].state = TaskInstanceState.FAILED
+    tis["t0"].state = TaskInstanceState.UPSTREAM_FAILED
+    session.commit()
+
+    with create_session() as other_session:
+        dag.set_task_instance_state(
+            task_id="fail_task",
+            run_id=dr.run_id,
+            state=TaskInstanceState.SUCCESS,
+            session=other_session,
+        )
+
+    decision = dr.task_instance_scheduling_decisions(session=session)
+    assert [ti.task_id for ti in decision.schedulable_tis] == ["t0"]
+    assert tis["fail_task"].state == TaskInstanceState.SUCCESS
+    assert tis["t0"].state is None
+    assert tis["t1"].state is None

Review Comment:
   These new regression tests mutate TI states by direct attribute assignment (e.g. `tis[...].state = ...`) rather than using `TaskInstance.set_state()`. The test intent is clearer and closer to real scheduler/API behavior if state changes go through `set_state()` or the same helper used in the API patch path, especially since `set_state()` also maintains timestamps/duration.



##########
airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py:
##########
@@ -214,3 +215,39 @@ def test_unmapped_parent_skip_mapped_downstream(session, dag_maker):
     assert len(list(dep.get_dep_statuses(tis["op2"], session, DepContext()))) == 1
     assert not dep.is_met(tis["op2"], session)
     assert tis["op2"].state == State.SKIPPED
+
+
+def test_cleared_skipmixin_parent_does_not_skip_with_stale_finished_tis(session, dag_maker):
+    start_date = pendulum.datetime(2020, 1, 1)
+    with dag_maker(
+        "test_cleared_skipmixin_parent_does_not_skip_with_stale_finished_tis",
+        schedule=None,
+        start_date=start_date,
+        session=session,
+    ):
+        op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op3")
+        op2 = EmptyOperator(task_id="op2")
+        op3 = EmptyOperator(task_id="op3")
+        op1 >> [op2, op3]
+
+    dagrun = dag_maker.create_dagrun(run_type=DagRunType.MANUAL, state=State.RUNNING)
+    tis = {ti.task_id: ti for ti in dagrun.task_instances}
+    run_task_instance(tis["op1"], op1)
+    session.commit()
+
+    with create_session() as other_session:
+        other_ti = other_session.scalar(
+            select(TaskInstance).where(
+                TaskInstance.dag_id == dagrun.dag_id,
+                TaskInstance.run_id == dagrun.run_id,
+                TaskInstance.task_id == "op1",
+            )
+        )
+        assert other_ti is not None
+        other_ti.state = State.NONE
+        other_session.flush()

Review Comment:
   This test simulates a cross-session update, but it updates `op1.state` to `State.NONE` without updating the `state` *and* without clearing related fields like `end_date`/`duration` that `TaskInstance.set_state(None, ...)` would normally manage. To better match real “clear” semantics (and avoid inconsistent rows that might affect other deps), consider using the same API used by clear operations (e.g., `ti.set_state(None, ...)` or the relevant clear helper) rather than assigning `state` directly.
   ```suggestion
           other_ti.set_state(None, session=other_session)
   ```



##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1238,7 +1247,7 @@ 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)
+        tis = self.get_task_instances(session=session, state=State.task_states, refresh_from_db=True)
         self.log.debug("number of tis tasks for %s: %s task(s)", self, len(tis))

Review Comment:
   `task_instance_scheduling_decisions()` now always loads TIs with `refresh_from_db=True`, which forces SQLAlchemy to overwrite any already-loaded instances in the session. This is on a scheduler hot path; please consider limiting refresh to the specific dependency-evaluation snapshot that needs it (e.g. only finished TIs via `DepContext.ensure_finished_tis`) or add a brief comment explaining why always-refreshing the full `State.task_states` set is required and acceptable for performance.



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org