shahar1 commented on code in PR #76669:
URL: https://github.com/apache/airflow/pull/76669#discussion_r4501460259
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -2873,6 +2873,72 @@ def
test_verifying_removed_map_index_does_not_allocate_attempt(dag_maker, sessio
)
[email protected]("airflow._shared.observability.metrics.stats._get_backend")
+def
test_restoring_shrink_removed_map_index_without_resolved_length_adds_no_attempt(
+ mock_get_backend, dag_maker, session
+):
+ """Restoring an index before the map length is known must not record a Dag
re-add or a try."""
+ mock_stats = mock.MagicMock(spec=StatsLogger)
+ mock_get_backend.return_value = mock_stats
+
+ with dag_maker(session=session) as dag:
+
+ @task
+ def producer(): ...
+
+ @task
+ def consumer(arg): ...
+
+ consumer.expand(arg=producer())
+
+ dr = dag_maker.create_dagrun()
+ producer_ti = dr.get_task_instance("producer", session=session)
+ producer_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(producer_ti, [1, 2, 3], session=session)
+ session.flush()
+ dr.task_instance_scheduling_decisions(session=session)
+
+ # A producer re-run returns a shorter list, so indexes 1 and 2 are removed
at try 1.
+ push_mapped_length(producer_ti, [1], session=session)
+ dag_version_id = DagVersion.get_latest_version(dag.dag_id,
session=session).id
+ dr.verify_integrity(dag_version_id=dag_version_id, session=session)
+ removed = dr.get_task_instance("consumer", map_index=1, session=session)
+ assert removed.state == TaskInstanceState.REMOVED
+ removed_id = removed.id
+ removed_try_number = removed.try_number
+
+ # Clearing only the producer puts the removed indexes back before the new
list is known.
+ producer_ti = producer_ti.prepare_db_for_next_try(session)
+ producer_ti.state = None
+ session.flush()
+ dr.verify_integrity(dag_version_id=dag_version_id, session=session)
+
+ restored = dr.get_task_instance("consumer", map_index=1, session=session)
+ assert restored.state is None
+ assert restored.id == removed_id
+ assert restored.try_number == removed_try_number
Review Comment:
**minor**: `removed_try_number` is `0` here. The index is created at try 0
and is never scheduled, so the `if ti.try_number > 0` guard on `main` already
keeps the same id and try. With `dagrun.py` reverted to `main`, only the
`task_restored_to_dag` assertion fails. AGENTS.md: "Every changed or added
behaviour must have a test; every test must fail without the PR's change."
Putting index 1 at try 1 (scheduled or run) before the shrink would cover the
scenario from the issue.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -2411,6 +2412,14 @@ def _check_for_removed_or_restored_tasks(
continue
if ti.state == TaskInstanceState.REMOVED:
+ if mapping_unresolved:
Review Comment:
**minor**: `mapping_unresolved` cannot tell a shrink-removal from a task
that left the Dag and was added back in a later version. Such a task is marked
REMOVED in the `AirflowException` branch, then restored here after a "run on
latest version" clear while its producer is pending. That re-add now loses the
`task_restored_to_dag` metric and the info log, while the comment gives the
shrink case as the reason. Either narrow the check, or say in the comment that
re-adds are included.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -2411,6 +2412,14 @@ def _check_for_removed_or_restored_tasks(
continue
if ti.state == TaskInstanceState.REMOVED:
+ if mapping_unresolved:
+ # The instance may have been removed because a producer
shrank the list, not
+ # because the task left the Dag, and the resolved length
is still unknown. Put
+ # it back without counting a Dag re-add or allocating an
attempt that may never
+ # run; _revise_map_indexes_if_mapped removes it again if
the list stays short.
+ self.log.debug("Restoring mapped instance '%s' while its
map length is unresolved", ti)
+ ti.state = None
Review Comment:
**major**: this skips `prepare_db_for_next_try` for every try number,
including a removed attempt that already ran. The issue's own repro removes the
indexes at try 1. The row goes back to `state=None` and keeps the first run's
`start_date` / `end_date` / `hostname`. `schedule_tis` keeps a non-zero try
(`case((TI.try_number == 0, 1), else_=TI.try_number)`), so an index that comes
back in range re-runs as the same try on the same row, and the first run's
record is overwritten. A local probe confirmed this: on this PR the row keeps
try 1 with the old dates, while on `main` try 1 is archived and try 2 starts
clean. Skipping the successor looks safe only when the attempt never started,
for example `ti.start_date is None`. Details are in the review body.
--
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]
