Skip to content

Fix Map zone gather race for async process-type tasks#776

Draft
GeigerJ2 wants to merge 7 commits into
aiidateam:mainfrom
GeigerJ2:fix/map-gather-race-async-tasks
Draft

Fix Map zone gather race for async process-type tasks#776
GeigerJ2 wants to merge 7 commits into
aiidateam:mainfrom
GeigerJ2:fix/map-gather-race-async-tasks

Conversation

@GeigerJ2

@GeigerJ2 GeigerJ2 commented Apr 13, 2026

Copy link
Copy Markdown
Contributor

TL;DR Stop cloning the pass-through gather_item; read the gathered values from the source clones directly

Mapping over an async process-type task (CalcJob, WorkChain, @task.graph) inside a Map zone could crash with a KeyError or hang the engine. At runtime a Map clones its body once per item into prefixed copies (key_0_<task>, key_1_<task>, ...), stored in source_task.mapped_tasks[prefix].

The zone then gathers each clone's output through a pass-through gather_item node, and update_map_task_state read each item's result out of a per-item clone of that gather_item. For synchronous source tasks that's fine: each runs inline and is done before the engine moves on. For async ones, it races: the awaitable cascade can finish the mapped source tasks and reach the gather phase before continue_workgraph has scheduled the gather_item clones, so their results are missing (KeyError) or the clones sit in PLANNED (hang).

Two changes:
generate_mapped_tasks stops cloning the gather_item: it's a pure pass-through, so its clones only ever held copies, and the stray PLANNED clone was the hang. update_map_task_state instead reads each value directly from the mapped source task's clones (source_task.mapped_tasks), which is the same value one link upstream and is race-free, because a source task's results are written by update_task_state before the cascade reaches the gather.

The new read also passes default=None, which matters when a mapped item fails: are_childen_finished counts FAILED/SKIPPED as finished, so the zone still reaches the gather. This doesn't hide the failure. Like any failed task, the run finishes with the engine's existing exit status 302 (is_finished_ok False, message naming the failed task); default=None only stops the gather from raising a KeyError on the missing result, which on main excepted the whole engine and surfaced as KeyError: 'sum1', burying the real cause (key_1 raised a ValueError). The failed item lands as None in the gathered namespace, which is debatable: the honest value is arguably the failed process node itself (it carries the exit status and traceback), not a None that reads like data. That is a Map-semantics decision, separate from this race fix.

The existing map tests only map over synchronous tasks that always succeed, so they passed on main and never covered this. Added two to tests/test_map.py (a @task.graph source, and a failing item) that fail on main with KeyError: 'sum1' and pass here. Also dropped a stale commented-out block from update_template_task_state.

#777 is stacked on this and exposes the gathered outputs to the client after wg.run(). I'd land this engine-side fix first; #777 re-touches update_map_task_state and can rebase onto it.

MWE
from typing import Annotated
from aiida import orm
from aiida_workgraph import Map, WorkGraph, dynamic, namespace, task

@task()
def generate_data(n: int) -> Annotated[dict, namespace(data=dynamic(int))]:
    return {'data': {f'key_{i}': i for i in range(n)}}

@task()
def add(x, y):
    return x + y

@task.graph  # process-type: runs as its own sub-process (the async case)
def add_workflow(x, y) -> int:
    return add(x=x, y=y).result

@task()
def calc_sum(data: Annotated[dict, dynamic(orm.Int)]) -> float:
    return sum(data.values())

with WorkGraph('mwe') as wg:
    data = generate_data(n=3).data
    with Map(data) as map_zone:
        out1 = add_workflow(x=map_zone.item.value, y=1).result
        map_zone.gather({'sum1': out1})
    out3 = calc_sum(data=map_zone.outputs.sum1).result
    wg.run()

print(out3.value)   # main: KeyError: 'sum1' (engine EXCEPTED) -- here: 6
Why reading the source clone gives the same value

The gather link runs source_clone.<from_socket> -> gather_item_clone.<to_socket>, and gather_item's executor is return_input (identity), so the value on the gather_item clone is literally the value already on the source clone.

old (gather_item clone) new (source clone)
which clones gather_item.mapped_tasks source_task.mapped_tasks
read key link.to_socket._name link.from_socket._scoped_name
accessor d[key] get_nested_dict(d, key, default=None)

Both loops use the same prefixes and write to the same _task_results[name][link.to_socket._name]. _scoped_name + get_nested_dict is used on the source side because a source output can be a nested namespace, whereas the gather_item input was always flat; for a flat socket they coincide.

`update_map_task_state` aggregated gathered outputs by iterating over
the gather_item's own clones. For async process-type source tasks
(CalcJob, WorkChain, @task.graph) the awaitable callback cascade can
finish all mapped source tasks and reach the gather phase before
`continue_workgraph` has scheduled the gather_item clones. The missing
clone results raised `KeyError` or left clones stuck in PLANNED, hanging
the engine's finalize path.

Two changes fix this:

- `generate_mapped_tasks` skips cloning the gather_item entirely. Its
  executor is a pure pass-through (`return_input`), so no per-item clone
  is needed, and skipping it removes the PLANNED-clone hang.

- `update_map_task_state` reads results directly from each mapped source
  task's clones (via `source_task.mapped_tasks`). This is race-free
  because source results are populated by `update_task_state` before the
  cascade can reach the gather phase. A `default=None` fallback also
  lets a failed mapped task finish gracefully (exit 302) instead of
  excepting the engine.

Add regression tests covering an async `@task.graph` source and a failed
mapped task; the existing test only maps over synchronous tasks that
always succeed and so passed before the fix.
GeigerJ2 added a commit to GeigerJ2/aiida-workgraph that referenced this pull request Jun 24, 2026
Task state and task action were encoded as bare strings checked against
hand-written `in [...]` lists copy-pasted across the engine: no single
source of truth, a typo silently never matched, and a new member had to
be remembered in every list.

Add `aiida_workgraph/enums.py` with `TaskState`/`TaskAction` (str-mixed
enums) and a `RuntimeInfoKey` `Literal` for the runtime-info dispatch
key, and migrate every call site to them. The duplicated terminal-state
check now lives once as `TERMINAL_TASK_STATES` (with `is_terminal`),
which is the kind of thing that made aiidateam#776 subtle.

The `str` mixin keeps each member equal to and JSON-serialising as the
same bare string, so persisted process-node state is byte-identical and
the migration is behaviour-preserving. A `__str__` returning the value
keeps reports and logs reading `RUNNING`, not `TaskState.RUNNING`, on
every supported Python version.

Fix the latent action casing bug: normalise once at the RPC boundary
(`apply_task_actions` validates the string into a `TaskAction` and
dispatches via `match`/`assert_never`), store the canonical value,
compare against members everywhere, and drop the scattered `.upper()`
calls. An unknown action now raises instead of being silently swallowed
(the engine's RPC layer captures and logs it without disturbing the
running workflow).

Leave the genuine `node.process_state` comparisons as plain strings
(aiida-core's `ProcessState` domain), and drop the dead `WAITING`/
`PAUSED` entries from the `kill_tasks` task-state membership list, where
a task state can never take those values. `enums.py` is importable but
not re-exported at the package top level.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant