You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/08/13 13:32:22 UTC
(airflow) 03/03: fix map_index is null dynamic task mapping bug (#54249)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 3251e68cf9c509cd08f0067a9e8adcc354dc3263
Author: Isaiah Iruoha <76...@users.noreply.github.com>
AuthorDate: Wed Aug 13 06:22:06 2025 -0700
fix map_index is null dynamic task mapping bug (#54249)
closes: #52881
When tasks within nested task groups try to access XCom values from sibling tasks in the same mapped parent task group instance, the system incorrectly uses `map_index=null` instead of the correct map index, causing XCom lookup failures.
<img width="1033" height="494" alt="Screenshot 2025-08-07 at 1 47 06 PM" src="https://github.com/user-attachments/assets/f6c1767b-a999-470d-b090-478643dcdc01" />
```javascript
No XCom value found; defaulting to None.: key="return_value": dag_id="test_dag": task_id="expandable_task_group.inner_task_group.alter_input": run_id="manual__2025-07-04T14:56:47.815002+00:00": map_index=null
```
The issue is in `airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py` in the `_get_upstream_map_indexes` function.
The logic only checks the immediate parent task group (upstream_task.task_group) but doesn't consider nested scenarios where tasks are within a regular task group that is itself within a mapped task group.
Replace the immediate parent check with a hierarchy-aware check using `get_closest_mapped_task_group()`. This ensures that both tasks recognize they share the same mapped task group ancestor and use the same map_index.
- Verified proposed solution resolves issue: #52881 when tested locally using `breeze start-airflow` against the latest main branch development environment.
- Added `test_nested_mapped_task_group_upstream_indexes` to verify that tasks in nested mapped task groups correctly resolve upstream map indexes.
(cherry picked from commit f685c61704017a256fbb4393b3f4d0e6722aa869)
---
.../execution_api/routes/task_instances.py | 11 ++--
.../versions/head/test_task_instances.py | 75 ++++++++++++++++++++++
2 files changed, 80 insertions(+), 6 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index e22d0a5f34d..ea56e2f37d4 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -66,7 +66,6 @@ from airflow.models.trigger import Trigger
from airflow.models.xcom import XComModel
from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated
from airflow.sdk.definitions.asset import Asset, AssetUniqueKey
-from airflow.sdk.definitions.taskgroup import MappedTaskGroup
from airflow.utils import timezone
from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
@@ -289,20 +288,20 @@ def ti_run(
def _get_upstream_map_indexes(
task: Operator, ti_map_index: int, run_id: str, session: SessionDep
) -> Iterator[tuple[str, int | list[int] | None]]:
+ task_mapped_group = task.get_closest_mapped_task_group()
for upstream_task in task.upstream_list:
+ upstream_mapped_group = upstream_task.get_closest_mapped_task_group()
map_indexes: int | list[int] | None
- if not isinstance(upstream_task.task_group, MappedTaskGroup):
+ if upstream_mapped_group is None:
# regular tasks or non-mapped task groups
map_indexes = None
- elif task.task_group == upstream_task.task_group:
- # tasks in the same mapped task group
- # the task should use the map_index as the previous task in the same mapped task group
+ elif task_mapped_group == upstream_mapped_group:
+ # tasks in the same mapped task group hierarchy
map_indexes = ti_map_index
else:
# tasks not in the same mapped task group
# the upstream mapped task group should combine the return xcom as a list and return it
mapped_ti_count: int
- upstream_mapped_group = upstream_task.task_group
try:
# for cases that does not need to resolve xcom
mapped_ti_count = upstream_mapped_group.get_parse_time_mapped_ti_count()
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index d18b18b673e..4017a04d82e 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -301,6 +301,81 @@ class TestTIRunState:
upstream_map_indexes = response.json()["upstream_map_indexes"]
assert upstream_map_indexes == expected_upstream_map_indexes[(ti.task_id, ti.map_index)]
+ def test_nested_mapped_task_group_upstream_indexes(self, client, dag_maker):
+ """
+ Test that upstream_map_indexes are correctly computed for tasks in nested mapped task groups.
+ """
+ with dag_maker("test_nested_mapped_tg", serialized=True):
+
+ @task
+ def alter_input(inp: str) -> str:
+ return f"{inp}_Altered"
+
+ @task
+ def print_task(orig_input: str, altered_input: str) -> str:
+ return f"orig:{orig_input},altered:{altered_input}"
+
+ @task_group
+ def inner_task_group(orig_input: str) -> None:
+ altered_input = alter_input(orig_input)
+ print_task(orig_input, altered_input)
+
+ @task_group
+ def expandable_task_group(param: str) -> None:
+ inner_task_group(param)
+
+ expandable_task_group.expand(param=["One", "Two", "Three"])
+
+ dr = dag_maker.create_dagrun()
+
+ # Set all alter_input tasks to success so print_task can run
+ for ti in dr.get_task_instances():
+ if "alter_input" in ti.task_id and ti.map_index >= 0:
+ ti.state = State.SUCCESS
+ elif "print_task" in ti.task_id and ti.map_index >= 0:
+ ti.set_state(State.QUEUED)
+ dag_maker.session.flush()
+
+ # Expected upstream_map_indexes for each print_task instance
+ expected_upstream_map_indexes = {
+ ("expandable_task_group.inner_task_group.print_task", 0): {
+ "expandable_task_group.inner_task_group.alter_input": 0
+ },
+ ("expandable_task_group.inner_task_group.print_task", 1): {
+ "expandable_task_group.inner_task_group.alter_input": 1
+ },
+ ("expandable_task_group.inner_task_group.print_task", 2): {
+ "expandable_task_group.inner_task_group.alter_input": 2
+ },
+ }
+
+ # Get only the expanded print_task instances (not the template)
+ print_task_tis = [
+ ti for ti in dr.get_task_instances() if "print_task" in ti.task_id and ti.map_index >= 0
+ ]
+
+ # Test each print_task instance
+ for ti in print_task_tis:
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/run",
+ json={
+ "state": "running",
+ "hostname": "random-hostname",
+ "unixname": "random-unixname",
+ "pid": 100,
+ "start_date": "2024-09-30T12:00:00Z",
+ },
+ )
+
+ assert response.status_code == 200
+ upstream_map_indexes = response.json()["upstream_map_indexes"]
+ expected = expected_upstream_map_indexes[(ti.task_id, ti.map_index)]
+
+ assert upstream_map_indexes == expected, (
+ f"Task {ti.task_id}[{ti.map_index}] should have upstream_map_indexes {expected}, "
+ f"but got {upstream_map_indexes}"
+ )
+
def test_dynamic_task_mapping_with_xcom(self, client, dag_maker, create_task_instance, session, run_task):
"""
Test that the Task Instance upstream_map_indexes is correctly fetched when to running the Task Instances with xcom