🐛 Map: fix gather race for async sources (#776)

`update_map_task_state` gathered from per-item clones of the
pass-through `gather_item`. For async process-type sources (CalcJob,
WorkChain, @task.graph) the awaitable cascade could reach the gather
before `continue_workgraph` had scheduled those clones, raising
`KeyError` or leaving them PLANNED and hanging the finalize path.

- `generate_mapped_tasks` no longer clones `gather_item`; the gather
  reads each value directly from the mapped source task's clones, which
  is race-free (source results are recorded before the cascade reaches
  the zone).

- Failure handling is fail-fast and keyed on the gather sources: the
  zone fails only when a gather source produced no result (its clone
  FAILED, or was SKIPPED in an iteration that errored), then skips the
  zone's downstream tasks. A deliberately-untaken `If` branch gathers
  `None` and finishes; an unrelated body-task failure leaves the gather
  complete and surfaces via the ordinary exit 302.

- `Map.gather` rejects a source produced outside the zone, validating
  all sources before mutating so a rejected build stays rebuildable.

- `map_info` keeps the source -> gather_item edges (as template edges)
  so the GUI can still draw the gather.

Also drop a stale commented-out gather block from
`update_template_task_state`, and add tests for the async source, the
failure / skip / If-in-Map paths, the out-of-zone reject, and map_info.

---

Co-authored-by: Edward Linscott <edwardlinscott@gmail.com>

Author: GeigerJ2

src/aiida_workgraph/engine/task_manager.py

@@ -305,7 +305,8 @@ def execute_map_task(self, task, kwargs):
             for prefix, value in source.items():
                 new_tasks, new_links = self.generate_mapped_tasks(task, prefix=prefix)
                 self.update_map_item_task_state(item_task, prefix, value)
-            map_info['children'] = list(new_tasks.keys())
+            # include the (uncloned) gather_item so its restored edges are not dangling
+            map_info['children'] = list(new_tasks.keys()) + [task.gather_item_task.name]
             map_info['links'] = new_links
         self.state_manager.set_task_runtime_info(name, 'map_info', map_info)
         # gather task finishes immediately
@@ -472,6 +473,15 @@ def generate_mapped_tasks(self, zone_task: Task, prefix: str) -> None:
         all_links = []
         child_tasks = self.get_all_children(zone_task.name)
         for child_task in child_tasks:
+            # The gather_item task is a pure pass-through aggregator
+            # (executor=return_input); the map zone reads directly from the
+            # mapped source tasks in `update_map_task_state`, so cloning
+            # gather_item would just create unused clones. Skipping the
+            # clone also avoids a race where, for async process-type source
+            # tasks (CalcJob, WorkChain, @task.graph), the gather_item
+            # clones stay PLANNED and hang the engine's finalize path.
+            if self.process.wg.tasks[child_task].identifier == 'workgraph.gather_item':
+                continue
             # since the child task is mapped, it should be skipped
             self.state_manager.set_task_runtime_info(child_task, 'state', TaskState.MAPPED)
             task = self.copy_task(child_task, prefix)
@@ -482,6 +492,15 @@ def generate_mapped_tasks(self, zone_task: Task, prefix: str) -> None:
         new_links = self._patch_cloned_tasks(new_tasks, all_links)
         # update process.wg.connectivity so the new tasks are recognized in child_node, zone references, etc.
         self._patch_connectivity(new_tasks)
+        # gather_item is intentionally not cloned, so `_patch_cloned_tasks` drops
+        # the source -> gather_item edges (its `to_task` is not in new_tasks) and
+        # the GUI loses them. Re-add each as its template edge, the same shape the
+        # other map_info links use (the GUI expands them per prefix); this is
+        # display-only, no engine link is created.
+        gather_item = zone_task.gather_item_task
+        for link in gather_item.inputs._all_links:
+            if link.from_task.name in new_tasks:
+                new_links.append(link.to_dict())
         return new_tasks, new_links
 
     def update_map_item_task_state(self, item_task, prefix, value: Any):

src/aiida_workgraph/engine/task_state.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
-from typing import Optional, Tuple, List, Any
+from typing import Optional, Tuple, List, Any, Iterator, TYPE_CHECKING
 from typing_extensions import assert_never
 from aiida.orm.utils.serialize import serialize
 from aiida_workgraph.orm.utils import deserialize_safe
 from aiida.orm import ProcessNode, Data
 from aiida_workgraph.enums import TERMINAL_TASK_STATES, RuntimeInfoKey, TaskState
 from node_graph.socket import BaseSocket, TaskSocketNamespace
 
+if TYPE_CHECKING:
+    from aiida_workgraph.task import Task
+
 
 class TaskStateManager:
     """
@@ -291,49 +294,104 @@ def update_zone_task_state(self, name: str) -> None:
             self.process.report(f'Task: {name} finished.')
             self.update_parent_task_state(name)
 
+    def _iter_zone_clones(self, zone: 'Task') -> Iterator[Tuple[str, 'Task']]:
+        """Yield ``(prefix, clone)`` for every mapped clone within a Map zone.
+
+        Clones live under each body template's ``mapped_tasks``, keyed by prefix
+        (the zone's child list holds templates, one clone per prefix), so walking
+        the template tree reaches all of them, including those in a nested `If`.
+        """
+        stack = list(getattr(zone, 'children', []))
+        while stack:
+            template = stack.pop()
+            yield from (template.mapped_tasks or {}).items()
+            stack.extend(getattr(template, 'children', []))
+
     def update_map_task_state(self, name: str) -> None:
         """Update the map task state.
         1) check if all child tasks are finished.
         2) gather the results of all the mapped tasks.
         3) update the parent task state.
         """
+        from aiida_workgraph.utils import get_nested_dict
+
         finished, _ = self.are_childen_finished(name)
-        if finished:
-            map_zone = self.process.wg.tasks[name]
-            # gather the results of all the mapped tasks
-            gather_task = map_zone.gather_item_task
-            for input in gather_task.inputs:
-                if input._name.startswith('_'):
-                    continue
-                results = {}
-                link = input._links[0]
-                for prefix, mapped_task in self.process.wg.tasks[gather_task.name].mapped_tasks.items():
-                    results[prefix] = self.ctx._task_results[mapped_task.name][link.to_socket._name]
-                self.ctx._task_results[name][link.to_socket._name] = results
-            self.set_task_runtime_info(name, 'state', TaskState.FINISHED)
-            # self.update_meta_tasks(name)
-            self.process.report(f'Task: {name} finished.')
-            self.update_meta_tasks(name)
+        if not finished:
+            return
+        map_zone = self.process.wg.tasks[name]
+        # Gather the results of all the mapped tasks.
+        #
+        # We aggregate directly from each mapped SOURCE task (the task whose
+        # output is linked into the template gather_item), not via the
+        # gather_item itself. The gather_item template is a pure pass-through
+        # aggregator (executor=return_input) and is intentionally not cloned
+        # per item in `generate_mapped_tasks`, so there are no gather_item
+        # clones to read from. Reading directly from the source's
+        # `_task_results` is also race-free: the source's results are
+        # populated by `update_task_state` before any cascade can reach here,
+        # which matters when the source is an async process-type task
+        # (CalcJob, WorkChain, or a @task.graph sub-workflow).
+        gather_task = map_zone.gather_item_task
+        gather_links = [
+            input_socket._links[0]
+            for input_socket in gather_task.inputs
+            if not input_socket._name.startswith('_') and input_socket._links
+        ]
+        # Fail the zone (fail-fast) only when a GATHER SOURCE produced no result:
+        # its clone FAILED, or was SKIPPED in an iteration that errored (something
+        # it depended on FAILED). This is keyed on the gather sources, not every
+        # clone, so a false `If` branch gathers None (SKIPPED with no FAILED clone
+        # in its iteration), and an unrelated body-task failure is left to the
+        # ordinary 302 path rather than discarding an otherwise-complete gather.
+        failed_prefixes = {
+            prefix
+            for prefix, clone in self._iter_zone_clones(map_zone)
+            if self.get_task_runtime_info(clone.name, 'state') == TaskState.FAILED
+        }
+        missing_prefixes = set()
+        for link in gather_links:
+            for prefix, clone in (self.process.wg.tasks[link.from_task.name].mapped_tasks or {}).items():
+                clone_state = self.get_task_runtime_info(clone.name, 'state')
+                if clone_state == TaskState.FAILED or (clone_state == TaskState.SKIPPED and prefix in failed_prefixes):
+                    missing_prefixes.add(prefix)
+        if missing_prefixes:
+            self.set_task_runtime_info(name, 'state', TaskState.FAILED)
+            # Skip the zone's downstream tasks, as `on_task_failed` does for an
+            # ordinary failure; otherwise they run on the missing gather output
+            # and pollute the report with their own consequent failures.
+            self.set_tasks_state(self.process.wg.connectivity['child_node'][name], TaskState.SKIPPED)
+            self.process.report(
+                f'Task: {name} failed, no result from mapped item(s): {", ".join(sorted(missing_prefixes))}.'
+            )
             self.update_parent_task_state(name)
+            return
+        for link in gather_links:
+            source_clones = self.process.wg.tasks[link.from_task.name].mapped_tasks or {}
+            results = {}
+            for prefix, clone in source_clones.items():
+                # A source with no recorded value for this socket gathers None: a
+                # false `If` branch (SKIPPED), or a FINISHED source that returned
+                # no value for it. They read the same here; telling them apart is
+                # the resilient-Map follow-up.
+                results[prefix] = get_nested_dict(
+                    self.ctx._task_results[clone.name],
+                    link.from_socket._scoped_name,
+                    default=None,
+                )
+            self.ctx._task_results[name][link.to_socket._name] = results
+        self.set_task_runtime_info(name, 'state', TaskState.FINISHED)
+        self.process.report(f'Task: {name} finished.')
+        self.update_meta_tasks(name)
+        self.update_parent_task_state(name)
 
     def update_template_task_state(self, name: str) -> None:
         """Update the template task state.
         1) check if all child tasks are finished.
-        2) gather the results of all the mapped tasks.
-        3) update the parent task state.
+        2) update the parent task state.
         """
         finished, _ = self.are_childen_finished(name)
         if finished:
-            # # gather the results of all the mapped tasks
-            # results = {}
-            # for prefix, mapped_task in self.process.wg.tasks[name].mapped_tasks.items():
-            #     for output in mapped_task.outputs:
-            #         if output._name in self.ctx._task_results[mapped_task.name]:
-            #             results.setdefault(output._name, {})
-            #             results[output._name][prefix] = self.ctx._task_results[mapped_task.name][output._name]
-            # self.ctx._task_results[name] = results
             self.set_task_runtime_info(name, 'state', TaskState.FINISHED)
-            # self.update_meta_tasks(name)
             self.process.report(f'Task: {name} finished.')
             self.update_parent_task_state(name)
 

src/aiida_workgraph/tasks/builtins.py

@@ -133,8 +133,29 @@ def gather_item_task(self) -> Task:
                 return cast(Task, child)
         return self.add_task('workgraph.gather_item')
 
+    def _contains_task(self, target: Task) -> bool:
+        """True if ``target`` is a descendant of this zone (recursing nested zones)."""
+        stack = list(self.children)
+        while stack:
+            task = stack.pop()
+            if task.name == target.name:
+                return True
+            stack.extend(getattr(task, 'children', []))
+        return False
+
     def gather(self, sockets: Dict[str, BaseSocket]) -> BaseSocket:
         """Collect per-entry results into the zone outputs, one namespace per name."""
+        # Validate every source before mutating anything (the loop below and even
+        # `self.gather_item_task` add state), so a rejection leaves the zone
+        # unchanged and rebuildable.
+        for name, socket in sockets.items():
+            source = socket._task
+            if not self._contains_task(source):
+                msg = (
+                    f"Map.gather() source '{name}' ('{source.name}') is outside the Map zone (one value, not "
+                    f'per-iteration). Move a per-iteration task inside the zone, or use a shared value (graph input) directly.'
+                )
+                raise ValueError(msg)
         gather_item = self.gather_item_task
         for name in sockets:
             gather_item.add_input_spec('workgraph.any', name=name)