🐛 Map: fix gather race for async source tasks

`update_map_task_state` gathered from the gather_item's own clones. For
async process-type sources (CalcJob, WorkChain, @task.graph) the
awaitable cascade can reach the gather before `continue_workgraph` has
scheduled those clones, raising `KeyError` or leaving them PLANNED and
hanging the finalize path.

`generate_mapped_tasks` no longer clones the gather_item: its executor
is a pure pass-through, so the clones only ever held copies. The gather
reads each value from the mapped source task's clones instead, which is
race-free because source results are recorded before the cascade
reaches the zone.

A failed iteration produces no result, so the zone now goes FAILED and
names the affected map keys, rather than gathering `None` and finishing
with a namespace that is silently one entry short.

Also drop a stale commented-out gather block from
`update_template_task_state`; the aggregation lives in
`update_map_task_state`.

Author: GeigerJ2

src/aiida_workgraph/engine/task_manager.py

@@ -472,6 +472,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)

src/aiida_workgraph/engine/task_state.py

@@ -297,43 +297,72 @@ def update_map_task_state(self, name: str) -> None:
         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
+        ]
+        # An iteration that raised is FAILED, and `on_task_failed` SKIPs its
+        # downstream children, so a broken item surfaces as either state on the
+        # source clone. Both mean it produced nothing to gather, so fail the
+        # whole zone instead of emitting a namespace that is silently missing
+        # that item: a FINISHED zone whose gathered dict is short reads as a
+        # complete result and quietly corrupts whatever aggregates it.
+        failed_prefixes = sorted(
+            prefix
+            for link in gather_links
+            for prefix, clone in (self.process.wg.tasks[link.from_task.name].mapped_tasks or {}).items()
+            if self.get_task_runtime_info(clone.name, 'state') in (TaskState.FAILED, TaskState.SKIPPED)
+        )
+        if failed_prefixes:
+            self.set_task_runtime_info(name, 'state', TaskState.FAILED)
+            self.process.report(f'Task: {name} failed, no result from mapped item(s): {", ".join(failed_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():
+                # Defensive only: every surviving clone is FINISHED with its
+                # result recorded, since the failure paths returned above.
+                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)
 

tests/test_map.py

@@ -22,6 +22,20 @@ def add(x, y):
     return x + y
 
 
+@task.graph
+def add_workflow(x, y) -> int:
+    """Async process-type source task (runs as its own sub-process) for use inside a Map zone."""
+    return add(x=x, y=y).result
+
+
+@task()
+def maybe_fail(x, y):
+    """Add two numbers, but fail for one specific item to exercise the failure path."""
+    if x == 1:
+        raise ValueError('boom on x==1')
+    return x + y
+
+
 @task()
 def calc_sum(data: Annotated[dict, dynamic(orm.Int)]) -> float:
     """Compute the sum of all provided values."""
@@ -71,3 +85,45 @@ def test_map_value_and_key():
     # values 0, 1 -> (0+10)+(1+10) = 21; keys are the user's own source keys
     assert total.value == 21
     assert joined.value == 'key_0,key_1'
+
+
+def test_map_zone_async_source():
+    """Map over an async process-type source task (`@task.graph`) and gather it.
+
+    Regression test for the gather race: when the mapped source is a process-type
+    task, the awaitable cascade can reach the gather phase before the gather_item
+    clones are scheduled. Before the fix this raised ``KeyError`` and excepted the
+    engine; the existing ``test_map_zone`` does not catch it because it maps over
+    plain synchronous ``@task`` functions that all succeed.
+    """
+    x = 1
+    n = 3
+    with WorkGraph('map_async_source') as wg:
+        data = generate_data(n=n).data
+        with Map(data) as map_zone:
+            out1 = add_workflow(x=map_zone.value, y=x).result
+            map_zone.gather({'sum1': out1})
+        out3 = calc_sum(data=map_zone.outputs.sum1).result
+        wg.run()
+        # values are 0+1, 1+1, 2+1 -> 1 + 2 + 3 = 6
+        assert out3.value == 6
+
+
+def test_map_zone_failed_iteration_fails_the_zone():
+    """A failed mapped iteration must fail the zone, not except and not under-report.
+
+    Before the fix the gather ``KeyError``-ed on the item that produced no
+    result, excepting the engine and burying the real cause. The zone now goes
+    FAILED and gathers nothing, rather than reporting FINISHED with a namespace
+    that is silently missing the failed item.
+    """
+    n = 3
+    with WorkGraph('map_fail') as wg:
+        data = generate_data(n=n).data
+        with Map(data) as map_zone:
+            out1 = maybe_fail(x=map_zone.value, y=10).result
+            map_zone.gather({'sum1': out1})
+        wg.run()
+    assert map_zone.state == 'FAILED'
+    assert wg.process.exit_status == 302
+    assert 'key_1_maybe_fail' in wg.process.exit_message