Bump aiida-shell to v0.8 (#326)

* Update parser signature for `aiida-shell` v0.8
* Extend ShellJob docs: split section with different Computer/Code options.
* Use `prepare_shell_job_inputs` to prepare the inputs

---------

Co-authored-by: superstar54 <[email protected]>

aiida_workgraph/engine/utils.py

@@ -129,36 +129,16 @@ def prepare_for_python_task(task: dict, kwargs: dict, var_kwargs: dict) -> dict:
     return inputs
 
 
-def prepare_for_shell_task(task: dict, kwargs: dict) -> dict:
+def prepare_for_shell_task(task: dict, inputs: dict) -> dict:
     """Prepare the inputs for ShellJob"""
-    from aiida_shell.launch import prepare_code, convert_nodes_single_file_data
-    from aiida.common import lang
-    from aiida.orm import AbstractCode
+    from aiida_shell.launch import prepare_shell_job_inputs
+    import inspect
 
-    command = kwargs.pop("command", None)
-    resolve_command = kwargs.pop("resolve_command", False)
-    metadata = kwargs.pop("metadata", {})
-    # setup code
-    if isinstance(command, str):
-        computer = (metadata or {}).get("options", {}).pop("computer", None)
-        code = prepare_code(command, computer, resolve_command)
-    else:
-        lang.type_check(command, AbstractCode)
-        code = command
-    # update the tasks with links
-    nodes = convert_nodes_single_file_data(kwargs.pop("nodes", {}))
-    # find all keys in kwargs start with "nodes."
-    for key in list(kwargs.keys()):
-        if key.startswith("nodes."):
-            nodes[key[6:]] = kwargs.pop(key)
-    metadata.update({"call_link_label": task["name"]})
-    inputs = {
-        "code": code,
-        "nodes": nodes,
-        "filenames": kwargs.pop("filenames", {}),
-        "arguments": kwargs.pop("arguments", []),
-        "outputs": kwargs.pop("outputs", []),
-        "parser": kwargs.pop("parser", None),
-        "metadata": metadata or {},
-    }
+    # Retrieve the signature of `prepare_shell_job_inputs` to determine expected input parameters.
+    signature = inspect.signature(prepare_shell_job_inputs)
+    kwargs = {key: inputs.pop(key, None) for key in signature.parameters.keys()}
+    inputs.update(prepare_shell_job_inputs(**kwargs))
+
+    inputs.setdefault("metadata", {})
+    inputs["metadata"].update({"call_link_label": task["name"]})
     return inputs

aiida_workgraph/utils/__init__.py

@@ -454,7 +454,7 @@ def get_or_create_code(
     try:
         return orm.load_code(f"{code_label}@{computer}")
     except NotExistent:
-        description = f"Python code on computer: {computer}"
+        description = f"Code on computer: {computer}"
         computer = orm.load_computer(computer)
         code_path = code_path or code_label
         code = InstalledCode(

docs/gallery/built-in/autogen/shelljob.py

@@ -24,7 +24,6 @@
 # Run a shell command without any arguments. Here we run the `date` command to show the date.
 #
 
-
 from aiida_workgraph import WorkGraph
 
 wg = WorkGraph(name="test_shell_date")
@@ -34,6 +33,54 @@
 # Print out the result:
 print("\nResult: ", date_task.outputs["stdout"].value.get_content())
 
+# %%
+# Under the hood, an AiiDA ``Code`` instance ``date`` will be created on the ``localhost`` computer. In addition, it is also
+# possible to specify a different, remote computer. For the sake of this example, we create a mock remote computer (you
+# would usually have this pre-configured already, e.g., your favorite HPC resource):
+
+from aiida import orm
+
+created, mock_remote_computer = orm.Computer.collection.get_or_create(
+    label="my-remote-computer",
+    description="A mock remote computer",
+    hostname="my-remote-computer",
+    workdir="/tmp/aiida",
+    transport_type="core.ssh",
+    scheduler_type="core.direct",
+)
+if created:
+    mock_remote_computer.store()
+
+# We can then specify the remote computer via the ``metadata``:
+
+wg = WorkGraph(name="test_shell_date_remote")
+date_task_remote = wg.add_task(
+    "ShellJob",
+    command="date",
+    metadata={"computer": orm.load_computer("my-remote-computer")},
+)
+
+#%%
+# In addition, it is possible to pass a pre-configured ``Code``. Let's again create a mock ``Code``:
+
+from aiida_workgraph.utils import get_or_create_code
+
+mock_remote_code = get_or_create_code(
+    code_label="remote-date",
+    computer="my-remote-computer",
+    code_path="/usr/bin/date",
+)
+
+# Which we can now directly pass to the ``command`` argument of ``add_task``:
+
+wg = WorkGraph(name="test_shell_date_preconfigured")
+preconf_task = wg.add_task(
+    "ShellJob", command=orm.load_code("remote-date@my-remote-computer")
+)
+
+# Note, we are not running or submitting the ``WorkGraph`` in the above two examples, as we are using mocked
+# ``Computer`` and ``Code`` entities for demonstration purposes.
+
 # %%
 # Running a shell command with arguments
 # =======================================
@@ -88,7 +135,7 @@
 load_profile()
 
 
-def parser(self, dirpath):
+def parser(dirpath):
     from aiida.orm import Int
 
     return {"result": Int((dirpath / "stdout").read_text().strip())}
@@ -137,5 +184,6 @@ def parser(self, dirpath):
 # %%
 # What's Next
 # ===========
-# For more examples of ``aiida-shell``, please refer to its `docs <https://aiida-shell.readthedocs.io/en/latest/howto.html#>`_.
+# Overall, WorkGraph's ``ShellJob`` builds on top of ``aiida-shell`` and mirrors its functionality. So features
+# implemented in ``aiida-shell`` should also be available for the ``ShellJob``. For more examples of ``aiida-shell``, please refer to its `docs <https://aiida-shell.readthedocs.io/en/latest/howto.html#>`_.
 #

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
     "node-graph>=0.1.0",
     "aiida-core>=2.3",
     "cloudpickle",
-    "aiida-shell==0.7.3",
+    "aiida-shell~=0.8",
     "fastapi",
     "uvicorn",
     "pydantic_settings",

tests/test_shell.py

@@ -75,7 +75,7 @@ def test_shell_graph_builder():
     def add_multiply(x, y):
         """Add two numbers and multiply the result by 2."""
         # define the parser function
-        def parser(self, dirpath):
+        def parser(dirpath):
             from aiida.orm import Int
 
             return {"result": Int((dirpath / "stdout").read_text().strip())}
@@ -99,7 +99,7 @@ def parser(self, dirpath):
             command="bc",
             arguments=["{expression}"],
             nodes={"expression": job1.outputs["stdout"]},
-            parser=(parser),
+            parser=parser,
             parser_outputs=[
                 {"identifier": "workgraph.any", "name": "result"}
             ],  # add a "result" output socket from the parser