Add scheduler commands and mpirun extra params at the `Code` level

[GeigerJ2]

ORIGINAL ISSUE BELOW

---

With a code that is set up like this:

❯ verdi code show 2
-----------------------  --------------------------------------
PK                       2
UUID                     ad20bb6d-66c2-4ec7-a2c1-cce26bb2e97f
Type                     core.code.installed
Label                    pw
Description              abc
Default calc job plugin  quantumespresso.pw
Use double quotes        False
With mpi
Prepend text             uenv start quantumespresso/v7.3.1:v1
Append text
Computer                 eiger (eiger.alps.cscs.ch), pk: 3
Filepath executable      /user-environment/env/default/bin/pw.x
-----------------------  --------------------------------------

On the new CSCS Alps infrastructure, even though the code actually works for running simulations, verdi code test fails, as the filepath_executable cannot be found. This is because uenv images are actually squashfs file systems that get mounted when running uenv start, meaning the file does not exist on the user's file system before. This is different from the common setup where the required module load commands are in the prepend_text and are executed at runtime to make the necessary libraries available, but the filepath_executable already points to an existing file, e.g., see here:
https://github.com/aiidateam/aiida-code-registry/blob/2532dfd5cf49bb92d9143f084a6b7b0239785450/daint.cscs.ch/codes/pw-7.2-hybrid.yaml#L5-L9

While this is a uenv, and therefore CSCS specific issue, as many of us use the CSCS machines, I thought it's still good to have it documented somewhere.

A simple fix could be to add a --run-prepend-text option to verdi code test, which should then make the test run successfully.

[ayushjariyal]

@GeigerJ2 , I want to work on this issue.

[GeigerJ2]

Hi @ayushjariyal, thanks for your interest in this issue. This is pretty specific, and probably not the best to work on. What about #6978 instead, as it relates again to the verdi storage info command where you already started working on?

[aman-coder03]

I'd like to work on this issue if it's not already taken.

[mikibonacci]

Hi @GeigerJ2, I think this is a critical issue, we might need to discuss to find a solution for this. (we probably need to open an issue not just for the testing).
From just the AiiDA side of things, the solution is easy: we can just add in the metadata.options.custom_scheduler_commands

#SBATCH --uenv=<name of the uenv>
#SBATCH --view=default

or, in the metadata.options.mpirun_extra_params (I hope it is the right one and the syntax is correct -- I will correct it in case)

['--uenv=<name of the uenv>', '--view=default'] 

and it will work. (the prepend_text of the code should contain uenv image pull <name of the uenv> to automatically pull the image if not there).
However, the real problem is when we want to pack in useful pre-compiled (and/or templated) yaml files. In other words, we may want to have this additional commands/params defined in a static way, so that the user does not need to do it manually. This is crucial for AiiDAlab QE app, for example, where the user cannot (and should not) define these inputs.
I see a couple of paths for this static case:

• Best for me: we allow to define custom scheduler commands and mpirun extra params in the code definition, adding two fields (we should not do that in the computer definition/commands/parameters, as different codes behave differently)
• Ok for me: we add, in the Mpirun command of the computer, a ${something}. Then, in the code prepend text, we define the export something=--uenv=<name of the uenv> --view=default (i.e. what I should put in mpirun_extra_params in the dynamic case). It's a bit involved and might not be working, not sure, but it can be a quick and zero-implementation-required fix (to be documented of course)

Additionally, we can add the uenv sbatch commands in the code prepend text, but the prepend text of the computer should not then contain any additional line which is not a sbatch command (e.g. export OMP...), otherwise it will not work

pinging also @giovannipizzi for comments on this.

[giovannipizzi]

Just a question to confirm. It is equivalent to put:

• Approach 1: add in the submission script:

  #SBATCH --uenv=<name of the uenv>
  #SBATCH --view=default
  

• Approach 2: uenv start quantumespresso/v7.3.1:v1 in the prepend_test
• Approach 3: --uenv=<name of the uenv> --view=default after mpirun/srun?

If they are really eqiuvalent, we might not need to support all three, but then we need to be very clear in docs, as in principle the user might want to do any of the three.

Approach 2 is already supported with the code-specific prepend-text, right?
Approach 1 requires code-specific custom scheduler commands, and approach 3 code-specific mpirun extra params, that is what Miki suggests. I think it's a good idea to support both.

So I suggest to implement these last two options (code-specific custom scheduler commands and mpirun extra params), and we support all three possibilities and we can provide YAML configuration files without user hassle, as Miki recommedns.

However, I think adding the prepend-text to verdi code test in in general a good idea, so we should add this option I think (and would - probably? - fix the main issue mentioned here if the user defines the code with approach 2. This might be enough to have support for these codes in the verdi code registry, since we can probably even do automated checks to see which approach it's used for any entry, or at least document it clearly (and then people will use other examples as templates for new coders, I think).

Approach 3 could also be supported somehow by verdi code test, with some caveats, but (like approach 1) would require a bit of custom logic, I fear, as it's not obvious which of the SBATCH commands (or of the mpirun-extra-params) should be put where. I don't know if we want to implement this, as the amount of code then becomes just too much (and it's also scheduler specific etc.). At most I would suggest detecting if specific keywords are used (as e.g. --uenv in the scheduler part) and, if the test command fails, showing an additional message pointing users to the docs.

[GeigerJ2]

Thanks @mikibonacci and @giovannipizzi for the discussion. And, sorry @mikibonacci for not getting back to it sooner. I originally read your comment, realized it required some thought, and then shelved it and forgot about it 😅

I've opened #7272 (still draft) that addresses the points raised here. Here's a summary of what's implemented:

As suggested by @giovannipizzi, both custom_scheduler_commands and mpirun_extra_params are now available as code-level properties on AbstractCode. This means all three approaches for uenv support can now be configured statically on the code definition:

• Approach 1 (#SBATCH --uenv=...): Set via code.custom_scheduler_commands
• Approach 2 (uenv start ...): Already supported via code.prepend_text
• Approach 3 (--uenv=... --view=default after srun): Set via code.mpirun_extra_params

These are additive with the existing CalcJob-level metadata.options—during presubmit, code-level values come first, then CalcJob-level values are appended. This follows the same layered pattern as prepend_text (computer → code → calcjob).

This means they can be defined in YAML configuration files for code registries and AiiDAlab, without requiring the user to set them manually on each calculation, @mikibonacci.

The PR also adds the --run-prepend-text flag on verdi code test, that executes both the Computer's and Code's prepend_text on the remote machine before checking if the executable exists. This fixes the original issue where verdi code test fails for uenv codes because the filesystem is only mounted after uenv start runs.

Note that this flag only helps for approach 2 (prepend_text). For approaches 1 and 3, the scheduler directives and mpirun params only take effect during actual job submission, so they can't be simulated in an SSH test session. verdi code test without the flag will still fail for those, but the actual job submission will work correctly. As mentioned by @giovannipizzi, this probably requires a bunch of custom hard-coded logic that is probably not worth adding.

---

As a side note:

One concern I always had with the possibility of setting those parameters at multiple places (computer, code, calcjob) is that it can be confusing for users, as it's not always immediately obvious where is the correct place to set them. That said, I the more I think about it, the layered approach actually does make sense as they aren't competing "sources of truth" that override each other, but they are additive at different scopes:

• Computer level: applies to all jobs on this machine (e.g. module load gcc)
• Code level: applies to all jobs using this specific code (e.g. #SBATCH --uenv=quantumespresso/...)
• CalcJob level (metadata.options): applies to this specific job only (e.g. one-off overrides)

So, possibly something we should be more clear about in the docs, telling users what each level does and when to use which.

[giovannipizzi]

Thanks Julian, and indeed the last point you make is exactly the point! If not clear let's clarify!

As you are at it: is there any other configurable property of computer (during setup or configure) that is not of code but could make sense? Anything else from things that can be set on the calcjob itself that it's worth considering adding? (maybe nothing, but as we are woejibg on this, a quick consideration of this point might help)

[GeigerJ2]

Good call, @giovannipizzi. I (Claude) did a systematic review of all runtime/scheduler options across Computer, Code, and CalcJob. I guess nobody took a global view on this infrastructure for a while, and instead options were added ad-hoc, as the need arose (same as we would be doing here). Instead, I do think it's good to now have this global look, and align everything, that's what is now done in #7272.

TL;DR

Now done in #7272

• custom_scheduler_commands on Code + Computer ✅ — full 3-level concatenation (Computer → Code → CalcJob)
• mpirun_extra_params on Code ✅ — concatenated (Code → CalcJob)
• environment_variables on Code + Computer ✅ — full 3-level dict merging (later overrides same key)

Worth discussing (not in #7272, but would be different PR)

• queue_name / account on Code — useful as defaults (e.g. GPU code always on gpu partition), but override semantics differ from the existing additive pattern

Not worth adding

• Everything else — either job-specific, hardware-specific, or already covered

---

Full overview: where each option lives and how it merges

Property	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	#7272	#7272	yes	Concatenated (Computer → Code → CalcJob)
environment_variables	#7272	#7272	yes	Dict merge (later overrides same key)
mpirun_extra_params	—	#7272	yes	Concatenated (Code → CalcJob)
use_double_quotes	yes	yes	—	Per code-info (shell quoting)
with_mpi	—	yes	yes (withmpi)	Must agree if both set
mpirun_command	yes	—	—	—
default_mpiprocs_per_machine	yes	—	—	Fallback for resources
default_memory_per_machine	yes	—	—	Fallback for resources
environment_variables_double_quotes	—	—	yes	—
queue_name	—	—	yes	—
account	—	—	yes	—
qos	—	—	yes	—
resources	—	—	yes	—
max_wallclock_seconds	—	—	yes	—
max_memory_kb	—	—	yes	—
priority	—	—	yes	—
rerunnable	—	—	yes	—
import_sys_environment	—	—	yes	—
shebang	yes	—	—	—
workdir	yes	—	—	—
minimum_job_poll_interval	yes	—	—	—
parser_name	—	—	yes	—
default_calc_job_plugin	—	yes	—	—
wrap_cmdline_params	—	yes	—	—
input_filename / output_filename	—	—	yes	—
submit_script_filename	—	—	yes	—
scheduler_stdout / scheduler_stderr	—	—	yes	—
additional_retrieve_list	—	—	yes	—
stash.* / unstash.*	—	—	yes	—