apache / apache/airflow

Task state store is keyed by positional map_index and survives a clear, so a re-expanded list silently inherits the wrong item's state

Open
#72,725 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 9h
Merged PRs (30d)
472

Description

## Summary

`task_state_store` rows are keyed by `map_index` — a **positional** index into the expanded upstream list — and they survive a clear-and-rerun by design.

If the upstream list changes between the first run and the clear-and-rerun (an insertion, a deletion, a reordering — routine for "list the new files since yesterday"), map index *N* on the second attempt is a **different item** than map index *N* on the first, and it silently inherits the earlier item's state.

The run reports success.

## Where

- `airflow-core/src/airflow/models/task_state_store.py` — unique on `(dag_run_id, task_id, map_index, key)`
- `airflow-core/src/airflow/state/metastore.py` — `_get` / `_store` / `_delete` match on `(dag_id, run_id, task_id, map_index, key)`; the async twins do the same
- `task-sdk/src/airflow/sdk/execution_time/task_runner.py` — builds `TaskScope(dag_id, run_id, task_id, map_index)`

Airflow's own clear path is careful about identity elsewhere: `clear_task_instances` regenerates the TaskInstance uuid (`taskinstance.py`, `self.id = uuid7()`), and a rerun task's XComs are deleted before re-execution (`task_runner.py`). Only the state store is addressed positionally.

## Reproduction

A mapped task that checkpoints per item, run once over `["a","b","c"]`, then **cleared and rerun in place** (what the UI "Clear task" button does, via `clear_task_instances`) with the upstream now returning `["z","a","b","c"]`:

```
--- ATTEMPT 1 (items=[a,b,c]) ---
OBSERVE map_index=0 item='a' checkpoint_read=None -> wrote 'a'
OBSERVE map_index=1 item='b' checkpoint_read=None -> wrote 'b'
OBSERVE map_index=2 item='c' checkpoint_read=None -> wrote 'c'

--- ATTEMPT 2 (clear-and-rerun in place, items=[z,a,b,c]) ---
OBSERVE map_index=0 item='z' checkpoint_read='a' [MISPAIRED]
OBSERVE map_index=1 item='a' checkpoint_read='b' [MISPAIRED]
OBSERVE map_index=2 item='b' checkpoint_read='c' [MISPAIRED]
OBSERVE map_index=3 item='c' checkpoint_read=None -> wrote 'c'
[rerun] final dagrun state: success
```

Every pre-existing checkpoint was handed to the wrong item. In the documented external-job-resumption pattern this means the instance processing `z` reattaches to the external job submitted for `a`.

The DAG:

```python
@dag(schedule=None, start_date=LOGICAL_DATE, catchup=False)
def d():
@task
def make_items():
return _items() # reads a file; changed between the two attempts

@task
def process(item, **context):
store = context["task_state_store"]
ck = store.get("checkpoint")
if ck is None:
store.set("checkpoint", item)
else:
print(f"map_index={context['ti'].map_index} item={item!r} checkpoint_read={ck!r}")

process.expand(item=make_items())
```

## Negative controls

**A — same mechanism, list unchanged** (the documented survive-a-clear case):

```
OBSERVE map_index=0 item='a' checkpoint_read='a' [OK]
OBSERVE map_index=1 item='b' checkpoint_read='b' [OK]
OBSERVE map_index=2 item='c' checkpoint_read='c' [OK]
```

**B — shifted list but a different DagRun** (the near miss that correctly does not trigger — `run_id` in the key does isolate):

```
OBSERVE map_index=0 item='z' checkpoint_read=None -> wrote 'z'
OBSERVE map_index=1 item='a' checkpoint_read=None -> wrote 'a'
```

## The silence

Driver exit code 0. All task instances `success`. DagRun `success`. Grepping the full run log for `[warning`, `[error` and `Traceback` yields nothing related. Nothing notices that the expansion width or contents changed, and nothing notices that a checkpoint written under a different input is being served.

## What the docs say

`airflow-core/docs/core-concepts/task-state-store.rst`:

> When a task is dynamically mapped (`task.expand(...)`), each map index has its own task state store namespace.

presented as an isolation guarantee, with no caveat that the index is positional or that it can name a different item after re-expansion. The same page states that surviving an operator-initiated clear is the intended use, so the persistence is not the problem — what persists is addressed by position.

Separately, `docs/authoring-and-scheduling/dynamic-task-mapping.rst` states that **the order of expansion is not guaranteed** and that mapped tasks are assigned an integer index. The two statements are never joined.

The class comment in `models/task_state_store.py` reasons about retries ("retries of the same task share the same rows — that is the point") and about different DAG runs having different `dag_run_id`. The re-expansion case is not considered.

The sibling `AssetScope` docstring **does** carry the equivalent warning, that `name` and `uri` are not guaranteed to be unique over time. `TaskScope` has none.

The documented custom-backend example builds the external object path as `f"airflow/task-store/{scope.dag_id}/{scope.run_id}/{scope.task_id}/{scope.map_index}/{key}"`, so the positional identity propagates into user object stores and the mispairing follows it out of the metadata database.

## Note on reproducing this

`dag.test()` cannot show it: `get_or_create_dagrun` **deletes** the existing DagRun, and the `ON DELETE CASCADE` foreign key wipes `task_state_store` with it, so every read comes back `None`. That looks like the store is safe and is not. The reproduction above uses `clear_task_instances` — the same function the public clear endpoint calls — and keeps the DagRun row.

## Not tested

- No scheduler/executor/API-server deployment. Tasks ran through Airflow's own `_run_task` in-process against SQLite; the Postgres/MySQL upsert dialects in `metastore.py` were not exercised.
- The re-run loop in the harness is mine, modelled on the body of `dag.test()`, because no public helper re-runs an existing DagRun in process. The **clear** and the **task execution** are Airflow's code.
- Custom (non-metastore) state backends and the external-ref serialization path.
- The async `a*` variants — identical code shape, presumed identical, unexercised.
- `map_index_template`-named indices, which change the UI label but not the stored `map_index`.
- The list-shrink direction: rows at now-nonexistent higher indices persist until expiry and would be picked up if the list grows again.
- `clear_on_success=True`, which would mask this for successful tasks but not for tasks cleared from a failed state.
- Backfill, and multi-argument `expand` / `expand_kwargs` / `.concat()`.

I have not checked whether this has been raised before; a pointer to an existing issue is welcome and I am happy to close this in favour of one.

## Version

Airflow `3.4.0`, source at `33d9915992980e1a082f665d33861f39171393da`, installed from source, Python 3.12, SQLite metastore.

Contributor guide

Open the contributing guide

Research direction

Start with task_state_store.py, the _get/_store/_delete paths in state/metastore.py, and TaskScope construction in task_runner.py; use the clear-and-rerun reproduction described in the issue. Trace how map_index reaches the database and documented custom backends, then define an identity approach that preserves same-item retries while preventing changed expansions from reusing another item's state. Done should include coverage for the changed-list, unchanged-list, and different-DagRun cases, plus aligned documentation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, databases
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.