Add option to use double quotes for Code and Computer CLI argumen…

…ts (#5478)

So far, the command line arguments of the code exeuction line in the
submit script written by the `Scheduler` plugin, would be wrapped by
single quotes to prevent special characters, such as spaces, from being
interpreted incorrectly. However, for certain use cases, the quoting has
an undesirable effect. For example, if an argument contains a bash
variable, indicated by a `$` sign, the single quotes will prevent it
from being dereferenced correctly.

A concrete use-case is the requested feature of being able to define a
Docker container as a `Code`, which needs certain arguments to be quoted
in double quotes. This PR adds a new property `use_double_quotes` to the
`Code` and `Computer` classes that determines whether arguments should
be quoted using single or double quotes.

The settings are communicated through the `cmdline_params` attribute of
a `JobTemplateCodeInfo` instance to the scheduler plugin. The plugin is
then responsible for respecting these booleans. To be able to
distinguish between the command line arguments of the computer and the
code, they are passed separately in `prepend_cmdline_params` and
`cmdline_params`, respectively.

Co-authored-by: Sebastiaan Huber <mail@sphuber.net>

aiida/cmdline/commands/cmd_code.py

@@ -72,6 +72,7 @@ def set_code_builder(ctx, param, value):
 @options_code.REMOTE_ABS_PATH()
 @options_code.FOLDER()
 @options_code.REL_PATH()
+@options_code.USE_DOUBLE_QUOTES()
 @options_code.PREPEND_TEXT()
 @options_code.APPEND_TEXT()
 @options.NON_INTERACTIVE()

aiida/cmdline/commands/cmd_computer.py

@@ -204,6 +204,7 @@ def set_computer_builder(ctx, param, value):
 @options_computer.MPI_RUN_COMMAND()
 @options_computer.MPI_PROCS_PER_MACHINE()
 @options_computer.DEFAULT_MEMORY_PER_MACHINE()
+@options_computer.USE_DOUBLE_QUOTES()
 @options_computer.PREPEND_TEXT()
 @options_computer.APPEND_TEXT()
 @options.NON_INTERACTIVE()

aiida/cmdline/params/options/commands/code.py

@@ -109,6 +109,15 @@ def validate_label_uniqueness(ctx, _, value):
     help='[if --store-in-db]: Relative path of the executable inside the code-folder.'
 )
 
+USE_DOUBLE_QUOTES = OverridableOption(
+    '--use-double-quotes/--not-use-double-quotes',
+    default=False,
+    cls=InteractiveOption,
+    prompt='Escape CLI arguments in double quotes',
+    help='Whether the executable and arguments of the code in the submission script should be escaped with single '
+    'or double quotes.'
+)
+
 LABEL = options.LABEL.clone(
     prompt='Label',
     callback=validate_label_uniqueness,

aiida/cmdline/params/options/commands/computer.py

@@ -117,6 +117,15 @@ def should_call_default_mpiprocs_per_machine(ctx):  # pylint: disable=invalid-na
     help='The default amount of RAM (kB) that should be allocated per machine (node), if not otherwise specified.'
 )
 
+USE_DOUBLE_QUOTES = OverridableOption(
+    '--use-double-quotes/--not-use-double-quotes',
+    default=False,
+    cls=InteractiveOption,
+    prompt='Escape CLI arguments in double quotes',
+    help='Whether the command line arguments before and after the executable in the submission script should be '
+    'escaped with single or double quotes.'
+)
+
 PREPEND_TEXT = OverridableOption(
     '--prepend-text',
     cls=TemplateInteractiveOption,

aiida/engine/processes/calcjobs/calcjob.py

@@ -708,16 +708,16 @@ def presubmit(self, folder: Folder) -> CalcInfo:
                 this_withmpi = code_info.withmpi
 
             if this_withmpi:
-                this_argv = (
-                    mpi_args + extra_mpirun_params + [this_code.get_execname()] +
-                    (code_info.cmdline_params if code_info.cmdline_params is not None else [])
-                )
+                prepend_cmdline_params = mpi_args + extra_mpirun_params
             else:
-                this_argv = [this_code.get_execname()
-                             ] + (code_info.cmdline_params if code_info.cmdline_params is not None else [])
+                prepend_cmdline_params = []
+
+            cmdline_params = [this_code.get_execname()] + (code_info.cmdline_params or [])
 
             tmpl_code_info = JobTemplateCodeInfo()
-            tmpl_code_info.cmdline_params = this_argv
+            tmpl_code_info.prepend_cmdline_params = prepend_cmdline_params
+            tmpl_code_info.cmdline_params = cmdline_params
+            tmpl_code_info.use_double_quotes = [computer.get_use_double_quotes(), this_code.get_use_double_quotes()]
             tmpl_code_info.stdin_name = code_info.stdin_name
             tmpl_code_info.stdout_name = code_info.stdout_name
             tmpl_code_info.stderr_name = code_info.stderr_name

aiida/orm/computers.py

@@ -439,6 +439,22 @@ def get_append_text(self) -> str:
     def set_append_text(self, val: str) -> None:
         self.set_property('append_text', str(val))
 
+    def get_use_double_quotes(self) -> bool:
+        """Return whether the command line parameters of this computer should be escaped with double quotes.
+
+        :returns: True if to escape with double quotes, False otherwise which is also the default.
+        """
+        return self.get_property('use_double_quotes', False)
+
+    def set_use_double_quotes(self, val: bool) -> None:
+        """Set whether the command line parameters of this computer should be escaped with double quotes.
+
+        :param use_double_quotes: True if to escape with double quotes, False otherwise.
+        """
+        from aiida.common.lang import type_check
+        type_check(val, bool)
+        self.set_property('use_double_quotes', val)
+
     def get_mpirun_command(self) -> List[str]:
         """
         Return the mpirun command. Must be a list of strings, that will be

aiida/orm/nodes/data/code.py

@@ -351,6 +351,22 @@ def get_input_plugin_name(self):
         """
         return self.base.attributes.get('input_plugin', None)
 
+    def set_use_double_quotes(self, use_double_quotes: bool):
+        """Set whether the command line invocation of this code should be escaped with double quotes.
+
+        :param use_double_quotes: True if to escape with double quotes, False otherwise.
+        """
+        from aiida.common.lang import type_check
+        type_check(use_double_quotes, bool)
+        self.base.attributes.set('use_double_quotes', use_double_quotes)
+
+    def get_use_double_quotes(self) -> bool:
+        """Return whether the command line invocation of this code should be escaped with double quotes.
+
+        :returns: True if to escape with double quotes, False otherwise which is also the default.
+        """
+        return self.base.attributes.get('use_double_quotes', False)
+
     def set_append_text(self, code):
         """
         Pass a string of code that will be put in the scheduler script after the

aiida/orm/utils/builders/code.py

@@ -63,6 +63,7 @@ def new(self):
         code.label = self._get_and_count('label', used)
         code.description = self._get_and_count('description', used)
         code.set_input_plugin_name(self._get_and_count('input_plugin', used))
+        code.set_use_double_quotes(self._get_and_count('use_double_quotes', used))
         code.set_prepend_text(self._get_and_count('prepend_text', used))
         code.set_append_text(self._get_and_count('append_text', used))
 
@@ -98,6 +99,7 @@ def get_code_spec(code):
         spec['label'] = code.label
         spec['description'] = code.description
         spec['input_plugin'] = code.get_input_plugin_name()
+        spec['use_double_quotes'] = code.get_use_double_quotes()
         spec['prepend_text'] = code.get_prepend_text()
         spec['append_text'] = code.get_append_text()
 

aiida/orm/utils/builders/computer.py

@@ -39,6 +39,7 @@ def get_computer_spec(computer):
         spec['hostname'] = computer.hostname
         spec['scheduler'] = computer.scheduler_type
         spec['transport'] = computer.transport_type
+        spec['use_double_quotes'] = computer.get_use_double_quotes()
         spec['prepend_text'] = computer.get_prepend_text()
         spec['append_text'] = computer.get_append_text()
         spec['work_dir'] = computer.get_workdir()
@@ -76,6 +77,7 @@ def new(self):
         computer.description = self._get_and_count('description', used)
         computer.scheduler_type = self._get_and_count('scheduler', used)
         computer.transport_type = self._get_and_count('transport', used)
+        computer.set_use_double_quotes(self._get_and_count('use_double_quotes', used))
         computer.set_prepend_text(self._get_and_count('prepend_text', used))
         computer.set_append_text(self._get_and_count('append_text', used))
         computer.set_workdir(self._get_and_count('work_dir', used))

aiida/schedulers/datastructures.py

@@ -370,13 +370,18 @@ class JobTemplateCodeInfo:
     `Scheduler.get_submit_script` will pass a list of these objects to `Scheduler._get_run_line` which
     should build up the code execution line based on the parameters specified in this dataclass.
 
+    :param preprend_cmdline_params: list of unescaped command line arguments that are to be prepended to the executable.
     :param cmdline_params: list of unescaped command line parameters.
+    :param use_double_quotes: list of two booleans. If true, use double quotes to escape command line arguments. The
+        first value applies to `prepend_cmdline_params` and the second to `cmdline_params`.
     :param stdin_name: filename of the the stdin file descriptor.
     :param stdout_name: filename of the the `stdout` file descriptor.
     :param stderr_name: filename of the the `stderr` file descriptor.
     :param join_files: boolean, if true, `stderr` should be redirected to `stdout`.
     """
+    prepend_cmdline_params: list[str] = field(default_factory=list)
     cmdline_params: list[str] = field(default_factory=list)
+    use_double_quotes: list[bool] = field(default_factory=lambda: [False, False])
     stdin_name: None | str = None
     stdout_name: None | str = None
     stderr_name: None | str = None

aiida/schedulers/scheduler.py

@@ -219,19 +219,28 @@ def _get_run_line(self, codes_info, codes_run_mode):
         list_of_runlines = []
 
         for code_info in codes_info:
+            computer_use_double_quotes = code_info.use_double_quotes[0]
+            code_use_double_quotes = code_info.use_double_quotes[1]
+
             command_to_exec_list = []
+            for arg in code_info.prepend_cmdline_params:
+                command_to_exec_list.append(escape_for_bash(arg, use_double_quotes=computer_use_double_quotes))
             for arg in code_info.cmdline_params:
-                command_to_exec_list.append(escape_for_bash(arg))
+                command_to_exec_list.append(escape_for_bash(arg, use_double_quotes=code_use_double_quotes))
             command_to_exec = ' '.join(command_to_exec_list)
 
-            stdin_str = f'< {escape_for_bash(code_info.stdin_name)}' if code_info.stdin_name else ''
-            stdout_str = f'> {escape_for_bash(code_info.stdout_name)}' if code_info.stdout_name else ''
+            escape_stdin_name = escape_for_bash(code_info.stdin_name, use_double_quotes=computer_use_double_quotes)
+            escape_stdout_name = escape_for_bash(code_info.stdout_name, use_double_quotes=computer_use_double_quotes)
+            escape_sterr_name = escape_for_bash(code_info.stderr_name, use_double_quotes=computer_use_double_quotes)
+
+            stdin_str = f'< {escape_stdin_name}' if code_info.stdin_name else ''
+            stdout_str = f'> {escape_stdout_name}' if code_info.stdout_name else ''
 
             join_files = code_info.join_files
             if join_files:
                 stderr_str = '2>&1'
             else:
-                stderr_str = f'2> {escape_for_bash(code_info.stderr_name)}' if code_info.stderr_name else ''
+                stderr_str = f'2> {escape_sterr_name}' if code_info.stderr_name else ''
 
             output_string = f'{command_to_exec} {stdin_str} {stdout_str} {stderr_str}'
 

tests/cmdline/commands/test_code.py

@@ -150,8 +150,13 @@ def test_noninteractive_remote(run_cli_command, aiida_localhost, non_interactive
     """Test non-interactive remote code setup."""
     label = 'noninteractive_remote'
     options = [
-        '--non-interactive', f'--label={label}', '--description=description', '--input-plugin=core.arithmetic.add',
-        '--on-computer', f'--computer={aiida_localhost.label}', '--remote-abs-path=/remote/abs/path'
+        '--non-interactive',
+        f'--label={label}',
+        '--description=description',
+        '--input-plugin=core.arithmetic.add',
+        '--on-computer',
+        f'--computer={aiida_localhost.label}',
+        '--remote-abs-path=/remote/abs/path',
     ]
     run_cli_command(cmd_code.setup_code, options)
     assert isinstance(load_code(label), Code)

tests/cmdline/commands/test_computer.py

@@ -351,6 +351,7 @@ def init_profile(self, aiida_profile_clean, run_cli_command):  # pylint: disable
         self.comp_builder.description = 'Test Computer'
         self.comp_builder.scheduler = 'core.direct'
         self.comp_builder.work_dir = '/tmp/aiida'
+        self.comp_builder.use_double_quotes = False
         self.comp_builder.prepend_text = ''
         self.comp_builder.append_text = ''
         self.comp_builder.mpiprocs_per_machine = 8

tests/engine/processes/calcjobs/test_calc_job.py

@@ -37,6 +37,30 @@ def raise_exception(exception, *args, **kwargs):
     raise exception()
 
 
+@pytest.mark.requires_rmq
+class DummyCalcJob(CalcJob):
+    """`DummyCalcJob` implementation to test the calcinfo with container code."""
+
+    @classmethod
+    def define(cls, spec):
+        super().define(spec)
+
+    def prepare_for_submission(self, folder):
+        from aiida.common.datastructures import CalcInfo, CodeInfo
+
+        codeinfo = CodeInfo()
+        codeinfo.code_uuid = self.inputs.code.uuid
+        codeinfo.cmdline_params = ['--version', '-c']
+        codeinfo.stdin_name = 'aiida.in'
+        codeinfo.stdout_name = 'aiida.out'
+        codeinfo.stderr_name = 'aiida.err'
+
+        calcinfo = CalcInfo()
+        calcinfo.codes_info = [codeinfo]
+
+        return calcinfo
+
+
 @pytest.mark.requires_rmq
 class FileCalcJob(CalcJob):
     """Example `CalcJob` implementation to test the `provenance_exclude_list` functionality.
@@ -136,6 +160,72 @@ def test_multi_codes_run_parallel(aiida_local_code_factory, file_regression, par
     file_regression.check(content, extension='.sh')
 
 
+@pytest.mark.requires_rmq
+@pytest.mark.usefixtures('clear_database_before_test', 'chdir_tmp_path')
+@pytest.mark.parametrize('computer_use_double_quotes', [True, False])
+def test_computer_double_quotes(aiida_local_code_factory, file_regression, computer_use_double_quotes):
+    """test that bash script quote escape behaviour can be controlled"""
+    computer = orm.Computer(
+        label='test-code-computer',
+        transport_type='core.local',
+        hostname='localhost',
+        scheduler_type='core.direct',
+    ).store()
+    computer.set_use_double_quotes(computer_use_double_quotes)
+
+    inputs = {
+        'code': aiida_local_code_factory('core.arithmetic.add', '/bin/bash', computer),
+        'metadata': {
+            'dry_run': True,
+            'options': {
+                'resources': {
+                    'num_machines': 1,
+                    'num_mpiprocs_per_machine': 1
+                },
+                'withmpi': True,
+            }
+        }
+    }
+
+    _, node = launch.run_get_node(DummyCalcJob, **inputs)
+    folder_name = node.dry_run_info['folder']
+    submit_script_filename = node.get_option('submit_script_filename')
+    with open(os.path.join(folder_name, submit_script_filename), encoding='utf8') as handle:
+        content = handle.read()
+
+    file_regression.check(content, extension='.sh')
+
+
+@pytest.mark.requires_rmq
+@pytest.mark.usefixtures('clear_database_before_test', 'chdir_tmp_path')
+@pytest.mark.parametrize('code_use_double_quotes', [True, False])
+def test_code_double_quotes(aiida_localhost, file_regression, code_use_double_quotes):
+    """test that bash script quote escape behaviour can be controlled"""
+    code = orm.Code(remote_computer_exec=(aiida_localhost, '/bin/bash'))
+    code.set_use_double_quotes(code_use_double_quotes)
+    inputs = {
+        'code': code.store(),
+        'metadata': {
+            'dry_run': True,
+            'options': {
+                'resources': {
+                    'num_machines': 1,
+                    'num_mpiprocs_per_machine': 1
+                },
+                'withmpi': True,
+            }
+        }
+    }
+
+    _, node = launch.run_get_node(DummyCalcJob, **inputs)
+    folder_name = node.dry_run_info['folder']
+    submit_script_filename = node.get_option('submit_script_filename')
+    with open(os.path.join(folder_name, submit_script_filename), encoding='utf8') as handle:
+        content = handle.read()
+
+    file_regression.check(content, extension='.sh')
+
+
 @pytest.mark.requires_rmq
 @pytest.mark.usefixtures('aiida_profile_clean', 'chdir_tmp_path')
 @pytest.mark.parametrize('calcjob_withmpi', [True, False])

tests/engine/processes/calcjobs/test_calc_job/test_code_double_quotes_False_.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+exec > _scheduler-stdout.txt
+exec 2> _scheduler-stderr.txt
+
+
+'mpirun' '-np' '1' '/bin/bash' '--version' '-c' < 'aiida.in' > 'aiida.out' 2> 'aiida.err'

tests/engine/processes/calcjobs/test_calc_job/test_code_double_quotes_True_.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+exec > _scheduler-stdout.txt
+exec 2> _scheduler-stderr.txt
+
+
+'mpirun' '-np' '1' "/bin/bash" "--version" "-c" < 'aiida.in' > 'aiida.out' 2> 'aiida.err'

tests/engine/processes/calcjobs/test_calc_job/test_computer_double_quotes_False_.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+exec > _scheduler-stdout.txt
+exec 2> _scheduler-stderr.txt
+
+
+'mpirun' '-np' '1' '/bin/bash' '--version' '-c' < 'aiida.in' > 'aiida.out' 2> 'aiida.err'

tests/engine/processes/calcjobs/test_calc_job/test_computer_double_quotes_True_.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+exec > _scheduler-stdout.txt
+exec 2> _scheduler-stderr.txt
+
+
+"mpirun" "-np" "1" '/bin/bash' '--version' '-c' < "aiida.in" > "aiida.out" 2> "aiida.err"

tests/orm/test_computers.py

@@ -99,6 +99,7 @@ def init_profile(self, aiida_profile_clean):  # pylint: disable=unused-argument
         self.comp_builder.work_dir = '/tmp/aiida'
         self.comp_builder.prepend_text = ''
         self.comp_builder.append_text = ''
+        self.comp_builder.use_double_quotes = False
         self.comp_builder.mpiprocs_per_machine = 8
         self.comp_builder.default_memory_per_machine = 1000000
         self.comp_builder.mpirun_command = 'mpirun'