Add support for Docker images to use as Code for CalcJobs (#5841)

The `ContainerizedCode` is currently not quite compatible with Docker's
containerization technology. The reason is that the executable,
including all of its command line arguments and the file descriptor
redirections need to be part of a single quoted string, passed to the
`bash -c` command that is run inside the container.

For example, to run `pw.x` the submit line needs to look like:

    docker run -i {image_name} sh -c "pw.x -input input.in"

To enable this, a new attribute is added to the `JobTemplateCodeInfo`
dataclass. This instructs the `Scheduler` to wrap the executable and all
its arguments in quotes. The `AbstractCode` has a new attribute with the
same name that is `False` by default to keep backwards-compatibility
but which can be set to `True` to enable compatiblity with Docker.

The attribute has an associated getter and setter, which is added to the
`AbstractCode` and not the `ContainerizedCode` because the `Scheduler`
plugin will attempt to retrieve this attribute for all code types. If
the properties would be added to `ContainerizedCode`, the scheduler
plugin would except with an `AttributeError`.

aiida/engine/processes/calcjobs/calcjob.py

@@ -938,6 +938,7 @@ def presubmit(self, folder: Folder) -> CalcInfo:
             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.use_double_quotes]
+            tmpl_code_info.wrap_cmdline_params = this_code.wrap_cmdline_params
             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/nodes/data/code/abstract.py

@@ -40,6 +40,7 @@ class AbstractCode(Data, metaclass=abc.ABCMeta):
     _KEY_ATTRIBUTE_APPEND_TEXT: str = 'append_text'
     _KEY_ATTRIBUTE_PREPEND_TEXT: str = 'prepend_text'
     _KEY_ATTRIBUTE_USE_DOUBLE_QUOTES: str = 'use_double_quotes'
+    _KEY_ATTRIBUTE_WRAP_CMDLINE_PARAMS: str = 'wrap_cmdline_params'
     _KEY_EXTRA_IS_HIDDEN: str = 'hidden'  # Should become ``is_hidden`` once ``Code`` is dropped
 
     def __init__(
@@ -49,6 +50,7 @@ def __init__(
         prepend_text: str = '',
         use_double_quotes: bool = False,
         is_hidden: bool = False,
+        wrap_cmdline_params: bool = False,
         **kwargs
     ):
         """Construct a new instance.
@@ -57,13 +59,16 @@ def __init__(
         :param append_text: The text that should be appended to the run line in the job script.
         :param prepend_text: The text that should be prepended to the run line in the job script.
         :param use_double_quotes: Whether the command line invocation of this code should be escaped with double quotes.
+        :param wrap_cmdline_params: Whether to wrap the executable and all its command line parameters into quotes to
+            form a single string. This is required to enable support for Docker with the ``ContainerizedCode``.
         :param is_hidden: Whether the code is hidden.
         """
         super().__init__(**kwargs)
         self.default_calc_job_plugin = default_calc_job_plugin
         self.append_text = append_text
         self.prepend_text = prepend_text
         self.use_double_quotes = use_double_quotes
+        self.wrap_cmdline_params = wrap_cmdline_params
         self.is_hidden = is_hidden
 
     @abc.abstractmethod
@@ -220,6 +225,25 @@ def use_double_quotes(self, value: bool) -> None:
         type_check(value, bool)
         self.base.attributes.set(self._KEY_ATTRIBUTE_USE_DOUBLE_QUOTES, value)
 
+    @property
+    def wrap_cmdline_params(self) -> bool:
+        """Return whether all command line parameters should be wrapped with double quotes to form a single argument.
+
+        ..note:: This is required to support certain containerization technologies, such as Docker.
+
+        :return: ``True`` if command line parameters should be wrapped, ``False`` otherwise.
+        """
+        return self.base.attributes.get(self._KEY_ATTRIBUTE_WRAP_CMDLINE_PARAMS, False)
+
+    @wrap_cmdline_params.setter
+    def wrap_cmdline_params(self, value: bool) -> None:
+        """Set whether all command line parameters should be wrapped with double quotes to form a single argument.
+
+        :param value: ``True`` if command line parameters should be wrapped, ``False`` otherwise.
+        """
+        type_check(value, bool)
+        self.base.attributes.set(self._KEY_ATTRIBUTE_WRAP_CMDLINE_PARAMS, value)
+
     @property
     def is_hidden(self) -> bool:
         """Return whether the code is hidden.

aiida/orm/nodes/data/code/containerized.py

@@ -58,15 +58,15 @@ def filepath_executable(self, value: str) -> None:
 
     @property
     def engine_command(self) -> str:
-        """Return the engine command with image as template field of the containerized code
+        """Return the engine command with image as template field of the containerized code.
 
         :return: The engine command of the containerized code
         """
         return self.base.attributes.get(self._KEY_ATTRIBUTE_ENGINE_COMMAND)
 
     @engine_command.setter
     def engine_command(self, value: str) -> None:
-        """Set the engine command of the containerized code
+        """Set the engine command of the containerized code.
 
         :param value: The engine command of the containerized code
         """
@@ -79,20 +79,19 @@ def engine_command(self, value: str) -> None:
 
     @property
     def image_name(self) -> str:
-        """The image name of container
+        """The image name of container.
 
         :return: The image name of container.
         """
         return self.base.attributes.get(self._KEY_ATTRIBUTE_IMAGE_NAME)
 
     @image_name.setter
     def image_name(self, value: str) -> None:
-        """Set the image name of container
+        """Set the image name of container.
 
         :param value: The image name of container.
         """
         type_check(value, str)
-
         self.base.attributes.set(self._KEY_ATTRIBUTE_IMAGE_NAME, value)
 
     def get_prepend_cmdline_params(
@@ -125,6 +124,13 @@ def _get_cli_options(cls) -> dict:
                 'prompt': 'Image name',
                 'help': 'Name of the image container in which to the run the executable.',
             },
+            'wrap_cmdline_params': {
+                'is_flag': True,
+                'default': False,
+                'help': 'Whether all command line parameters to be passed to the engine command should be wrapped in '
+                'a double quotes to form a single argument. This should be set to `True` for Docker.',
+                'prompt': 'Wrap command line parameters',
+            }
         }
         options.update(**super()._get_cli_options())
 

aiida/schedulers/datastructures.py

@@ -379,6 +379,10 @@ class JobTemplateCodeInfo:
     :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 wrap_cmdline_params: Boolean, by default ``False``. If set to ``True``, all the command line arguments,
+        which includes the ``cmdline_params`` but also all file descriptor redirections (stdin, stderr and stdoout),
+        should be wrapped in double quotes, turning it into a single command line argument. This is necessary to enable
+        support for certain containerization technologies such as Docker.
     :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.
@@ -387,6 +391,7 @@ class JobTemplateCodeInfo:
     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])
+    wrap_cmdline_params: bool = False
     stdin_name: None | str = None
     stdout_name: None | str = None
     stderr_name: None | str = None

aiida/schedulers/scheduler.py

@@ -222,18 +222,20 @@ def _get_run_line(self, codes_info: list[JobTemplateCodeInfo], codes_run_mode: C
             to launch the multiple codes.
         :return: string with format: [executable] [args] {[ < stdin ]} {[ < stdout ]} {[2>&1 | 2> stderr]}
         """
+        # pylint: disable=too-many-locals
         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 = []
+            prepend_cmdline_params = []
             for arg in code_info.prepend_cmdline_params:
-                command_to_exec_list.append(escape_for_bash(arg, use_double_quotes=computer_use_double_quotes))
+                prepend_cmdline_params.append(escape_for_bash(arg, use_double_quotes=computer_use_double_quotes))
+
+            cmdline_params = []
             for arg in code_info.cmdline_params:
-                command_to_exec_list.append(escape_for_bash(arg, use_double_quotes=code_use_double_quotes))
-            command_to_exec = ' '.join(command_to_exec_list)
+                cmdline_params.append(escape_for_bash(arg, use_double_quotes=code_use_double_quotes))
 
             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)
@@ -248,9 +250,17 @@ def _get_run_line(self, codes_info: list[JobTemplateCodeInfo], codes_run_mode: C
             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}'
+            cmdline_params.extend([stdin_str, stdout_str, stderr_str])
+
+            prepend_cmdline_params_string = ' '.join(prepend_cmdline_params)
+            cmdline_params_string = ' '.join(cmdline_params)
+
+            if code_info.wrap_cmdline_params:
+                cmdline_params_string = escape_for_bash(cmdline_params_string, use_double_quotes=True)
+
+            run_line = f'{prepend_cmdline_params_string} {cmdline_params_string}'.strip()
 
-            list_of_runlines.append(output_string)
+            list_of_runlines.append(run_line)
 
         self.logger.debug(f'_get_run_line output: {list_of_runlines}')
 

docs/source/topics/data_types.rst

@@ -610,29 +610,53 @@ The ``ContainerizedCode`` is compatible with a variety of containerization techn
 
 .. tab-set::
 
-    .. tab-item:: Singularity
+    .. tab-item:: Docker
 
-        To use `Singularity <https://singularity-docs.readthedocs.io/en/latest/>`__ use the following ``engine_command`` when setting up the code:
+        To use `Docker <https://www.docker.com/>`_ ``aiida-core==2.3.0`` or higher is required in order to be able to set ``wrap_cmdline_params = True``.
+        When setting up a code for a Docker container, use the following ``engine_command`` when setting up the code:
 
         .. code-block:: console
 
-            singularity exec --bind $PWD:$PWD {image_name}
+            docker run -i -v $PWD:/workdir:rw -w /workdir {image_name} sh -c
 
-    .. tab-item:: Sarus
+        .. note:: Currently running with MPI is not yet supported, as it needs to be called inside of the container which is currently not possible.
+            The associated computer should also be configured to have the setting ``use_double_quotes = False``.
+            This can be set from the Python API using ``load_computer('idenfitier').set_use_double_quotes(False)``.
+
+        The following configuration provides an example to setup Quantum ESPRESSO's ``pw.x`` to be run by Docker on the local host
+
+        .. code-block:: yaml
 
-        To use `Sarus <https://sarus.readthedocs.io/en/stable/>`__ use the following ``engine_command`` when setting up the code:
+            label: qe-pw-on-docker
+            computer: localhost
+            engine_command: docker run -i -v $PWD:/workdir:rw -w /workdir {image_name} sh -c
+            image_name: haya4kun/quantum_espresso
+            filepath_executable: pw.x
+            default_calc_job_plugin: quantumespresso.pw
+            use_double_quotes: false
+            wrap_cmdline_params: true
+
+        Save the configuration to ``code.yml`` and create the code using the ``verdi`` CLI:
 
         .. code-block:: console
 
-            sarus run --mount=src=$PWD,dst=/workdir,type=bind --workdir=/workdir {image_name}
+            verdi code create core.code.containerized -n --config=code.yml
 
+    .. tab-item:: Singularity
 
-Using `Docker <https://www.docker.com/>`__ directly is currently not supported because:
+        To use `Singularity <https://singularity-docs.readthedocs.io/en/latest/>`_ use the following ``engine_command`` when setting up the code:
 
-* The Docker daemon always runs as the root user and the files created in the working directory inside the container will usually be owned by root if uid is not specified in the image, which prevents AiiDA from deleting those files after execution.
-* Docker cannot be launched as a normal MPI program to propagate execution context to the container application.
+        .. code-block:: console
 
-Support may be added at a later time.
+            singularity exec --bind $PWD:$PWD {image_name}
+
+    .. tab-item:: Sarus
+
+        To use `Sarus <https://sarus.readthedocs.io/en/stable/>`_ use the following ``engine_command`` when setting up the code:
+
+        .. code-block:: console
+
+            sarus run --mount=src=$PWD,dst=/workdir,type=bind --workdir=/workdir {image_name}
 
 
 

tests/engine/processes/calcjobs/test_calc_job.py

@@ -269,7 +269,44 @@ def test_containerized_code(file_regression, aiida_localhost):
     _, node = launch.run_get_node(DummyCalcJob, **inputs)
     folder_name = node.dry_run_info['folder']
     submit_script_filename = node.get_option('submit_script_filename')
-    content = (pathlib.Path(folder_name) / submit_script_filename).read_bytes().decode('utf-8')
+    content = (pathlib.Path(folder_name) / submit_script_filename).read_text()
+
+    file_regression.check(content, extension='.sh')
+
+
+@pytest.mark.requires_rmq
+@pytest.mark.usefixtures('chdir_tmp_path')
+def test_containerized_code_wrap_cmdline_params(file_regression, aiida_localhost):
+    """Test :class:`~aiida.orm.nodes.data.code.containerized.ContainerizedCode` with ``wrap_cmdline_params = True``."""
+    aiida_localhost.set_use_double_quotes(False)
+    engine_command = """docker run -i -v $PWD:/workdir:rw -w /workdir {image_name} sh -c"""
+    containerized_code = orm.ContainerizedCode(
+        default_calc_job_plugin='core.arithmetic.add',
+        filepath_executable='/bin/bash',
+        engine_command=engine_command,
+        image_name='ubuntu',
+        computer=aiida_localhost,
+        wrap_cmdline_params=True,
+    ).store()
+
+    inputs = {
+        'code': containerized_code,
+        'metadata': {
+            'dry_run': True,
+            'options': {
+                'resources': {
+                    'num_machines': 1,
+                    'num_mpiprocs_per_machine': 1
+                },
+                'withmpi': False,
+            }
+        }
+    }
+
+    _, node = launch.run_get_node(DummyCalcJob, **inputs)
+    folder_name = node.dry_run_info['folder']
+    submit_script_filename = node.get_option('submit_script_filename')
+    content = (pathlib.Path(folder_name) / submit_script_filename).read_text()
 
     file_regression.check(content, extension='.sh')
 
@@ -305,7 +342,7 @@ def test_containerized_code_withmpi_true(file_regression, aiida_localhost):
     _, node = launch.run_get_node(DummyCalcJob, **inputs)
     folder_name = node.dry_run_info['folder']
     submit_script_filename = node.get_option('submit_script_filename')
-    content = (pathlib.Path(folder_name) / submit_script_filename).read_bytes().decode('utf-8')
+    content = (pathlib.Path(folder_name) / submit_script_filename).read_text()
 
     file_regression.check(content, extension='.sh')
 
@@ -379,9 +416,9 @@ def test_portable_code(tmp_path, aiida_localhost):
     for filename in code.base.repository.list_object_names():
         assert filename in uploaded_files
 
-    content = (pathlib.Path(folder_name) / code.filepath_executable).read_bytes().decode('utf-8')
-    subcontent = (pathlib.Path(folder_name) / 'sub' / 'dummy').read_bytes().decode('utf-8')
-    subsubcontent = (pathlib.Path(folder_name) / 'sub' / 'sub' / 'sub-dummy').read_bytes().decode('utf-8')
+    content = (pathlib.Path(folder_name) / code.filepath_executable).read_text()
+    subcontent = (pathlib.Path(folder_name) / 'sub' / 'dummy').read_text()
+    subsubcontent = (pathlib.Path(folder_name) / 'sub' / 'sub' / 'sub-dummy').read_text()
 
     assert content == 'bash implementation'
     assert subcontent == 'dummy'

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

@@ -0,0 +1,6 @@
+#!/bin/bash
+exec > _scheduler-stdout.txt
+exec 2> _scheduler-stderr.txt
+
+
+'docker' 'run' '-i' '-v' '$PWD:/workdir:rw' '-w' '/workdir' 'ubuntu' 'sh' '-c' "'/bin/bash' '--version' '-c' < 'aiida.in' > 'aiida.out' 2> 'aiida.err'"

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

@@ -3,6 +3,6 @@ exec > _scheduler-stdout.txt
 exec 2> _scheduler-stderr.txt
 
 
-'/bin/bash'   
+'/bin/bash'
 
-'/bin/bash'   
+'/bin/bash'

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

@@ -3,9 +3,9 @@ exec > _scheduler-stdout.txt
 exec 2> _scheduler-stderr.txt
 
 
-'/bin/bash'    &
+'/bin/bash' &
 
-'/bin/bash'    &
+'/bin/bash' &
 
 wait
 

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

@@ -3,6 +3,6 @@ exec > _scheduler-stdout.txt
 exec 2> _scheduler-stderr.txt
 
 
-'/bin/bash'   
+'/bin/bash'
 
-'/bin/bash'   
+'/bin/bash'

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

@@ -3,6 +3,6 @@ exec > _scheduler-stdout.txt
 exec 2> _scheduler-stderr.txt
 
 
-'/bin/bash'   
+'/bin/bash'
 
-'mpirun' '-np' '1' '/bin/bash'   
+'mpirun' '-np' '1' '/bin/bash'