Fix test and update doc

Author: superstar54

docs/gallery/advanced/autogen/context_manager.py

@@ -485,7 +485,9 @@ def generate_add_multiply_workgraph():
 # ``Map`` zone
 # ============
 # .. warning::
-#   **This feature is experimental.** The API for ``Map`` zone is subject to change in future releases. We welcome your feedback on its functionality.
+#   **This feature is experimental.** The API for ``Map`` zone is subject to change in future releases.
+#   We welcome your feedback on its functionality.
+#   **ctx** does not work inside a ``Map`` zone yet.
 #
 # The ``Map`` context manager works similarly to Python's built-in ``map`` function.
 # By accessing the ``item`` member of the ``Map`` context, we can pass each individual element (e.g. a dictionary entry) to tasks.
@@ -522,10 +524,13 @@ def get_value(data, key):
 
 with WorkGraph('AddMap') as wg:
     with Map(data) as map_zone:
-        wg.outputs.result = add(
-            x=get_value(map_zone.item, 'x').result,
-            y=get_value(map_zone.item, 'y').result,
+        result = add(
+            x=get_value(map_zone.item.value, 'x').result,
+            y=get_value(map_zone.item.value, 'y').result,
         ).result
+        map_zone.gather({'result': result})
+        wg.ctx.result = map_zone.outputs.result
+        wg.outputs.result = wg.ctx.result
 
 wg.run()
 
@@ -561,17 +566,18 @@ def get_value(data, key):
 
 
 @task
-def aggregate_sum(data):
+def aggregate_sum(data: spec.dynamic(Any)) -> int:
     return sum(data.values())
 
 
 with WorkGraph('AddAggregate') as wg:
     with Map(data) as map_zone:
         added_numbers = add(
-            x=get_value(map_zone.item, 'x').result,
-            y=get_value(map_zone.item, 'y').result,
+            x=get_value(map_zone.item.value, 'x').result,
+            y=get_value(map_zone.item.value, 'y').result,
         ).result
-    wg.outputs.result = aggregate_sum(added_numbers).result
+        map_zone.gather({'result': added_numbers})
+    wg.outputs.result = aggregate_sum(map_zone.outputs.result).result
 
 wg.run()
 
@@ -598,6 +604,10 @@ def aggregate_sum(data):
 #
 # Let's run an add-multiply workflow with a hardcoded multiplication factor:
 
+
+# In order to restart a workflow, the tasks should be importable from a module (i.e. not defined on-the-fly).
+from aiida_workgraph.tasks.tests import add, multiply
+
 with WorkGraph('AddMultiplyToBeContinued', inputs=spec.namespace(x=Any, y=Any)) as wg1:
     the_sum = add(
         x=wg1.inputs.x,

docs/gallery/advanced/autogen/node_graph_programming.py

@@ -195,9 +195,9 @@ def generate_data(N) -> spec.namespace(result=spec.dynamic(Any)):
 
 
 @task
-def calc_sum(**kwargs):
+def calc_sum(data: spec.dynamic(Any)) -> int:
     """Calculates the sum of all keyword arguments' values."""
-    return sum(kwargs.values())
+    return sum(data.values())
 
 
 # %%
@@ -215,11 +215,13 @@ def calc_sum(**kwargs):
 map_task = wg.add_task('workgraph.map_zone', source=data_task.outputs.result)
 
 # Inside the Map Zone, add 1 to each item
-add_task_in_map = map_task.add_task(add, x=map_task.item, y=1)
+add_task_in_map = map_task.add_task(add, x=map_task.item.value, y=1)
+
+map_task.gather({'result': add_task_in_map.outputs.result})
 
 # After the Map Zone, sum all the results from the add_task_in_map
 # The 'kwargs' input allows collecting all dynamic outputs from the mapped tasks.
-sum_task = wg.add_task(calc_sum, kwargs=add_task_in_map.outputs.result)
+sum_task = wg.add_task(calc_sum, data=map_task.outputs.result)
 
 # Set the final output of the workgraph
 wg.outputs.result = sum_task.outputs.result

pyproject.toml

@@ -20,12 +20,12 @@ classifiers = [
 dependencies = [
   "numpy",
   "scipy",
-  "node-graph==0.3.12",
+  "node-graph==0.3.13",
   "node-graph-widget>=0.0.5",
   "aiida-core~=2.7.1",
   "cloudpickle",
   "aiida-shell~=0.8",
-  "aiida-pythonjob==0.4.4",
+  "aiida-pythonjob==0.4.5",
   "jsonschema"
 ]
 description = "Design flexible node-based workflow for AiiDA calculation."

src/aiida_workgraph/engine/task_manager.py

@@ -303,6 +303,9 @@ def execute_map_task(self, task, kwargs):
             map_info['children'] = list(new_tasks.keys())
             map_info['links'] = new_links
         self.state_manager.set_task_runtime_info(name, 'map_info', map_info)
+        # gather task finishes immediately
+        gather_task = task.gather_item_task
+        self.state_manager.set_task_runtime_info(gather_task.name, 'state', 'FINISHED')
 
         self.continue_workgraph()
 
@@ -472,7 +475,8 @@ def generate_mapped_tasks(self, zone_task: Task, prefix: str) -> None:
 
     def update_map_item_task_state(self, item_task, prefix, value: Any):
         new_name = f'{prefix}_{item_task.name}'
-        self.ctx._task_results[new_name]['item'] = value
+        self.ctx._task_results[new_name]['key'] = prefix
+        self.ctx._task_results[new_name]['value'] = value
         self.state_manager.set_task_runtime_info(new_name, 'state', 'FINISHED')
 
     def copy_task(self, name: str, prefix: str) -> 'Task':

src/aiida_workgraph/engine/task_state.py

@@ -249,8 +249,10 @@ def update_parent_task_state(self, name: str) -> None:
             node_type = parent_task.node_type.upper()
             if node_type == 'WHILE':
                 self.update_while_task_state(parent_task.name)
-            elif node_type in ['IF', 'ZONE', 'MAP']:
+            elif node_type in ['IF', 'ZONE']:
                 self.update_zone_task_state(parent_task.name)
+            elif node_type == 'MAP':
+                self.update_map_task_state(parent_task.name)
 
         # If the task is a mapped child, update its parent's "template" (the original map node)
         if self.process.wg.tasks[name].map_data:
@@ -283,24 +285,49 @@ def update_zone_task_state(self, name: str) -> None:
             self.process.report(f'Task: {name} finished.')
             self.update_parent_task_state(name)
 
-    def update_template_task_state(self, name: str) -> None:
-        """Update the template task state.
+    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.
         """
         finished, _ = self.are_childen_finished(name)
         if finished:
+            map_zone = self.process.wg.tasks[name]
             # 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
+            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', 'FINISHED')
+            # self.update_meta_tasks(name)
+            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.
+        """
+        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', 'FINISHED')
+            # self.update_meta_tasks(name)
             self.process.report(f'Task: {name} finished.')
             self.update_parent_task_state(name)
 

src/aiida_workgraph/executors/builtins.py

@@ -28,7 +28,6 @@ def get_item(data: dict, key: str) -> Any:
     return data.get(key, None)
 
 
-def return_inputs(**inputs: dict) -> dict:
-    """Return the inputs"""
-    inputs['items'] = inputs.pop('item', {})
-    return inputs
+def return_input(**kwargs: Any) -> dict:
+    """Return the input"""
+    return kwargs

src/aiida_workgraph/tasks/builtins.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import Any, Dict
 from aiida_workgraph.task import ChildTaskSet, Task
 from aiida_workgraph import task, namespace, meta
@@ -8,7 +9,8 @@
 from node_graph.node_spec import NodeSpec
 from node_graph.socket_spec import SocketSpec, SocketSpecMeta
 from typing import Annotated
-from aiida_workgraph.executors.builtins import get_item, update_ctx, get_context, select
+from aiida_workgraph.executors.builtins import update_ctx, get_context, select, return_input
+from node_graph.node import BuiltinPolicy
 
 
 class GraphLevelTask(_GraphIOSharedMixin, Task):
@@ -41,7 +43,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.children = ChildTaskSet(parent=self)
 
-    def add_task(self, *args, **kwargs):
+    def add_task(self, *args, **kwargs) -> Task:
         """Syntactic sugar to add a task to the zone."""
         task = self.graph.add_task(*args, **kwargs)
         self.children.add(task)
@@ -92,8 +94,9 @@ class Map(Zone):
         node_type='MAP',
         catalog='Control',
         inputs=namespace(
-            source=Annotated[Any, SocketSpec('workgraph.any', link_limit=100000)],
+            source=SocketSpec('workgraph.any', link_limit=100000),
         ),
+        outputs=namespace(),
         base_class_path='aiida_workgraph.tasks.builtins.Map',
     )
 
@@ -109,14 +112,29 @@ def item(self):
         map_item_task = self.add_task('workgraph.map_item')
         return map_item_task.outputs
 
-      def gather(self, socket: BaseSocket) -> None:
-        gather_item = self.graph.add_task('workgraph.gather_item')
-        self.graph.add_link(socket, gather_item.inputs.value)
-        return gather_item.outputs.values
+    @property
+    def gather_item_task(self) -> Task | None:
+        for child in self.children:
+            if child.identifier == 'workgraph.gather_item':
+                return child
+        gather_item = self.add_task('workgraph.gather_item')
+        return gather_item
+
+    def gather(self, sockets: Dict[str, BaseSocket]) -> None:
+        gather_item = self.gather_item_task
+        for name in sockets:
+            gather_item.add_input_spec('workgraph.any', name=name)
+            self.add_output_spec('workgraph.any', name=name)
+        gather_item.set_inputs(sockets)
+        return gather_item.outputs
+
 
 class MapItem(Task):
     """MapItem"""
 
+    # turn off framework builtins for these graph-level nodes
+    _BUILTINS_POLICY = BuiltinPolicy(input_wait=False, output_wait=False, default_output=False)
+
     _default_spec = NodeSpec(
         identifier='workgraph.map_item',
         node_type='Normal',
@@ -125,32 +143,26 @@ class MapItem(Task):
             source=SocketSpec('workgraph.any', link_limit=100000, meta=SocketSpecMeta(required=False)),
             key=SocketSpec('workgraph.string', meta=SocketSpecMeta(required=False)),
         ),
-        outputs=namespace(item=SocketSpec('workgraph.any')),
-        executor=RuntimeExecutor.from_callable(get_item),
+        outputs=namespace(key=str, value=any),
         base_class_path='aiida_workgraph.tasks.builtins.MapItem',
     )
 
 
 class GatherItem(Task):
     """GatherItem"""
 
-    identifier = 'workgraph.gather_item'
-    name = 'GatherItem'
-    node_type = 'Normal'
-    catalog = 'Control'
-
-    def create_sockets(self) -> None:
-        self.inputs._clear()
-        self.outputs._clear()
+    # turn off framework builtins for these graph-level nodes
+    _BUILTINS_POLICY = BuiltinPolicy(input_wait=True, output_wait=False, default_output=False)
 
-        self.add_input('workgraph.any', 'value')
-        self.add_output('workgraph.namespace', 'values')
-        self.add_output('workgraph.any', '_wait')
-
-    def get_executor(self):
-        from aiida_workgraph.executors.builtins import return_inputs
-
-        return RuntimeExecutor.from_callable(return_inputs)
+    _default_spec = NodeSpec(
+        identifier='workgraph.gather_item',
+        node_type='Normal',
+        catalog='Control',
+        inputs=namespace(),
+        outputs=namespace(),
+        executor=RuntimeExecutor.from_callable(return_input),
+        base_class_path='aiida_workgraph.tasks.builtins.GatherItem',
+    )
 
 
 class SetContext(Task):

src/aiida_workgraph/tasks/tests.py

@@ -0,0 +1,13 @@
+from aiida_workgraph import task
+
+
+@task
+def add(x, y):
+    """Add two numbers."""
+    return x + y
+
+
+@task
+def multiply(x, y):
+    """Multiply two numbers."""
+    return x * y