aiida-pythonjob 0.2.1__tar.gz → 0.2.2__tar.gz

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiida-pythonjob
-Version: 0.2.1
+Version: 0.2.2
 Summary: Run Python functions on a remote computer.
 Project-URL: Source, https://github.com/aiidateam/aiida-pythonjob
 Author-email: Xing Wang <xingwang1991@gmail.com>

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/src/aiida_pythonjob/__init__.py

@@ -1,6 +1,6 @@
 """AiiDA plugin that run Python function on remote computers."""
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 from .calculations import PythonJob
 from .decorator import pyfunction

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/src/aiida_pythonjob/calculations/pyfunction.py

@@ -38,20 +38,18 @@ class PyFunction(Process):
 
     @property
     def func(self) -> t.Callable[..., t.Any]:
+        import cloudpickle
+
         if self._func is None:
-            self._func = self.inputs.function_data.pickled_function.value
+            self._func = cloudpickle.loads(self.inputs.function_data.pickled_function)
         return self._func
 
     @classmethod
     def define(cls, spec: ProcessSpec) -> None:  # type: ignore[override]
         """Define the process specification, including its inputs, outputs and known exit codes."""
         super().define(spec)
-        spec.input_namespace("function_data")
-        spec.input("function_data.name", valid_type=Str, serializer=to_aiida_type)
-        spec.input("function_data.source_code", valid_type=Str, serializer=to_aiida_type, required=False)
+        spec.input_namespace("function_data", dynamic=True, required=True)
         spec.input("function_data.outputs", valid_type=List, serializer=to_aiida_type, required=False)
-        spec.input("function_data.pickled_function", valid_type=Data, required=False)
-        spec.input("function_data.mode", valid_type=Str, serializer=to_aiida_type, required=False)
         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(
@@ -88,10 +86,10 @@ class PyFunction(Process):
     def get_function_name(self) -> str:
         """Return the name of the function to run."""
         if "name" in self.inputs.function_data:
-            name = self.inputs.function_data.name.value
+            name = self.inputs.function_data.name
         else:
             try:
-                name = self.inputs.function_data.pickled_function.value.__name__
+                name = self.func.__name__
             except AttributeError:
                 # If a user doesn't specify name, fallback to something generic
                 name = "anonymous_function"
@@ -202,8 +200,10 @@ class PyFunction(Process):
                 if exit_code.status != 0:
                     return exit_code
             if len(top_level_output_list) == 1:
-                # If output name in results, use it
-                if top_level_output_list[0]["name"] in results:
+                # User returned a single (nested) dict with AiiDA data nodes as values
+                if self.already_serialized(results):
+                    top_level_output_list = [{"name": key, "value": value} for key, value in results.items()]
+                elif top_level_output_list[0]["name"] in results:
                     top_level_output_list[0]["value"] = self.serialize_output(
                         results.pop(top_level_output_list[0]["name"]),
                         top_level_output_list[0],
@@ -228,8 +228,14 @@ class PyFunction(Process):
                 if len(results) > 0:
                     self.logger.warning(f"Found extra results that are not included in the output: {results.keys()}")
         elif len(top_level_output_list) == 1:
-            # Single top-level output, single result
-            top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
+            # Single top-level output
+            # There are two cases:
+            # 1. The output as a whole will be serialized as the single output
+            # 2. The output is a mapping with already AiiDA data nodes as values, no need to serialize
+            if self.already_serialized(results):
+                top_level_output_list[0]["value"] = results
+            else:
+                top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
         else:
             return self.exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
         # Store the outputs
@@ -238,6 +244,20 @@ class PyFunction(Process):
 
         return ExitCode()
 
+    def already_serialized(self, results):
+        """Check if the results are already serialized."""
+        import collections
+
+        if isinstance(results, Data):
+            return True
+        elif isinstance(results, collections.abc.Mapping):
+            for value in results.values():
+                if not self.already_serialized(value):
+                    return False
+            return True
+        else:
+            return False
+
     def find_output(self, name):
         """Find the output spec with the given name."""
         for output in self.output_list:

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/src/aiida_pythonjob/calculations/pythonjob.py

@@ -7,6 +7,7 @@ import typing as t
 
 from aiida.common.datastructures import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
+from aiida.common.lang import override
 from aiida.engine import CalcJob, CalcJobProcessSpec
 from aiida.orm import (
     Data,
@@ -37,17 +38,14 @@ class PythonJob(CalcJob):
     _DEFAULT_INPUT_FILE = "script.py"
     _DEFAULT_OUTPUT_FILE = "aiida.out"
     _DEFAULT_PARENT_FOLDER_NAME = "./parent_folder/"
+    _SOURCE_CODE_KEY = "source_code"
 
     @classmethod
     def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
         """Define the process specification, including its inputs, outputs and known exit codes."""
         super().define(spec)
-        spec.input_namespace("function_data")
-        spec.input("function_data.name", valid_type=Str, serializer=to_aiida_type)
-        spec.input("function_data.source_code", valid_type=Str, serializer=to_aiida_type, required=False)
+        spec.input_namespace("function_data", dynamic=True, required=True)
         spec.input("function_data.outputs", valid_type=List, serializer=to_aiida_type, required=False)
-        spec.input("function_data.pickled_function", valid_type=Data, required=False)
-        spec.input("function_data.mode", valid_type=Str, serializer=to_aiida_type, required=False)
         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(
@@ -175,13 +173,9 @@ class PythonJob(CalcJob):
     def get_function_name(self) -> str:
         """Return the name of the function to run."""
         if "name" in self.inputs.function_data:
-            name = self.inputs.function_data.name.value
+            name = self.inputs.function_data.name
         else:
-            try:
-                name = self.inputs.function_data.pickled_function.value.__name__
-            except AttributeError:
-                # If a user doesn't specify name, fallback to something generic
-                name = "anonymous_function"
+            name = "anonymous_function"
         return name
 
     def _build_process_label(self) -> str:
@@ -192,6 +186,13 @@ class PythonJob(CalcJob):
             name = self.get_function_name()
             return f"PythonJob<{name}>"
 
+    @override
+    def _setup_db_record(self) -> None:
+        """Set up the database record for the process."""
+        super()._setup_db_record()
+        if "source_code" in self.inputs.function_data:
+            self.node.base.attributes.set(self._SOURCE_CODE_KEY, self.inputs.function_data.source_code)
+
     def on_create(self) -> None:
         """Called when a Process is created."""
         super().on_create()
@@ -223,19 +224,13 @@ class PythonJob(CalcJob):
         else:
             parent_folder_name = self._DEFAULT_PARENT_FOLDER_NAME
 
-        function_data = self.inputs.function_data
-
         # Build the Python script
-        source_code = function_data.get("source_code")
-        if "pickled_function" in self.inputs.function_data:
-            pickled_function = self.inputs.function_data.pickled_function.get_serialized_value()
-        else:
-            pickled_function = None
-        # Generate script.py content
+        source_code = self.node.base.attributes.get(self._SOURCE_CODE_KEY, None)
+        pickled_function = self.inputs.function_data.pickled_function
         function_name = self.get_function_name()  # or some user-defined name
         script_content = generate_script_py(
             pickled_function=pickled_function,
-            source_code=source_code.value if source_code else None,
+            source_code=source_code,
             function_name=function_name,
         )
 

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/src/aiida_pythonjob/utils.py

@@ -54,8 +54,6 @@ def inspect_function(
     # the source code is not saved in the pickle file
     import cloudpickle
 
-    from aiida_pythonjob.data.pickled_data import PickledData
-
     if inspect_source:
         try:
             source_code = inspect.getsource(func)
@@ -70,10 +68,10 @@ def inspect_function(
     if register_pickle_by_value:
         module = importlib.import_module(func.__module__)
         cloudpickle.register_pickle_by_value(module)
-        pickled_function = PickledData(value=func)
+        pickled_function = cloudpickle.dumps(func)
         cloudpickle.unregister_pickle_by_value(module)
     else:
-        pickled_function = PickledData(value=func)
+        pickled_function = cloudpickle.dumps(func)
 
     return {"source_code": source_code, "mode": "use_pickled_function", "pickled_function": pickled_function}
 

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/tests/test_pyfunction.py

@@ -1,3 +1,4 @@
+from aiida import orm
 from aiida.engine import run_get_node
 from aiida_pythonjob import pyfunction
 
@@ -118,3 +119,15 @@ def test_override_outputs():
     assert result["add_multiply"]["add"]["order1"].value == 3
     assert result["add_multiply"]["add"]["order2"].value == 5
     assert result["add_multiply"]["multiply"].value == 2
+
+
+def test_aiida_node_as_inputs_outputs():
+    """Test function with AiiDA nodes as inputs and outputs."""
+
+    @pyfunction()
+    def add(x, y):
+        return {"sum": orm.Int(x + y), "diff": orm.Int(x - y)}
+
+    result, node = run_get_node(add, x=orm.Int(1), y=orm.Int(2))
+    assert set(result.keys()) == {"sum", "diff"}
+    assert result["sum"].value == 3

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/tests/test_pythonjob.py

@@ -310,3 +310,22 @@ def test_local_function(fixture_localhost):
     )
     result, node = run_get_node(PythonJob, **inputs)
     assert result["result"].value == 8
+
+
+@pytest.mark.usefixtures("started_daemon_client")
+def test_submit(fixture_localhost):
+    """Test decorator."""
+    from aiida.engine import submit
+
+    def add(x, y):
+        return x + y
+
+    inputs = prepare_pythonjob_inputs(
+        add,
+        function_inputs={"x": 1, "y": 2},
+        process_label="add",
+    )
+    node = submit(PythonJob, **inputs, wait=True)
+
+    assert node.outputs.result.value == 3
+    assert node.process_label == "add"

{aiida_pythonjob-0.2.1 → aiida_pythonjob-0.2.2}/tests/test_utils.py

@@ -12,19 +12,12 @@ def test_build_function_data():
     assert function_data["name"] == "build_function_data"
     assert "source_code" in function_data
     assert "pickled_function" in function_data
-    node = function_data["pickled_function"]
-    with node.base.repository.open(node.FILENAME, mode="rb") as f:
-        text = f.read()
-        assert b"cloudpickle" not in text
-
+    assert b"cloudpickle" not in function_data["pickled_function"]
     function_data = build_function_data(build_function_data, register_pickle_by_value=True)
     assert function_data["name"] == "build_function_data"
     assert "source_code" in function_data
     assert "pickled_function" in function_data
-    node = function_data["pickled_function"]
-    with node.base.repository.open(node.FILENAME, mode="rb") as f:
-        text = f.read()
-        assert b"cloudpickle" in text
+    assert b"cloudpickle" in function_data["pickled_function"]
 
     def local_function(x, y):
         return x + y