Merge pull request #1803 from rsarm/feat/preload-cmds

[feat] Add a new `prepare_cmds` partition configuration parameter which allows to emit commands before loading any environment

Author: Vasileios Karakasis

docs/config_reference.rst

@@ -308,6 +308,16 @@ System Partition Configuration
    This option is relevant only when ReFrame executes with the `asynchronous execution policy <pipeline.html#execution-policies>`__.
 
 
+.. js:attribute:: .systems[].partitions[].prepare_cmds
+
+   :required: No
+   :default: ``[]``
+
+   List of shell commands to be emitted before any environment loading commands are emitted.
+
+   .. versionadded:: 3.5.0
+
+
 .. js:attribute:: .systems[].partitions[].resources
 
    :required: No

docs/tutorial_advanced.rst

@@ -389,13 +389,16 @@ Generally, ReFrame generates the job shell scripts using the following pattern:
 
    #!/bin/bash -l
    {job_scheduler_preamble}
-   {test_environment}
+   {prepare_cmds}
+   {env_load_cmds}
    {prerun_cmds}
    {parallel_launcher} {executable} {executable_opts}
    {postrun_cmds}
 
 The ``job_scheduler_preamble`` contains the backend job scheduler directives that control the job allocation.
-The ``test_environment`` are the necessary commands for setting up the environment of the test.
+The ``prepare_cmds`` are commands that can be emitted before the test environment commands.
+These can be specified with the :js:attr:`prepare_cmds <.systems[].partitions[].prepare_cmds>` partition configuration option.
+The ``env_load_cmds`` are the necessary commands for setting up the environment of the test.
 These include any modules or environment variables set at the `system partition level <config_reference.html#system-partition-configuration>`__ or any `modules <regression_test_api.html#reframe.core.pipeline.RegressionTest.modules>`__ or `environment variables <regression_test_api.html#reframe.core.pipeline.RegressionTest.variables>`__ set at the test level.
 Then the commands specified in :attr:`prerun_cmds <reframe.core.pipeline.RegressionTest.prerun_cmds>` follow, while those specified in the :attr:`postrun_cmds <reframe.core.pipeline.RegressionTest.postrun_cmds>` come after the launch of the parallel job.
 The parallel launch itself consists of three parts:

reframe/core/pipeline.py

@@ -1220,6 +1220,7 @@ def compile(self):
             try:
                 self._build_job.prepare(
                     build_commands, environs,
+                    self._current_partition.prepare_cmds,
                     login=rt.runtime().get_option('general/0/use_login_shell'),
                     trap_errors=True
                 )
@@ -1345,6 +1346,7 @@ def run(self):
                 self.logger.debug('Generating the run script')
                 self._job.prepare(
                     commands, environs,
+                    self._current_partition.prepare_cmds,
                     login=rt.runtime().get_option('general/0/use_login_shell'),
                     trap_errors=rt.runtime().get_option(
                         'general/0/trap_job_errors'

reframe/core/schedulers/__init__.py

@@ -330,7 +330,7 @@ def nodelist(self):
     def submit_time(self):
         return self._submit_time
 
-    def prepare(self, commands, environs=None, **gen_opts):
+    def prepare(self, commands, environs=None, prepare_cmds=None, **gen_opts):
         environs = environs or []
         if self.num_tasks <= 0:
             getlogger().debug(f'[F] Flexible node allocation requested')
@@ -357,6 +357,10 @@ def prepare(self, commands, environs=None, **gen_opts):
         with shell.generate_script(self.script_filename,
                                    **gen_opts) as builder:
             builder.write_prolog(self.scheduler.emit_preamble(self))
+            prepare_cmds = prepare_cmds or []
+            for c in prepare_cmds:
+                builder.write_body(c)
+
             builder.write(runtime.emit_loadenv_commands(*environs))
             for c in commands:
                 builder.write_body(c)

reframe/core/systems.py

@@ -22,7 +22,7 @@ class SystemPartition(jsonext.JSONSerializable):
 
     def __init__(self, parent, name, sched_type, launcher_type,
                  descr, access, container_environs, resources,
-                 local_env, environs, max_jobs):
+                 local_env, environs, max_jobs, prepare_cmds):
         getlogger().debug(f'Initializing system partition {name!r}')
         self._parent_system = parent
         self._name = name
@@ -35,6 +35,7 @@ def __init__(self, parent, name, sched_type, launcher_type,
         self._local_env = local_env
         self._environs = environs
         self._max_jobs = max_jobs
+        self._prepare_cmds = prepare_cmds
         self._resources = {r['name']: r['options'] for r in resources}
 
     @property
@@ -98,6 +99,14 @@ def max_jobs(self):
         '''
         return self._max_jobs
 
+    @property
+    def prepare_cmds(self):
+        '''Commands to be emitted before loading the modules.
+
+        :type: :class:`List[str]`
+        '''
+        return self._prepare_cmds
+
     @property
     def name(self):
         '''The name of this partition.
@@ -317,7 +326,8 @@ def create(cls, site_config):
                         modules=site_config.get(f'{partid}/modules'),
                         variables=site_config.get(f'{partid}/variables')
                     ),
-                    max_jobs=site_config.get(f'{partid}/max_jobs')
+                    max_jobs=site_config.get(f'{partid}/max_jobs'),
+                    prepare_cmds=site_config.get(f'{partid}/prepare_cmds')
                 )
             )
 

reframe/schemas/config.json

@@ -220,6 +220,10 @@
                                 "modules": {"$ref": "#/defs/modules_list"},
                                 "variables": {"$ref": "#/defs/envvar_list"},
                                 "max_jobs": {"type": "number"},
+                                "prepare_cmds": {
+                                    "type": "array",
+                                    "items": {"type": "string"}
+                                 },
                                 "resources": {
                                     "type": "array",
                                     "items": {
@@ -472,6 +476,7 @@
         "systems/partitions/modules": [],
         "systems/partitions/variables": [],
         "systems/partitions/environs": [],
-        "systems/partitions/max_jobs": 8
+        "systems/partitions/max_jobs": 8,
+        "systems/partitions/prepare_cmds": []
     }
 }

unittests/test_schedulers.py

@@ -115,27 +115,32 @@ def fake_job(make_job):
     return ret
 
 
-def prepare_job(job, command='hostname', pre_run=None, post_run=None):
+def prepare_job(job, command='hostname',
+                pre_run=None, post_run=None,
+                prepare_cmds=None):
     environs = [Environment(name='foo', modules=['testmod_foo'])]
     pre_run = pre_run or ['echo prerun']
     post_run = post_run or ['echo postrun']
+    prepare_cmds = prepare_cmds or ['echo prepare']
     with rt.module_use('unittests/modules'):
         job.prepare(
             [
                 *pre_run,
                 job.launcher.run_command(job) + ' ' + command,
                 post_run
             ],
-            environs
+            environs,
+            prepare_cmds
         )
 
 
 def assert_job_script_sanity(job):
     '''Assert the sanity of the produced script file.'''
     with open(job.script_filename) as fp:
-        matches = re.findall(r'echo prerun|echo postrun|hostname',
+        matches = re.findall(r'echo prepare|echo prerun|echo postrun|hostname',
                              fp.read())
-        assert ['echo prerun', 'hostname', 'echo postrun'] == matches
+        assert ['echo prepare', 'echo prerun', 'hostname',
+                'echo postrun'] == matches
 
 
 def _expected_slurm_directives(job):
@@ -508,7 +513,8 @@ def test_cancel_with_grace(minimal_job, scheduler, local_only):
     prepare_job(minimal_job,
                 command='sleep 5 &',
                 pre_run=['trap -- "" TERM'],
-                post_run=['echo $!', 'wait'])
+                post_run=['echo $!', 'wait'],
+                prepare_cmds=[''])
     minimal_job.submit()
 
     # Stall a bit here to let the the spawned process start and install its
@@ -552,7 +558,8 @@ def test_cancel_term_ignore(minimal_job, scheduler, local_only):
                 command=os.path.join(fixtures.TEST_RESOURCES_CHECKS,
                                      'src', 'sleep_deeply.sh'),
                 pre_run=[''],
-                post_run=[''])
+                post_run=[''],
+                prepare_cmds=[''])
     minimal_job.submit()
 
     # Stall a bit here to let the the spawned process start and install its