abhijeets25012-tech commented on code in PR #76669:
URL: https://github.com/apache/airflow/pull/76669#discussion_r4511865337
##########
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:
Good point. The existing test only covers try 0, which doesn't reproduce the
reported issue. I've updated the regression test to schedule or execute map
index 1 before shrinking the mapping, so it covers restoration after try 1 and
verifies that the previous attempt isn't reused.
##########
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:
Thanks for catching this. You're right that restoring an already-executed
attempt without preparing its successor can leave stale attempt metadata on the
row. I've updated the logic to distinguish attempts that have actually started
from those that have never started, so previously executed attempts follow the
normal next-attempt preparation path.
--
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]
