From f03ff5597cc1d9ec84c3f9ccb7cc2a1cb6d2ae49 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Tue, 19 Nov 2024 18:43:04 +0100 Subject: [PATCH] Avoid passing `None` to arguments of `prepare_shell_job_inputs` (#351) Originally noted in #348, and fixes #348. The reason for the bug was that in this line: https://github.com/aiidateam/aiida-workgraph/blob/8f6260b44595f79471f6ea70e8efe0f0e158cdab/aiida_workgraph/engine/utils.py#L139 iteration was done over all input arguments of the `prepare_shell_job_inputs` function, using `None` as the default for the ones that weren't contained in the `inputs` dictionary passed to the `prepare_for_shell_task` function of WorkGraph. This would lead to `None` values being passed explicitly to the arguments of `prepare_shell_job_inputs`, in particular `resolve_command`, effectively overriding its default `True` value. Inside `aiida-shell` the check for the absolute path of the executable inside the `prepare_code` function (see [here](https://github.com/sphuber/aiida-shell/blob/14866d1450aa252ec414373e86e98611e2eae9db/src/aiida_shell/launch.py#L176-L184)) would thus be skipped, and the WorkGraph would just fail silently at a later stage. With this PR, only the input arguments that are also contained in the `inputs` of the `prepare_for_shell_task` are being explicitly passed, while the rest is left unset. --- aiida_workgraph/engine/utils.py | 22 ++++++++++++++++++++-- tests/test_shell.py | 10 ++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/aiida_workgraph/engine/utils.py b/aiida_workgraph/engine/utils.py index aefb56d4..c14b4915 100644 --- a/aiida_workgraph/engine/utils.py +++ b/aiida_workgraph/engine/utils.py @@ -136,8 +136,26 @@ def prepare_for_shell_task(task: dict, inputs: dict) -> dict: # 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)) + aiida_shell_input_keys = signature.parameters.keys() + + # Iterate over all WorkGraph `inputs`, and extract the ones which are expected by `prepare_shell_job_inputs` + inputs_aiida_shell_subset = { + key: inputs[key] for key in inputs.keys() if key in aiida_shell_input_keys + } + + try: + aiida_shell_inputs = prepare_shell_job_inputs(**inputs_aiida_shell_subset) + except ValueError: + raise + + # We need to remove the original input-keys, as they might be offending for the call to `launch_shell_job` + # E.g., `inputs` originally can contain `command`, which gets, however, transformed to # + # `code` by `prepare_shell_job_inputs` + for key in inputs_aiida_shell_subset.keys(): + inputs.pop(key) + + # Finally, we update the original `inputs` with the modified ones from the call to `prepare_shell_job_inputs` + inputs = {**inputs, **aiida_shell_inputs} inputs.setdefault("metadata", {}) inputs["metadata"].update({"call_link_label": task["name"]}) diff --git a/tests/test_shell.py b/tests/test_shell.py index 6b836da1..3160e71f 100644 --- a/tests/test_shell.py +++ b/tests/test_shell.py @@ -4,6 +4,16 @@ from aiida.orm import SinglefileData, load_computer +def test_prepare_for_shell_task_nonexistent(): + """Check that the `ValueError` raised by `aiida-shell` for a non-extistent executable is captured by WorkGraph.""" + from aiida_workgraph.engine.utils import prepare_for_shell_task + + task = {"name": "test"} + inputs = {"command": "abc42"} + with pytest.raises(ValueError, match="failed to determine the absolute path"): + prepare_for_shell_task(task=task, inputs=inputs) + + @pytest.mark.usefixtures("started_daemon_client") def test_shell_command(fixture_localhost): """Test the ShellJob with command as a string."""