Update parser signature for aiida-shell v0.8, extend ShellJob docs,…

… use `prepare_shell_job_inputs`

aiida_workgraph/engine/utils.py

@@ -131,34 +131,22 @@ def prepare_for_python_task(task: dict, kwargs: dict, var_kwargs: dict) -> dict:
 
 def prepare_for_shell_task(task: dict, kwargs: 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 (
+ convert_nodes_single_file_data,
+ prepare_shell_job_inputs,
+ )
 
- 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", {}))
+ metadata.update({"call_link_label": task["name"]})
+
  # find all keys in kwargs start with "nodes."
+ nodes = convert_nodes_single_file_data(kwargs.pop("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 {},
- }
+
+ kwargs.update({"nodes": nodes, "metadata": metadata})
+
+ inputs = prepare_shell_job_inputs(**kwargs)
+
  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

@@ -34,6 +34,52 @@
 # 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 +134,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 +183,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",