aiida-pythonjob 0.1.0__py3-none-any.whl

aiida_pythonjob/__init__.py

@@ -0,0 +1,16 @@
+"""AiiDA plugin that run Python function on remote computers."""
+
+__version__ = "0.1.0"
+
+from .calculations import PythonJob
+from .data import PickledData, PickledFunction
+from .launch import prepare_pythonjob_inputs
+from .parsers import PythonJobParser
+
+__all__ = (
+    "PythonJob",
+    "PickledData",
+    "PickledFunction",
+    "prepare_pythonjob_inputs",
+    "PythonJobParser",
+)

aiida_pythonjob/calculations/__init__.py

@@ -0,0 +1,3 @@
+from .pythonjob import PythonJob
+
+__all__ = ("PythonJob",)

aiida_pythonjob/calculations/pythonjob.py

@@ -0,0 +1,306 @@
+"""Calcjob to run a Python function on a remote computer."""
+
+from __future__ import annotations
+
+import pathlib
+import typing as t
+
+from aiida.common.datastructures import CalcInfo, CodeInfo
+from aiida.common.extendeddicts import AttributeDict
+from aiida.common.folders import Folder
+from aiida.engine import CalcJob, CalcJobProcessSpec
+from aiida.orm import (
+    Data,
+    FolderData,
+    List,
+    RemoteData,
+    SinglefileData,
+    Str,
+    to_aiida_type,
+)
+
+from aiida_pythonjob.data.pickled_function import PickledFunction, to_pickled_function
+
+__all__ = ("PythonJob",)
+
+
+class PythonJob(CalcJob):
+    """Calcjob to run a Python function on a remote computer."""
+
+    _internal_retrieve_list = []
+    _retrieve_singlefile_list = []
+    _retrieve_temporary_list = []
+
+    _DEFAULT_INPUT_FILE = "script.py"
+    _DEFAULT_OUTPUT_FILE = "aiida.out"
+    _DEFAULT_PARENT_FOLDER_NAME = "./parent_folder/"
+
+    @classmethod
+    def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
+        """Define the process specification, including its inputs, outputs and known exit codes.
+
+        :param spec: the calculation job process spec to define.
+        """
+        super().define(spec)
+        spec.input(
+            "function",
+            valid_type=PickledFunction,
+            serializer=to_pickled_function,
+            required=False,
+        )
+        spec.input(
+            "function_source_code",
+            valid_type=Str,
+            serializer=to_aiida_type,
+            required=False,
+        )
+        spec.input("function_name", 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
+        )  # , serializer=serialize_to_aiida_nodes)
+        spec.input(
+            "function_outputs",
+            valid_type=List,
+            default=lambda: List(),
+            required=False,
+            serializer=to_aiida_type,
+            help="The information of the output ports",
+        )
+        spec.input(
+            "parent_folder",
+            valid_type=(RemoteData, FolderData, SinglefileData),
+            required=False,
+            help="Use a local or remote folder as parent folder (for restarts and similar)",
+        )
+        spec.input(
+            "parent_folder_name",
+            valid_type=Str,
+            required=False,
+            serializer=to_aiida_type,
+            help="""Default name of the subfolder that you want to create in the working directory,
+            in which you want to place the files taken from parent_folder""",
+        )
+        spec.input(
+            "parent_output_folder",
+            valid_type=Str,
+            default=None,
+            required=False,
+            serializer=to_aiida_type,
+            help="Name of the subfolder inside 'parent_folder' from which you want to copy the files",
+        )
+        spec.input_namespace(
+            "upload_files",
+            valid_type=(FolderData, SinglefileData),
+            required=False,
+            help="The folder/files to upload",
+        )
+        spec.input_namespace(
+            "copy_files",
+            valid_type=(RemoteData,),
+            required=False,
+            help="The folder/files to copy from the remote computer",
+        )
+        spec.input(
+            "additional_retrieve_list",
+            valid_type=List,
+            default=None,
+            required=False,
+            serializer=to_aiida_type,
+            help="The names of the files to retrieve",
+        )
+        spec.outputs.dynamic = True
+        # set default options (optional)
+        spec.inputs["metadata"]["options"]["parser_name"].default = "pythonjob.pythonjob"
+        spec.inputs["metadata"]["options"]["input_filename"].default = "script.py"
+        spec.inputs["metadata"]["options"]["output_filename"].default = "aiida.out"
+        spec.inputs["metadata"]["options"]["resources"].default = {
+            "num_machines": 1,
+            "num_mpiprocs_per_machine": 1,
+        }
+        # start exit codes - marker for docs
+        spec.exit_code(
+            310,
+            "ERROR_READING_OUTPUT_FILE",
+            invalidates_cache=True,
+            message="The output file could not be read.",
+        )
+        spec.exit_code(
+            320,
+            "ERROR_INVALID_OUTPUT",
+            invalidates_cache=True,
+            message="The output file contains invalid output.",
+        )
+        spec.exit_code(
+            321,
+            "ERROR_RESULT_OUTPUT_MISMATCH",
+            invalidates_cache=True,
+            message="The number of results does not match the number of outputs.",
+        )
+
+    def _build_process_label(self) -> str:
+        """Use the function name as the process label.
+
+        :returns: The process label to use for ``ProcessNode`` instances.
+        """
+        if "process_label" in self.inputs:
+            return self.inputs.process_label.value
+        else:
+            data = self.get_function_data()
+            return f"PythonJob<{data['name']}>"
+
+    def on_create(self) -> None:
+        """Called when a Process is created."""
+
+        super().on_create()
+        self.node.label = self._build_process_label()
+
+    def get_function_data(self) -> dict[str, t.Any]:
+        """Get the function data.
+
+        :returns: The function data.
+        """
+        if "function" in self.inputs:
+            metadata = self.inputs.function.metadata
+            metadata["source_code"] = metadata["import_statements"] + "\n" + metadata["source_code_without_decorator"]
+            return metadata
+        else:
+            return {
+                "source_code": self.inputs.function_source_code.value,
+                "name": self.inputs.function_name.value,
+            }
+
+    def prepare_for_submission(self, folder: Folder) -> CalcInfo:
+        """Prepare the calculation for submission.
+
+        1) Write the python script to the folder.
+        2) Write the inputs to a pickle file and save it to the folder.
+
+        :param folder: A temporary folder on the local file system.
+        :returns: A :class:`aiida.common.datastructures.CalcInfo` instance.
+        """
+        import cloudpickle as pickle
+
+        dirpath = pathlib.Path(folder._abspath)
+        inputs: dict[str, t.Any]
+
+        if self.inputs.function_inputs:
+            inputs = dict(self.inputs.function_inputs)
+        else:
+            inputs = {}
+        if "parent_folder_name" in self.inputs:
+            parent_folder_name = self.inputs.parent_folder_name.value
+        else:
+            parent_folder_name = self._DEFAULT_PARENT_FOLDER_NAME
+        function_data = self.get_function_data()
+        # create python script to run the function
+        script = f"""
+import pickle
+
+# define the function
+{function_data["source_code"]}
+
+# load the inputs from the pickle file
+with open('inputs.pickle', 'rb') as handle:
+    inputs = pickle.load(handle)
+
+# run the function
+result = {function_data["name"]}(**inputs)
+# save the result as a pickle file
+with open('results.pickle', 'wb') as handle:
+    pickle.dump(result, handle)
+"""
+        # write the script to the folder
+        with folder.open(self.options.input_filename, "w", encoding="utf8") as handle:
+            handle.write(script)
+        # symlink = settings.pop('PARENT_FOLDER_SYMLINK', False)
+        symlink = True
+
+        remote_copy_list = []
+        local_copy_list = []
+        remote_symlink_list = []
+        remote_list = remote_symlink_list if symlink else remote_copy_list
+
+        source = self.inputs.get("parent_folder", None)
+
+        if source is not None:
+            if isinstance(source, RemoteData):
+                dirpath = pathlib.Path(source.get_remote_path())
+                if self.inputs.parent_output_folder is not None:
+                    dirpath = pathlib.Path(source.get_remote_path()) / self.inputs.parent_output_folder.value
+                remote_list.append(
+                    (
+                        source.computer.uuid,
+                        str(dirpath),
+                        parent_folder_name,
+                    )
+                )
+            elif isinstance(source, FolderData):
+                dirname = self.inputs.parent_output_folder.value if self.inputs.parent_output_folder is not None else ""
+                local_copy_list.append((source.uuid, dirname, parent_folder_name))
+            elif isinstance(source, SinglefileData):
+                local_copy_list.append((source.uuid, source.filename, source.filename))
+        if "upload_files" in self.inputs:
+            upload_files = self.inputs.upload_files
+            for key, source in upload_files.items():
+                # replace "_dot_" with "." in the key
+                new_key = key.replace("_dot_", ".")
+                if isinstance(source, FolderData):
+                    local_copy_list.append((source.uuid, "", new_key))
+                elif isinstance(source, SinglefileData):
+                    local_copy_list.append((source.uuid, source.filename, source.filename))
+                else:
+                    raise ValueError(
+                        f"""Input folder/file: {source} is not supported.
+Only AiiDA SinglefileData and FolderData are allowed."""
+                    )
+        if "copy_files" in self.inputs:
+            copy_files = self.inputs.copy_files
+            for key, source in copy_files.items():
+                # replace "_dot_" with "." in the key
+                new_key = key.replace("_dot_", ".")
+                dirpath = pathlib.Path(source.get_remote_path())
+                remote_list.append((source.computer.uuid, str(dirpath), new_key))
+        # create pickle file for the inputs
+        input_values = {}
+        for key, value in inputs.items():
+            if isinstance(value, Data) and hasattr(value, "value"):
+                # get the value of the pickled data
+                input_values[key] = value.value
+            # TODO: should check this recursively
+            elif isinstance(value, (AttributeDict, dict)):
+                # if the value is an AttributeDict, use recursively
+                input_values[key] = {k: v.value for k, v in value.items()}
+            else:
+                raise ValueError(
+                    f"Input data {value} is not supported. Only AiiDA data Node with a value attribute is allowed. "
+                )
+        # save the value as a pickle file, the path is absolute
+        filename = "inputs.pickle"
+        dirpath = pathlib.Path(folder._abspath)
+        with folder.open(filename, "wb") as handle:
+            pickle.dump(input_values, handle)
+        # create a singlefiledata object for the pickled data
+        file_data = SinglefileData(file=f"{dirpath}/{filename}")
+        file_data.store()
+        local_copy_list.append((file_data.uuid, file_data.filename, filename))
+
+        codeinfo = CodeInfo()
+        codeinfo.stdin_name = self.options.input_filename
+        codeinfo.stdout_name = self.options.output_filename
+        codeinfo.code_uuid = self.inputs.code.uuid
+
+        calcinfo = CalcInfo()
+        calcinfo.codes_info = [codeinfo]
+        calcinfo.local_copy_list = local_copy_list
+        calcinfo.remote_copy_list = remote_copy_list
+        calcinfo.remote_symlink_list = remote_symlink_list
+        calcinfo.retrieve_list = ["results.pickle", self.options.output_filename]
+        if self.inputs.additional_retrieve_list is not None:
+            calcinfo.retrieve_list += self.inputs.additional_retrieve_list.get_list()
+        calcinfo.retrieve_list += self._internal_retrieve_list
+
+        calcinfo.retrieve_temporary_list = self._retrieve_temporary_list
+        calcinfo.retrieve_singlefile_list = self._retrieve_singlefile_list
+
+        return calcinfo

aiida_pythonjob/config.py

@@ -0,0 +1,14 @@
+import json
+
+from aiida.manage.configuration.settings import AIIDA_CONFIG_FOLDER
+
+
+def load_config() -> dict:
+    """Load the configuration from the config file."""
+    config_file_path = AIIDA_CONFIG_FOLDER / "pythonjob.json"
+    try:
+        with config_file_path.open("r") as f:
+            config = json.load(f)
+    except FileNotFoundError:
+        config = {}
+    return config

aiida_pythonjob/data/__init__.py

@@ -0,0 +1,4 @@
+from .pickled_data import PickledData
+from .pickled_function import PickledFunction
+
+__all__ = ("PickledData", "PickledFunction")

aiida_pythonjob/data/pickled_data.py

@@ -0,0 +1,87 @@
+"""`Data` sub class to represent any data using pickle."""
+
+import sys
+from pickle import UnpicklingError
+
+import cloudpickle
+from aiida import orm
+
+
+class Dict(orm.Dict):
+    @property
+    def value(self):
+        return self.get_dict()
+
+
+class List(orm.List):
+    @property
+    def value(self):
+        return self.get_list()
+
+
+class PickledData(orm.Data):
+    """Data to represent a pickled value using cloudpickle."""
+
+    FILENAME = "value.pkl"  # Class attribute to store the filename
+
+    def __init__(self, value=None, **kwargs):
+        """Initialize a `PickledData` node instance.
+
+        :param value: raw Python value to initialize the `PickledData` node from.
+        """
+        super().__init__(**kwargs)
+        self.set_value(value)
+
+    def __str__(self):
+        return f"{super().__str__()} : {self.get_value()}"
+
+    @property
+    def value(self):
+        """Return the contents of this node.
+
+        :return: The unpickled value.
+        """
+        return self.get_value()
+
+    @value.setter
+    def value(self, value):
+        self.set_value(value)
+
+    def get_value(self):
+        """Return the contents of this node, unpickling the stored value.
+
+        :return: The unpickled value.
+        """
+        return self._get_value_from_file()
+
+    def _get_value_from_file(self):
+        """Read the pickled value from file and return it."""
+        try:
+            with self.base.repository.open(self.FILENAME, mode="rb") as f:
+                return cloudpickle.loads(f.read())  # Deserialize the value
+        except (UnpicklingError, ValueError) as e:
+            raise ImportError(
+                "Failed to load the pickled value. This may be due to an incompatible pickle protocol. "
+                "Please ensure that the correct environment and cloudpickle version are being used."
+            ) from e
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                "Failed to load the pickled value. This may be due to a missing module. "
+                "Please ensure that the correct environment and cloudpickle version are being used."
+            ) from e
+
+    def set_value(self, value):
+        """Set the contents of this node by pickling the provided value.
+
+        :param value: The Python value to pickle and store.
+        """
+        # Serialize the value and store it
+        serialized_value = cloudpickle.dumps(value)
+        self.base.repository.put_object_from_bytes(serialized_value, self.FILENAME)
+
+        # Store relevant metadata
+        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        self.base.attributes.set("python_version", python_version)
+        self.base.attributes.set("serializer_module", cloudpickle.__name__)
+        self.base.attributes.set("serializer_version", cloudpickle.__version__)
+        self.base.attributes.set("pickle_protocol", cloudpickle.DEFAULT_PROTOCOL)

aiida_pythonjob/data/pickled_function.py

@@ -0,0 +1,145 @@
+import inspect
+import textwrap
+from typing import Any, Callable, Dict, _SpecialForm, get_type_hints
+
+from .pickled_data import PickledData
+
+
+class PickledFunction(PickledData):
+    """Data class to represent a pickled Python function."""
+
+    def __init__(self, value=None, **kwargs):
+        """Initialize a PickledFunction node instance.
+
+        :param value: a Python function
+        """
+        super().__init__(**kwargs)
+        if not callable(value):
+            raise ValueError("value must be a callable Python function")
+        self.set_value(value)
+        self.set_attribute(value)
+
+    def __str__(self):
+        return f"PickledFunction<{self.base.attributes.get('function_name')}> pk={self.pk}"
+
+    @property
+    def metadata(self):
+        """Return a dictionary of metadata."""
+        return {
+            "name": self.base.attributes.get("name"),
+            "import_statements": self.base.attributes.get("import_statements"),
+            "source_code": self.base.attributes.get("source_code"),
+            "source_code_without_decorator": self.base.attributes.get("source_code_without_decorator"),
+            "type": "function",
+            "is_pickle": True,
+        }
+
+    @classmethod
+    def build_callable(cls, func):
+        """Return the executor for this node."""
+        import cloudpickle as pickle
+
+        executor = {
+            "executor": pickle.dumps(func),
+            "type": "function",
+            "is_pickle": True,
+        }
+        executor.update(cls.inspect_function(func))
+        return executor
+
+    def set_attribute(self, value):
+        """Set the contents of this node by pickling the provided function.
+
+        :param value: The Python function to pickle and store.
+        """
+        # Serialize the function and extract metadata
+        serialized_data = self.inspect_function(value)
+
+        # Store relevant metadata
+        self.base.attributes.set("name", serialized_data["name"])
+        self.base.attributes.set("import_statements", serialized_data["import_statements"])
+        self.base.attributes.set("source_code", serialized_data["source_code"])
+        self.base.attributes.set(
+            "source_code_without_decorator",
+            serialized_data["source_code_without_decorator"],
+        )
+
+    @classmethod
+    def inspect_function(cls, func: Callable) -> Dict[str, Any]:
+        """Serialize a function for storage or transmission."""
+        try:
+            # we need save the source code explicitly, because in the case of jupyter notebook,
+            # the source code is not saved in the pickle file
+            source_code = inspect.getsource(func)
+            # Split the source into lines for processing
+            source_code_lines = source_code.split("\n")
+            function_source_code = "\n".join(source_code_lines)
+            # Find the first line of the actual function definition
+            for i, line in enumerate(source_code_lines):
+                if line.strip().startswith("def "):
+                    break
+            function_source_code_without_decorator = "\n".join(source_code_lines[i:])
+            function_source_code_without_decorator = textwrap.dedent(function_source_code_without_decorator)
+            # we also need to include the necessary imports for the types used in the type hints.
+            try:
+                required_imports = cls.get_required_imports(func)
+            except Exception as e:
+                required_imports = {}
+                print(f"Failed to get required imports for function {func.__name__}: {e}")
+            # Generate import statements
+            import_statements = "\n".join(
+                f"from {module} import {', '.join(types)}" for module, types in required_imports.items()
+            )
+        except Exception as e:
+            print(f"Failed to inspect function {func.__name__}: {e}")
+            function_source_code = ""
+            function_source_code_without_decorator = ""
+            import_statements = ""
+        return {
+            "name": func.__name__,
+            "source_code": function_source_code,
+            "source_code_without_decorator": function_source_code_without_decorator,
+            "import_statements": import_statements,
+        }
+
+    @classmethod
+    def get_required_imports(cls, func: Callable) -> Dict[str, set]:
+        """Retrieve type hints and the corresponding modules."""
+        type_hints = get_type_hints(func)
+        imports = {}
+
+        def add_imports(type_hint):
+            if isinstance(type_hint, _SpecialForm):  # Handle special forms like Any, Union, Optional
+                module_name = "typing"
+                type_name = type_hint._name or str(type_hint)
+            elif hasattr(type_hint, "__origin__"):  # This checks for higher-order types like List, Dict
+                module_name = type_hint.__module__
+                type_name = getattr(type_hint, "_name", None) or getattr(type_hint.__origin__, "__name__", None)
+                for arg in getattr(type_hint, "__args__", []):
+                    if arg is type(None):
+                        continue
+                    add_imports(arg)  # Recursively add imports for each argument
+            elif hasattr(type_hint, "__module__"):
+                module_name = type_hint.__module__
+                type_name = type_hint.__name__
+            else:
+                return  # If no module or origin, we can't import it, e.g., for literals
+
+            if type_name is not None:
+                if module_name not in imports:
+                    imports[module_name] = set()
+                imports[module_name].add(type_name)
+
+        for _, type_hint in type_hints.items():
+            add_imports(type_hint)
+
+        return imports
+
+
+def to_pickled_function(value):
+    """Convert a Python function to a `PickledFunction` instance."""
+    return PickledFunction(value)
+
+
+class PickledLocalFunction(PickledFunction):
+    """PickledFunction subclass for local functions."""

aiida_pythonjob/data/serializer.py

@@ -0,0 +1,122 @@
+import sys
+from importlib.metadata import entry_points
+from typing import Any
+
+from aiida import common, orm
+
+from aiida_pythonjob.config import load_config
+
+from .pickled_data import PickledData
+
+
+def get_serializer_from_entry_points() -> dict:
+    """Retrieve the serializer from the entry points."""
+    # import time
+
+    # ts = time.time()
+    configs = load_config()
+    serializers = configs.get("serializers", {})
+    excludes = serializers.get("excludes", [])
+    # Retrieve the entry points for 'aiida.data' and store them in a dictionary
+    eps = entry_points()
+    if sys.version_info >= (3, 10):
+        group = eps.select(group="aiida.data")
+    else:
+        group = eps.get("aiida.data", [])
+    eps = {}
+    for ep in group:
+        # split the entry point name by first ".", and check the last part
+        key = ep.name.split(".", 1)[-1]
+        # skip key without "." because it is not a module name for a data type
+        if "." not in key or key in excludes:
+            continue
+        eps.setdefault(key, [])
+        eps[key].append(ep)
+
+    # print("Time to load entry points: ", time.time() - ts)
+    # check if there are duplicates
+    for key, value in eps.items():
+        if len(value) > 1:
+            if key in serializers:
+                [ep for ep in value if ep.name == serializers[key]]
+                eps[key] = [ep for ep in value if ep.name == serializers[key]]
+                if not eps[key]:
+                    raise ValueError(f"Entry point {serializers[key]} not found for {key}")
+            else:
+                msg = f"Duplicate entry points for {key}: {[ep.name for ep in value]}"
+                raise ValueError(msg)
+    return eps
+
+
+eps = get_serializer_from_entry_points()
+
+
+def serialize_to_aiida_nodes(inputs: dict | None = None) -> dict:
+    """Serialize the inputs to a dictionary of AiiDA data nodes.
+
+    Args:
+        inputs (dict): The inputs to be serialized.
+
+    Returns:
+        dict: The serialized inputs.
+    """
+    new_inputs = {}
+    # save all kwargs to inputs port
+    for key, data in inputs.items():
+        new_inputs[key] = general_serializer(data)
+    return new_inputs
+
+
+def clean_dict_key(data):
+    """Replace "." with "__dot__" in the keys of a dictionary."""
+    if isinstance(data, dict):
+        return {k.replace(".", "__dot__"): clean_dict_key(v) for k, v in data.items()}
+    return data
+
+
+def general_serializer(data: Any, check_value=True) -> orm.Node:
+    """Serialize the data to an AiiDA data node."""
+    if isinstance(data, orm.Data):
+        if check_value and not hasattr(data, "value"):
+            raise ValueError("Only AiiDA data Node with a value attribute is allowed.")
+        return data
+    elif isinstance(data, common.extendeddicts.AttributeDict):
+        # if the data is an AttributeDict, use it directly
+        return data
+    # if is string with syntax {{}}, this is a port will read data from ctx
+    elif isinstance(data, str) and data.startswith("{{") and data.endswith("}}"):
+        return data
+    # if data is a class instance, get its __module__ and class name as a string
+    # for example, an Atoms will have ase.atoms.Atoms
+    else:
+        data = clean_dict_key(data)
+        # try to get the serializer from the entry points
+        data_type = type(data)
+        ep_key = f"{data_type.__module__}.{data_type.__name__}"
+        # search for the key in the entry points
+        if ep_key in eps:
+            try:
+                new_node = eps[ep_key][0].load()(data)
+            except Exception as e:
+                raise ValueError(f"Error in serializing {ep_key}: {e}")
+            finally:
+                # try to save the node to da
+                try:
+                    new_node.store()
+                    return new_node
+                except Exception:
+                    # try to serialize the value as a PickledData
+                    try:
+                        new_node = PickledData(data)
+                        new_node.store()
+                        return new_node
+                    except Exception as e:
+                        raise ValueError(f"Error in serializing {ep_key}: {e}")
+        else:
+            # try to serialize the data as a PickledData
+            try:
+                new_node = PickledData(data)
+                new_node.store()
+                return new_node
+            except Exception as e:
+                raise ValueError(f"Error in serializing {ep_key}: {e}")

aiida_pythonjob/launch.py

@@ -0,0 +1,67 @@
+from typing import Any, Callable
+
+from aiida.orm import AbstractCode, Computer, FolderData, List, SinglefileData, Str
+
+from .data.pickled_function import PickledFunction
+from .data.serializer import serialize_to_aiida_nodes
+from .utils import get_or_create_code
+
+
+def prepare_pythonjob_inputs(
+    function: Callable[..., Any],
+    function_inputs: dict[str, Any] | None = None,
+    function_outputs: dict[str, Any] | None = None,
+    code: AbstractCode | None = None,
+    command_info: dict[str, str] | None = None,
+    computer: str | Computer = "localhost",
+    metadata: dict[str, Any] | None = None,
+    upload_files: dict[str, str] = {},
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Prepare the inputs for PythonJob"""
+    import os
+
+    # get the names kwargs for the PythonJob, which are the inputs before _wait
+    executor = PickledFunction.build_callable(function)
+    new_upload_files = {}
+    # change the string in the upload files to SingleFileData, or FolderData
+    for key, source in upload_files.items():
+        # only alphanumeric and underscores are allowed in the key
+        # replace all "." with "_dot_"
+        new_key = key.replace(".", "_dot_")
+        if isinstance(source, str):
+            if os.path.isfile(source):
+                new_upload_files[new_key] = SinglefileData(file=source)
+            elif os.path.isdir(source):
+                new_upload_files[new_key] = FolderData(tree=source)
+        elif isinstance(source, (SinglefileData, FolderData)):
+            new_upload_files[new_key] = source
+        else:
+            raise ValueError(f"Invalid upload file type: {type(source)}, {source}")
+    #
+    if code is None:
+        command_info = command_info or {}
+        code = get_or_create_code(computer=computer, **command_info)
+    # get the source code of the function
+    function_name = executor["name"]
+    if executor.get("is_pickle", False):
+        function_source_code = executor["import_statements"] + "\n" + executor["source_code_without_decorator"]
+    else:
+        function_source_code = f"from {executor['module']} import {function_name}"
+
+    # serialize the kwargs into AiiDA Data
+    function_inputs = function_inputs or {}
+    function_inputs = serialize_to_aiida_nodes(function_inputs)
+    # transfer the args to kwargs
+    inputs = {
+        "process_label": "PythonJob<{}>".format(function_name),
+        "function_source_code": Str(function_source_code),
+        "function_name": Str(function_name),
+        "code": code,
+        "function_inputs": function_inputs,
+        "upload_files": new_upload_files,
+        "function_outputs": List(function_outputs),
+        "metadata": metadata or {},
+        **kwargs,
+    }
+    return inputs

aiida_pythonjob/parsers/__init__.py

@@ -0,0 +1,3 @@
+from .pythonjob import PythonJobParser
+
+__all__ = ("PythonJobParser",)

aiida_pythonjob/parsers/pythonjob.py

@@ -0,0 +1,111 @@
+"""Parser for an `PythonJob` job."""
+
+from aiida.engine import ExitCode
+from aiida.parsers.parser import Parser
+
+from aiida_pythonjob.data.serializer import general_serializer
+
+
+class PythonJobParser(Parser):
+    """Parser for an `PythonJob` job."""
+
+    def parse(self, **kwargs):
+        """Parse the contents of the output files stored in the `retrieved` output node.
+
+        The function_outputs could be a namespce, e.g.,
+        function_outputs=[
+            {"identifier": "namespace", "name": "add_multiply"},
+            {"name": "add_multiply.add"},
+            {"name": "add_multiply.multiply"},
+            {"name": "minus"},
+        ]
+        """
+        import pickle
+
+        function_outputs = self.node.inputs.function_outputs.get_list()
+        if len(function_outputs) == 0:
+            function_outputs = [{"name": "result"}]
+        self.output_list = function_outputs
+        # first we remove nested outputs, e.g., "add_multiply.add"
+        top_level_output_list = [output for output in self.output_list if "." not in output["name"]]
+        exit_code = 0
+        try:
+            with self.retrieved.base.repository.open("results.pickle", "rb") as handle:
+                results = pickle.load(handle)
+                if isinstance(results, tuple):
+                    if len(top_level_output_list) != len(results):
+                        self.exit_codes.ERROR_RESULT_OUTPUT_MISMATCH
+                    for i in range(len(top_level_output_list)):
+                        top_level_output_list[i]["value"] = self.serialize_output(results[i], top_level_output_list[i])
+                elif isinstance(results, dict) and len(top_level_output_list) > 1:
+                    # pop the exit code if it exists
+                    exit_code = results.pop("exit_code", 0)
+                    for output in top_level_output_list:
+                        if output.get("required", False):
+                            if output["name"] not in results:
+                                self.exit_codes.ERROR_MISSING_OUTPUT
+                        output["value"] = self.serialize_output(results.pop(output["name"]), output)
+                    # if there are any remaining results, raise an warning
+                    if results:
+                        self.logger.warning(
+                            f"Found extra results that are not included in the output: {results.keys()}"
+                        )
+                elif isinstance(results, dict) and len(top_level_output_list) == 1:
+                    exit_code = results.pop("exit_code", 0)
+                    # if output name in results, use it
+                    if top_level_output_list[0]["name"] in results:
+                        top_level_output_list[0]["value"] = self.serialize_output(
+                            results[top_level_output_list[0]["name"]],
+                            top_level_output_list[0],
+                        )
+                    # otherwise, we assume the results is the output
+                    else:
+                        top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
+                elif len(top_level_output_list) == 1:
+                    # otherwise, we assume the results is the output
+                    top_level_output_list[0]["value"] = self.serialize_output(results, top_level_output_list[0])
+                else:
+                    raise ValueError("The number of results does not match the number of outputs.")
+                for output in top_level_output_list:
+                    self.out(output["name"], output["value"])
+                if exit_code:
+                    if isinstance(exit_code, dict):
+                        exit_code = ExitCode(exit_code["status"], exit_code["message"])
+                    elif isinstance(exit_code, int):
+                        exit_code = ExitCode(exit_code)
+                    return exit_code
+        except OSError:
+            return self.exit_codes.ERROR_READING_OUTPUT_FILE
+        except ValueError as exception:
+            self.logger.error(exception)
+            return self.exit_codes.ERROR_INVALID_OUTPUT
+
+    def find_output(self, name):
+        """Find the output with the given name."""
+        for output in self.output_list:
+            if output["name"] == name:
+                return output
+        return None
+
+    def serialize_output(self, result, output):
+        """Serialize outputs."""
+
+        name = output["name"]
+        if output.get("identifier", "Any").upper() in ["NAMESPACE", "WORKGRAPH.NAMESPACE"]:
+            if isinstance(result, dict):
+                serialized_result = {}
+                for key, value in result.items():
+                    full_name = f"{name}.{key}"
+                    full_name_output = self.find_output(full_name)
+                    if full_name_output and full_name_output.get("identifier", "Any").upper() in [
+                        "NAMESPACE",
+                        "WORKGRAPH.NAMESPACE",
+                    ]:
+                        serialized_result[key] = self.serialize_output(value, full_name_output)
+                    else:
+                        serialized_result[key] = general_serializer(value)
+                return serialized_result
+            else:
+                self.exit_codes.ERROR_INVALID_OUTPUT
+        else:
+            return general_serializer(result)

aiida_pythonjob/utils.py

@@ -0,0 +1,31 @@
+from typing import Optional
+
+from aiida.common.exceptions import NotExistent
+from aiida.orm import Computer, InstalledCode, load_code, load_computer
+
+
+def get_or_create_code(
+    label: str = "python3",
+    computer: Optional[str | Computer] = "localhost",
+    filepath_executable: Optional[str] = None,
+    prepend_text: str = "",
+):
+    """Try to load code, create if not exit."""
+
+    try:
+        return load_code(f"{label}@{computer}")
+    except NotExistent:
+        description = f"Code on computer: {computer}"
+        computer = load_computer(computer)
+        filepath_executable = filepath_executable or label
+        code = InstalledCode(
+            computer=computer,
+            label=label,
+            description=description,
+            filepath_executable=filepath_executable,
+            default_calc_job_plugin="pythonjob.pythonjob",
+            prepend_text=prepend_text,
+        )
+
+        code.store()
+        return code

aiida_pythonjob-0.1.0.dist-info/METADATA

@@ -0,0 +1,82 @@
+Metadata-Version: 2.3
+Name: aiida-pythonjob
+Version: 0.1.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>
+License: MIT License
+        
+        Copyright (c) 2024 AiiDA team
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+Keywords: aiida,plugin
+Classifier: Development Status :: 3 - Alpha
+Classifier: Framework :: AiiDA
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python
+Requires-Python: >=3.9
+Requires-Dist: aiida-core<3,>=2.3
+Requires-Dist: cloudpickle
+Requires-Dist: voluptuous
+Provides-Extra: docs
+Requires-Dist: furo; extra == 'docs'
+Requires-Dist: markupsafe<2.1; extra == 'docs'
+Requires-Dist: nbsphinx; extra == 'docs'
+Requires-Dist: sphinx; extra == 'docs'
+Requires-Dist: sphinx-gallery; extra == 'docs'
+Requires-Dist: sphinx-rtd-theme; extra == 'docs'
+Requires-Dist: sphinxcontrib-contentui; extra == 'docs'
+Requires-Dist: sphinxcontrib-details-directive; extra == 'docs'
+Provides-Extra: pre-commit
+Requires-Dist: pre-commit~=3.5; extra == 'pre-commit'
+Description-Content-Type: text/markdown
+
+# 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)
+[![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/)
+
+Efficiently design and manage flexible workflows with AiiDA, featuring an interactive GUI, checkpoints, provenance tracking, error-resistant, and remote execution capabilities.
+
+
+
+## Installation
+
+```console
+    pip install aiida-pythonjob
+```
+
+To install the latest version from source, first clone the repository and then install using `pip`:
+
+```console
+git clone https://github.com/aiidateam/aiida-pythonjob
+cd aiida-pythonjob
+pip install -e .
+```
+
+
+## Documentation
+Explore the comprehensive [documentation](https://aiida-pythonjob.readthedocs.io/en/latest/) to discover all the features and capabilities of AiiDA Workgraph.
+
+
+## License
+[MIT](http://opensource.org/licenses/MIT)

aiida_pythonjob-0.1.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+aiida_pythonjob/__init__.py,sha256=8PyvN-j-1t8OpA30SINJmQrtBHcfce_V47Y1clUu1e8,383
+aiida_pythonjob/config.py,sha256=ZHeS_7qRmkea4ltWpju5vsOHvYfi9mVdpklWfUu2LRs,384
+aiida_pythonjob/launch.py,sha256=Ch9TtgJA6IXlvapg8sspC_QDgni0xJWj1OIISimbMs0,2709
+aiida_pythonjob/utils.py,sha256=HqxmYD76xFwqwBR3jezHNM0zw1uB1WPCn6dHH3NLJDM,947
+aiida_pythonjob/calculations/__init__.py,sha256=-xxZoADS038pH5T_ssINO-Ql5q7SY31nPYHlua6f0Ns,59
+aiida_pythonjob/calculations/pythonjob.py,sha256=ktE67rIx0yitf7wek4XMGB8nhNAg1omgh2juNqlRMJo,11826
+aiida_pythonjob/data/__init__.py,sha256=fgl_e1ZbQ4It8l1zfQUM1p3dcIWitG-A8RHpy6N7wgk,130
+aiida_pythonjob/data/pickled_data.py,sha256=x4ZQOlExzdbww2_ojg5BxKGr4-didDMjXRUXLGUgUD8,2901
+aiida_pythonjob/data/pickled_function.py,sha256=VNpbXKxgntXKKhfxrvdJKs9NtV6ormAwcJ5BQvlSMmM,5975
+aiida_pythonjob/data/serializer.py,sha256=a0Gyx3SH6J9C7Vra1DKSRyDpBHAEn8vtu1En-mlemfs,4454
+aiida_pythonjob/parsers/__init__.py,sha256=-gwCpe5K-lMR1Yonp0AKST_z3y63xsfDoFRpTnHSfHw,71
+aiida_pythonjob/parsers/pythonjob.py,sha256=WM_nNKUXUcnWgoqZbbosI_AKBCphlt2ZXsV7XGnilz8,5386
+aiida_pythonjob-0.1.0.dist-info/METADATA,sha256=avZe5Bw0XFo0SGsWGhIGYjrdaQHgsDwaqtj8gVJeNFc,3606
+aiida_pythonjob-0.1.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+aiida_pythonjob-0.1.0.dist-info/entry_points.txt,sha256=F7WrhW9wyC4WIu5FYxozTGuzrJGcaY1HHk_WeVU9PIM,349
+aiida_pythonjob-0.1.0.dist-info/licenses/LICENSE,sha256=YHoQimZwJx-5oMGat60FgIuvY6rTfVSJN_NPe66nlPc,1067
+aiida_pythonjob-0.1.0.dist-info/RECORD,,

aiida_pythonjob-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.26.3
+Root-Is-Purelib: true
+Tag: py3-none-any

aiida_pythonjob-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,9 @@
+[aiida.calculations]
+pythonjob.pythonjob = aiida_pythonjob.calculations.pythonjob:PythonJob
+
+[aiida.data]
+pythonjob.pickled_data = aiida_pythonjob.data.pickled_data:PickledData
+pythonjob.pickled_function = aiida_pythonjob.data.pickled_function:PickledFunction
+
+[aiida.parsers]
+pythonjob.pythonjob = aiida_pythonjob.parsers.pythonjob:PythonJobParser

aiida_pythonjob-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 AiiDA team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.