aiida-pythonjob 0.3.3__tar.gz → 0.3.5__tar.gz

{aiida_pythonjob-0.3.3/src/aiida_pythonjob.egg-info → aiida_pythonjob-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiida-pythonjob
-Version: 0.3.3
+Version: 0.3.5
 Summary: Run Python functions on a remote computer.
 Author-email: Xing Wang <xingwang1991@gmail.com>
 License: MIT License
@@ -61,7 +61,7 @@ Dynamic: license-file
 
 # AiiDA-PythonJob
 [![PyPI version](https://badge.fury.io/py/aiida-pythonjob.svg)](https://badge.fury.io/py/aiida-pythonjob)
-[![Unit test](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci.yml/badge.svg)](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci.yml)
+[![Unit test](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci-tests.yml/badge.svg)](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci-tests.yml)
 [![codecov](https://codecov.io/gh/aiidateam/aiida-pythonjob/branch/main/graph/badge.svg)](https://codecov.io/gh/aiidateam/aiida-pythonjob)
 [![Docs status](https://readthedocs.org/projects/aiida-pythonjob/badge)](http://aiida-pythonjob.readthedocs.io/)
 

{aiida_pythonjob-0.3.3 → aiida_pythonjob-0.3.5}/README.md

@@ -1,6 +1,6 @@
 # AiiDA-PythonJob
 [![PyPI version](https://badge.fury.io/py/aiida-pythonjob.svg)](https://badge.fury.io/py/aiida-pythonjob)
-[![Unit test](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci.yml/badge.svg)](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci.yml)
+[![Unit test](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci-tests.yml/badge.svg)](https://github.com/aiidateam/aiida-pythonjob/actions/workflows/ci-tests.yml)
 [![codecov](https://codecov.io/gh/aiidateam/aiida-pythonjob/branch/main/graph/badge.svg)](https://codecov.io/gh/aiidateam/aiida-pythonjob)
 [![Docs status](https://readthedocs.org/projects/aiida-pythonjob/badge)](http://aiida-pythonjob.readthedocs.io/)
 

{aiida_pythonjob-0.3.3 → aiida_pythonjob-0.3.5}/pyproject.toml

@@ -60,13 +60,13 @@ Source = "https://github.com/aiidateam/aiida-pythonjob"
 "pythonjob.builtins.float" = "aiida.orm.nodes.data.float:Float"
 "pythonjob.builtins.str" = "aiida.orm.nodes.data.str:Str"
 "pythonjob.builtins.bool" = "aiida.orm.nodes.data.bool:Bool"
-"pythonjob.builtins.list" = "aiida_pythonjob.data.data_wrapper:List"
-"pythonjob.builtins.dict" = "aiida_pythonjob.data.data_wrapper:Dict"
+"pythonjob.builtins.list" = "aiida.orm.nodes.data.list:List"
+"pythonjob.builtins.dict" = "aiida.orm.nodes.data.dict:Dict"
 "pythonjob.numpy.float32" = "aiida.orm.nodes.data.float:Float"
 "pythonjob.numpy.float64" = "aiida.orm.nodes.data.float:Float"
 "pythonjob.numpy.int64" = "aiida.orm.nodes.data.int:Int"
 "pythonjob.numpy.bool_" = "aiida.orm.nodes.data.bool:Bool"
-"pythonjob.numpy.ndarray" = "aiida_pythonjob.data.data_wrapper:ArrayData"
+"pythonjob.numpy.ndarray" = "aiida.orm.nodes.data.array.array.ArrayData"
 
 [project.entry-points."aiida.calculations"]
 "pythonjob.pythonjob" = "aiida_pythonjob.calculations.pythonjob:PythonJob"
@@ -98,12 +98,15 @@ filterwarnings = [
 source = ["src/aiida_pythonjob"]
 
 [tool.ruff]
-line-length = 120
+line-length = 121
 
 [tool.ruff.lint]
 ignore = [
   "F403",
   "F405",
+  "F821",
+  "F722",
+  "E203",
   "PLR0911",
   "PLR0912",
   "PLR0913",

{aiida_pythonjob-0.3.3 → aiida_pythonjob-0.3.5}/src/aiida_pythonjob/__init__.py

@@ -1,20 +1,21 @@
 """AiiDA plugin that run Python function on remote computers."""
 
-__version__ = "0.3.3"
+__version__ = "0.3.5"
 
 from node_graph import socket_spec as spec
 
 from .calculations import PyFunction, PythonJob
 from .decorator import pyfunction
-from .launch import prepare_pythonjob_inputs
+from .launch import prepare_pyfunction_inputs, prepare_pythonjob_inputs
 from .parsers import PythonJobParser
 
 __all__ = (
-    "PythonJob",
-    "PyFunction",
-    "pyfunction",
     "PickledData",
-    "prepare_pythonjob_inputs",
+    "PyFunction",
+    "PythonJob",
     "PythonJobParser",
+    "prepare_pyfunction_inputs",
+    "prepare_pythonjob_inputs",
+    "pyfunction",
     "spec",
 )

{aiida_pythonjob-0.3.3 → aiida_pythonjob-0.3.5}/src/aiida_pythonjob/calculations/__init__.py

@@ -1,4 +1,4 @@
 from .pyfunction import PyFunction
 from .pythonjob import PythonJob
 
-__all__ = ("PythonJob", "PyFunction")
+__all__ = ("PyFunction", "PythonJob")

aiida_pythonjob-0.3.5/src/aiida_pythonjob/calculations/common.py

@@ -0,0 +1,121 @@
+"""Common helpers and mixins shared by PyFunction and PythonJob."""
+
+from __future__ import annotations
+
+from typing import Any, MutableMapping
+
+from aiida.orm import Data, Str, to_aiida_type
+
+from aiida_pythonjob.data.deserializer import deserialize_to_raw_python_data
+
+# Attribute keys stored on ProcessNode.base.attributes
+ATTR_OUTPUTS_SPEC = "outputs_spec"
+ATTR_USE_PICKLE = "use_pickle"
+ATTR_SERIALIZERS = "serializers"
+ATTR_DESERIALIZERS = "deserializers"
+
+
+def add_common_function_io(spec) -> None:
+    """Attach inputs common to both in-process and remote Python execution.
+
+    Works with both :class:`~aiida.engine.ProcessSpec` and
+    :class:`~aiida.engine.CalcJobProcessSpec`.
+    """
+    spec.input_namespace("function_data", dynamic=True, required=True)
+    spec.input(
+        "metadata.outputs_spec",
+        valid_type=dict,
+        required=False,
+        help="Specification for the outputs.",
+    )
+    spec.input(
+        "metadata.use_pickle",
+        valid_type=bool,
+        required=False,
+        help="Allow pickling of function inputs and outputs.",
+    )
+    spec.input("process_label", valid_type=Str, serializer=to_aiida_type, required=False)
+
+    spec.input_namespace("function_inputs", valid_type=Data, required=False)
+
+    spec.input(
+        "metadata.deserializers",
+        valid_type=dict,
+        required=False,
+        help="Deserializers to convert input AiiDA nodes to raw Python data.",
+    )
+    spec.input(
+        "metadata.serializers",
+        valid_type=dict,
+        required=False,
+        help="Serializers to convert raw Python data to AiiDA nodes.",
+    )
+    spec.exit_code(
+        320,
+        "ERROR_DESERIALIZE_INPUTS_FAILED",
+        invalidates_cache=True,
+        message="Failed to unpickle inputs.\n{exception}\n{traceback}",
+    )
+    spec.exit_code(
+        321,
+        "ERROR_INVALID_OUTPUT",
+        invalidates_cache=True,
+        message="The output file contains invalid output.",
+    )
+    spec.exit_code(
+        322,
+        "ERROR_RESULT_OUTPUT_MISMATCH",
+        invalidates_cache=True,
+        message="The number of results does not match the number of outputs.",
+    )
+
+    spec.inputs.validator = validate_function_inputs
+
+
+def validate_function_inputs(inputs: MutableMapping[str, Any], _):
+    """Validate that ``function_inputs`` can be deserialized.
+
+    Uses ``metadata.deserializers`` if provided. Raises if invalid.
+    """
+    deserializers = inputs.get("metadata", {}).get("deserializers", {})
+    function_inputs = inputs.get("function_inputs", {})
+    # This should raise if any datum cannot be deserialized.
+    deserialize_to_raw_python_data(function_inputs, deserializers=deserializers, dry_run=True)
+
+
+class FunctionProcessMixin:
+    """Mixin providing common metadata handling and labeling logic.
+
+    Place this mixin **before** :class:`~aiida.engine.Process`/``CalcJob`` in the MRO.
+    """
+
+    label_template: str = "{name}"
+    default_name: str = "anonymous_function"
+
+    def _extract_declared_name(self) -> str | None:  # pragma: no cover - trivial
+        """Try to read a user-declared function name from inputs.
+
+        Subclasses may extend this (e.g. by inspecting a pickled function).
+        """
+        try:
+            if "name" in self.inputs.function_data:
+                return self.inputs.function_data.name
+        except Exception:
+            pass
+        return None
+
+    def get_function_name(self) -> str:  # used by both PyFunction and PythonJob
+        return self._extract_declared_name() or self.default_name
+
+    def _build_process_label(self) -> str:  # called by AiiDA engine
+        if "process_label" in self.inputs:
+            return self.inputs.process_label.value
+        return self.label_template.format(name=self.get_function_name())
+
+    def _setup_metadata(self, metadata: dict) -> None:  # type: ignore[override]
+        """Store common metadata on the ProcessNode and forward the rest."""
+        self.node.base.attributes.set(ATTR_OUTPUTS_SPEC, metadata.pop("outputs_spec", {}))
+        self.node.base.attributes.set(ATTR_USE_PICKLE, metadata.pop("use_pickle", False))
+        self.node.base.attributes.set(ATTR_SERIALIZERS, metadata.pop("serializers", {}))
+        self.node.base.attributes.set(ATTR_DESERIALIZERS, metadata.pop("deserializers", {}))
+        super()._setup_metadata(metadata)

aiida_pythonjob-0.3.5/src/aiida_pythonjob/calculations/pyfunction.py

@@ -0,0 +1,133 @@
+"""Proess to run a Python function locally"""
+
+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
+from aiida.orm import CalcFunctionNode
+from node_graph.socket_spec import SocketSpec
+
+from aiida_pythonjob.calculations.common import (
+    ATTR_DESERIALIZERS,
+    ATTR_OUTPUTS_SPEC,
+    ATTR_SERIALIZERS,
+    ATTR_USE_PICKLE,
+    FunctionProcessMixin,
+    add_common_function_io,
+)
+from aiida_pythonjob.data.deserializer import deserialize_to_raw_python_data
+from aiida_pythonjob.parsers.utils import parse_outputs
+
+__all__ = ("PyFunction",)
+
+
+class PyFunction(FunctionProcessMixin, Process):
+    """Run a Python function in-process, using :class:`SocketSpec` for I/O."""
+
+    _node_class = CalcFunctionNode
+    label_template = "{name}"
+    default_name = "anonymous_function"
+
+    def __init__(self, *args, **kwargs) -> None:
+        if kwargs.get("enable_persistence", False):
+            raise RuntimeError("Cannot persist a function 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 (restore pickled function)."""
+        super().load_instance_state(saved_state, load_context)
+        self._func = cloudpickle.loads(self.inputs.function_data.pickled_function)
+
+    @property
+    def func(self) -> t.Callable[..., t.Any]:
+        if self._func is None:
+            self._func = cloudpickle.loads(self.inputs.function_data.pickled_function)
+        return self._func
+
+    def _extract_declared_name(self) -> str | None:
+        name = super()._extract_declared_name()
+        if name:
+            return name
+        try:
+            return self.func.__name__
+        except Exception:
+            return None
+
+    @classmethod
+    def define(cls, spec: ProcessSpec) -> None:  # type: ignore[override]
+        """Define inputs/outputs and exit codes."""
+        super().define(spec)
+        add_common_function_io(spec)
+        spec.inputs.dynamic = True
+        spec.outputs.dynamic = True
+        spec.exit_code(
+            323,
+            "ERROR_FUNCTION_EXECUTION_FAILED",
+            invalidates_cache=True,
+            message="Function execution failed.\n{exception}\n{traceback}",
+        )
+
+    @override
+    def _setup_db_record(self) -> None:
+        super()._setup_db_record()
+        self.node.store_source_info(self.func)
+
+    def execute(self) -> dict[str, t.Any] | None:
+        """Mirror calcfunction behavior: unwrap single-output dicts to a bare value."""
+        result = super().execute()
+        if result and len(result) == 1 and self.SINGLE_OUTPUT_LINKNAME in result:
+            return result[self.SINGLE_OUTPUT_LINKNAME]
+        return result
+
+    @override
+    def run(self) -> ExitCode | None:
+        # Respect caching semantics (from aiida-core calcfunction implementation)
+        if self.node.exit_status is not None:
+            return ExitCode(self.node.exit_status, self.node.exit_message)
+
+        # Deserialize inputs
+        try:
+            inputs = dict(self.inputs.function_inputs or {})
+            deserializers = self.node.base.attributes.get(ATTR_DESERIALIZERS, {})
+            inputs = deserialize_to_raw_python_data(inputs, deserializers=deserializers)
+        except Exception as exception:
+            return self.exit_codes.ERROR_DESERIALIZE_INPUTS_FAILED.format(
+                exception=str(exception), traceback=traceback.format_exc()
+            )
+
+        # Execute function
+        try:
+            results = self.func(**inputs)
+        except Exception as exception:
+            return self.exit_codes.ERROR_FUNCTION_EXECUTION_FAILED.format(
+                exception=str(exception), traceback=traceback.format_exc()
+            )
+
+        # Parse & attach outputs
+        outputs_spec = SocketSpec.from_dict(self.node.base.attributes.get(ATTR_OUTPUTS_SPEC) or {})
+        use_pickle = self.node.base.attributes.get(ATTR_USE_PICKLE, False)
+        serializers = self.node.base.attributes.get(ATTR_SERIALIZERS, {})
+        outputs, exit_code = parse_outputs(
+            results,
+            output_spec=outputs_spec,
+            exit_codes=self.exit_codes,
+            logger=self.logger,
+            serializers=serializers,
+            use_pickle=use_pickle,
+        )
+        if exit_code:
+            return exit_code
+
+        for name, value in (outputs or {}).items():
+            self.out(name, value)
+        return ExitCode()