Resolve live producers by region and read their XCom by attempt

Once a loop body or a mapped region can hold several live task instances with
the same task_id and map_index, a consumer can no longer find its producer from
(dag_id, run_id, task_id, map_index) alone, so the lookup has to know where the
caller sits: in the same iteration, in the previous one (what a loop body means
by "the last result"), outside the loop, or in an explicitly named region. When
more than one live candidate still fits we have no sensible option to raise
an error.

Only live (non-archived or superceded) try are candidates, and their data is
read by UUID. An archived try keeps its XCom under its own UUID, so reading by
the resolved try is exact and cannot revive data from work a clear replaced.

Callers that predate regions must see what they saw before, so the default read
scope stays the sentinel region and regional rows appear only when a caller asks
for them. Lookups of earlier runs refuse a non-sentinel region because a region
belongs to a single Dag run, so the producer has to be resolved again for each
run.

The scheduler detected changes to upstream state and tracked map-length
revisions by TaskInstanceKey, which cannot tell two regions apart. One region's
expansion would have marked another's as changed, so both now key on the attempt
and on (task, region).

Author: ashb

airflow-core/src/airflow/models/dagrun.py

@@ -124,7 +124,6 @@
         TaskInstance as TIDataModel,
     )
     from airflow.models.dag_version import DagVersion
-    from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.sdk import DAG as SDKDAG
     from airflow.serialization.definitions.dag import SerializedDAG
     from airflow.serialization.definitions.mappedoperator import Operator
@@ -1064,6 +1063,7 @@ def get_task_instance(
         task_id: str,
         *,
         map_index: int = -1,
+        region_id: UUID = SENTINEL_REGION_ID,
         session: Session = NEW_SESSION,
     ) -> TI | None:
         """
@@ -1078,6 +1078,7 @@ def get_task_instance(
             task_id=task_id,
             session=session,
             map_index=map_index,
+            region_id=region_id,
         )
 
     @staticmethod
@@ -1088,6 +1089,7 @@ def fetch_task_instance(
         task_id: str,
         *,
         map_index: int = -1,
+        region_id: UUID = SENTINEL_REGION_ID,
         session: Session = NEW_SESSION,
     ) -> TI | None:
         """
@@ -1099,7 +1101,9 @@ def fetch_task_instance(
         :param session: Sqlalchemy ORM Session
         """
         return session.scalars(
-            select(TI).filter_by(dag_id=dag_id, run_id=dag_run_id, task_id=task_id, map_index=map_index)
+            select(TI).filter_by(
+                dag_id=dag_id, run_id=dag_run_id, task_id=task_id, map_index=map_index, region_id=region_id
+            )
         ).one_or_none()
 
     def get_dag(self) -> SerializedDAG:
@@ -1683,7 +1687,7 @@ def _get_ready_tis(
         finished_tis: list[TI],
         session: Session,
     ) -> tuple[list[TI], bool, bool]:
-        old_states: dict[TaskInstanceKey, Any] = {}
+        old_states: dict[UUID, Any] = {}
         ready_tis: list[TI] = []
         changed_tis = False
 
@@ -1734,14 +1738,14 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
 
         # Check dependencies.
         expansion_happened = False
-        # Set of task ids for which was already done _revise_map_indexes_if_mapped
-        revised_map_index_task_ids: set[str] = set()
+        # (task id, region id) pairs for which _revise_map_indexes_if_mapped was already done
+        revised_expansion_keys: set[tuple[str, UUID]] = set()
         for schedulable in itertools.chain(schedulable_tis, additional_tis):
             if TYPE_CHECKING:
                 assert isinstance(schedulable.task, Operator)
             old_state = schedulable.state
             if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
-                old_states[schedulable.key] = old_state
+                old_states[schedulable.id] = old_state
                 continue
             # If schedulable is not yet expanded, try doing it now. This is
             # called in two places: First and ideally in the mini scheduler at
@@ -1760,14 +1764,18 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
                     # later in this same pass must see the post-expansion count).
                     dep_context.invalidate_upstream_task_id_counts()
             if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
-                # It's enough to revise map index once per task id,
+                # It's enough to revise map index once per task and region,
                 # checking the map index for each mapped task significantly slows down scheduling
-                if schedulable.task.task_id not in revised_map_index_task_ids:
+                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(
-                        schedulable.task, dag_version_id=schedulable.dag_version_id, session=session
+                        schedulable.task,
+                        dag_version_id=schedulable.dag_version_id,
+                        region_id=schedulable.region_id,
+                        session=session,
                     )
                     ready_tis.extend(revised_tis)
-                    revised_map_index_task_ids.add(schedulable.task.task_id)
+                    revised_expansion_keys.add(expansion_key)
                     if revised_tis:
                         # Revising a mapped task can add new instances, growing its instance count
                         # the same way expansion does. Drop the upstream-count memo so a downstream
@@ -1781,10 +1789,9 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
                     ready_tis.append(schedulable)
 
         # Check if any ti changed state
-        tis_filter = TI.filter_for_tis(old_states)
-        if tis_filter is not None:
-            fresh_tis = session.scalars(select(TI).where(tis_filter)).all()
-            changed_tis = any(ti.state != old_states[ti.key] for ti in fresh_tis)
+        if old_states:
+            fresh_tis = session.scalars(select(TI).where(TI.id.in_(old_states))).all()
+            changed_tis = any(ti.state != old_states[ti.id] for ti in fresh_tis)
 
         return ready_tis, changed_tis, expansion_happened
 
@@ -2156,7 +2163,12 @@ def _create_task_instances(
             session.rollback()
 
     def _revise_map_indexes_if_mapped(
-        self, task: Operator, *, dag_version_id: UUID | None, session: Session
+        self,
+        task: Operator,
+        *,
+        dag_version_id: UUID | None,
+        session: Session,
+        region_id: UUID = SENTINEL_REGION_ID,
     ) -> list[TI]:
         """
         Check if task increased or reduced in length and handle appropriately.
@@ -2181,7 +2193,7 @@ def _revise_map_indexes_if_mapped(
                 TI.dag_id == self.dag_id,
                 TI.task_id == task.task_id,
                 TI.run_id == self.run_id,
-                TI.region_id == SENTINEL_REGION_ID,
+                TI.region_id == region_id,
             )
         )
         existing_indexes = set(query)
@@ -2194,7 +2206,7 @@ def _revise_map_indexes_if_mapped(
                     TI.dag_id == self.dag_id,
                     TI.task_id == task.task_id,
                     TI.run_id == self.run_id,
-                    TI.region_id == SENTINEL_REGION_ID,
+                    TI.region_id == region_id,
                     TI.map_index.in_(removed_indexes),
                 )
                 .values(state=TaskInstanceState.REMOVED)
@@ -2209,13 +2221,20 @@ def _revise_map_indexes_if_mapped(
             task_id=task.task_id,
             run_id=self.run_id,
             map_indexes=missing_indexes,
-            region_id=SENTINEL_REGION_ID,
+            region_id=region_id,
             session=session,
         )
 
         new_tis: list[TI] = []
         for index in missing_indexes:
-            ti = TI(task, run_id=self.run_id, map_index=index, state=None, dag_version_id=dag_version_id)
+            ti = TI(
+                task,
+                run_id=self.run_id,
+                map_index=index,
+                region_id=region_id,
+                state=None,
+                dag_version_id=dag_version_id,
+            )
             ti.try_number = last_tries.get(index, -1) + 1
             ti.max_tries += ti.try_number
             self.log.debug("Expanding TIs upserted %s", ti)