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

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