apache / apache/airflow

Mismatched Xcom Map Index when Dynamic Mapping over TaskGroup and not all mapped tasks have run

Open
#40,321 6 comments 3 reactions 0 assignees View on GitHub
affected_version:3.0 area:core area:dynamic-task-mapping kind:bug priority:medium
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 10h
Merged PRs (30d)
483

Description

### Apache Airflow version

Other Airflow 2 version (please specify below)

### If "Other Airflow 2 version" selected, which one?

2.8.3

### What happened?

If a TaskGroup is dynamically mapped, the Xcom map index for downstream tasks to pull from is mismatched. This is due to the Xcom length not being the same length as the number of mapped tasks resulting in the following error:
` File "/home/airflow/.local/lib/python3.11/site-packages/airflow/models/xcom.py", line 795, in __getitem__
raise IndexError(key) from None
IndexError: 5`

This is a simplified DAG that reproduces the problem. One mapped task is caused to reschedule instead of complete causing the shortened Xcom dict:

```from datetime import datetime, timedelta

from airflow.decorators import task, task_group
from airflow.operators.python import get_current_context
from airflow import DAG
from airflow.operators.empty import EmptyOperator
from airflow.sensors.base import PokeReturnValue

default_args = {
'owner': 'dev',
'depends_on_past': False,
'start_date': datetime(2019, 1, 1),
'provide_context': True,
'retries': 0,
'retry_delay': timedelta(seconds=30)
}

@task
def parse_csv_schedule():
items_dict = [{'A': 1}, {'B': 2}, {'C': 3}, {'D': 4}, {'E': 5}, {'F': 6}]
return items_dict

@task_group(group_id="process_items")
def process_items(items_dict: dict):
@task.sensor(poke_interval=90, mode="reschedule")
def retrieve_item(item: dict) -> PokeReturnValue:
for letter, integer in item.items():
if integer == 3:
print("Number is 3 - reschedule")
return PokeReturnValue(is_done=False)
else:
print("Number is not 3 - Let's go")
return PokeReturnValue(is_done=True, xcom_value=item)

@task
def process_item(item: dict):
context = get_current_context()
ti = context["ti"]
map_index = ti.map_index
print(f"Taskflow XCOM: {item}")
xcom_dict = ti.xcom_pull("process_items.retrieve_item")[map_index]
print(xcom_dict)

item_dict = retrieve_item(item=items_dict)
process_item(item=item_dict)

with DAG(dag_id='af041_Debugger', default_args=default_args, max_active_runs=1, schedule_interval=None):
end = EmptyOperator(task_id='end', trigger_rule='none_failed_min_one_success')
items_dict = parse_csv_schedule()
items_dict >> process_items.expand(items_dict=items_dict) >> end
```

Resulting in:
![image](https://github.com/apache/airflow/assets/39799836/79af125a-44a8-4360-ab82-f266935b5303)

![image](https://github.com/apache/airflow/assets/39799836/8511db9f-f05f-49da-a9b2-6d2d068aa147)

The last task (map_index 5) fails as the Xcom list length only has 5 items instead of 6. What is more of an issue is the task with map_index=3, has pulled the xcom for map_index=4!

This issues appears to happen with both Taskflow API and traditional operators.

### What you think should happen instead?

- Mapped tasks should run independently of each other.
- Downstream tasks should always be dependent on their associated upstream tasks
- If a mapped task does not complete then downstream tasks should not be able to access said task's Xcom
- Downstream tasks should only be able to access Xcoms from their related upstream tasks

Ideally a downstream task could just perform a xcom pull from a upstream task with the same map_index without needing to call ti.map_index. If the Xcom must be a list of all mapped tasks then there should be an explicit reference to the map_index or at the very least, the length of the dict should be the same as the number of mapped tasks.

### How to reproduce

Run the provided DAG. You will notice that the last mapped task fails as the map_index for the upstream Xcom does not exist. This is due to the length of the Xcom dict being too short due to one of the mapped tasks not completing.

If there is a better way to access Xcoms from upstream tasks (not just the direct upstream via Taskflow) please let me know as it feels like I'm not using TaskGroups and Xcoms the way it's intended.

### Operating System

Ubuntu 22.04

### Versions of Apache Airflow Providers

_No response_

### Deployment

Official Apache Airflow Helm Chart

### Deployment details

AWS EKS

### Anything else?

100% reproducible

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Contributor guide

Open the contributing guide

Research direction

Start with the XCom handling shown in airflow/models/xcom.py around line 795, then run the provided reproducing DAG on Airflow 2.8.3. Trace mapped TaskGroup XCom pulls when one sensor reschedules, and verify that downstream tasks do not receive shifted or missing values for their associated map_index.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.