Refactor workgraph into core

environment.yml

@@ -20,6 +20,7 @@ dependencies:
 - jedi<0.19
 - jinja2~=3.0
 - kiwipy[rmq]~=0.9.0
+- node-graph~=0.6.5
 - numpy<3,>=1.21
 - paramiko~=3.0
 - pgsu~=0.3.0

pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
   'jedi<0.19',
   'jinja2~=3.0',
   'kiwipy[rmq]~=0.9.0',
+  'node-graph~=0.6.5',
   'numpy>=1.21,<3',
   'paramiko~=3.0',
   'pgsu~=0.3.0',
@@ -117,13 +118,16 @@ requires-python = '>=3.10'
 'core.code.containerized' = 'aiida.orm.nodes.data.code.containerized:ContainerizedCode'
 'core.code.installed' = 'aiida.orm.nodes.data.code.installed:InstalledCode'
 'core.code.portable' = 'aiida.orm.nodes.data.code.portable:PortableCode'
+'core.datetime' = 'aiida.orm.nodes.data.datetime:DateTimeData'
 'core.dict' = 'aiida.orm.nodes.data.dict:Dict'
 'core.enum' = 'aiida.orm.nodes.data.enum:EnumData'
 'core.float' = 'aiida.orm.nodes.data.float:Float'
 'core.folder' = 'aiida.orm.nodes.data.folder:FolderData'
+'core.function' = 'aiida.orm.nodes.data.function:FunctionData'
 'core.int' = 'aiida.orm.nodes.data.int:Int'
 'core.jsonable' = 'aiida.orm.nodes.data.jsonable:JsonableData'
 'core.list' = 'aiida.orm.nodes.data.list:List'
+'core.none' = 'aiida.orm.nodes.data.none:NoneData'
 'core.numeric' = 'aiida.orm.nodes.data.numeric:NumericType'
 'core.orbital' = 'aiida.orm.nodes.data.orbital:OrbitalData'
 'core.remote' = 'aiida.orm.nodes.data.remote.base:RemoteData'
@@ -151,6 +155,7 @@ requires-python = '>=3.10'
 'process.workflow' = 'aiida.orm.nodes.process.workflow.workflow:WorkflowNode'
 'process.workflow.workchain' = 'aiida.orm.nodes.process.workflow.workchain:WorkChainNode'
 'process.workflow.workfunction' = 'aiida.orm.nodes.process.workflow.workfunction:WorkFunctionNode'
+'process.workflow.workgraph' = 'aiida.orm.nodes.process.workflow.workgraph:WorkGraphNode'
 
 [project.entry-points.'aiida.orm']
 'core.auth_info' = 'aiida.orm.authinfos:AuthInfo'
@@ -391,6 +396,7 @@ module = [
   'graphviz.*',
   'kiwipy.*',
   'mayavi.*',
+  'node_graph.*',
   'pgsu.*',
   'pgtest.*',
   'trogon.*',

src/aiida/engine/processes/workchains/workchain.py

@@ -131,6 +131,9 @@ def __init__(
 
         self._stepper: Stepper | None = None
         self._awaitables: list[Awaitable] = []
+        # The pks of awaitables whose completion callback is already registered. This is runtime state, callbacks
+        # do not survive a checkpoint, so it is not persisted and is reset in `load_instance_state`.
+        self._registered_awaitable_pks: set[int] = set()
         self._context = AttributeDict()
 
     @classmethod
@@ -174,13 +177,51 @@ def load_instance_state(self, saved_state, load_context):
         self._stepper = None
         stepper_state = saved_state.get(self._STEPPER_STATE, None)
         if stepper_state is not None:
-            self._stepper = self.spec().get_outline().recreate_stepper(stepper_state, self)  # type: ignore[arg-type]
+            self._stepper = self._recreate_stepper(stepper_state)
 
         self.set_logger(self.node.logger)
 
+        # Callbacks do not survive the checkpoint, so nothing is registered yet on the reloaded process.
+        self._registered_awaitable_pks = set()
         if self._awaitables:
             self._action_awaitables()
 
+    def _create_stepper(self) -> Stepper:
+        """Return the stepper that drives this work chain.
+
+        This is the seam for supplying a different execution strategy. The default steps through the outline declared
+        on the spec, but a subclass may return any :class:`plumpy.workchains.Stepper`, for example one that derives the
+        order of execution from a graph of data dependencies instead of a static outline.
+
+        A subclass that overrides this should also override :meth:`_recreate_stepper`, otherwise its processes cannot
+        be restored from a checkpoint.
+        """
+        return self.spec().get_outline().create_stepper(self)  # type: ignore[arg-type]
+
+    def _recreate_stepper(self, saved_state: t.Any) -> Stepper:
+        """Restore the stepper from the state it wrote to the checkpoint.
+
+        The counterpart of :meth:`_create_stepper`, called when a process is loaded from a checkpoint rather than
+        started fresh.
+
+        :param saved_state: the state previously returned by ``Stepper.save()``
+        """
+        return self.spec().get_outline().recreate_stepper(saved_state, self)  # type: ignore[arg-type]
+
+    @property
+    def _awaitable_barrier(self) -> bool:
+        """Whether each step waits for everything it launched before the next one begins.
+
+        This is the difference between the two execution models, and it is a property of the stepping strategy, so
+        the value is taken from the stepper. ``True``, the default, is the outline model: :meth:`_do_step` clears
+        the awaitables at the start of every step, so a step forms a barrier over the children it launched and the
+        process only resumes once all of them have finished. A stepper that schedules by data dependencies wants
+        ``False``: the awaitables persist across steps and the process resumes as each child finishes, so
+        independent branches stay in flight together. A stepper opts into the streaming model by defining
+        ``awaitable_barrier = False`` on itself.
+        """
+        return getattr(self._stepper, 'awaitable_barrier', True)
+
     @Protect.final
     def on_run(self):
         super().on_run()
@@ -299,7 +340,7 @@ def _update_process_status(self) -> None:
     @override
     @Protect.final
     async def run(self) -> t.Any:
-        self._stepper = self.spec().get_outline().create_stepper(self)  # type: ignore[arg-type]
+        self._stepper = self._create_stepper()
         return await run_with_portal(self._do_step)
 
     def _do_step(self) -> t.Any:
@@ -312,7 +353,11 @@ def _do_step(self) -> t.Any:
         """
         from .context import ToContext
 
-        self._awaitables = []
+        # Under the barrier model the awaitables belong to a single step and are cleared before the next one, which
+        # is what forces every step to wait for all the children it launched. A streaming stepper keeps them, so
+        # children launched in earlier steps stay in flight while later steps run.
+        if self._awaitable_barrier:
+            self._awaitables = []
         result: t.Any = None
 
         try:
@@ -380,24 +425,41 @@ def on_wait(self, awaitables: t.Sequence[t.Awaitable]):
             self.call_soon(self.resume)
 
     def _action_awaitables(self) -> None:
-        """Handle the awaitables that are currently registered with the work chain.
+        """Register the completion callback for each awaitable that does not already have one.
 
         Depending on the class type of the awaitable's target a different callback
         function will be bound with the awaitable and the runner will be asked to
-        call it when the target is completed
+        call it when the target is completed.
+
+        The registration is guarded against duplicates: under the barrier model the awaitables are cleared each
+        step so the same one is never seen twice, but a streaming stepper keeps its awaitables across steps and
+        would otherwise register a further callback for the same awaitable on every pass through the waiting state.
         """
         for awaitable in self._awaitables:
+            if awaitable.pk in self._registered_awaitable_pks:
+                continue
             if awaitable.target == AwaitableTarget.PROCESS:
                 callback = functools.partial(self.call_soon, self._on_awaitable_finished, awaitable)
                 self.runner.call_on_process_finish(awaitable.pk, callback)
+                self._registered_awaitable_pks.add(awaitable.pk)
             else:
                 raise AssertionError(f"invalid awaitable target '{awaitable.target}'")
 
+    def _on_awaitable_resolved(self, awaitable: Awaitable) -> None:
+        """Hook called once a finished awaitable has been resolved onto the context, before the resume decision.
+
+        Defaults to doing nothing. A subclass can use it to run bookkeeping that must see the resolved value and
+        must happen before the process is resumed, without having to reimplement :meth:`_on_awaitable_finished`.
+
+        :param awaitable: the awaitable that has just been resolved
+        """
+
     def _on_awaitable_finished(self, awaitable: Awaitable) -> None:
         """Callback function, for when an awaitable process instance is completed.
 
-        The awaitable will be effectuated on the context of the work chain and removed from the internal list. If all
-        awaitables have been dealt with, the work chain process is resumed.
+        The awaitable will be effectuated on the context of the work chain and removed from the internal list. The
+        process is then resumed: under the barrier model only once every awaitable has finished, and under the
+        streaming model as soon as this one does, so a finished child can unblock its dependents while others run.
 
         :param awaitable: an Awaitable instance
         """
@@ -414,6 +476,8 @@ def _on_awaitable_finished(self, awaitable: Awaitable) -> None:
             value = node  # type: ignore[assignment]
 
         self._resolve_awaitable(awaitable, value)
+        self._registered_awaitable_pks.discard(awaitable.pk)
+        self._on_awaitable_resolved(awaitable)
 
-        if self.state == ProcessState.WAITING and not self._awaitables:
+        if self.state == ProcessState.WAITING and (not self._awaitable_barrier or not self._awaitables):
             self.resume()

src/aiida/orm/__init__.py

@@ -52,13 +52,15 @@
     'ComputerEntityLoader',
     'ContainerizedCode',
     'Data',
+    'DateTimeData',
     'Dict',
     'Entity',
     'EntityExtras',
     'EntityTypes',
     'EnumData',
     'Float',
     'FolderData',
+    'FunctionData',
     'Group',
     'GroupEntityLoader',
     'ImportGroup',
@@ -77,6 +79,7 @@
     'NodeEntityLoader',
     'NodeLinksManager',
     'NodeRepository',
+    'NoneData',
     'NumericType',
     'OrbitalData',
     'OrderSpecifier',
@@ -104,10 +107,13 @@
     'User',
     'WorkChainNode',
     'WorkFunctionNode',
+    'WorkGraphNode',
     'WorkflowNode',
     'XyData',
     'cif_from_ase',
+    'deserialize_to_raw_python_data',
     'find_bandgap',
+    'general_serializer',
     'get_loader',
     'get_query_type_from_type_string',
     'get_type_string_from_class',
@@ -119,6 +125,7 @@
     'load_node',
     'load_node_class',
     'pycifrw_from_cif',
+    'serialize_to_aiida_nodes',
     'to_aiida_type',
     'validate_link',
 )

src/aiida/orm/nodes/__init__.py

@@ -31,10 +31,12 @@
     'Code',
     'ContainerizedCode',
     'Data',
+    'DateTimeData',
     'Dict',
     'EnumData',
     'Float',
     'FolderData',
+    'FunctionData',
     'InstalledCode',
     'Int',
     'JsonableData',
@@ -44,6 +46,7 @@
     'Node',
     'NodeAttributes',
     'NodeRepository',
+    'NoneData',
     'NumericType',
     'OrbitalData',
     'PortableCode',
@@ -62,12 +65,16 @@
     'UpfData',
     'WorkChainNode',
     'WorkFunctionNode',
+    'WorkGraphNode',
     'WorkflowNode',
     'XyData',
     'cif_from_ase',
+    'deserialize_to_raw_python_data',
     'find_bandgap',
+    'general_serializer',
     'has_pycifrw',
     'pycifrw_from_cif',
+    'serialize_to_aiida_nodes',
     'to_aiida_type',
 )
 

src/aiida/orm/nodes/data/__init__.py

@@ -18,16 +18,21 @@
 from .cif import *
 from .code import *
 from .data import *
+from .datetime import *
+from .deserializer import *
 from .dict import *
 from .enum import *
 from .float import *
 from .folder import *
+from .function import *
 from .int import *
 from .jsonable import *
 from .list import *
+from .none import *
 from .numeric import *
 from .orbital import *
 from .remote import *
+from .serializer import *
 from .singlefile import *
 from .str import *
 from .structure import *
@@ -43,16 +48,19 @@
     'Code',
     'ContainerizedCode',
     'Data',
+    'DateTimeData',
     'Dict',
     'EnumData',
     'Float',
     'FolderData',
+    'FunctionData',
     'InstalledCode',
     'Int',
     'JsonableData',
     'Kind',
     'KpointsData',
     'List',
+    'NoneData',
     'NumericType',
     'OrbitalData',
     'PortableCode',
@@ -70,9 +78,12 @@
     'UpfData',
     'XyData',
     'cif_from_ase',
+    'deserialize_to_raw_python_data',
     'find_bandgap',
+    'general_serializer',
     'has_pycifrw',
     'pycifrw_from_cif',
+    'serialize_to_aiida_nodes',
     'to_aiida_type',
 )
 

src/aiida/orm/nodes/data/datetime.py

@@ -0,0 +1,45 @@
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""`Data` sub class to represent a :class:`datetime.datetime` value."""
+
+from __future__ import annotations
+
+import datetime
+
+from .base import to_aiida_type
+from .data import Data
+
+__all__ = ('DateTimeData',)
+
+
+@to_aiida_type.register(datetime.datetime)
+def _(value):
+    return DateTimeData(value)
+
+
+class DateTimeData(Data):
+    """`Data` sub class to store a :class:`datetime.datetime` object.
+
+    The value is stored as an ISO-8601 string for portability across backends and reconstructed on access.
+    """
+
+    def __init__(self, value, **kwargs):
+        if not isinstance(value, datetime.datetime):
+            msg = f'expected a datetime.datetime, got {type(value)}'
+            raise TypeError(msg)
+        super().__init__(**kwargs)
+        self.base.attributes.set('datetime', value.isoformat())
+
+    @property
+    def value(self) -> datetime.datetime:
+        """Return the stored value as a :class:`datetime.datetime`."""
+        return datetime.datetime.fromisoformat(self.base.attributes.get('datetime'))
+
+    def __str__(self) -> str:
+        return str(self.value)