use run_remotely to set PythonJob (#84)

When adding a new node, one can set `run_remotely=True` to ask the engine to run the function on a remote computer.
In the inputs, user can set `"computer": "merlin6"` to tell which remote computer to run this job, the engine will try to load `python3@{computer}` code, it will create a new one if not exit.

One can also set the `code_label`, `python_path`, `prepend_text` to specific the code.

In the `metadata.custom_scheduler_commands`, one can add commands to load modules and activate the conda environment and so on.

Author: superstar54

aiida_workgraph/calculations/python.py

@@ -14,6 +14,7 @@
     List,
     FolderData,
     RemoteData,
+    to_aiida_type,
 )
 
 
@@ -45,13 +46,21 @@ def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
         :param spec: the calculation job process spec to define.
         """
         super().define(spec)
-        spec.input("function_source_code", required=False)
-        spec.input("function_name", 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_namespace("kwargs", valid_type=Data, required=False)
         spec.input(
             "output_name_list",
             valid_type=List,
             required=False,
+            serializer=to_aiida_type,
             help="The names of the output ports",
         )
         spec.input(
@@ -62,16 +71,18 @@ def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
         )
         spec.input(
             "parent_output_folder",
-            valid_type=(Str),
+            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(
             "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

aiida_workgraph/collection.py

@@ -24,6 +24,9 @@ def new(
         # build the node on the fly if the identifier is a callable
         if callable(identifier):
             identifier = build_node_from_callable(identifier)
+            if kwargs.pop("run_remotely", False):
+                # this is a PythonJob
+                identifier = build_PythonJob_node(identifier)
         if isinstance(identifier, str) and identifier.upper() == "PYTHONJOB":
             # copy the inputs and outputs from the function node to the PythonJob node
             identifier = build_PythonJob_node(kwargs.pop("function"))

aiida_workgraph/decorator.py

@@ -236,13 +236,22 @@ def build_node_from_AiiDA(
 def build_PythonJob_node(func: Callable) -> Node:
     """Build PythonJob node from function."""
     from aiida_workgraph.calculations.python import PythonJob
+    from copy import deepcopy
 
     ndata = {"executor": PythonJob, "node_type": "CALCJOB"}
     _, ndata_py = build_node_from_AiiDA(ndata)
-    ndata = func.ndata
+    ndata = deepcopy(func.ndata)
     # merge the inputs and outputs from the PythonJob node to the function node
     # skip the already existed inputs and outputs
     inputs = ndata["inputs"]
+    inputs.extend(
+        [
+            ["String", "computer"],
+            ["String", "code_label"],
+            ["String", "code_path"],
+            ["String", "prepend_text"],
+        ]
+    )
     outputs = ndata["outputs"]
     for input in ndata_py["inputs"]:
         if input not in inputs:
@@ -252,6 +261,7 @@ def build_PythonJob_node(func: Callable) -> Node:
             outputs.append(output)
     # append the kwargs of the PythonJob node to the function node
     kwargs = ndata["kwargs"]
+    kwargs.extend(["computer", "code_label", "code_path", "prepend_text"])
     kwargs.extend(ndata_py["kwargs"])
     ndata["inputs"] = inputs
     ndata["outputs"] = outputs

aiida_workgraph/engine/workgraph.py

@@ -826,10 +826,23 @@ def run_nodes(self, names: t.List[str], continue_workgraph: bool = True) -> None
             elif node["metadata"]["node_type"].upper() in ["PYTHONJOB"]:
                 from aiida_workgraph.calculations.python import PythonJob
                 from aiida_workgraph.calculations.general_data import GeneralData
+                from aiida_workgraph.utils import get_or_create_code
 
                 print("node  type: Python.")
                 # normal function does not have a process
-                code = kwargs.pop("code")
+                code = kwargs.pop("code", None)
+                computer = kwargs.pop("computer", None)
+                code_label = kwargs.pop("code_label", None)
+                code_path = kwargs.pop("code_path", None)
+                prepend_text = kwargs.pop("prepend_text", None)
+                #
+                if code is None:
+                    code = get_or_create_code(
+                        computer=computer or "localhost",
+                        code_label=code_label or "python3",
+                        code_path=code_path,
+                        prepend_text=prepend_text,
+                    )
                 parent_folder = kwargs.pop("parent_folder", None)
                 metadata = kwargs.pop("metadata", {})
                 metadata.update({"call_link_label": name})
@@ -847,7 +860,6 @@ def run_nodes(self, names: t.List[str], continue_workgraph: bool = True) -> None
                         inputs[key] = value
                     else:
                         inputs[key] = GeneralData(value)
-                print("inputs: ", inputs)
                 # outputs
                 output_name_list = [output["name"] for output in node["outputs"]]
 

aiida_workgraph/utils/__init__.py

@@ -1,6 +1,7 @@
 from typing import Any, Dict, Optional, Union
 from aiida.engine.processes import Process
 from aiida import orm
+from aiida.common.exceptions import NotExistent
 
 
 def get_executor(data: Dict[str, Any]) -> Union[Process, Any]:
@@ -236,3 +237,31 @@ def get_processes_latest(pk: int) -> Dict[str, Dict[str, Union[int, str]]]:
                         "mtime": nodes.mtime,
                     }
     return nodes
+
+
+def get_or_create_code(
+    computer: str = "localhost",
+    code_label: str = "python3",
+    code_path: str = None,
+    prepend_text: str = "",
+):
+    """Try to load code, create if not exit."""
+    from aiida.orm.nodes.data.code.installed import InstalledCode
+
+    try:
+        return orm.load_code(f"{code_label}@{computer}")
+    except NotExistent:
+        description = f"Python code on computer: {computer}"
+        computer = orm.load_computer(computer)
+        code_path = code_path or code_label
+        code = InstalledCode(
+            computer=computer,
+            label=code_label,
+            description=description,
+            filepath_executable=code_path,
+            default_calc_job_plugin="workgraph.python",
+            prepend_text=prepend_text,
+        )
+
+        code.store()
+        return code

docs/source/built-in/aiida_python.ipynb

@@ -65,14 +65,13 @@
     "    return x*y\n",
     "\n",
     "wg = WorkGraph(\"first_workflow\")\n",
-    "wg.nodes.new(\"PythonJob\", function=add, name=\"add\")\n",
-    "wg.nodes.new(\"PythonJob\", function=multiply, name=\"multiply\", x=wg.nodes[\"add\"].outputs[0])\n",
+    "wg.nodes.new(add, name=\"add\", run_remotely=True)\n",
+    "wg.nodes.new(multiply, name=\"multiply\", x=wg.nodes[\"add\"].outputs[0], run_remotely=True)\n",
     "\n",
-    "code = load_code(\"python@localhost\")\n",
     "\n",
     "#------------------------- Submit the calculation -------------------\n",
-    "wg.submit(inputs = {\"add\": {\"x\": 2, \"y\": 3, \"code\": code},\n",
-    "                    \"multiply\": {\"y\": 4, \"code\": code}},\n",
+    "wg.submit(inputs = {\"add\": {\"x\": 2, \"y\": 3, \"computer\": \"localhost\"},\n",
+    "                    \"multiply\": {\"y\": 4, \"computer\": \"localhost\"}},\n",
     "          wait=True)\n",
     "print(\"\\nResult of multiply is {} \\n\\n\".format(wg.nodes[\"multiply\"].outputs['result'].value))\n",
     "```\n",