add register_pickle_by_value

Author: superstar54

docs/gallery/autogen/how_to.py

@@ -347,6 +347,33 @@ def add(x, y):
 print("exit_status:", node.exit_status)
 print("exit_message:", node.exit_message)
 
+######################################################################
+# Using `register_pickle_by_value`
+# --------------------------------
+#
+# If the function is defined inside an external module that is **not installed** on
+# the remote computer, this can cause import errors during execution.
+#
+# **Solution:**
+# By enabling `register_pickle_by_value=True`, the function is serialized **by value**
+# instead of being referenced by its module path. This embeds the function unpickled
+# even if the original module is unavailable on the remote computer.
+#
+# **Example:**
+#
+# .. code-block:: python
+#
+#     inputs = prepare_pythonjob_inputs(
+#         my_function,
+#         function_inputs={"x": 1, "y": 2},
+#         computer="localhost",
+#         register_pickle_by_value=True,  # Ensures function is embedded
+#     )
+#
+# **Important Considerations:**: If the function **contains import statements**,
+# the imported modules **must still be installed** on the remote computer.
+#
+
 
 ######################################################################
 # Define your data serializer and deserializer

pyproject.toml

@@ -28,6 +28,9 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+dev = [
+  "hatch",
+]
 pre-commit = [
   'pre-commit~=3.5',
 ]

src/aiida_pythonjob/launch.py

@@ -22,6 +22,7 @@ def prepare_pythonjob_inputs(
     function_data: dict | None = None,
     deserializers: dict | None = None,
     serializers: dict | None = None,
+    register_pickle_by_value: bool = False,
     **kwargs: Any,
 ) -> Dict[str, Any]:
     """Prepare the inputs for PythonJob"""
@@ -33,7 +34,7 @@ def prepare_pythonjob_inputs(
         raise ValueError("Only one of function or function_data should be provided")
     # if function is a function, inspect it and get the source code
     if function is not None and inspect.isfunction(function):
-        function_data = build_function_data(function)
+        function_data = build_function_data(function, register_pickle_by_value=register_pickle_by_value)
     new_upload_files = {}
     # change the string in the upload files to SingleFileData, or FolderData
     for key, source in upload_files.items():

src/aiida_pythonjob/utils.py

@@ -1,3 +1,4 @@
+import importlib
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, _SpecialForm, get_type_hints
 
@@ -6,8 +7,6 @@
 
 
 def import_from_path(path: str) -> Any:
-    import importlib
-
     module_name, object_name = path.rsplit(".", 1)
     module = importlib.import_module(module_name)
     try:
@@ -47,24 +46,39 @@ def add_imports(type_hint):
     return imports
 
 
-def inspect_function(func: Callable) -> Dict[str, Any]:
+def inspect_function(
+    func: Callable, inspect_source: bool = False, register_pickle_by_value: bool = False
+) -> Dict[str, Any]:
     """Serialize a function for storage or transmission."""
     # we need save the source code explicitly, because in the case of jupyter notebook,
     # the source code is not saved in the pickle file
+    import cloudpickle
+
     from aiida_pythonjob.data.pickled_data import PickledData
 
-    try:
-        source_code = inspect.getsource(func)
-        # Split the source into lines for processing
-        source_code_lines = source_code.split("\n")
-        source_code = "\n".join(source_code_lines)
-    except OSError:
-        source_code = "Failed to retrieve source code."
+    if inspect_source:
+        try:
+            source_code = inspect.getsource(func)
+            # Split the source into lines for processing
+            source_code_lines = source_code.split("\n")
+            source_code = "\n".join(source_code_lines)
+        except OSError:
+            source_code = "Failed to retrieve source code."
+    else:
+        source_code = ""
+
+    if register_pickle_by_value:
+        module = importlib.import_module(func.__module__)
+        cloudpickle.register_pickle_by_value(module)
+        pickled_function = PickledData(value=func)
+        cloudpickle.unregister_pickle_by_value(module)
+    else:
+        pickled_function = PickledData(value=func)
 
-    return {"source_code": source_code, "mode": "use_pickled_function", "pickled_function": PickledData(value=func)}
+    return {"source_code": source_code, "mode": "use_pickled_function", "pickled_function": pickled_function}
 
 
-def build_function_data(func: Callable) -> Dict[str, Any]:
+def build_function_data(func: Callable, register_pickle_by_value: bool = False) -> Dict[str, Any]:
     """Inspect the function and return a dictionary with the function data."""
     import types
 
@@ -73,15 +87,10 @@ def build_function_data(func: Callable) -> Dict[str, Any]:
         function_data = {"name": func.__name__}
         if func.__module__ == "__main__" or "." in func.__qualname__.split(".", 1)[-1]:
             # Local or nested callable, so pickle the callable
-            function_data.update(inspect_function(func))
+            function_data.update(inspect_function(func, inspect_source=True))
         else:
             # Global callable (function/class), store its module and name for reference
-            function_data.update(
-                {
-                    "mode": "use_module_path",
-                    "source_code": f"from {func.__module__} import {func.__name__}",
-                }
-            )
+            function_data.update(inspect_function(func, register_pickle_by_value=register_pickle_by_value))
     else:
         raise TypeError("Provided object is not a callable function or class.")
     return function_data

tests/test_utils.py

@@ -1,17 +1,46 @@
+import pytest
 from aiida_pythonjob.utils import build_function_data
 
 
 def test_build_function_data():
-    from math import sqrt
-
-    function_data = build_function_data(sqrt)
-    assert function_data == {
-        "name": "sqrt",
-        "mode": "use_module_path",
-        "source_code": "from math import sqrt",
-    }
-    #
-    try:
-        function_data = build_function_data(1)
-    except Exception as e:
-        assert str(e) == "Provided object is not a callable function or class."
+    """Test the build_function_data function behavior."""
+
+    with pytest.raises(TypeError, match="Provided object is not a callable function or class."):
+        build_function_data(1)
+
+    function_data = build_function_data(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
+
+    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
+
+    def local_function(x, y):
+        return x + y
+
+    function_data = build_function_data(local_function)
+    assert function_data["name"] == "local_function"
+    assert "source_code" in function_data
+    assert function_data["mode"] == "use_pickled_function"
+
+    def outer_function():
+        def nested_function(x, y):
+            return x + y
+
+        return nested_function
+
+    nested_func = outer_function()
+    function_data = build_function_data(nested_func)
+    assert function_data["name"] == "nested_function"
+    assert function_data["mode"] == "use_pickled_function"