✨ Add Code-level scheduler cmds & mpirun params

docs/source/howto/tricks_real_world_runs.rst

@@ -10,15 +10,21 @@ Setting runtime configuration of a calculation job
 ==================================================
 
 When submitting a job, you can provide specific instructions which should be put in the job submission script (e.g. for the Slurm scheduler).
-They can be provided via the ``metadata.options`` dictionary of the calculation job.
+These can be configured at multiple levels, and are **additive** — they get concatenated during submission, not overridden:
+
+- **Computer level**: applies to *all* jobs on this machine (e.g. ``module load gcc``). Set during ``verdi computer setup``.
+- **Code level**: applies to all jobs using *this specific code* (e.g. ``#SBATCH --uenv=quantumespresso/...``). Set during ``verdi code create``.
+- **CalcJob level**: applies to *this specific job* only (e.g. one-off overrides). Set via the ``metadata.options`` dictionary of the calculation job.
+
+During submission, values from all three levels are joined together (computer first, then code, then CalcJob).
 
 In particular, you can provide:
 
 - custom scheduler directives, like additional ``#SBATCH`` commands for the Slurm scheduler
 - prepend text to the job script (e.g. if you need to load specific modules or set specific environment variables which are not already specified in the computer/code setup)
 - additional mpirun parameters (e.g. if you need to bind processes to cores, etc.)
 
-The full list of available options can be found in the `AiiDA CalcJob options documentation <https://aiida--7048.org.readthedocs.build/projects/aiida-core/en/7048/topics/calculations/usage.html#options>`__.
+For a complete overview of which options are available at which level and how they are merged, see :ref:`topics:calculations:usage:calcjobs:scheduler-and-runtime-options`.
 
 
 Basic pattern
@@ -61,7 +67,9 @@ In case you are submitting a ``PwBaseWorkChain``, these options should be set un
 Custom scheduler directives (e.g. extra ``#SBATCH``)
 ----------------------------------------------------
 
-Use ``custom_scheduler_commands`` to inject raw scheduler lines near the top of the submit script (before any non-scheduler command):
+Use ``custom_scheduler_commands`` to inject raw scheduler lines near the top of the submit script (before any non-scheduler command).
+
+**On a per-job basis** via ``metadata.options``:
 
 .. code-block:: python
 
@@ -71,16 +79,35 @@ Use ``custom_scheduler_commands`` to inject raw scheduler lines near the top of
 	#SBATCH --hint=nomultithread
 	""".strip()
 
+**On the code** (applies to all jobs using this code):
+
+.. code-block:: yaml
+
+	# In the code YAML configuration:
+	custom_scheduler_commands: |
+	    #SBATCH --uenv=quantumespresso/v7.3.1:v1
+	    #SBATCH --view=default
+
+This is particularly useful for CSCS Alps ``uenv`` images or similar setups where the scheduler must be told to mount a specific environment before the job starts.
+
 Notes:
 
 - Keep the lines valid for your scheduler (Slurm here; adapt to PBS/LSF/etc.).
 - Use this when a directive is not covered by a dedicated option.
+- Can also be set at the **Computer** level (during ``verdi computer setup``) for directives that apply to all jobs on that machine.
+- All three levels are concatenated (Computer first, then Code, then CalcJob).
 
 
 Prepend/append shell code to the job script
 -------------------------------------------
 
-Use ``prepend_text`` to add shell commands immediately before launching the code, and ``append_text`` for commands executed right after the code finishes:
+Use ``prepend_text`` to add shell commands immediately before launching the code, and ``append_text`` for commands executed right after the code finishes.
+
+``prepend_text`` can be set at three levels (all are concatenated in the submission script):
+
+- **Computer**: set during ``verdi computer setup`` — applies to all jobs on this machine (e.g. ``module load gcc``)
+- **Code**: set during ``verdi code create`` — applies to all jobs using this code (e.g. ``uenv start quantumespresso/v7.3.1:v1``)
+- **CalcJob**: set via ``metadata.options`` — applies to this specific job only
 
 .. code-block:: python
 
@@ -92,20 +119,43 @@ Use ``prepend_text`` to add shell commands immediately before launching the code
 	echo "Run finished on $(hostname) at $(date)"
 	""".strip()
 
-.. tip:: for simple environment variables you can also use ``environment_variables`` (AiiDA will export them for you):
+Environment variables
+---------------------
+
+Use ``environment_variables`` to set environment variables that AiiDA will export for you in the submit script.
+This can be set at three levels (all are merged, with more specific levels overriding):
+
+- **Computer**: set during ``verdi computer setup`` — applies to all jobs on this machine
+- **Code**: set during ``verdi code create`` — applies to all jobs using this code
+- **CalcJob**: set via ``metadata.options`` — applies to this specific job only
+
+**On a per-job basis** via ``metadata.options``:
 
 .. code-block:: python
 
 	builder.metadata.options.environment_variables = {
 		 'OMP_NUM_THREADS': '1',
 	}
 
+**On the code** (applies to all jobs using this code):
+
+.. code-block:: yaml
+
+	# In the code YAML configuration:
+	environment_variables:
+	    OMP_NUM_THREADS: '1'
+	    LD_LIBRARY_PATH: '/opt/lib'
+
+.. note::
+	``environment_variables`` is a dictionary of strings. Computer-level, code-level, and CalcJob-level values are merged (later levels override earlier ones for the same key).
 
 
 Extra parameters to mpirun (or equivalent)
 ------------------------------------------
 
-Set ``mpirun_extra_params`` to pass flags to the MPI launcher in addition to the computer's configured ``mpirun_command``:
+Set ``mpirun_extra_params`` to pass flags to the MPI launcher in addition to the computer's configured ``mpirun_command``.
+
+**On a per-job basis** via ``metadata.options``:
 
 .. code-block:: python
 
@@ -114,8 +164,17 @@ Set ``mpirun_extra_params`` to pass flags to the MPI launcher in addition to the
 		 '--bind-to', 'core', '--map-by', 'socket:PE=2',
 	]
 
+**On the code** (applies to all jobs using this code):
+
+.. code-block:: yaml
+
+	# In the code YAML configuration:
+	mpirun_extra_params:
+	    - '--uenv=quantumespresso/v7.3.1:v1'
+	    - '--view=default'
+
 .. note::
-	``mpirun_extra_params`` is a list/tuple of strings; AiiDA will join them with spaces. Keep launcher-specific flags consistent with your cluster (OpenMPI, MPICH, srun, etc.).
+	``mpirun_extra_params`` is a list of strings; AiiDA will join them with spaces. Code-level and CalcJob-level values are concatenated (code first, then CalcJob). Keep launcher-specific flags consistent with your cluster (OpenMPI, MPICH, srun, etc.).
 
 
 Full list of metadata available
@@ -139,7 +198,7 @@ Here is the full list of options that can be set in ``builder.metadata``.
       - additional_retrieve_list (list | tuple | None): Relative file paths to retrieve in addition to what the plugin specifies.
       - append_text (str): Text appended to the scheduler-job script just after the code execution.
       - custom_scheduler_commands (str): Raw scheduler directives inserted before any non-scheduler command (e.g. extra ``#SBATCH`` lines).
-      - environment_variables (dict): Environment variables to export for this calculation.
+      - environment_variables (dict): Environment variables to export for this calculation. Merged with computer-level and code-level values (CalcJob overrides Code overrides Computer).
       - environment_variables_double_quotes (bool): If True, use double quotes instead of single quotes to escape ``environment_variables``.
       - import_sys_environment (bool): If True, the submission script will load the system environment variables.
       - input_filename (str): Name of the main input file written to the remote working directory.

docs/source/topics/calculations/usage.rst

@@ -986,6 +986,95 @@ Best Practices and Considerations
 
    - **Error handling**: In case of custom stashing, it's your sole responsibility to manage the failures and log errors.
 
+.. _topics:calculations:usage:calcjobs:scheduler-and-runtime-options:
+
+Scheduler and runtime options
+-----------------------------
+
+Several options that control the submission script can be set at up to three levels: **Computer**, **Code**, and **CalcJob**.
+During submission, AiiDA merges them automatically — text options are concatenated (Computer first, then Code, then CalcJob), while dictionary options are merged (later levels override earlier ones for the same key).
+
+The table below shows where each option can be set and how values from different levels are combined.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 12 12 12 34
+
+   * - Option
+     - Computer
+     - Code
+     - CalcJob
+     - Merging strategy
+   * - ``prepend_text``
+     - yes
+     - yes
+     - yes
+     - Concatenated (Computer → Code → CalcJob)
+   * - ``append_text``
+     - yes
+     - yes
+     - yes
+     - Concatenated (CalcJob → Code → Computer)
+   * - ``custom_scheduler_commands``
+     - yes
+     - yes
+     - yes
+     - Concatenated (Computer → Code → CalcJob)
+   * - ``environment_variables``
+     - yes
+     - yes
+     - yes
+     - Dict merge (Computer → Code → CalcJob, later overrides same key)
+   * - ``mpirun_extra_params``
+     -
+     - yes
+     - yes
+     - Concatenated (Code → CalcJob)
+   * - ``with_mpi``
+     -
+     - yes
+     - yes (``withmpi``)
+     - Must agree if set at multiple levels; error on conflict
+   * - ``use_double_quotes``
+     - yes
+     - yes
+     -
+     - Per code-info
+   * - ``mpirun_command``
+     - yes
+     -
+     -
+     - —
+   * - ``default_mpiprocs_per_machine``
+     - yes
+     -
+     -
+     - Fallback for ``resources``
+   * - ``default_memory_per_machine``
+     - yes
+     -
+     -
+     - Fallback for ``resources``
+   * - ``resources``
+     -
+     -
+     - yes
+     - —
+   * - ``queue_name``, ``account``, ``qos``
+     -
+     -
+     - yes
+     - —
+   * - ``max_wallclock_seconds``, ``max_memory_kb``
+     -
+     -
+     - yes
+     - —
+
+Options that appear only at the CalcJob level (``resources``, ``queue_name``, ``max_wallclock_seconds``, etc.) are job-specific and do not support layered configuration.
+
+For practical examples of how to use these options, see :ref:`how-to:real-world-tricks`.
+
 .. _topics:calculations:usage:calcjobs:options:
 
 Options

src/aiida/cmdline/__init__.py

@@ -30,6 +30,7 @@
     'GroupParamType',
     'HostnameType',
     'IdentifierParamType',
+    'JSONDictParamType',
     'LabelStringType',
     'LazyChoice',
     'MpirunCommandParamType',

src/aiida/cmdline/commands/cmd_code.py

@@ -147,8 +147,16 @@ def setup_code(ctx, non_interactive, **kwargs):
 
 @verdi_code.command('test')
 @arguments.CODE(callback=set_code_builder)
+@click.option(
+    '--run-prepend-text',
+    is_flag=True,
+    default=False,
+    help="Execute the computer's and code's prepend text on the remote computer before checking the executable. "
+    'This is useful when the executable only becomes available after running setup commands '
+    '(e.g. `module load ...` or `uenv start ...`).',
+)
 @with_dbenv()
-def code_test(code):
+def code_test(code, run_prepend_text):
     """Run tests for the given code to check whether it is usable.
 
     For remote codes the following checks are performed:
@@ -160,7 +168,7 @@ def code_test(code):
 
     if isinstance(code, InstalledCode):
         try:
-            code.validate_filepath_executable()
+            code.validate_filepath_executable(run_prepend_text=run_prepend_text)
         except exceptions.ValidationError as exception:
             echo.echo_critical(f'validation failed: {exception}')
 

src/aiida/cmdline/commands/cmd_computer.py

@@ -284,8 +284,10 @@ def set_computer_builder(ctx, param, value):
 @options_computer.MPI_PROCS_PER_MACHINE()
 @options_computer.DEFAULT_MEMORY_PER_MACHINE()
 @options_computer.USE_DOUBLE_QUOTES()
+@options_computer.CUSTOM_SCHEDULER_COMMANDS()
 @options_computer.PREPEND_TEXT()
 @options_computer.APPEND_TEXT()
+@options_computer.ENVIRONMENT_VARIABLES()
 @options.NON_INTERACTIVE()
 @options.CONFIG_FILE()
 @click.pass_context
@@ -337,8 +339,12 @@ def computer_setup(ctx, non_interactive, **kwargs):
 @options_computer.DEFAULT_MEMORY_PER_MACHINE(
     contextual_default=partial(get_parameter_default, 'default_memory_per_machine')
 )
+@options_computer.CUSTOM_SCHEDULER_COMMANDS(
+    contextual_default=partial(get_parameter_default, 'custom_scheduler_commands')
+)
 @options_computer.PREPEND_TEXT(contextual_default=partial(get_parameter_default, 'prepend_text'))
 @options_computer.APPEND_TEXT(contextual_default=partial(get_parameter_default, 'append_text'))
+@options_computer.ENVIRONMENT_VARIABLES(contextual_default=partial(get_parameter_default, 'environment_variables'))
 @options.NON_INTERACTIVE()
 @click.pass_context
 @with_dbenv()
@@ -490,8 +496,10 @@ def computer_show(computer):
         ['Mpirun command', ' '.join(computer.get_mpirun_command())],
         ['Default #procs/machine', computer.get_default_mpiprocs_per_machine()],
         ['Default memory (kB)/machine', computer.get_default_memory_per_machine()],
+        ['Custom scheduler commands', computer.get_custom_scheduler_commands()],
         ['Prepend text', computer.get_prepend_text()],
         ['Append text', computer.get_append_text()],
+        ['Environment variables', computer.get_environment_variables()],
     ]
     echo_tabulate(table)
 
@@ -790,8 +798,10 @@ def computer_export_setup(computer, output_file, overwrite, sort):
         'mpiprocs_per_machine': computer.get_default_mpiprocs_per_machine(),
         'default_memory_per_machine': computer.get_default_memory_per_machine(),
         'use_double_quotes': computer.get_use_double_quotes(),
+        'custom_scheduler_commands': computer.get_custom_scheduler_commands(),
         'prepend_text': computer.get_prepend_text(),
         'append_text': computer.get_append_text(),
+        'environment_variables': computer.get_environment_variables(),
     }
 
     if output_file is None:

src/aiida/cmdline/groups/dynamic.py

@@ -176,12 +176,19 @@ def list_options(self, entry_point: str) -> list[t.Callable[[FC], FC]]:
 
             default = field_info.default_factory if field_info.default is PydanticUndefined else field_info.default
 
+            # Check if the type is a list or tuple (e.g. list[str], tuple[str, ...]) which should be handled as
+            # Click ``multiple=True`` options, allowing the option to be specified multiple times on the CLI.
+            # Dict types (e.g. dict[str, str]) are handled with a custom JSON-based Click ParamType.
+            origin = t.get_origin(field_info.annotation)
+            is_multiple = origin in (list, tuple)
+            is_dict = origin is dict
+
             # If the annotation has the ``__args__`` attribute it is an instance of a type from ``typing`` and the real
             # type can be gotten from the arguments. For example it could be ``typing.Union[str, None]`` calling
             # ``typing.Union[str, None].__args__`` will return the tuple ``(str, NoneType)``. So to get the real type,
             # we simply remove all ``NoneType`` and the remaining type should be the type of the option.
             if hasattr(field_info.annotation, '__args__'):
-                args = list(filter(lambda e: e is not type(None), field_info.annotation.__args__))
+                args = list(filter(lambda e: e is not type(None) and e is not Ellipsis, field_info.annotation.__args__))
                 # Click parameters only support specifying a single type, so we default to the first one even if the
                 # pydantic model defines multiple.
                 field_type = args[0]
@@ -196,6 +203,15 @@ def list_options(self, entry_point: str) -> list[t.Callable[[FC], FC]]:
                 'default': default,
                 'help': field_info.description,
             }
+
+            if is_multiple:
+                options_spec[key]['multiple'] = True
+
+            if is_dict:
+                from aiida.cmdline.params.types import JSONDictParamType
+
+                options_spec[key]['type'] = JSONDictParamType()
+
             for metadata in field_info.metadata:
                 for metadata_key, metadata_value in metadata.items():
                     if metadata_key in ('priority', 'short_name', 'option_cls'):

src/aiida/cmdline/params/__init__.py

@@ -27,6 +27,7 @@
     'GroupParamType',
     'HostnameType',
     'IdentifierParamType',
+    'JSONDictParamType',
     'LabelStringType',
     'LazyChoice',
     'MpirunCommandParamType',