kaxil commented on code in PR #74407:
URL: https://github.com/apache/airflow/pull/74407#discussion_r4269807555
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -2581,6 +2581,20 @@ def expand_mapped_task(self, *, session: Session) ->
tuple[Sequence[TaskInstance
else:
dr = unmapped_ti.dag_run
state = unmapped_ti.state
+ if unmapped_ti.try_number > 0:
+ # Archived attempts only exist once the placeholder has
been retried. They would have
Review Comment:
These aren't always placeholder attempts.
`_check_for_removed_or_restored_tasks` keeps a sole -1 TI for expansion when a
task becomes mapped, so an unmapped task that ran and was then removed and
restored as mapped (restore goes through `prepare_db_for_next_try`), or was
`up_for_retry` when the DAG changed, reaches here with archived attempts that
really ran. This deletes them along with their XCom and rendered fields, where
3.3 moved those history rows to index 0 through the `ON UPDATE CASCADE`. I
haven't run either path. It's probably rare enough to accept, but then the
comment should say so rather than calling them placeholders.
##########
RELEASE_NOTES.rst:
##########
@@ -24,6 +24,25 @@
.. towncrier release notes start
+Airflow 3.4.0 (unreleased)
+--------------------------
+
+.. note::
+ Downgrading to any version earlier than 3.4.0 can be slow, depending on the
number of
+ XCom rows that have been written after upgrading. Upgrading was tested on a
Review Comment:
This still reads as if the cost scales with rows written after the upgrade.
On Postgres the downgrade re-adds `xcom_task_instance_fkey` and `rtif_ti_fkey`
as validated constraints through `_redirect_legacy`, which scans every row of
the legacy tables, pre-upgrade ones included, and SQLite's batch path copies
them. On the 300m row table mentioned here that's the slow part of the
rollback. Could the note say the cost on Postgres and SQLite depends on the
total size of `xcom` and `rendered_task_instance_fields`?
##########
airflow-core/tests/unit/models/test_mappedoperator.py:
##########
@@ -228,6 +228,66 @@ def test_expand_mapped_task_failed_state_in_db(dag_maker,
session):
queue_workload.assert_called_once()
[email protected]("zero_index_exists", [False, True])
+def test_expand_mapped_task_deletes_archived_attempts_of_the_placeholder(
+ dag_maker, session, zero_index_exists
+):
+ with dag_maker(session=session, serialized=True) as dag:
+ task1 = BaseOperator(task_id="op1")
+ other = BaseOperator(task_id="other")
+ mapped =
MockOperator.partial(task_id="task_2").expand(arg2=task1.output)
+ dr = dag_maker.create_dagrun()
+ other_dr = dag_maker.create_dagrun(run_id="other_run",
logical_date=DEFAULT_DATE + timedelta(days=1))
+ mapped_deser = dag.task_dict[mapped.task_id]
+ dr.get_task_instance(mapped.task_id,
session=session).prepare_db_for_next_try(
+ session
+ ).prepare_db_for_next_try(session)
+
+ zero_index = TaskInstance(
+ mapped_deser,
+ run_id=dr.run_id,
+ map_index=0,
+ state=TaskInstanceState.SUCCESS,
+ dag_version_id=DagVersion.get_latest_version(dr.dag_id).id,
+ )
+ session.add(zero_index)
+ session.flush()
+ if zero_index_exists:
+ zero_index.prepare_db_for_next_try(session)
+ else:
+ zero_index.archive(reason="retry", session=session)
+ other_task_ti = dr.get_task_instance(other.task_id, session=session)
+ other_task_ti.prepare_db_for_next_try(session)
+ other_run_ti = other_dr.get_task_instance(mapped.task_id, session=session)
+ other_run_ti.prepare_db_for_next_try(session)
+ survivors = {zero_index.id, other_task_ti.id, other_run_ti.id}
Review Comment:
Nothing here binds the `dag_id` filter on the new delete. `other_task_ti`
and `other_run_ti` cover `task_id` and `run_id`, but no other DAG with the same
task_id and run_id is seeded. Run IDs like `scheduled__...` collide across DAGs
all the time, so losing that filter would delete another DAG's attempts. Could
a second DAG with the same mapped task_id and run_id go into `survivors`?
##########
airflow-core/src/airflow/migrations/versions/0142_3_4_0_unify_task_attempt_ownership.py:
##########
@@ -370,24 +370,244 @@ def upgrade():
)
+def _build_task_instance_table():
+ return sa.table(
+ "task_instance",
+ sa.column("id"),
+ *(sa.column(c) for c in _COPY_COLUMNS),
+ sa.column("working_set", sa.Boolean()),
+ sa.column("dag_version_id"),
+ )
+
+
+def _build_legacy_owner_table():
+ return sa.table(
+ "legacy_task_data_owner", *(sa.column(c) for c in _COORDINATES),
sa.column("task_instance_id")
+ )
+
+
+def _select_latest_attempts(ti):
+ newer = ti.alias("newer")
+ # An anti-join, because grouping by max(try_number) aggregates the whole
table however few rows need copying.
+ return (
+ sa.select(ti.c.id, *(ti.c[c] for c in _COORDINATES))
+ .where(
+ ~sa.exists().where(
+ *(newer.c[c] == ti.c[c] for c in _COORDINATES),
newer.c.try_number > ti.c.try_number
+ )
+ )
+ .subquery()
+ )
+
+
+def _has_duplicate_attempts(bind, ti) -> bool:
+ coordinates = [ti.c[c] for c in _COORDINATES]
+ for where, group_by in (
+ (ti.c.working_set.is_not(None), coordinates),
+ (sa.true(), [*coordinates, ti.c.try_number]),
+ ):
+ duplicated =
sa.select(1).select_from(ti).where(where).group_by(*group_by).having(sa.func.count()
> 1)
+ if bind.scalar(duplicated.limit(1)):
+ return True
+ return False
+
+
+def _select_archived_ids(ti):
+ return sa.select(ti.c.id).where(ti.c.working_set.is_(None))
+
+
+def _count_orphaned_archived_attempts(bind, ti) -> int:
+ current = ti.alias("current")
+ orphaned = (
+ sa.select(sa.func.count())
+ .select_from(ti)
+ .where(
+ ti.c.working_set.is_(None),
+ ~sa.exists().where(
+ *(current.c[c] == ti.c[c] for c in _COORDINATES),
current.c.working_set.is_not(None)
+ ),
+ )
+ )
+ return bind.scalar(orphaned)
+
+
+def _has_moved_owners_with_legacy_rows(bind, ti) -> bool:
+ owner = _build_legacy_owner_table()
+ moved = sa.or_(*(owner.c[c] != ti.c[c] for c in _COORDINATES))
+ for name in ("xcom_v1", "rtif_v1"):
+ legacy = sa.table(name, *(sa.column(c) for c in _COORDINATES))
+ found = (
+ sa.select(1)
+ .select_from(owner.join(ti, ti.c.id == owner.c.task_instance_id))
+ .where(moved, sa.exists().where(*(legacy.c[c] == owner.c[c] for c
in _COORDINATES)))
+ )
+ if bind.scalar(found.limit(1)):
+ return True
+ return False
+
+
+def _delete_moved_owners(ti):
+ owner = _build_legacy_owner_table()
+ op.execute(
+ owner.delete().where(
+ sa.exists().where(
+ ti.c.id == owner.c.task_instance_id,
+ sa.or_(*(ti.c[c] != owner.c[c] for c in _COORDINATES)),
+ )
+ )
+ )
+
+
+def _delete_legacy_rows_of_archived_attempts(ti):
+ owner = _build_legacy_owner_table()
+ archived_ids = _select_archived_ids(ti)
+ # Legacy rows owned by an archived attempt are hidden from its successor,
so keeping them would expose them again.
+ for name in ("xcom_v1", "rtif_v1"):
+ legacy = sa.table(name, *(sa.column(c) for c in _COORDINATES))
+ op.execute(
+ legacy.delete().where(
+ sa.exists().where(
+ *(owner.c[c] == legacy.c[c] for c in _COORDINATES),
+ owner.c.task_instance_id.in_(archived_ids),
+ )
+ )
+ )
+
+
+def _move_owners_to_current_attempts(ti):
+ owner = _build_legacy_owner_table()
+ current = ti.alias("current")
+ archived_ids = _select_archived_ids(ti)
+ xcom_v2 = sa.table("xcom_v2", sa.column("task_instance_id"))
+ rtif_v2 = sa.table("rtif_v2", sa.column("task_instance_id"))
+ # Deleting archived attempts must not cascade into legacy rows through an
owner that points at one.
+ op.execute(
+ owner.update()
+ .where(owner.c.task_instance_id.in_(archived_ids))
+ .values(
+ task_instance_id=sa.select(current.c.id)
+ .where(*(current.c[c] == owner.c[c] for c in _COORDINATES),
current.c.working_set.is_not(None))
+ .scalar_subquery()
+ )
+ )
+ # Legacy tables reference owners by coordinates, so attempts created after
the upgrade need an owner row first.
+ op.execute(
+ owner.insert().from_select(
+ [*_COORDINATES, "task_instance_id"],
+ sa.select(*(ti.c[c] for c in _COORDINATES), ti.c.id).where(
+ ti.c.working_set.is_not(None),
+ ti.c.id.in_(sa.select(xcom_v2.c.task_instance_id))
+ | ti.c.id.in_(sa.select(rtif_v2.c.task_instance_id)),
+ ~sa.exists().where(*(owner.c[c] == ti.c[c] for c in
_COORDINATES)),
+ ),
+ )
+ )
+
+
+def _delete_replaced_legacy_rows(v1, v2, latest, *matching):
+ op.execute(
+ v1.delete().where(
+ sa.exists().where(
+ v2.c.task_instance_id == latest.c.id,
+ *(latest.c[c] == v1.c[c] for c in _COORDINATES),
+ *matching,
+ )
+ )
+ )
+
+
+def _copy_xcom_to_legacy(latest):
+ columns = ("key", "value", "timestamp", "dag_result", "mapped_length")
+ v2 = sa.table("xcom_v2", sa.column("task_instance_id"), *(sa.column(c) for
c in columns))
+ v1 = sa.table("xcom_v1", *(sa.column(c) for c in (*_COORDINATES, *columns,
"dag_run_id")))
+ dag_run = sa.table("dag_run", sa.column("id"), sa.column("dag_id"),
sa.column("run_id"))
+ _delete_replaced_legacy_rows(v1, v2, latest, v2.c.key == v1.c.key)
+ source = v2.join(latest, v2.c.task_instance_id == latest.c.id).join(
+ dag_run, sa.and_(dag_run.c.dag_id == latest.c.dag_id, dag_run.c.run_id
== latest.c.run_id)
+ )
+ op.execute(
+ v1.insert().from_select(
+ [*_COORDINATES, *columns, "dag_run_id"],
+ sa.select(
+ *(latest.c[c] for c in _COORDINATES), *(v2.c[c] for c in
columns), dag_run.c.id
+ ).select_from(source),
+ )
+ )
+
+
+def _copy_rendered_fields_to_legacy(latest):
+ columns = ("rendered_fields", "k8s_pod_yaml")
+ v2 = sa.table("rtif_v2", sa.column("task_instance_id"), *(sa.column(c) for
c in columns))
+ v1 = sa.table("rtif_v1", *(sa.column(c) for c in (*_COORDINATES,
*columns)))
+ _delete_replaced_legacy_rows(v1, v2, latest)
+ op.execute(
+ v1.insert().from_select(
+ [*_COORDINATES, *columns],
+ sa.select(*(latest.c[c] for c in _COORDINATES), *(v2.c[c] for c in
columns)).select_from(
+ v2.join(latest, v2.c.task_instance_id == latest.c.id)
+ ),
+ )
+ )
+
+
+def _move_archived_attempts_to_history(ti):
+ history = sa.table(
+ "task_instance_history",
+ sa.column("task_instance_id"),
+ *(sa.column(c) for c in _COPY_COLUMNS),
+ sa.column("dag_version_id"),
+ )
+ archived = ti.c.working_set.is_(None)
+ archived_ids = _select_archived_ids(ti)
+ op.execute(
+ history.insert().from_select(
+ ["task_instance_id", *_COPY_COLUMNS, "dag_version_id"],
+ sa.select(ti.c.id, *(ti.c[c] for c in _COPY_COLUMNS),
ti.c.dag_version_id).where(archived),
+ )
+ )
+ hitl = sa.table("hitl_detail", sa.column("ti_id"), *(sa.column(c) for c in
_HITL_COLUMNS))
+ hitl_history = sa.table(
+ "hitl_detail_history", sa.column("ti_history_id"), *(sa.column(c) for
c in _HITL_COLUMNS)
+ )
+ op.execute(
+ hitl_history.insert().from_select(
+ ["ti_history_id", *_HITL_COLUMNS],
+ sa.select(hitl.c.ti_id, *(hitl.c[c] for c in _HITL_COLUMNS))
+ .join(ti, ti.c.id == hitl.c.ti_id)
+ .where(archived),
+ )
+ )
+ # SQLite runs this with foreign keys off, so cascades cannot be relied on.
+ for child in ("hitl_detail", "task_instance_note", "task_reschedule"):
+ child_table = sa.table(child, sa.column("ti_id"))
+
op.execute(child_table.delete().where(child_table.c.ti_id.in_(archived_ids)))
+ op.execute(ti.delete().where(archived))
+
+
def downgrade():
- """Refuse downgrade when the predecessor cannot represent retained
ownership."""
+ """Fold the latest attempt's data back into the legacy tables and restore
archived attempts to history."""
if op.get_context().as_sql:
raise RuntimeError("Offline downgrade cannot verify retained attempt
ownership")
bind = op.get_bind()
- if bind.scalar(sa.text("SELECT 1 FROM task_instance WHERE working_set IS
NULL LIMIT 1")):
- raise RuntimeError("Cannot downgrade attempt ownership with historical
attempts")
- for name in ("xcom_v2", "rtif_v2"):
- if bind.scalar(sa.text(f"SELECT 1 FROM {name} LIMIT 1")):
- raise RuntimeError(f"Cannot downgrade attempt ownership with data
in {name}")
- if bind.scalar(
- sa.text(
- "SELECT 1 FROM legacy_task_data_owner o JOIN task_instance t ON
t.id=o.task_instance_id "
- "WHERE o.dag_id<>t.dag_id OR o.task_id<>t.task_id OR
o.run_id<>t.run_id OR o.map_index<>t.map_index LIMIT 1"
+ ti = _build_task_instance_table()
+ if _has_duplicate_attempts(bind, ti):
+ raise RuntimeError("Cannot downgrade attempt ownership with multiple
attempts sharing coordinates")
+ if orphans := _count_orphaned_archived_attempts(bind, ti):
+ raise RuntimeError(
+ f"Cannot downgrade attempt ownership: {orphans} archived attempts
(task_instance rows with "
+ "working_set IS NULL) have no current attempt at their
coordinates; delete them first"
)
- ):
+ if _has_moved_owners_with_legacy_rows(bind, ti):
raise RuntimeError("Cannot downgrade attempt ownership after legacy
owners changed coordinates")
with _sqlite_rebuilds():
+ _create_history_tables()
+ _delete_moved_owners(ti)
+ _delete_legacy_rows_of_archived_attempts(ti)
+ _move_owners_to_current_attempts(ti)
+ latest = _select_latest_attempts(ti)
+ _copy_xcom_to_legacy(latest)
+ _copy_rendered_fields_to_legacy(latest)
+ _move_archived_attempts_to_history(ti)
Review Comment:
This keeps the current row's `try_number`, but 3.4 has already bumped it for
a pending retry while 3.3.2 bumps it again in `schedule_tis` (`TI.try_number +
1`). With `retries=2` and try 1 failed on 3.4, the TI comes back as try 2 in
`up_for_retry`, 3.3 runs it as try 3, and `3 <= 2` fails the task one retry
early. Same for a TI that was cleared but not scheduled yet. Could the
downgrade decrement `try_number` on current rows in `up_for_retry` or NULL
state that have an archived attempt at `try_number - 1`, plus a test for each?
--
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]
