Container code

aiida/cmdline/commands/cmd_code.py

@@ -64,6 +64,9 @@ def set_code_builder(ctx, param, value):
 # guaranteed only for the interactive case, however. For the non-interactive case, the callback is called explicitly
 # once more in the command body to cover the case when the label is specified before the computer.
 @verdi_code.command('setup')
+@options_code.ON_CONTAINER()
+@options_code.CONTAINER_ENGINE_COMMAND()
+@options_code.IMAGE()
 @options_code.ON_COMPUTER()
 @options_code.COMPUTER()
 @options_code.LABEL()
@@ -72,6 +75,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()
@@ -84,7 +88,10 @@ def setup_code(ctx, non_interactive, **kwargs):
 
     options_code.validate_label_uniqueness(ctx, None, kwargs['label'])
 
-    if kwargs.pop('on_computer'):
+    if kwargs.pop('on_container'):
+        kwargs['code_type'] = CodeBuilder.CodeType.ON_CONTAINER
+        kwargs.pop('on_computer')
+    elif kwargs.pop('on_computer'):
         kwargs['code_type'] = CodeBuilder.CodeType.ON_COMPUTER
     else:
         kwargs['code_type'] = CodeBuilder.CodeType.STORE_AND_UPLOAD

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

@@ -15,6 +15,14 @@
 from aiida.cmdline.params.options.overridable import OverridableOption
 
 
+def is_containerized(ctx):
+    return bool(ctx.params.get('on_container'))
+
+
+def is_not_containerized(ctx):
+    return bool(not is_containerized(ctx))
+
+
 def is_on_computer(ctx):
     return bool(ctx.params.get('on_computer'))
 
@@ -23,6 +31,10 @@ def is_not_on_computer(ctx):
     return bool(not is_on_computer(ctx))
 
 
+def is_on_computer_or_container(ctx):
+    return is_on_computer(ctx) or is_containerized(ctx)
+
+
 def validate_label_uniqueness(ctx, _, value):
     """Validate the uniqueness of the label of the code.
 
@@ -68,11 +80,38 @@ def validate_label_uniqueness(ctx, _, value):
     return value
 
 
+ON_CONTAINER = OverridableOption(
+    '--on-container/--not-on-container',
+    is_eager=False,
+    default=False,
+    help='Whether the code is installed on the container.'
+)
+
+IMAGE = OverridableOption(
+    '--image',
+    prompt='image name or absolute path to image file',
+    required_fn=is_containerized,
+    prompt_fn=is_containerized,
+    cls=InteractiveOption,
+    help='image name or absolute path to image file.'
+)
+
+CONTAINER_ENGINE_COMMAND = OverridableOption(
+    '--container-engine-command',
+    prompt='container cmdline params tmpl',
+    required_fn=is_containerized,
+    prompt_fn=is_containerized,
+    cls=InteractiveOption,
+    help='container cmdline params tmpl.'
+)
+
 ON_COMPUTER = OverridableOption(
     '--on-computer/--store-in-db',
     is_eager=False,
     default=True,
     cls=InteractiveOption,
+    required_fn=is_not_containerized,
+    prompt_fn=is_not_containerized,
     prompt='Installed on target computer?',
     help='Whether the code is installed on the target computer, or should be copied to the target computer each time '
     'from a local path.'
@@ -81,8 +120,8 @@ def validate_label_uniqueness(ctx, _, value):
 REMOTE_ABS_PATH = OverridableOption(
     '--remote-abs-path',
     prompt='Remote absolute path',
-    required_fn=is_on_computer,
-    prompt_fn=is_on_computer,
+    required_fn=is_on_computer_or_container,
+    prompt_fn=is_on_computer_or_container,
     type=types.AbsolutePathParamType(dir_okay=False),
     cls=InteractiveOption,
     help='[if --on-computer]: Absolute path to the executable on the target computer.'
@@ -109,6 +148,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',
+    is_eager=False,
+    default=False,
+    cls=InteractiveOption,
+    prompt='Whether use double quotes for code cmdline params?',
+    help='Whether the code executable parameters in script using the double quotes to escape.'
+)
+
 LABEL = options.LABEL.clone(
     prompt='Label',
     callback=validate_label_uniqueness,
@@ -132,8 +180,8 @@ def validate_label_uniqueness(ctx, _, value):
 COMPUTER = options.COMPUTER.clone(
     prompt='Computer',
     cls=InteractiveOption,
-    required_fn=is_on_computer,
-    prompt_fn=is_on_computer,
+    required_fn=is_on_computer_or_container,
+    prompt_fn=is_on_computer_or_container,
     help='Name of the computer, on which the code is installed.'
 )
 

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

@@ -117,6 +117,16 @@ 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',
+    is_eager=False,
+    default=False,
+    cls=InteractiveOption,
+    prompt='Whether use double quotes for prepend cmdline params?',
+    help='Whether the cmdline parameters before the code executable parameters in '
+    'script using the double quotes to escape.',
+)
+
 PREPEND_TEXT = OverridableOption(
     '--prepend-text',
     cls=TemplateInteractiveOption,

aiida/common/datastructures.py

@@ -99,6 +99,11 @@ class CodeInfo(DefaultFieldsAttributeDict):
     This attribute-dictionary contains the information needed to execute a code.
     Possible attributes are:
 
+    * ``prepend_cmdline_params``: (optional) a list of strings, containing parameters to be
+      written in front of cmdline_params::
+
+        prepend_cmdline_params cmdline_params ...
+
     * ``cmdline_params``: a list of strings, containing parameters to be written on
       the command line right after the call to the code, as for example::
 
@@ -136,16 +141,24 @@ class CodeInfo(DefaultFieldsAttributeDict):
 
     * ``withmpi``: if True, executes the code with mpirun (or another MPI installed
       on the remote computer)
+    * ``use_double_quotes``: is a tuple contain two bool values which decide whether use double quotes for
+      prepend_cmdline_params, cmdline_params respectively. If True, the code cmdline parameters of
+      bash script escaped by quotes.
     * ``code_uuid``: the uuid of the code associated to the CodeInfo
+    * ``custom_cmdline_string``: If this field provided, only the content of this string will print and
+      the string will print as it is without any escape
     """
     _default_fields = (
+        'prepend_cmdline_params',
         'cmdline_params',  # as a list of strings
         'stdin_name',
         'stdout_name',
         'stderr_name',
         'join_files',
         'withmpi',
-        'code_uuid'
+        'use_double_quotes',  # tuple to control double quotes of (prepend_cmdline_params, cmdline_params)
+        'code_uuid',
+        'custom_cmdline_string',
     )
 
 

aiida/common/escaping.py

@@ -12,7 +12,7 @@
 import re
 
 
-def escape_for_bash(str_to_escape):
+def escape_for_bash(str_to_escape, use_double_quotes=False):
     """
     This function takes any string and escapes it in a way that
     bash will interpret it as a single string.
@@ -34,14 +34,23 @@ def escape_for_bash(str_to_escape):
     Finally, note that for python I have to enclose the string '"'"'
     within triple quotes to make it work, getting finally: the complicated
     string found below.
+
+    :param str_to_escape: the string to escape.
+    :param use_double_quotes: boolean, if ``True``, use double quotes instead of single quotes.
+    :return: the escaped string.
     """
     if str_to_escape is None:
         return ''
 
     str_to_escape = str(str_to_escape)
-
-    escaped_quotes = str_to_escape.replace("'", """'"'"'""")
-    return f"'{escaped_quotes}'"
+    if use_double_quotes:
+        escaped_quotes = str_to_escape.replace('"', '''"'"'"''')
+        escaped = f'"{escaped_quotes}"'
+    else:
+        escaped_quotes = str_to_escape.replace("'", """'"'"'""")
+        escaped = f"'{escaped_quotes}'"
+
+    return escaped
 
 
 # Mapping of "SQL" tokens into corresponding regex expressions

aiida/engine/daemon/execmanager.py

@@ -27,7 +27,7 @@
 from aiida.common.datastructures import CalcInfo
 from aiida.common.folders import SandboxFolder
 from aiida.common.links import LinkType
-from aiida.orm import CalcJobNode, Code, FolderData, Node, RemoteData, load_node
+from aiida.orm import CalcJobNode, FolderData, Node, RemoteData, load_code, load_node
 from aiida.orm.utils.log import get_dblogger_extra
 from aiida.repository.common import FileType
 from aiida.schedulers.datastructures import JobState
@@ -86,7 +86,7 @@ def upload_calculation(
     computer = node.computer
 
     codes_info = calc_info.codes_info
-    input_codes = [load_node(_.code_uuid, sub_classes=(Code,)) for _ in codes_info]
+    input_codes = [load_code(_.code_uuid) for _ in codes_info]
 
     logger_extra = get_dblogger_extra(node)
     transport.set_logger_extra(logger_extra)

aiida/engine/processes/calcjobs/calcjob.py

@@ -22,6 +22,7 @@
 from aiida.common.folders import Folder
 from aiida.common.lang import classproperty, override
 from aiida.common.links import LinkType
+from aiida.orm.nodes.data.containerized_code import ContainerizedCode
 
 from ..exit_code import ExitCode
 from ..ports import PortNamespace
@@ -244,6 +245,9 @@ def define(cls, spec: CalcJobProcessSpec) -> None:  # type: ignore[override]
             help='If set to true, the submission script will load the system environment variables',)
         spec.input('metadata.options.environment_variables', valid_type=dict, default=lambda: {},
             help='Set a dictionary of custom environment variables for this calculation',)
+        spec.input('metadata.options.envvar_double_quotes', valid_type=bool, default=False,
+            help='If set to True, use double quotes instead of single quotes to escape the environment variables '
+                 'specified in ``environment_variables``.',)
         spec.input('metadata.options.priority', valid_type=str, required=False,
             help='Set the priority of the job to be queued')
         spec.input('metadata.options.max_memory_kb', valid_type=int, required=False,
@@ -583,7 +587,7 @@ def presubmit(self, folder: Folder) -> CalcInfo:
         from aiida.common.datastructures import CodeInfo, CodeRunMode
         from aiida.common.exceptions import InputValidationError, InvalidOperation, PluginInternalError, ValidationError
         from aiida.common.utils import validate_list_of_string_tuples
-        from aiida.orm import Code, Computer, load_node
+        from aiida.orm import Code, Computer, load_code
         from aiida.schedulers.datastructures import JobTemplate
 
         inputs = self.node.get_incoming(link_type=LinkType.INPUT_CALC)
@@ -687,7 +691,51 @@ def presubmit(self, folder: Folder) -> CalcInfo:
 
             if code_info.code_uuid is None:
                 raise PluginInternalError('CalcInfo should have the information of the code to be launched')
-            this_code = load_node(code_info.code_uuid, sub_classes=(Code,))
+            this_code = load_code(code_info.code_uuid)
+
+            # Set code_info use_double_quotes based on computer and code setup if code_info not set by plugin
+            if code_info.use_double_quotes is None:
+                code_info.use_double_quotes = (computer.get_use_double_quotes(), this_code.get_use_double_quotes())
+
+            # set this_argv only for container code with custom_cmdline_string to override other parameters
+            if isinstance(this_code, ContainerizedCode):
+                from string import Template
+
+                from aiida.common.escaping import escape_for_bash
+
+                # This part copy from scheduler _get_run_line
+                stdin_str = f'< {escape_for_bash(code_info.stdin_name, code_info.use_double_quotes[1])}' if code_info.stdin_name else ''
+                stdout_str = f'> {escape_for_bash(code_info.stdout_name, code_info.use_double_quotes[1])}' 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, code_info.use_double_quotes[1])}' if code_info.stderr_name else ''
+
+                append_cmdline_list = [stdin_str, stdout_str, stderr_str]
+
+                command_to_exec_list = [
+                    escape_for_bash(
+                        this_code.get_container_exec_path(), use_double_quotes=code_info.use_double_quotes[1]
+                    )
+                ]
+
+                if code_info.cmdline_params is not None:
+                    for arg in code_info.cmdline_params:
+                        command_to_exec_list.append(
+                            escape_for_bash(arg, use_double_quotes=code_info.use_double_quotes[1])
+                        )
+
+                command_to_exec = ' '.join(command_to_exec_list + append_cmdline_list)
+
+                # containerized code run line always use its own mpi setting rather config from prepend
+                code_info.prepend_cmdline_params = None
+
+                temp_obj = Template(this_code.get_container_engine_command())
+
+                # use safe_substitute since $ also be placeholder for ENV_VAR
+                code_info.custom_cmdline_string = temp_obj.safe_substitute(exec_str=command_to_exec)
 
             # To determine whether this code should be run with MPI enabled, we get the value that was set in the inputs
             # of the entire process, which can then be overwritten by the value from the `CodeInfo`. This allows plugins
@@ -702,18 +750,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 [])
-                )
+                code_info.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 [])
+                code_info.prepend_cmdline_params = None
 
-            # overwrite the old cmdline_params and add codename and mpirun stuff
+            this_argv = [this_code.get_execname()
+                         ] + (code_info.cmdline_params if code_info.cmdline_params is not None else [])
             code_info.cmdline_params = this_argv
 
             codes_info.append(code_info)
+
         job_tmpl.codes_info = codes_info
 
         # set the codes execution mode, default set to `SERIAL`
@@ -731,6 +777,7 @@ def presubmit(self, folder: Folder) -> CalcInfo:
         job_tmpl.import_sys_environment = self.node.get_option('import_sys_environment')
 
         job_tmpl.job_environment = self.node.get_option('environment_variables')
+        job_tmpl.envvar_double_quotes = self.node.get_option('envvar_double_quotes')
 
         queue_name = self.node.get_option('queue_name')
         account = self.node.get_option('account')

aiida/manage/tests/pytest_fixtures.py

@@ -164,7 +164,15 @@ def test_1(aiida_local_code_factory):
     :rtype: object
     """
 
-    def get_code(entry_point, executable, computer=aiida_localhost, label=None, prepend_text=None, append_text=None):
+    def get_code(
+        entry_point,
+        executable,
+        computer=aiida_localhost,
+        label=None,
+        use_double_quotes=False,
+        prepend_text=None,
+        append_text=None
+    ):
         """Get local code.
 
         Sets up code for given entry point on given computer.
@@ -201,6 +209,8 @@ def get_code(entry_point, executable, computer=aiida_localhost, label=None, prep
         code.label = label
         code.description = label
 
+        code.set_use_double_quotes(use_double_quotes)
+
         if prepend_text is not None:
             code.set_prepend_text(prepend_text)
 

aiida/orm/__init__.py

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