aiida-pythonjob 0.2.5__tar.gz → 0.3.0__tar.gz

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/.readthedocs.yml

@@ -12,6 +12,7 @@ build:
       - rabbitmq-server -detached
       - sleep 5
       - rabbitmq-diagnostics status
+      - pip list
       - verdi presto
       - verdi daemon start
       - verdi status

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiida-pythonjob
-Version: 0.2.5
+Version: 0.3.0
 Summary: Run Python functions on a remote computer.
 Project-URL: Source, https://github.com/aiidateam/aiida-pythonjob
 Author-email: Xing Wang <xingwang1991@gmail.com>
@@ -37,6 +37,7 @@ Requires-Python: >=3.9
 Requires-Dist: aiida-core<3,>=2.3
 Requires-Dist: ase
 Requires-Dist: cloudpickle
+Requires-Dist: node-graph==0.2.22
 Provides-Extra: dev
 Requires-Dist: hatch; extra == 'dev'
 Provides-Extra: docs

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/docs/environment.yml

@@ -3,5 +3,5 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - aiida-core
+  - aiida-core~=2.6.3
   - aiida-core.services

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/docs/gallery/autogen/pyfunction.py

@@ -6,14 +6,14 @@ PyFunction
 
 ######################################################################
 # Default outputs
-# --------------
+# -----------------
 #
 # The default output of the function is `result`. The `pyfunction` task
 # will store the result as one node in the database with the key `result`.
 #
 from aiida import load_profile
 from aiida.engine import run_get_node
-from aiida_pythonjob import pyfunction
+from aiida_pythonjob import pyfunction, spec
 
 load_profile()
 
@@ -35,7 +35,7 @@ print("result: ", result)
 #
 
 
-@pyfunction(outputs=[{"name": "sum"}, {"name": "diff"}])
+@pyfunction(outputs=spec.namespace(sum=any, diff=any))
 def add(x, y):
     return {"sum": x + y, "diff": x - y}
 
@@ -48,7 +48,7 @@ print("diff: ", result["diff"])
 
 ######################################################################
 # Namespace Output
-# --------------
+# -----------------
 #
 # The `pyfunction` allows users to define namespace outputs. A namespace output
 # is a dictionary with keys and values returned by a function. Each value in
@@ -70,7 +70,7 @@ from ase import Atoms  # noqa: E402
 from ase.build import bulk  # noqa: E402
 
 
-@pyfunction(outputs=[{"name": "scaled_structures", "identifier": "namespace"}])
+@pyfunction(outputs=spec.dynamic(Atoms))
 def generate_structures(structure: Atoms, factor_lst: list) -> dict:
     """Scale the structure by the given factor_lst."""
     scaled_structures = {}
@@ -78,12 +78,12 @@ def generate_structures(structure: Atoms, factor_lst: list) -> dict:
         atoms = structure.copy()
         atoms.set_cell(atoms.cell * factor_lst[i], scale_atoms=True)
         scaled_structures[f"s_{i}"] = atoms
-    return {"scaled_structures": scaled_structures}
+    return scaled_structures
 
 
 result, node = run_get_node(generate_structures, structure=bulk("Al"), factor_lst=[0.95, 1.0, 1.05])
 print("scaled_structures: ")
-for key, value in result["scaled_structures"].items():
+for key, value in result.items():
     print(key, value)
 
 
@@ -115,7 +115,7 @@ print("exit_message:", node.exit_message)
 
 ######################################################################
 # Define your data serializer and deserializer
-# --------------
+# ----------------------------------------------
 #
 # PythonJob search data serializer from the `aiida.data` entry point by the
 # module name and class name (e.g., `ase.atoms.Atoms`).

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/docs/gallery/autogen/pythonjob.py

@@ -63,14 +63,14 @@ PythonJob
 
 ######################################################################
 # Default outputs
-# --------------
+# ----------------
 #
 # The default output of the function is `result`. The `PythonJob` task
 # will store the result as one node in the database with the key `result`.
 #
 from aiida import load_profile
 from aiida.engine import run_get_node
-from aiida_pythonjob import PythonJob, prepare_pythonjob_inputs
+from aiida_pythonjob import PythonJob, prepare_pythonjob_inputs, spec
 
 load_profile()
 
@@ -91,7 +91,7 @@ print("result: ", result["result"])
 # Custom outputs
 # --------------
 # If the function return a dictionary with fixed number of keys, and you
-# want to store the values as separate outputs, you can specify the `output_ports` parameter.
+# want to store the values as separate outputs, you can specify the `outputs_spec` parameter.
 # For a dynamic number of outputs, you can use the namespace output, which is explained later.
 #
 
@@ -103,10 +103,7 @@ def add(x, y):
 inputs = prepare_pythonjob_inputs(
     add,
     function_inputs={"x": 1, "y": 2},
-    output_ports=[
-        {"name": "sum"},
-        {"name": "diff"},
-    ],
+    outputs_spec=spec.namespace(sum=any, diff=any),
 )
 result, node = run_get_node(PythonJob, **inputs)
 
@@ -117,7 +114,7 @@ print("diff: ", result["diff"])
 
 ######################################################################
 # Using parent folder
-# --------------
+# -----------------------
 # The parent_folder parameter allows a task to access the output files of
 # a parent task. This feature is particularly useful when you want to reuse
 # data generated by a previous computation in subsequent computations. In
@@ -142,7 +139,6 @@ def multiply(x, y):
 inputs1 = prepare_pythonjob_inputs(
     add,
     function_inputs={"x": 1, "y": 2},
-    output_ports=[{"name": "sum"}],
 )
 
 result1, node1 = run_get_node(PythonJob, inputs=inputs1)
@@ -150,7 +146,6 @@ result1, node1 = run_get_node(PythonJob, inputs=inputs1)
 inputs2 = prepare_pythonjob_inputs(
     multiply,
     function_inputs={"x": 1, "y": 2},
-    output_ports=[{"name": "product"}],
     parent_folder=result1["remote_folder"],
 )
 
@@ -160,7 +155,7 @@ print("result: ", result2)
 
 ######################################################################
 # Upload files or folders to the remote computer
-# --------------
+# -------------------------------------------------
 # The `upload_files` parameter allows users to upload files or folders to
 # the remote computer. The files will be uploaded to the working directory of the remote computer.
 #
@@ -202,7 +197,7 @@ print("result: ", result["result"])
 
 ######################################################################
 # Retrieve additional files from the remote computer
-# --------------
+# ----------------------------------------------------
 # Sometimes, one may want to retrieve additional files from the remote
 # computer after the job has finished. For example, one may want to retrieve
 # the output files generated by the `pw.x` calculation in Quantum ESPRESSO.
@@ -235,7 +230,7 @@ print("retrieved files: ", result["retrieved"].list_object_names())
 
 ######################################################################
 # Namespace Output
-# --------------
+# ------------------
 #
 # The `PythonJob` allows users to define namespace outputs. A namespace output
 # is a dictionary with keys and values returned by a function. Each value in
@@ -264,21 +259,55 @@ def generate_structures(structure: Atoms, factor_lst: list) -> dict:
         atoms = structure.copy()
         atoms.set_cell(atoms.cell * factor_lst[i], scale_atoms=True)
         scaled_structures[f"s_{i}"] = atoms
-    return {"scaled_structures": scaled_structures}
+    return scaled_structures
 
 
 inputs = prepare_pythonjob_inputs(
     generate_structures,
     function_inputs={"structure": bulk("Al"), "factor_lst": [0.95, 1.0, 1.05]},
-    output_ports=[{"name": "scaled_structures", "identifier": "namespace"}],
+    outputs_spec=spec.dynamic(Atoms),
 )
 
 result, node = run_get_node(PythonJob, inputs=inputs)
 print("scaled_structures: ")
-for key, value in result["scaled_structures"].items():
+for key, value in result.items():
     print(key, value)
 
 
+######################################################################
+# --------------------------
+# Nested Namespace
+# --------------------------
+#
+# One can also define nested namespace outputs by specifying the "ports" parameter.
+
+
+def generate_structures(structure: Atoms, factor_lst: list) -> dict:
+    """Scale the structure by the given factor_lst."""
+    scaled_structures = {}
+    volumes = {}
+    for i in range(len(factor_lst)):
+        atoms = structure.copy()
+        atoms.set_cell(atoms.cell * factor_lst[i], scale_atoms=True)
+        scaled_structures[f"s_{i}"] = atoms
+        volumes[f"v_{i}"] = atoms.get_volume()
+    return {
+        "scaled_structures": scaled_structures,
+        "volume": volumes,
+    }
+
+
+inputs = prepare_pythonjob_inputs(
+    generate_structures,
+    function_inputs={"structure": bulk("Al"), "factor_lst": [0.95, 1.0, 1.05]},
+    outputs_spec=spec.namespace(scaled_structures=spec.dynamic(Atoms), volume=spec.dynamic(float)),
+)
+
+result, node = run_get_node(PythonJob, inputs=inputs)
+print("result: ", result["scaled_structures"])
+print("volumes: ", result["volume"])
+
+
 ######################################################################
 # What if my calculation fails?
 # --------------------------------
@@ -375,7 +404,7 @@ print("exit_message:", node.exit_message)
 
 ######################################################################
 # Define your data serializer and deserializer
-# --------------
+# ----------------------------------------------
 #
 # PythonJob search data serializer from the `aiida.data` entry point by the
 # module name and class name (e.g., `ase.atoms.Atoms`).

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/docs/requirements.txt

@@ -1,4 +1,6 @@
-sphinx_rtd_theme==1.2.2
+sphinx_rtd_theme~=3.0
+click~=8.1.0
+pydantic~=2.11.0
 sphinx-gallery
 nbsphinx==0.9.2
 ipython

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "aiida-core>=2.3,<3",
     "ase",
     "cloudpickle",
+    "node-graph==0.2.22",
 ]
 
 [project.optional-dependencies]
@@ -163,3 +164,6 @@ features = ["docs"]
 build = [
   "make -C docs"
 ]
+
+[tool.hatch.metadata]
+allow-direct-references = true

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/src/aiida_pythonjob/__init__.py

@@ -1,6 +1,8 @@
 """AiiDA plugin that run Python function on remote computers."""
 
-__version__ = "0.2.5"
+__version__ = "0.3.0"
+
+from node_graph import spec
 
 from .calculations import PythonJob
 from .decorator import pyfunction
@@ -13,4 +15,5 @@ __all__ = (
     "PickledData",
     "prepare_pythonjob_inputs",
     "PythonJobParser",
+    "spec",
 )

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/src/aiida_pythonjob/calculations/pyfunction.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import traceback
 import typing as t
 
+import cloudpickle
+import plumpy
 from aiida.common.lang import override
 from aiida.engine import Process, ProcessSpec
 from aiida.engine.processes.exit_code import ExitCode
@@ -30,10 +32,18 @@ class PyFunction(Process):
         super().__init__(enable_persistence=False, *args, **kwargs)  # type: ignore[misc]
         self._func = None
 
+    @override
+    def load_instance_state(
+        self, saved_state: t.MutableMapping[str, t.Any], load_context: plumpy.persistence.LoadSaveContext
+    ) -> None:
+        """Load the instance state from the saved state."""
+
+        super().load_instance_state(saved_state, load_context)
+        # Restore the function from the pickled data
+        self._func = cloudpickle.loads(self.inputs.function_data.pickled_function)
+
     @property
     def func(self) -> t.Callable[..., t.Any]:
-        import cloudpickle
-
         if self._func is None:
             self._func = cloudpickle.loads(self.inputs.function_data.pickled_function)
         return self._func
@@ -189,7 +199,8 @@ class PyFunction(Process):
         if exit_code:
             return exit_code
         # Store the outputs
-        for output in self.output_ports["ports"]:
-            self.out(output["name"], output["value"])
+        for name, port in self.output_ports["ports"].items():
+            if "value" in port:
+                self.out(name, port["value"])
 
         return ExitCode()

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/src/aiida_pythonjob/decorator.py

@@ -60,8 +60,10 @@ def pyfunction(
             manager = get_manager()
             runner = manager.get_runner()
             # # Remove all the known inputs from the kwargs
-            output_ports = kwargs.pop("output_ports", None) or outputs
-            input_ports = kwargs.pop("input_ports", None) or inputs
+            outputs_spec = kwargs.pop("outputs_spec", None) or outputs
+            inputs_spec = kwargs.pop("inputs_spec", None) or inputs
+            input_ports = kwargs.pop("input_ports", None)
+            output_ports = kwargs.pop("output_ports", None)
             metadata = kwargs.pop("metadata", None)
             function_data = kwargs.pop("function_data", None)
             deserializers = kwargs.pop("deserializers", None)
@@ -73,6 +75,8 @@ def pyfunction(
             process_inputs = prepare_pyfunction_inputs(
                 function=function,
                 function_inputs=function_inputs,
+                inputs_spec=inputs_spec,
+                outputs_spec=outputs_spec,
                 input_ports=input_ports,
                 output_ports=output_ports,
                 metadata=metadata,

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/src/aiida_pythonjob/launch.py

@@ -2,23 +2,22 @@ from __future__ import annotations
 
 import inspect
 import os
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, Optional, Union
 
 from aiida import orm
+from node_graph.nodes.utils import generate_input_sockets, generate_output_sockets
 
-from .utils import (
-    build_function_data,
-    build_input_port_definitions,
-    format_input_output_ports,
-    get_or_create_code,
-)
+from .ports_adapter import inputs_sockets_to_ports, outputs_sockets_to_ports
+from .utils import build_function_data, get_or_create_code, serialize_ports
 
 
 def prepare_pythonjob_inputs(
     function: Optional[Callable[..., Any]] = None,
     function_inputs: Optional[Dict[str, Any]] = None,
-    input_ports: Optional[List[str | dict]] = None,
-    output_ports: Optional[List[str | dict]] = None,
+    inputs_spec: Optional[type] = None,
+    outputs_spec: Optional[type] = None,
+    output_ports: Optional[Dict[str, Any]] = None,
+    input_ports: Optional[Dict[str, Any]] = None,
     code: Optional[orm.AbstractCode] = None,
     command_info: Optional[Dict[str, str]] = None,
     computer: Union[str, orm.Computer] = "localhost",
@@ -32,7 +31,6 @@ def prepare_pythonjob_inputs(
     **kwargs: Any,
 ) -> Dict[str, Any]:
     """Prepare the inputs for PythonJob"""
-    from .utils import serialize_ports
 
     if function is None and function_data is None:
         raise ValueError("Either function or function_data must be provided")
@@ -62,14 +60,18 @@ def prepare_pythonjob_inputs(
     if code is None:
         command_info = command_info or {}
         code = get_or_create_code(computer=computer, **command_info)
-    output_ports = output_ports or [{"name": "result"}]
-    output_ports = {"name": "outputs", "identifier": "namespace", "ports": output_ports or [{"name": "result"}]}
-    input_ports = {"name": "inputs", "identifier": "namespace", "ports": input_ports or []}
-    input_ports = format_input_output_ports(input_ports)
-    input_ports = build_input_port_definitions(func=function, input_ports=input_ports)
-    function_data["output_ports"] = format_input_output_ports(output_ports)
+    # outputs
+    if not output_ports:
+        node_outputs = generate_output_sockets(function or (lambda **_: None), outputs=outputs_spec)
+        output_ports = outputs_sockets_to_ports(node_outputs)
+    # inputs
+    if not input_ports:
+        node_inputs = generate_input_sockets(function or (lambda **_: None), inputs=inputs_spec)
+        input_ports = inputs_sockets_to_ports(node_inputs)
+
+    function_data["output_ports"] = output_ports
     function_data["input_ports"] = input_ports
-    # serialize the kwargs into AiiDA Data
+    # serialize kwargs against the (nested) input schema
     function_inputs = function_inputs or {}
     function_inputs = serialize_ports(python_data=function_inputs, port_schema=input_ports, serializers=serializers)
     # replace "." with "__dot__" in the keys of a dictionary
@@ -112,8 +114,10 @@ def create_inputs(func, *args: Any, **kwargs: Any) -> dict[str, Any]:
 def prepare_pyfunction_inputs(
     function: Optional[Callable[..., Any]] = None,
     function_inputs: Optional[Dict[str, Any]] = None,
-    input_ports: Optional[List[str | dict]] = None,
-    output_ports: Optional[List[str | dict]] = None,
+    inputs_spec: Optional[type] = None,
+    outputs_spec: Optional[type] = None,
+    output_ports: Optional[Dict[str, Any]] = None,
+    input_ports: Optional[Dict[str, Any]] = None,
     metadata: Optional[Dict[str, Any]] = None,
     process_label: Optional[str] = None,
     function_data: dict | None = None,
@@ -125,8 +129,6 @@ def prepare_pyfunction_inputs(
     """Prepare the inputs for PythonJob"""
     import types
 
-    from .utils import serialize_ports
-
     if function is None and function_data is None:
         raise ValueError("Either function or function_data must be provided")
     if function is not None and function_data is not None:
@@ -139,11 +141,16 @@ def prepare_pyfunction_inputs(
             raise NotImplementedError("Built-in functions are not supported yet")
         else:
             raise ValueError("Invalid function type")
-    output_ports = {"name": "outputs", "identifier": "namespace", "ports": output_ports or [{"name": "result"}]}
-    input_ports = {"name": "inputs", "identifier": "namespace", "ports": input_ports or []}
-    input_ports = format_input_output_ports(input_ports)
-    input_ports = build_input_port_definitions(func=function, input_ports=input_ports)
-    function_data["output_ports"] = format_input_output_ports(output_ports)
+    # outputs
+    if not output_ports:
+        node_outputs = generate_output_sockets(function or (lambda **_: None), outputs=outputs_spec)
+        output_ports = outputs_sockets_to_ports(node_outputs)
+    # inputs
+    if not input_ports:
+        node_inputs = generate_input_sockets(function or (lambda **_: None), inputs=inputs_spec)
+        input_ports = inputs_sockets_to_ports(node_inputs)
+
+    function_data["output_ports"] = output_ports
     function_data["input_ports"] = input_ports
     # serialize the kwargs into AiiDA Data
     function_inputs = function_inputs or {}

{aiida_pythonjob-0.2.5 → aiida_pythonjob-0.3.0}/src/aiida_pythonjob/parsers/pythonjob.py

@@ -72,11 +72,14 @@ class PythonJobParser(Parser):
                     return exit_code
 
                 # Store the outputs
-                for output in self.output_ports["ports"]:
-                    self.out(output["name"], output["value"])
+                for name, port in self.output_ports["ports"].items():
+                    if "value" in port:
+                        self.out(name, port["value"])
 
         except OSError:
             return self.exit_codes.ERROR_READING_OUTPUT_FILE
         except ValueError as exception:
-            self.logger.error(exception)
+            self.logger.error(
+                f"An error occurred when attempting to parse the output of the calculation: ValueError: {exception!s}"
+            )
             return self.exit_codes.ERROR_INVALID_OUTPUT

aiida_pythonjob-0.3.0/src/aiida_pythonjob/ports_adapter.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+
+def _is_namespace_identifier(identifier: str, *, ns_id: str = "node_graph.namespace") -> bool:
+    return identifier == ns_id or identifier.endswith(".namespace")
+
+
+def _socket_meta_required(sock: Dict[str, Any]) -> Optional[bool]:
+    meta = sock.get("metadata", {}) or {}
+    if "required" in meta and meta["required"] is not None:
+        return bool(meta["required"])
+    if sock.get("property", {}).get("default", None) is not None:
+        return False
+    return None
+
+
+def _socket_meta_help(sock: Dict[str, Any]) -> Optional[str]:
+    return (sock.get("metadata") or {}).get("help")
+
+
+def _socket_to_port_object(sock: Dict[str, Any]) -> Dict[str, Any]:
+    ident = sock.get("identifier", "")
+    meta = sock.get("metadata", {}) or {}
+
+    if _is_namespace_identifier(ident):
+        ports: Dict[str, Dict[str, Any]] = {}
+        for name, child in (sock.get("sockets") or {}).items():
+            ports[name] = _socket_to_port_object(child)
+
+        obj: Dict[str, Any] = {"identifier": "NAMESPACE", "ports": ports}
+
+        if meta.get("dynamic"):
+            obj["dynamic"] = True
+            # Prefer full nested item namespace when present
+            if isinstance(meta.get("item"), dict):
+                obj["item"] = _socket_to_port_object(meta["item"])
+            else:
+                # fallback: identifier only (treat namespace id as empty namespace)
+                item_ident = meta.get("item_identifier", "node_graph.any")
+                obj["item"] = (
+                    {"identifier": "NAMESPACE", "ports": {}}
+                    if _is_namespace_identifier(item_ident)
+                    else {"identifier": "ANY"}
+                )
+
+        h = _socket_meta_help(sock)
+        if h:
+            obj["help"] = h
+        req = _socket_meta_required(sock)
+        if req is not None:
+            obj["required"] = req
+        return obj
+
+    # Leaf
+    obj = {"identifier": "ANY"}
+    h = _socket_meta_help(sock)
+    if h:
+        obj["help"] = h
+    req = _socket_meta_required(sock)
+    if req is not None:
+        obj["required"] = req
+    return obj
+
+
+def inputs_sockets_to_ports(node_inputs: Dict[str, Any]) -> Dict[str, Any]:
+    ports_map: Dict[str, Dict[str, Any]] = {}
+    for name, sock in (node_inputs.get("sockets") or {}).items():
+        ports_map[name] = _socket_to_port_object(sock)
+
+    schema: Dict[str, Any] = {"name": "inputs", "identifier": "NAMESPACE", "ports": ports_map}
+    meta = node_inputs.get("metadata", {}) or {}
+    if meta.get("dynamic"):
+        schema["dynamic"] = True
+        if isinstance(meta.get("item"), dict):
+            schema["item"] = _socket_to_port_object(meta["item"])
+        else:
+            item_ident = meta.get("item_identifier", "node_graph.any")
+            schema["item"] = (
+                {"identifier": "NAMESPACE", "ports": {}}
+                if _is_namespace_identifier(item_ident)
+                else {"identifier": "ANY"}
+            )
+
+    return schema
+
+
+def outputs_sockets_to_ports(node_outputs: Dict[str, Any]) -> Dict[str, Any]:
+    meta = node_outputs.get("metadata", {}) or {}
+    sockets = node_outputs.get("sockets") or {}
+
+    ports_map = {name: _socket_to_port_object(sock) for name, sock in sockets.items()}
+    schema: Dict[str, Any] = {"name": "outputs", "identifier": "NAMESPACE", "ports": ports_map}
+    if meta.get("dynamic"):
+        schema["dynamic"] = True
+        if isinstance(meta.get("item"), dict):
+            schema["item"] = _socket_to_port_object(meta["item"])
+        else:
+            item_ident = meta.get("item_identifier", "node_graph.any")
+            schema["item"] = (
+                {"identifier": "NAMESPACE", "ports": {}}
+                if _is_namespace_identifier(item_ident)
+                else {"identifier": "ANY"}
+            )
+    return schema