This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch changes-3.4.0rc3
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 1d06ae4012f3928d65b770e195901a9264228965
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Nov 4 22:43:31 2026 +0530
[v3-4-test] Keep mapped tasks when clearing with run on latest version
(#76420) (#76468)
* Keep mapped tasks when clearing with run on latest version
Clearing on the latest Dag version re-checks the run's task instances
against
that version. When the producer of a mapped task's input is cleared in the
same
request, its new attempt has no XCom yet, so the mapped length cannot be
resolved and every expanded instance was marked removed. Nothing expands a
removed index again, so the mapped task never re-ran and its downstream
tasks
ran as if it had finished. The same happened when a new version switched a
task from literal to XCom mapping before the producer had run.
The scheduler already trims or extends a mapped task once its upstreams have
finished, so keeping the instances lets it settle the count from the
producer's new value.
* Test that the scheduler keeps cleared mapped tasks on Dag change
The removal is not limited to clearing with "run on latest version". For a
run without a bundle version, the scheduler checks the run's integrity
whenever the Dag's structure changes, and before the fix that check removed
the mapped instances of a cleared producer that had not run again yet.
* Restore removed mapped tasks when their producer re-runs on latest version
A mapped index removed earlier, by a shrink or because the task left the
Dag, stayed removed after a clear on the latest version re-ran its
producer: the unresolved-count branch skipped the restore block, and
_revise_map_indexes_if_mapped never brings a removed index back.
(cherry picked from commit cd3ae408a1a0570920407ebe3833fe7da4d666d6)
Co-authored-by: Ephraim Anierobi <[email protected]>
---
airflow-core/src/airflow/models/dagrun.py | 11 ++--
airflow-core/tests/unit/jobs/test_scheduler_job.py | 48 +++++++++++++-
airflow-core/tests/unit/models/test_cleartasks.py | 77 ++++++++++++++++++++++
airflow-core/tests/unit/models/test_dagrun.py | 19 ++++--
4 files changed, 141 insertions(+), 14 deletions(-)
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index 810e308fb59..dbb0c76f55a 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -2211,13 +2211,10 @@ class DagRun(Base, LoggingMixin):
producer_contexts=coordinates.producer_contexts(ti),
)
except NotFullyPopulated:
- # Not all upstreams finished, so we can't tell what should
be here. Remove everything.
- if ti.region_index >= 0:
- self.log.debug(
- "Removing the unmapped TI '%s' as the mapping
can't be resolved yet", ti
- )
- ti.state = TaskInstanceState.REMOVED
- continue
+ # Not all upstreams finished, so the count isn't known
yet. Keep the expanded
+ # instances and restore removed ones below;
_revise_map_indexes_if_mapped trims
+ # or extends them once the upstreams finish, but never
restores a removed index.
+ pass
else:
# Upstreams finished, check there aren't any extras
if ti.region_index >= total_length:
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 268dbfb2ce2..6ed8be40735 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -186,8 +186,9 @@ from tests_common.test_utils.db import (
clear_db_triggers,
set_default_pool_slots,
)
+from tests_common.test_utils.mapping import push_mapped_length
from tests_common.test_utils.mock_executor import MockExecutor
-from tests_common.test_utils.mock_operators import CustomOperator
+from tests_common.test_utils.mock_operators import CustomOperator, MockOperator
from tests_common.test_utils.taskinstance import create_task_instance,
run_task_instance
from unit.listeners import dag_listener
from unit.models import TEST_DAGS_FOLDER
@@ -6190,6 +6191,51 @@ class TestSchedulerJob:
assert downstream_ti.state == State.SCHEDULED
assert downstream_ti.dag_version_id == latest_version_id
+ @pytest.mark.parametrize(
+ ("new_length", "expected_states"),
+ [
+ (1, {0: TaskInstanceState.SCHEDULED, 1:
TaskInstanceState.REMOVED}),
+ (3, dict.fromkeys(range(3), TaskInstanceState.SCHEDULED)),
+ ],
+ )
+ def test_verify_integrity_keeps_mapped_tis_while_cleared_producer_reruns(
+ self, dag_maker, session, new_length, expected_states
+ ):
+ with dag_maker("test_cleared_mapped_producer", session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ assert dr.bundle_version is None
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, [1, 2], session=session)
+ dr.task_instance_scheduling_decisions(session=session)
+ for ti in dr.get_task_instances(session=session):
+ ti.state = TaskInstanceState.SUCCESS
+ session.flush()
+ clear_task_instances(dr.get_task_instances(session=session), session)
+
+ with dag_maker("test_cleared_mapped_producer", session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ EmptyOperator(task_id="added")
+
+ self.job_runner = SchedulerJobRunner(job=Job())
+ self.job_runner._schedule_dag_run(dr, session)
+
+ tis = dr.get_task_instances(session=session)
+ assert "added" in {ti.task_id for ti in tis}
+ assert {ti.map_index: ti.state for ti in tis if ti.task_id ==
"mapped"} == {0: None, 1: None}
+
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, list(range(new_length)), session=session)
+ self.job_runner._schedule_dag_run(dr, session)
+ session.expire_all()
+
+ tis = dr.get_task_instances(session=session)
+ assert {ti.map_index: ti.state for ti in tis if ti.task_id ==
"mapped"} == expected_states
+
def test_verify_integrity_not_called_for_versioned_bundles(self,
dag_maker, session):
with dag_maker("test_verify_integrity_if_dag_not_changed") as dag:
BashOperator(task_id="dummy", bash_command="echo hi")
diff --git a/airflow-core/tests/unit/models/test_cleartasks.py
b/airflow-core/tests/unit/models/test_cleartasks.py
index 5ed43c424ec..5a616a340d0 100644
--- a/airflow-core/tests/unit/models/test_cleartasks.py
+++ b/airflow-core/tests/unit/models/test_cleartasks.py
@@ -65,6 +65,7 @@ from airflow.utils.types import DagRunTriggeredByType,
DagRunType
from tests_common.test_utils import db
from tests_common.test_utils.asserts import count_queries
from tests_common.test_utils.dag import sync_dag_to_db
+from tests_common.test_utils.mapping import push_mapped_length
from tests_common.test_utils.mock_operators import MockOperator
from tests_common.test_utils.taskinstance import run_task_instance
from unit.models import DEFAULT_DATE
@@ -1578,6 +1579,82 @@ class TestClearTasks:
assert attempt.working_set is None
assert attempt.dag_version_id == old_version_id
+ @pytest.mark.parametrize("dr_state", [DagRunState.SUCCESS,
DagRunState.RUNNING])
+ @pytest.mark.parametrize("new_length", [1, 3])
+ def test_clear_latest_with_producer_keeps_mapped_task_for_revision(
+ self, dag_maker, session, dr_state, new_length
+ ):
+ with dag_maker("test_clear_latest_mapped", bundle_version="v1",
session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ dr = dag_maker.create_dagrun(state=dr_state)
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, [1, 2], session=session)
+ dr.task_instance_scheduling_decisions(session=session)
+ for ti in dr.get_task_instances(session=session):
+ ti.state = TaskInstanceState.SUCCESS
+ session.flush()
+
+ with dag_maker("test_clear_latest_mapped", bundle_version="v2",
session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ EmptyOperator(task_id="added")
+
+ clear_task_instances(dr.get_task_instances(session=session), session,
run_on_latest_version=True)
+
+ mapped_states = {
+ ti.map_index: ti.state for ti in
dr.get_task_instances(session=session) if ti.task_id == "mapped"
+ }
+ assert mapped_states == {0: None, 1: None}
+
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, list(range(new_length)), session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+
+ schedulable = sorted(ti.map_index for ti in decision.schedulable_tis
if ti.task_id == "mapped")
+ assert schedulable == list(range(new_length))
+
+ @pytest.mark.parametrize(("new_length", "expected_schedulable"), [(1, []),
(3, [1, 2])])
+ def test_clear_latest_producer_restores_removed_mapped_tis(
+ self, dag_maker, session, new_length, expected_schedulable
+ ):
+ with dag_maker("test_clear_latest_restore", bundle_version="v1",
session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, [1, 2, 3], session=session)
+ dr.task_instance_scheduling_decisions(session=session)
+ # As left by an earlier producer re-run that shrank the mapping to one.
+ for ti in dr.get_task_instances(session=session):
+ ti.state = TaskInstanceState.REMOVED if ti.map_index > 0 else
TaskInstanceState.SUCCESS
+ session.flush()
+
+ with dag_maker("test_clear_latest_restore", bundle_version="v2",
session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ EmptyOperator(task_id="added")
+
+ clear_task_instances(
+ [dr.get_task_instance("size", session=session)], session,
run_on_latest_version=True
+ )
+
+ mapped_states = {
+ ti.map_index: ti.state for ti in
dr.get_task_instances(session=session) if ti.task_id == "mapped"
+ }
+ assert mapped_states == {0: TaskInstanceState.SUCCESS, 1: None, 2:
None}
+
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, list(range(new_length)), session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+
+ schedulable = sorted(ti.map_index for ti in decision.schedulable_tis
if ti.task_id == "mapped")
+ assert schedulable == expected_schedulable
+
def test_complete_restart_uses_latest_version_for_unpinned_run(self,
dag_maker, session):
with dag_maker("test_restart_unpinned", session=session):
EmptyOperator(task_id="work")
diff --git a/airflow-core/tests/unit/models/test_dagrun.py
b/airflow-core/tests/unit/models/test_dagrun.py
index 7ffdf347472..1c3f1208916 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -2625,7 +2625,7 @@ def test_mapped_literal_verify_integrity(dag_maker,
session):
def test_mapped_literal_to_xcom_arg_verify_integrity(dag_maker, session):
- """Test that when we change from literal to a XComArg the TIs are
removed"""
+ """Changing from a literal to an unresolved XComArg keeps the TIs until
the XComArg resolves."""
@task
def task_2(arg2): ...
@@ -2643,15 +2643,22 @@ def
test_mapped_literal_to_xcom_arg_verify_integrity(dag_maker, session):
dag_version_id = DagVersion.get_latest_version(dag_id=dr.dag_id,
session=session).id
dr.verify_integrity(dag_version_id=dag_version_id, session=session)
- indices = session.execute(
+ query = (
select(TI.map_index, TI.state)
.where(TI.task_id == "task_2", TI.dag_id == dr.dag_id, TI.run_id ==
dr.run_id)
.order_by(TI.map_index)
- ).all()
+ )
+ assert session.execute(query).all() == [(0, None), (1, None), (2, None),
(3, None)]
- assert indices == [
- (0, TaskInstanceState.REMOVED),
- (1, TaskInstanceState.REMOVED),
+ upstream = dr.get_task_instance(task_id="task_1", session=session)
+ upstream.state = TaskInstanceState.SUCCESS
+ push_mapped_length(upstream, [1, 2], session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+
+ assert sorted(ti.map_index for ti in decision.schedulable_tis if
ti.task_id == "task_2") == [0, 1]
+ assert session.execute(query).all() == [
+ (0, None),
+ (1, None),
(2, TaskInstanceState.REMOVED),
(3, TaskInstanceState.REMOVED),
]
