Fix Map zone gather race for async process-type tasks#776
Draft
GeigerJ2 wants to merge 7 commits into
Draft
Conversation
`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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 aMapzone could crash with aKeyErroror hang the engine. At runtime aMapclones its body once per item into prefixed copies (key_0_<task>,key_1_<task>, ...), stored insource_task.mapped_tasks[prefix].The zone then gathers each clone's output through a pass-through
gather_itemnode, andupdate_map_task_stateread each item's result out of a per-item clone of thatgather_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 beforecontinue_workgraphhas scheduled the gather_item clones, so their results are missing (KeyError) or the clones sit inPLANNED(hang).Two changes:
generate_mapped_tasksstops cloning thegather_item: it's a pure pass-through, so its clones only ever held copies, and the strayPLANNEDclone was the hang.update_map_task_stateinstead 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 byupdate_task_statebefore the cascade reaches the gather.The new read also passes
default=None, which matters when a mapped item fails:are_childen_finishedcountsFAILED/SKIPPEDas 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_okFalse, message naming the failed task);default=Noneonly stops the gather from raising aKeyErroron the missing result, which onmainexcepted the whole engine and surfaced asKeyError: 'sum1', burying the real cause (key_1raised aValueError). The failed item lands asNonein the gathered namespace, which is debatable: the honest value is arguably the failed process node itself (it carries the exit status and traceback), not aNonethat 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
mainand never covered this. Added two totests/test_map.py(a@task.graphsource, and a failing item) that fail onmainwithKeyError: 'sum1'and pass here. Also dropped a stale commented-out block fromupdate_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-touchesupdate_map_task_stateand can rebase onto it.MWE
Why reading the source clone gives the same value
The gather link runs
source_clone.<from_socket>->gather_item_clone.<to_socket>, andgather_item's executor isreturn_input(identity), so the value on the gather_item clone is literally the value already on the source clone.gather_item.mapped_taskssource_task.mapped_taskslink.to_socket._namelink.from_socket._scoped_named[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_dictis 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.