Support arbitrary nested dynamic socket (#712)

* allow arbitrary nested dynamic socket
* Bump node-graph to 0.4.2
* Bump aiida-pythonjob to 0.4.6
* Update docs

Author: superstar54

docs/gallery/howto/autogen/annotate_inputs_outputs.py

@@ -171,6 +171,11 @@ def SquareNumbersGenerator(n: int):
 # The graph shows that the ``generate_square_numbers`` task has multiple output nodes, one for each entry in the dynamically generated dictionary.
 # The ``dynamic(typing.Any)`` specification instructs the workgraph to treat each value in the returned dictionary as a separate output node of any type.
 #
+# .. note::
+#
+#    If no item type is specified (i.e., just ``dynamic()``), the namespace becomes fully dynamic,
+#    enabling arbitrary nested structures (e.g., dictionaries). Each value will then be recursively treated as a leaf node.
+#
 # Nested namespaces
 # ~~~~~~~~~~~~~~~~~
 #

pyproject.toml

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

src/aiida_workgraph/config.py

@@ -15,17 +15,6 @@
     WorkChain: 'WORKCHAIN',
 }
 
-builtin_inputs = {
-    '_wait': {
-        'link_limit': 1e6,
-        'metadata': {'arg_type': 'none', 'builtin_socket': True},
-    }
-}
-builtin_outputs = {
-    '_wait': {'metadata': {'arg_type': 'none', 'builtin_socket': True}},
-    '_outputs': {'metadata': {'arg_type': 'none', 'builtin_socket': True}},
-}
-
 
 def load_config() -> dict:
     """Load the configuration from the config file."""

src/aiida_workgraph/engine/task_state.py

@@ -108,7 +108,6 @@ def update_normal_task_state(self, name, results, success=True):
         """Set the results of a normal task.
         A normal task is created by decorating a function with @task().
         """
-        from aiida_workgraph.config import builtin_outputs
 
         if success:
             task = self.process.wg.tasks[name]
@@ -117,7 +116,9 @@ def update_normal_task_state(self, name, results, success=True):
                 if len(task.outputs) - 2 != len(results):
                     self.on_task_failed(name)
                     return self.process.exit_codes.OUTPUS_NOT_MATCH_RESULTS
-                output_names = [output._name for output in task.outputs if output._name not in builtin_outputs]
+                output_names = [
+                    output._name for output in task.outputs if output._metadata.extra.get('builtin_socket') is not True
+                ]
                 for i, output_name in enumerate(output_names):
                     self.ctx._task_results[name][output_name] = results[i]
             elif isinstance(results, dict):

src/aiida_workgraph/socket_spec.py

@@ -2,7 +2,7 @@
 from dataclasses import replace
 from typing import Any, Tuple, Dict
 from node_graph.socket_spec import (
-    SocketSpecMeta,
+    SocketMeta,
     SocketSpecSelect,
     SocketSpec,
     SocketSpecAPI as _SocketSpecAPI,
@@ -70,7 +70,7 @@ def _from_port(cls, port: Port | PortNamespace, *, parent_required: bool, role:
             ns = SocketSpec(
                 identifier=cls.NAMESPACE,
                 fields=fields,
-                meta=SocketSpecMeta(
+                meta=SocketMeta(
                     required=required_here,
                     is_metadata=getattr(port, 'is_metadata', False),
                     call_role=('kwargs' if role == 'input' else None),
@@ -81,8 +81,11 @@ def _from_port(cls, port: Port | PortNamespace, *, parent_required: bool, role:
             is_dyn = bool(getattr(port, 'dynamic', False))
             if is_dyn:
                 valid_type = getattr(port, 'valid_type', None)
-                item_ident = cls._identifier_from_valid_type(valid_type)
-                ns = replace(ns, dynamic=True, item=SocketSpec(identifier=item_ident))
+                if valid_type:
+                    item_ident = cls._identifier_from_valid_type(valid_type)
+                    ns = replace(ns, meta=replace(ns.meta, dynamic=True), item=SocketSpec(identifier=item_ident))
+                else:
+                    ns = replace(ns, meta=replace(ns.meta, dynamic=True), item=None)
             return ns
 
         # Leaf Port (InputPort/OutputPort)
@@ -91,7 +94,7 @@ def _from_port(cls, port: Port | PortNamespace, *, parent_required: bool, role:
         ident = cls._identifier_from_valid_type(valid_type)
         return SocketSpec(
             identifier=ident,
-            meta=SocketSpecMeta(
+            meta=SocketMeta(
                 required=required_here,
                 is_metadata=getattr(port, 'is_metadata', False),
                 call_role=('kwargs' if role == 'input' else None),

src/aiida_workgraph/task.py

@@ -74,16 +74,6 @@ def set_from_builder(self, builder: Any) -> None:
         data = get_dict_from_builder(builder)
         self.set_inputs(data)
 
-    def set_from_protocol(self, *args: Any, **kwargs: Any) -> None:
-        """Set the task inputs from protocol data."""
-
-        executor = self.get_executor().callable
-        # check if the executor has the get_builder_from_protocol method
-        if not hasattr(executor, 'get_builder_from_protocol'):
-            raise AttributeError(f'Executor {executor.__name__} does not have the get_builder_from_protocol method.')
-        builder = executor.get_builder_from_protocol(*args, **kwargs)
-        self.set_from_builder(builder)
-
     @classmethod
     def new(cls, identifier: Union[str, Callable], name: Optional[str] = None) -> 'Task':
         """Create a task from a identifier."""
@@ -177,21 +167,6 @@ def to_widget_value(self):
         wgdata = workgraph_to_short_json(wgdata)
         return wgdata
 
-    def _repr_mimebundle_(self, *args: Any, **kwargs: Any) -> any:
-        # if ipywdigets > 8.0.0, use _repr_mimebundle_ instead of _ipython_display_
-        self.widget.value = self.to_widget_value()
-        if hasattr(self.widget, '_repr_mimebundle_'):
-            return self.widget._repr_mimebundle_(*args, **kwargs)
-        else:
-            return self.widget._ipython_display_(*args, **kwargs)
-
-    def to_html(self, output: str = None, show_socket_depth: Optional[int] = None, **kwargs):
-        """Write a standalone html file to visualize the task."""
-        if show_socket_depth is None:
-            show_socket_depth = self.show_socket_depth
-        self.widget.value = self.to_widget_value()
-        return self.widget.to_html(output=output, **kwargs)
-
 
 class TaskSet:
     """

src/aiida_workgraph/tasks/builtins.py

@@ -7,7 +7,7 @@
 from node_graph import RuntimeExecutor
 from aiida import orm
 from node_graph.node_spec import NodeSpec
-from node_graph.socket_spec import SocketSpec, SocketSpecMeta
+from node_graph.socket_spec import SocketSpec, SocketMeta
 from typing import Annotated
 from aiida_workgraph.executors.builtins import update_ctx, get_context, select, return_input
 from node_graph.node import BuiltinPolicy
@@ -140,8 +140,8 @@ class MapItem(Task):
         node_type='Normal',
         catalog='Control',
         inputs=namespace(
-            source=SocketSpec('workgraph.any', link_limit=100000, meta=SocketSpecMeta(required=False)),
-            key=SocketSpec('workgraph.string', meta=SocketSpecMeta(required=False)),
+            source=SocketSpec('workgraph.any', link_limit=100000, meta=SocketMeta(required=False)),
+            key=SocketSpec('workgraph.string', meta=SocketMeta(required=False)),
         ),
         outputs=namespace(key=str, value=any),
         base_class_path='aiida_workgraph.tasks.builtins.MapItem',
@@ -173,7 +173,7 @@ class SetContext(Task):
         node_type='Normal',
         catalog='Control',
         inputs=namespace(
-            context=SocketSpec('workgraph.any', meta=SocketSpecMeta(required=False)),
+            context=SocketSpec('workgraph.any', meta=SocketMeta(required=False)),
             key=any,
             value=any,
         ),
@@ -189,7 +189,7 @@ class GetContext(Task):
         identifier='workgraph.get_context',
         node_type='Normal',
         catalog='Control',
-        inputs=namespace(context=SocketSpec('workgraph.any', meta=SocketSpecMeta(required=False)), key=any),
+        inputs=namespace(context=SocketSpec('workgraph.any', meta=SocketMeta(required=False)), key=any),
         outputs=namespace(result=any),
         executor=RuntimeExecutor.from_callable(get_context),
         base_class_path='aiida_workgraph.tasks.builtins.GetContext',

src/aiida_workgraph/tasks/pythonjob_tasks.py

@@ -8,7 +8,7 @@
 from aiida_pythonjob import pyfunction, PythonJob, PyFunction, MonitorPyFunction
 from aiida_pythonjob.utils import serialize_ports
 from aiida_workgraph.task import Task
-from node_graph.socket_spec import SocketSpec, SocketSpecSelect, SocketSpecMeta
+from node_graph.socket_spec import SocketSpec, SocketSpecSelect, SocketMeta
 from node_graph.node_spec import NodeSpec
 from aiida_workgraph.socket_spec import namespace
 from .function_task import build_callable_nodespec
@@ -291,9 +291,9 @@ def build_pythonjob_nodespec(
 
     # additions specific to PythonJob
     add_in = namespace(
-        computer=Annotated[str, SocketSpecMeta(required=False)],
-        command_info=Annotated[dict, SocketSpecMeta(required=False)],
-        register_pickle_by_value=Annotated[bool, SocketSpecMeta(required=False)],
+        computer=Annotated[str, SocketMeta(required=False)],
+        command_info=Annotated[dict, SocketMeta(required=False)],
+        register_pickle_by_value=Annotated[bool, SocketMeta(required=False)],
     )
 
     return build_callable_nodespec(

src/aiida_workgraph/tasks/shelljob_task.py

@@ -7,7 +7,7 @@
 from aiida_shell.launch import prepare_shell_job_inputs
 from node_graph.node_spec import NodeSpec
 from node_graph.executor import RuntimeExecutor
-from node_graph.socket_spec import SocketSpec, merge_specs, SocketSpecMeta
+from node_graph.socket_spec import SocketSpec, merge_specs, SocketMeta
 from aiida_workgraph.socket_spec import from_aiida_process, namespace
 from aiida_workgraph.task import Task, TaskHandle
 from aiida import orm
@@ -127,7 +127,7 @@ def _build_shelljob_nodespec(
     in_spec = replace(in_spec, fields={**in_spec.fields, 'code': patched_code})
 
     # Add additional inputs
-    additions_in = namespace(command=Any, resolve_command=Annotated[bool, SocketSpecMeta(required=False)])
+    additions_in = namespace(command=Any, resolve_command=Annotated[bool, SocketMeta(required=False)])
     in_spec = merge_specs(in_spec, additions_in)
 
     # Ensure stdout/stderr outputs

src/aiida_workgraph/workgraph.py

@@ -331,9 +331,6 @@ def update(self) -> None:
     def pk(self) -> Optional[int]:
         return self.process.pk if self.process else None
 
-    def update_ctx(self, value: Dict[str, Any]) -> None:
-        self.ctx._set_socket_value(value, link_limit=100000)
-
     @classmethod
     def from_dict(cls, wgdata: Dict[str, Any]) -> 'WorkGraph':
         if 'tasks' in wgdata:
@@ -603,19 +600,6 @@ def to_widget_value(self) -> Dict[str, Any]:
         wgdata = workgraph_to_short_json(wgdata)
         return wgdata
 
-    def _repr_mimebundle_(self, *args, **kwargs):
-        # if ipywdigets > 8.0.0, use _repr_mimebundle_ instead of _ipython_display_
-        self.widget.value = self.to_widget_value()
-        if hasattr(self.widget, '_repr_mimebundle_'):
-            return self.widget._repr_mimebundle_(*args, **kwargs)
-        else:
-            return self.widget._ipython_display_(*args, **kwargs)
-
-    def to_html(self, output: str = None, **kwargs):
-        """Write a standalone html file to visualize the workgraph."""
-        self.widget.value = self.to_widget_value()
-        return self.widget.to_html(output=output, **kwargs)
-
     def generate_provenance_graph(self):
         """Generate the provenance graph of the workgraph process."""
         from aiida_workgraph.utils import generate_provenance_graph