Add the ContainerizedCode data plugin (#5667)

This implementation of the `AbstractCode` interface allows running a
`CalcJob` within a container on a target computer. The data plugin
stores the name of the container image, the executable within that
container, the command with which the container should be launched and
the `Computer` on which the container can be run.

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

.github/config/add-singularity.yaml

@@ -0,0 +1,10 @@
+---
+label: add-singularity
+description: Bash run in Docker image through Singularity
+default_calc_job_plugin: core.arithmetic.add
+computer: localhost
+filepath_executable: /bin/sh
+image_name: docker://alpine:3
+engine_command: singularity exec --bind $PWD:$PWD {image_name}
+prepend_text: ' '
+append_text: ' '

.github/config/add.yaml

@@ -2,7 +2,6 @@
 label: add
 description: add
 default_calc_job_plugin: core.arithmetic.add
-on_computer: true
 computer: localhost
 filepath_executable: /bin/bash
 prepend_text: ' '

.github/config/doubler.yaml

@@ -2,7 +2,6 @@
 label: doubler
 description: doubler
 default_calc_job_plugin: core.templatereplacer
-on_computer: true
 computer: localhost
 filepath_executable: PLACEHOLDER_REMOTE_ABS_PATH_DOUBLER
 prepend_text: ' '

.github/config/localhost.yaml

@@ -6,6 +6,7 @@ transport: core.local
 scheduler: core.direct
 shebang: '#!/usr/bin/env bash'
 work_dir: PLACEHOLDER_WORK_DIR
+use_double_quotes: true
 mpirun_command: ' '
 mpiprocs_per_machine: 1
 prepend_text: ' '

.github/system_tests/test_containerized_code.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Test running a :class:`~aiida.orm.nodes.data.codes.containerized.ContainerizedCode` code."""
+from aiida import orm
+from aiida.engine import run_get_node
+
+
+def test_add_singularity():
+    """Test installed containerized code by add plugin"""
+    builder = orm.load_code('add-singularity@localhost').get_builder()
+    builder.x = orm.Int(4)
+    builder.y = orm.Int(6)
+    builder.metadata.options.resources = {'num_machines': 1, 'num_mpiprocs_per_machine': 1}
+
+    results, node = run_get_node(builder)
+
+    assert node.is_finished_ok
+    assert 'sum' in results
+    assert 'remote_folder' in results
+    assert 'retrieved' in results
+    assert results['sum'].value == 10
+
+
+def main():
+    test_add_singularity()
+
+
+if __name__ == '__main__':
+    main()

.github/workflows/nightly.yml

@@ -44,6 +44,9 @@ jobs:
 
         steps:
         -   uses: actions/checkout@v2
+        -   uses: eWaterCycle/setup-singularity@v7  # for containerized code test
+            with:
+                singularity-version: 3.8.7
 
         -   name: Cache Python dependencies
             uses: actions/cache@v1

.github/workflows/setup.sh

@@ -23,6 +23,7 @@ verdi computer configure core.local localhost --config "${CONFIG}/localhost-conf
 verdi computer test localhost
 verdi code create core.code.installed --non-interactive --config "${CONFIG}/doubler.yaml"
 verdi code create core.code.installed --non-interactive --config "${CONFIG}/add.yaml"
+verdi code create core.code.containerized --non-interactive --config "${CONFIG}/add-singularity.yaml"
 
 # set up slurm-ssh computer
 verdi computer setup --non-interactive --config "${CONFIG}/slurm-ssh.yaml"

.github/workflows/tests_nightly.sh

@@ -9,6 +9,7 @@ export PYTHONPATH="${PYTHONPATH}:${SYSTEM_TESTS}:${MODULE_POLISH}"
 
 verdi daemon start 4
 verdi -p test_aiida run ${SYSTEM_TESTS}/test_daemon.py
+verdi -p test_aiida run ${SYSTEM_TESTS}/test_containerized_code.py
 bash ${SYSTEM_TESTS}/test_polish_workchains.sh
 verdi daemon stop
 

aiida/orm/__init__.py

@@ -49,6 +49,7 @@
     'Comment',
     'Computer',
     'ComputerEntityLoader',
+    'ContainerizedCode',
     'DESCENDING',
     'Data',
     'Dict',

aiida/orm/nodes/__init__.py

@@ -31,6 +31,7 @@
     'CalculationNode',
     'CifData',
     'Code',
+    'ContainerizedCode',
     'Data',
     'Dict',
     'EnumData',

aiida/orm/nodes/data/__init__.py

@@ -43,6 +43,7 @@
     'Bool',
     'CifData',
     'Code',
+    'ContainerizedCode',
     'Data',
     'Dict',
     'EnumData',

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

@@ -7,13 +7,15 @@
 # pylint: disable=wildcard-import
 
 from .abstract import *
+from .containerized import *
 from .installed import *
 from .legacy import *
 from .portable import *
 
 __all__ = (
     'AbstractCode',
     'Code',
+    'ContainerizedCode',
     'InstalledCode',
     'PortableCode',
 )

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

@@ -0,0 +1,134 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Data plugin representing an executable code inside a container.
+
+The containerized code allows specifying a container image and the executable filepath within that container that is to
+be executed when a calculation job is run with this code.
+"""
+from __future__ import annotations
+
+import pathlib
+
+import click
+
+from aiida.common.lang import type_check
+
+from .installed import InstalledCode
+
+__all__ = ('ContainerizedCode',)
+
+
+class ContainerizedCode(InstalledCode):
+    """Data plugin representing an executable code in container on a remote computer."""
+    _KEY_ATTRIBUTE_ENGINE_COMMAND: str = 'engine_command'
+    _KEY_ATTRIBUTE_IMAGE_NAME: str = 'image_name'
+
+    def __init__(self, engine_command: str, image_name: str, **kwargs):
+        super().__init__(**kwargs)
+        self.engine_command = engine_command
+        self.image_name = image_name
+
+    @property
+    def filepath_executable(self) -> pathlib.PurePath:
+        """Return the filepath of the executable that this code represents.
+
+        .. note:: This is overridden from the base class since the path does not have to be absolute.
+
+        :return: The filepath of the executable.
+        """
+        return super().filepath_executable
+
+    @filepath_executable.setter
+    def filepath_executable(self, value: str) -> None:
+        """Set the filepath of the executable that this code represents.
+
+        .. note:: This is overridden from the base class since the path does not have to be absolute.
+
+        :param value: The filepath of the executable.
+        """
+        type_check(value, str)
+        self.base.attributes.set(self._KEY_ATTRIBUTE_FILEPATH_EXECUTABLE, value)
+
+    @property
+    def engine_command(self) -> str:
+        """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
+
+        :param value: The engine command of the containerized code
+        """
+        type_check(value, str)
+
+        if '{image_name}' not in value:
+            raise ValueError("the '{image_name}' template field should be in engine command.")
+
+        self.base.attributes.set(self._KEY_ATTRIBUTE_ENGINE_COMMAND, value)
+
+    @property
+    def image_name(self) -> str:
+        """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
+
+        :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(
+        self, mpi_args: list[str] | None = None, extra_mpirun_params: list[str] | None = None
+    ) -> list[str]:
+        """Return the list of prepend cmdline params for mpi seeting
+
+        :return: list of prepend cmdline parameters."""
+        engine_cmdline = self.engine_command.format(image_name=self.image_name)
+        engine_cmdline_params = engine_cmdline.split()
+
+        return (mpi_args or []) + (extra_mpirun_params or []) + engine_cmdline_params
+
+    @classmethod
+    def _get_cli_options(cls) -> dict:
+        """Return the CLI options that would allow to create an instance of this class."""
+        options = {
+            'engine_command': {
+                'required':
+                True,
+                'prompt':
+                'Engine command',
+                'help': (
+                    'The command to run the container. It must contain the placeholder '
+                    '{image_name} that will be replaced with the `image_name`.'
+                ),
+                'type':
+                click.STRING,
+            },
+            'image_name': {
+                'required': True,
+                'type': click.STRING,
+                'prompt': 'Image name',
+                'help': 'Name of the image container in which to the run the executable.',
+            },
+        }
+        options.update(**super()._get_cli_options())
+
+        return options