kaxil commented on code in PR #74346:
URL: https://github.com/apache/airflow/pull/74346#discussion_r4222252882
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1930,12 +1929,81 @@ def task_filter(task: Operator) -> bool:
)
# Create the missing tasks, including mapped tasks
- tis_to_create = self._create_tasks(
- (task for task in dag.task_dict.values() if task_filter(task)),
- task_creator,
- session=session,
+ missing_tasks = [task for task in dag.task_dict.values() if
task_filter(task)]
+ tis_to_create = list(
+ self._create_tasks(
+ missing_tasks,
+ task_creator,
+ session=session,
+ expand_literals=True,
+ )
)
self._create_task_instances(self.dag_id, tis_to_create,
created_counts, hook_is_noop, session=session)
+ empty_literal_tasks = [
+ task.task_id for task in missing_tasks if
self._get_literal_expansion_count(task) == 0
+ ]
+ if empty_literal_tasks:
+ session.execute(
+ update(TI)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ TI.task_id.in_(empty_literal_tasks),
+ TI.region_index == -1,
+ )
+ .values(state=TaskInstanceState.SKIPPED)
+ .execution_options(synchronize_session=False)
+ )
+
+ @staticmethod
+ def _get_literal_expansion_count(task: Operator) -> int | None:
+ """Return how many instances a task whose inputs are all literals
expands to, else None."""
+ from airflow.models.expandinput import NotFullyPopulated
+
+ try:
+ return task.get_parse_time_mapped_ti_count()
+ except (NotMapped, NotFullyPopulated):
+ return None
+
+ def _return_unmapped_instance_to_sentinel(self, ti: TI, *, session:
Session) -> bool:
+ """
+ Undo the placement of an unfinished instance whose task is no longer
mapped.
+
+ A mapped task's placeholder lives in the task's own region. When a
later Dag version stops mapping the
+ task, nothing else would move the placeholder back, and a plain task
is looked up in the sentinel region.
+ Returns whether the instance was marked removed.
+
+ A mapped task only turns plain mid-run when the run's definition is
replaced: an unversioned bundle
+ was reparsed, or a clear re-pinned the run to the latest version. Runs
pinned to a versioned bundle
+ never reach this.
+ """
+ if ti.region_id == SENTINEL_REGION_ID or (
+ ti.state is not None
+ and ti.state not in State.unfinished
+ and ti.state != TaskInstanceState.REMOVED
+ ):
+ return False
+ region = session.get(DynamicRegion, ti.region_id)
+ if region is None or region.node_id != ti.task_id or
region.parent_region_id is not None:
+ return False
+ sentinel_slot_taken = session.scalar(
+ select(TI.id)
+ .where(
+ TI.dag_id == ti.dag_id,
+ TI.run_id == ti.run_id,
+ TI.task_id == ti.task_id,
+ TI.region_id == SENTINEL_REGION_ID,
+ TI.region_index == -1,
+ TI.working_set.is_(True),
+ )
+ .limit(1)
+ )
+ if ti.region_index >= 0 or sentinel_slot_taken is not None:
Review Comment:
When a mapped task turns plain after it has expanded, every unfinished slot
is `>= 0`, so they all land here and go to REMOVED. There's no `-1` left to
move back (expansion promotes the placeholder to index 0 in place), and
`task_ids` already contains the task id, so `task_filter` in `verify_integrity`
never creates a plain instance. The task doesn't run again in that run, and
downstream trigger rules just see a removed upstream. Could the lowest
unfinished slot move to `(SENTINEL_REGION_ID, -1)` when that slot is free?
`test_task_no_longer_mapped_removes_unfinished_expanded_slots` currently
asserts exactly this outcome.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1930,12 +1929,81 @@ def task_filter(task: Operator) -> bool:
)
# Create the missing tasks, including mapped tasks
- tis_to_create = self._create_tasks(
- (task for task in dag.task_dict.values() if task_filter(task)),
- task_creator,
- session=session,
+ missing_tasks = [task for task in dag.task_dict.values() if
task_filter(task)]
+ tis_to_create = list(
+ self._create_tasks(
+ missing_tasks,
+ task_creator,
+ session=session,
+ expand_literals=True,
+ )
)
self._create_task_instances(self.dag_id, tis_to_create,
created_counts, hook_is_noop, session=session)
+ empty_literal_tasks = [
Review Comment:
This marks a zero-length literal mapping SKIPPED at run creation, before its
upstreams have run. With `A >> M.expand(x=[]) >> D` and `D` on `all_done` or
`none_failed`, `D` is runnable straight away and runs before `A`, and if `A`
fails `M` stays SKIPPED instead of going to UPSTREAM_FAILED. Before this change
the placeholder stayed `None` and `expand_mapped_task` skipped it only after
`are_dependencies_met`. Could this be limited to tasks with no upstream, or
dropped so the scheduler path handles it?
##########
airflow-core/src/airflow/models/task_coordinates.py:
##########
@@ -245,6 +275,41 @@ def get_loop_pass(self, ti: TaskCoordinate, *,
dag_version_id: UUID | None = Non
position = loop_position(regions, ti.region_id, ti.region_index,
loop.group_id)
return None if position is None else position[1]
+ def producer_contexts(
+ self,
+ caller: TaskInstance,
+ producer_task_ids: Collection[str] | None = None,
+ ) -> dict[str, ProducerContext]:
+ task = caller.task or self.get_task(
+ caller.dag_id, caller.run_id, caller.task_id,
dag_version_id=caller.dag_version_id
+ )
+ if producer_task_ids is None:
+ producer_task_ids = (
+ {op.task_id for op in task.iter_mapped_dependencies()} if
is_mapped(task) else set()
+ )
+ group = task.task_group
+ while group is not None:
+ if isinstance(group, SerializedMappedTaskGroup):
+ producer_task_ids.update(op.task_id for op in
group.iter_mapped_dependencies())
+ group = group.parent_group
+ caller_loop = enclosing_loop(task)
+ contexts = {}
+ for task_id in producer_task_ids:
+ producer = (
+ task.dag.get_task(task_id)
+ if caller.task is not None and task.dag is not None
+ else self.get_task(
+ caller.dag_id, caller.run_id, task_id,
dag_version_id=caller.dag_version_id
+ )
+ )
+ loop = enclosing_loop(producer)
+ if loop is None:
+ continue
+ if caller_loop is None or caller_loop.group_id != loop.group_id:
+ raise ValueError("A loop producer requires a consumer inside
the loop")
Review Comment:
This `ValueError` comes out of `expand_mapped_task` and
`_expand_mapped_task_if_needed` only catches `NotMapped`, so a mapped task
outside a loop that expands over a loop task's output raises on every scheduler
pass and the run never finishes. Could the Dag be rejected at parse time for
this shape, or the consumer go to UPSTREAM_FAILED here instead?
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py:
##########
@@ -256,44 +257,90 @@ class DagRunWaiter:
interval: float
result_task_ids: list[str] | None
+ def validate_result_selection(self, dag_run: DagRun, dag_bag: DagBagDep,
*, session: Session) -> None:
+ if self.result_task_ids == []:
+ return
+ query = select(TaskInstance.task_id,
TaskInstance.dag_version_id).where(
+ TaskInstance.dag_id == self.dag_id, TaskInstance.run_id ==
self.run_id
+ )
+ if self.result_task_ids is not None:
+ query = query.where(TaskInstance.task_id.in_(self.result_task_ids))
+ versions: dict[UUID | None, set[str]] = {}
+ for task_id, version in session.execute(query.distinct()):
+ versions.setdefault(version or dag_run.created_dag_version_id,
set()).add(task_id)
+ current_task_ids = {task_id for task_ids in versions.values() for
task_id in task_ids}
+ created_dag = (
+ dag_bag.get_dag(dag_run.created_dag_version_id, session=session)
+ if dag_run.created_dag_version_id is not None
+ else None
+ )
+ if created_dag is not None:
+ requested = (
+ set(self.result_task_ids) if self.result_task_ids is not None
else set(created_dag.task_ids)
+ )
+ versions.setdefault(dag_run.created_dag_version_id,
set()).update(requested - current_task_ids)
+ for version, task_ids in versions.items():
+ if version is None:
+ continue
+ dag = (
+ created_dag
+ if version == dag_run.created_dag_version_id
+ else dag_bag.get_dag(version, session=session)
+ )
+ if dag is None:
+ continue
+ for task_id in task_ids:
+ if not dag.has_task(task_id):
+ continue
+ task = dag.get_task(task_id)
+ if self.result_task_ids is None and not
task.returns_dag_result:
+ continue
+ if enclosing_loop(task) is not None:
Review Comment:
With this, every default `/wait` on a Dag whose result task sits inside a
loop returns 409, and nothing stops `dag.add_result()` on a loop-body task when
the Dag is defined. Would a parse-time check be better, with this as the
backstop?
##########
airflow-core/src/airflow/migrations/versions/0144_3_4_0_rename_stored_map_index.py:
##########
@@ -0,0 +1,93 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Rename the stored map index to region_index and add the dynamic region slot
key.
+
+Revision ID: 7f8c9a2d410e
+Revises: 54a27b6f9d01
+Create Date: 2026-09-27 23:00:00.000000
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = "7f8c9a2d410e"
+down_revision = "54a27b6f9d01"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+# ``legacy_task_data_owner``, ``xcom_v1`` and ``rtif_v1`` keep ``map_index``:
they record
+# the coordinates their rows had before regions existed and are never
rewritten.
+_TABLES = ("task_instance", "task_state_store")
+
+_SLOT_KEY_CONSTRAINT = "dynamic_region_slot_key_uq"
+
+
+def _rename(old: str, new: str) -> None:
+ for table in _TABLES:
+ op.alter_column(
+ table,
+ old,
+ new_column_name=new,
+ existing_type=sa.Integer(),
+ existing_nullable=False,
+ existing_server_default="-1",
+ )
+
+
+def _alter_dynamic_region(apply) -> None:
+ """Alter the table in batch mode; SQLite recreates it, so keep its
self-referencing cascade from firing."""
+ context = op.get_context()
+ if context.dialect.name != "sqlite":
+ with op.batch_alter_table("dynamic_region") as batch_op:
+ apply(batch_op)
+ return
+ with context.autocommit_block():
Review Comment:
On SQLite `autocommit_block()` commits the open transaction first, so the
rename and `add_column` in `upgrade()` are already committed when the rebuild
runs. If the rebuild fails, rerunning the upgrade stops at `_rename` because
`map_index` no longer exists. `_sqlite_rebuilds()` in 0142 keeps only the
PRAGMAs in autocommit and runs the rebuild in a nested transaction; could this
reuse it?
##########
airflow-core/tests/unit/migrations/test_0144_rename_stored_map_index.py:
##########
@@ -0,0 +1,301 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Tests for the stored map index rename migration, written against the schema as
it stood at that revision.
+
+The tables are declared here instead of taken from ``airflow.models``, and
alembic moves the database
+to the revision, so later migrations and model changes cannot change what
these tests exercise.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+from importlib import import_module
+from uuid import UUID, uuid4
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy import event, inspect
+
+from airflow import settings
+from airflow._shared.timezones import timezone
+from airflow.utils.db import downgrade, upgradedb
+from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime
+
+pytestmark = pytest.mark.db_test
+
+_migration =
import_module("airflow.migrations.versions.0144_3_4_0_rename_stored_map_index")
+REVISION = _migration.revision
+PREVIOUS_REVISION = _migration.down_revision
+DAG_ID = "migration_0144"
+RUN_ID = "run"
+REGION_ID = UUID(int=7)
+
+metadata = sa.MetaData()
+dag_run = sa.Table(
+ "dag_run",
+ metadata,
+ sa.Column("id", sa.Integer, primary_key=True),
+ sa.Column("dag_id", sa.String(250)),
+ sa.Column("run_id", sa.String(250)),
+ sa.Column("run_type", sa.String(50)),
+ sa.Column("run_after", UtcDateTime),
+)
+task_instance = sa.Table(
+ "task_instance",
+ metadata,
+ sa.Column("id", sa.Uuid, primary_key=True),
+ sa.Column("dag_id", sa.String(250)),
+ sa.Column("task_id", sa.String(250)),
+ sa.Column("run_id", sa.String(250)),
+ sa.Column("region_index", sa.Integer),
+ sa.Column("region_id", CompactUUID),
+ sa.Column("working_set", sa.Boolean),
+ sa.Column("try_number", sa.Integer),
+ sa.Column("pool", sa.String(256)),
+ sa.Column("pool_slots", sa.Integer),
+)
+task_state_store = sa.Table(
+ "task_state_store",
+ metadata,
+ sa.Column("id", sa.Integer, primary_key=True),
+ sa.Column("dag_run_id", sa.Integer),
+ sa.Column("dag_id", sa.String(250)),
+ sa.Column("task_id", sa.String(250)),
+ sa.Column("run_id", sa.String(250)),
+ sa.Column("region_index", sa.Integer),
+ sa.Column("region_id", CompactUUID),
+ sa.Column("key", sa.String(512)),
+ sa.Column("value", sa.Text),
+ sa.Column("updated_at", UtcDateTime),
+)
+dynamic_region = sa.Table(
Review Comment:
This `dynamic_region` table has no `parent_region_*` or `slot_key` columns
and the fixture seeds a single top-level region, so nothing here would notice
the self-referencing cascade that the `PRAGMA foreign_keys=OFF` guard is there
to stop, and nothing checks `slot_key` after the upgrade. A nested region row
that must survive the round trip, plus a `slot_key` column assertion, would
cover both.
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1467,11 +1488,35 @@ def field(name: str) -> Any:
return "<deferred>"
prefix = f"<TaskInstance: {field('dag_id')}.{field('task_id')}
{field('run_id')} "
- map_index = field("map_index")
- if map_index != -1:
- prefix += f"map_index={map_index} "
+ try:
+ coordinates = self.get_display_coordinates()
+ except (DetachedInstanceError, ObjectDeletedError):
+ coordinates = {}
+ prefix += "".join(f"{name}={value} " for name, value in
coordinates.items())
return prefix + f"[{field('state')}] ti_id={field('id')}>"
+ def get_display_coordinates(
+ self, *, include_unmapped: bool = False, task: Operator | None = None
+ ) -> dict[str, Any]:
+ """
+ Return the coordinates that identify this instance in messages shown
to users.
+
+ An instance outside a loop shows its ``map_index`` whatever region
holds it, because every mapped
+ expansion has its own region and the region id means nothing to a
user. An unmapped loop-body
+ instance shows its iteration. Any other instance, including one whose
task definition is not loaded,
+ shows its region and index, which is the only unambiguous form. The
definition is never loaded here;
+ a caller that holds it for an instance loaded without one passes it as
``task``.
+ """
+ from airflow.models.task_coordinates import enclosing_loop
+
+ index = self.region_index
+ task = task if task is not None else getattr(self, "task", None)
+ if self.region_id == SENTINEL_REGION_ID or (task is not None and
enclosing_loop(task) is None):
Review Comment:
`__repr__`, `_log_state` and `TaskInstanceNote.__repr__` call this without
`task`, and a TI loaded in the scheduler usually has none attached, so an
ordinary mapped TI still shows up as `region_id=<uuid> region_index=2` in those
messages. Whether the instance sits in its own task's top-level region is
knowable from the region row. Could the branch check that instead of whether
the definition happens to be attached?
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py:
##########
@@ -182,7 +182,9 @@ def
test_prior_dates_resolves_mapped_region_separately_for_each_run(
for dr in (old, current):
if dr is current and not current_regional:
Review Comment:
Now that `create_dagrun` already puts every mapped TI in its own region,
`current_regional=False` leaves the current run's TI regional too, so the two
cases test the same thing. Should the `False` case move it back to
`SENTINEL_REGION_ID` explicitly?
--
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]
