Merge pull request #2887 from teojgo/feat/prepare_cmds_env

[feat] Support `prepare_cmds` at the environment level

docs/config_reference.rst

@@ -786,6 +786,16 @@ They are associated with `system partitions <#system-partition-configuration>`__
    .. versionadded:: 3.9.1
 
 
+.. py:attribute:: environments.prepare_cmds
+
+   :required: No
+   :default: ``[]``
+
+   List of shell commands to be emitted before any commands that load the environment.
+
+   .. versionadded:: 4.3.0
+
+
 .. py:attribute:: environments.cc
 
    :required: No

reframe/core/environments.py

@@ -39,7 +39,7 @@ class Environment(jsonext.JSONSerializable):
     '''
 
     def __init__(self, name, modules=None, env_vars=None,
-                 extras=None, features=None):
+                 extras=None, features=None, prepare_cmds=None):
         modules = modules or []
         env_vars = env_vars or []
         self._name = name
@@ -56,6 +56,7 @@ def __init__(self, name, modules=None, env_vars=None,
 
         self._extras = extras or {}
         self._features = features or []
+        self._prepare_cmds = prepare_cmds or []
 
     @property
     def name(self):
@@ -135,6 +136,16 @@ def features(self):
         '''
         return self._features
 
+    @property
+    def prepare_cmds(self):
+        '''The prepare commands associated with this environment.
+
+        .. versionadded:: 4.3.0
+
+        :type: :class:`List[str]`
+        '''
+        return util.SequenceView(self._prepare_cmds)
+
     def __eq__(self, other):
         if not isinstance(other, type(self)):
             return NotImplemented
@@ -220,6 +231,7 @@ def __init__(self,
                  env_vars=None,
                  extras=None,
                  features=None,
+                 prepare_cmds=None,
                  cc='cc',
                  cxx='CC',
                  ftn='ftn',
@@ -230,7 +242,8 @@ def __init__(self,
                  fflags=None,
                  ldflags=None,
                  **kwargs):
-        super().__init__(name, modules, env_vars, extras, features)
+        super().__init__(name, modules, env_vars, extras, features,
+                         prepare_cmds)
         self._cc = cc
         self._cxx = cxx
         self._ftn = ftn

reframe/core/runtime.py

@@ -242,6 +242,9 @@ def _load_cmds_tracked(**module):
     env_snapshot = snapshot()
     commands = []
     for env in environs:
+        for cmd in env.prepare_cmds:
+            commands.append(cmd)
+
         for mod in env.modules_detailed:
             if runtime().get_option('general/0/resolve_module_conflicts'):
                 commands += _load_cmds_tracked(**mod)

reframe/core/systems.py

@@ -531,6 +531,9 @@ def create(cls, site_config):
                     env_vars=site_config.get(f'environments/@{e}/env_vars'),
                     extras=site_config.get(f'environments/@{e}/extras'),
                     features=site_config.get(f'environments/@{e}/features'),
+                    prepare_cmds=site_config.get(
+                        f'environments/@{e}/prepare_cmds'
+                    ),
                     cc=site_config.get(f'environments/@{e}/cc'),
                     cxx=site_config.get(f'environments/@{e}/cxx'),
                     ftn=site_config.get(f'environments/@{e}/ftn'),

reframe/schemas/config.json

@@ -388,6 +388,10 @@
                         "type": "array",
                         "items": {"type": "string"}
                     },
+                    "prepare_cmds": {
+                        "type": "array",
+                        "items": {"type": "string"}
+                    },
                     "extras": {
                         "type": "object",
                         "propertyNames": {

unittests/resources/config/settings.py

@@ -192,6 +192,7 @@ def hostname():
             'modules': [
                 {'name': 'PrgEnv-gnu', 'collection': False, 'path': None}
             ],
+            'prepare_cmds': ['echo prepare1', 'echo prepare2'],
             'extras': {
                 'foo': 2,
                 'bar': 'y'

unittests/test_environments.py

@@ -53,7 +53,9 @@ def env0():
     return env.Environment(
         'TestEnv1', ['testmod_foo'],
         [('_var0', 'val1'), ('_var2', '$_var0'), ('_var3', '${_var1}')],
-        {'foo': 1, 'bar': 2}
+        {'foo': 1, 'bar': 2},
+        ['feat1', 'feat2'],
+        ['echo prep1', 'echo prep2']
     )
 
 
@@ -81,16 +83,24 @@ def test_env_construction(base_environ, env0):
     # Assert extras
     assert env0.extras == {'foo': 1, 'bar': 2}
 
+    # Assert features
+    assert env0.features == ['feat1', 'feat2']
+
+    # Assert prepare_cmds
+    assert env0.prepare_cmds == ['echo prep1', 'echo prep2']
+
 
 def test_progenv_construction():
     environ = env.ProgEnvironment('myenv',
                                   modules=['modfoo'],
                                   env_vars=[('var', 'val')],
-                                  extras={'foo': 'bar'})
+                                  extras={'foo': 'bar'},
+                                  prepare_cmds=['echo prep1', 'echo prep2'])
     assert environ.name == 'myenv'
     assert environ.modules == ['modfoo']
     assert environ.env_vars == {'var': 'val'}
     assert environ.extras == {'foo': 'bar'}
+    assert environ.prepare_cmds == ['echo prep1', 'echo prep2']
 
 
 def test_env_snapshot(base_environ, env0, env1):
@@ -248,6 +258,7 @@ def test_emit_loadenv_commands(base_environ, user_runtime,
                                modules_system, env0):
     ms = rt.runtime().modules_system
     expected_commands = [
+        *env0.prepare_cmds,
         ms.emit_load_commands('testmod_foo')[0],
         'export _var0=val1',
         'export _var2=$_var0',
@@ -256,6 +267,16 @@ def test_emit_loadenv_commands(base_environ, user_runtime,
     assert expected_commands == rt.emit_loadenv_commands(env0)
 
 
+def test_emit_loadenv_nomod_commands(base_environ, env0):
+    expected_commands = [
+        *env0.prepare_cmds,
+        'export _var0=val1',
+        'export _var2=$_var0',
+        'export _var3=${_var1}',
+    ]
+    assert expected_commands == rt.emit_loadenv_commands(env0)
+
+
 def test_emit_loadenv_commands_ignore_confict(base_environ,
                                               make_exec_ctx, env0):
     if not test_util.has_sane_modules_system():
@@ -272,6 +293,7 @@ def test_emit_loadenv_commands_ignore_confict(base_environ,
     with ms.change_module_path(test_util.TEST_MODULES):
         ms.load_module('testmod_bar')
         expected_commands = [
+            *env0.prepare_cmds,
             ms.emit_load_commands('testmod_foo')[0],
             'export _var0=val1',
             'export _var2=$_var0',
@@ -286,6 +308,7 @@ def test_emit_loadenv_commands_with_confict(base_environ, user_runtime,
     modules_system.load_module('testmod_bar')
     ms = rt.runtime().modules_system
     expected_commands = [
+        *env0.prepare_cmds,
         ms.emit_unload_commands('testmod_bar')[0],
         ms.emit_load_commands('testmod_foo')[0],
         'export _var0=val1',