diff --git a/aiida_workgraph/engine/utils.py b/aiida_workgraph/engine/utils.py index 54e4b558..aefb56d4 100644 --- a/aiida_workgraph/engine/utils.py +++ b/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 diff --git a/aiida_workgraph/utils/__init__.py b/aiida_workgraph/utils/__init__.py index 4e958658..465e2f82 100644 --- a/aiida_workgraph/utils/__init__.py +++ b/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( diff --git a/docs/gallery/built-in/autogen/shelljob.py b/docs/gallery/built-in/autogen/shelljob.py index 50076ec6..7632ef71 100644 --- a/docs/gallery/built-in/autogen/shelljob.py +++ b/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 `_. +# 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 `_. # diff --git a/pyproject.toml b/pyproject.toml index 88ce268e..584970f2 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/test_shell.py b/tests/test_shell.py index 81246d75..6b836da1 100644 --- a/tests/test_shell.py +++ b/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