Do not store pickled function data (#24)

Do not save the pickled_function data anymore; only save the `source_code` as an attribute.

Author: superstar54

src/aiida_pythonjob/calculations/pyfunction.py

@@ -38,20 +38,18 @@ def __init__(self, *args, **kwargs) -> None:
 
     @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 @@ def define(cls, spec: ProcessSpec) -> None:  # type: ignore[override]
     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"

src/aiida_pythonjob/calculations/pythonjob.py

@@ -7,6 +7,7 @@
 
 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 @@ def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
     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 @@ def _build_process_label(self) -> str:
             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 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         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,
         )
 

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}
 

tests/test_pythonjob.py

@@ -310,3 +310,22 @@ def add(x, y):
     )
     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"

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