Merge pull request #3152 from ekouts/feat/env_sched_bind_options_2

[feat] Add support for associating scheduler resources with an environment

Author: ekouts

docs/config_reference.rst

@@ -689,7 +689,7 @@ ReFrame can launch containerized applications, but you need to configure properl
 Custom Job Scheduler Resources
 ==============================
 
-ReFrame allows you to define custom scheduler resources for each partition that you can then transparently access through the :attr:`~reframe.core.pipeline.RegressionTest.extra_resources` attribute of a regression test.
+ReFrame allows you to define custom scheduler resources for each partition that can then be transparently accessed through the :attr:`~reframe.core.pipeline.RegressionTest.extra_resources` attribute of a test or from an environment.
 
 .. py:attribute:: systems.partitions.resources.name
 
@@ -953,6 +953,18 @@ They are associated with `system partitions <#system-partition-configuration>`__
    It first looks for definitions for the current partition, then for the containing system and, finally, for global definitions (the ``*`` pseudo-system).
 
 
+.. py:attribute:: environments.resources
+
+   :required: No
+   :default: ``{}``
+
+   Scheduler resources associated with this environments.
+   
+   This is the equivalent of a test's :attr:`~reframe.core.pipeline.RegressionTest.extra_resources`.
+
+   .. versionadded:: 4.6
+
+
 .. _logging-config-reference:
 
 Logging Configuration

reframe/core/environments.py

@@ -242,6 +242,7 @@ def __init__(self,
                  cxxflags=None,
                  fflags=None,
                  ldflags=None,
+                 resources=None,
                  **kwargs):
         super().__init__(name, modules, env_vars, extras, features,
                          prepare_cmds)
@@ -254,6 +255,7 @@ def __init__(self,
         self._cxxflags = cxxflags or []
         self._fflags   = fflags   or []
         self._ldflags  = ldflags  or []
+        self._resources = resources or {}
 
     @property
     def cc(self):
@@ -326,3 +328,13 @@ def nvcc(self):
         :type: :class:`str`
         '''
         return self._nvcc
+
+    @property
+    def resources(self):
+        '''The scheduler resources associated with this environment.
+
+        .. versionadded:: 4.6.0
+
+        :type: :class:`Dict[str, object]`
+        '''
+        return self._resources

reframe/core/pipeline.py

@@ -2036,7 +2036,13 @@ def _get_cp_env():
 
     def _map_resources_to_jobopts(self):
         resources_opts = []
-        for r, v in self.extra_resources.items():
+        # Combine all resources into one dictionary, where extra_resources
+        # can overwrite _current_environ.resources
+        combined_resources = {
+            **self._current_environ.resources,
+            **self.extra_resources
+        }
+        for r, v in combined_resources.items():
             resources_opts += self._current_partition.get_resource(r, **v)
 
         return resources_opts

reframe/core/systems.py

@@ -549,7 +549,8 @@ def create(cls, site_config):
                     cflags=site_config.get(f'environments/@{e}/cflags'),
                     cxxflags=site_config.get(f'environments/@{e}/cxxflags'),
                     fflags=site_config.get(f'environments/@{e}/fflags'),
-                    ldflags=site_config.get(f'environments/@{e}/ldflags')
+                    ldflags=site_config.get(f'environments/@{e}/ldflags'),
+                    resources=site_config.get(f'environments/@{e}/resources')
                 ) for e in site_config.get(f'{partid}/environs')
                 if any(re.match(pattern, e) for pattern in env_patt)
             ]

reframe/schemas/config.json

@@ -410,6 +410,16 @@
                         "type": "array",
                         "items": {"$ref": "#/defs/alphanum_string"}
                     },
+                    "resources": {
+                        "type": "object",
+                        "propertyNames": {
+                            "pattern": "^[a-zA-Z_][a-zA-Z0-9_]*$"
+                        },
+                        "additionalProperties": {
+                            "type": "object",
+                            "additionalProperties": true
+                        }
+                    },
                     "target_systems": {"$ref": "#/defs/system_ref"}
                 },
                 "required": ["name"],

unittests/resources/config/settings.py

@@ -61,6 +61,13 @@ def hostname():
                                 '#DW jobdw capacity={capacity}',
                                 '#DW stage_in source={stagein_src}'
                             ]
+                        },
+                        {
+                            'name': 'uenv',
+                            'options': [
+                                '--mount={mount}',
+                                '--file={file}'
+                            ]
                         }
                     ],
                     'features': ['cuda', 'mpi'],

unittests/test_config.py

@@ -293,7 +293,7 @@ def test_select_subconfig(site_config):
     assert (site_config.get('systems/0/partitions/0/environs') ==
             ['PrgEnv-gnu', 'builtin'])
     assert site_config.get('systems/0/partitions/0/descr') == 'GPU partition'
-    assert len(site_config.get('systems/0/partitions/0/resources')) == 2
+    assert len(site_config.get('systems/0/partitions/0/resources')) == 3
     assert (site_config.get('systems/0/partitions/0/resources/@gpu/name') ==
             'gpu')
     assert site_config.get('systems/0/partitions/0/modules') == [
@@ -461,6 +461,12 @@ def test_system_create(site_config):
     assert resources_spec == ['#DW jobdw capacity=100GB',
                               '#DW stage_in source=/foo']
 
+    env_resources_spec = partition.get_resource(
+        'uenv', mount='mount_point', file='file_path'
+    )
+    assert env_resources_spec == ['--mount=mount_point',
+                                  '--file=file_path']
+
     # Check processor info
     assert partition.processor.info is not None
     assert partition.processor.topology is not None