Merge pull request #3573 from ekouts/feat/maint_nodes

[feat] Treat nodes in `MAINTENANCE` state as available for flexible allocations on Slurm `MAINT`-flagged reservations

Author: vkarak

docs/manpage.rst

@@ -1032,6 +1032,7 @@ The way the tests are generated and how they interact with the test filtering op
    - ``all``: Tests will run on all the nodes of their respective valid partitions regardless of the node state.
    - ``avail``: Tests will run on all the nodes of their respective valid partitions that are available for running jobs.
      Note that if a node is currently allocated to another job it is still considered as "available."
+     Also, for ReFrame partitions using the Slurm backends, if this option is used on a reservation with the ``MAINT`` flag set, then nodes in ``MAINTENANCE`` state will also be considered as available.
    - ``NODESTATE``: Tests will run on all the nodes of their respective valid partitions that are exclusively in state ``NODESTATE``.
      If ``NODESTATE`` is not specified, ``idle`` is assumed.
    - ``NODESTATE*``: Tests will run on all the nodes of their respective valid partitions that are at least in state ``NODESTATE``.
@@ -1060,8 +1061,13 @@ The way the tests are generated and how they interact with the test filtering op
       To achieve the previous behaviour, you should use ``--distribute=idle*``.
 
    .. versionchanged:: 4.9
+
       ``--distribute=NODESTATE`` now allows you to specify multiple valid states using the ``|`` character.
 
+   .. versionchanged:: 4.10
+
+      Nodes in ``MAINTENANCE`` state are considered available, if this option is run on a Slurm reservation with the ``MAINT`` flag set.
+
 .. option:: -P, --parameterize=[TEST.]VAR=VAL0,VAL1,...
 
    Parameterize a test on an existing variable or parameter.

reframe/core/schedulers/__init__.py

@@ -100,6 +100,57 @@ def filternodes(self, job, nodes):
         :meta private:
         '''
 
+    def filternodes_by_state(self, nodelist, state):
+        '''Filter nodes by their state
+
+        :arg nodelist: List of :class:`Node` instances to filter.
+        :arg state: The state of the nodes.
+            If ``all``, the initial list is returned untouched.
+            If ``avail``, only the available nodes will be returned.
+            All other values are interpreted as a state string.
+            The pipe character ``|`` can be used as to specify multiple
+            alternative node states.
+            State match is exclusive unless the ``*`` is added at the end of the
+            state string.
+            When defining multiple states using ``|``, ``*`` has to be added at
+            the end of each alternative state for which a non-exclusive match is
+            required.
+
+        :returns: the filtered node list
+
+        .. versionchanged:: 4.9
+
+           Support the ``|`` character to filter according to alternative states.
+
+        .. versionchanged:: 4.10
+
+           Moved inside the :class:`JobShceduler` class.
+        '''
+        if '|' in state:
+            allowed_states = state.split('|')
+            final_nodelist = set()
+            for s in allowed_states:
+                final_nodelist.update(
+                    self.filternodes_by_state(nodelist, s)
+                )
+
+            nodelist = final_nodelist
+        elif state == 'avail':
+            nodelist = {n for n in nodelist if self.is_node_avail(n)}
+        elif state != 'all':
+            if state.endswith('*'):
+                # non-exclusive state match
+                state = state[:-1]
+                nodelist = {
+                    n for n in nodelist if n.in_state(state)
+                }
+            else:
+                nodelist = {
+                    n for n in nodelist if n.in_statex(state)
+                }
+
+        return nodelist
+
     @abc.abstractmethod
     def submit(self, job):
         '''Submit a job.
@@ -153,51 +204,6 @@ def log(self, message, level=DEBUG2):
         getlogger().log(level, f'[S] {self.registered_name}: {message}')
 
 
-def filter_nodes_by_state(nodelist, state):
-    '''Filter nodes by their state
-
-    :arg nodelist: List of :class:`Node` instances to filter.
-    :arg state: The state of the nodes.
-        If ``all``, the initial list is returned untouched.
-        If ``avail``, only the available nodes will be returned.
-        All other values are interpreted as a state string.
-        The pipe character ``|`` can be used as to specify multiple
-        alternative node states.
-        State match is exclusive unless the ``*`` is added at the end of the
-        state string.
-        When defining multiple states using ``|``, ``*`` has to be added at
-        the end of each alternative state for which a non-exclusive match is
-        required.
-
-    :returns: the filtered node list
-
-    .. versionchanged:: 4.9
-       Support the ``|`` character to filter according to alternative states.
-    '''
-    if '|' in state:
-        allowed_states = state.split('|')
-        final_nodelist = set()
-        for s in allowed_states:
-            final_nodelist.update(filter_nodes_by_state(nodelist, s))
-
-        nodelist = final_nodelist
-    elif state == 'avail':
-        nodelist = {n for n in nodelist if n.is_avail()}
-    elif state != 'all':
-        if state.endswith('*'):
-            # non-exclusive state match
-            state = state[:-1]
-            nodelist = {
-                n for n in nodelist if n.in_state(state)
-            }
-        else:
-            nodelist = {
-                n for n in nodelist if n.in_statex(state)
-            }
-
-    return nodelist
-
-
 class Job(jsonext.JSONSerializable, metaclass=JobMeta):
     '''A job descriptor.
 
@@ -618,19 +624,21 @@ def guess_num_tasks(self):
             f'[F] Total available nodes: {len(available_nodes)}'
         )
 
+        available_nodes = self.scheduler.filternodes(self, available_nodes)
+        getlogger().debug(
+            f'[F] Total available after scheduler filter: '
+            f'{len(available_nodes)}'
+        )
+
         # Try to guess the number of tasks now
-        available_nodes = filter_nodes_by_state(
-            available_nodes, self.sched_flex_alloc_nodes.lower()
+        available_nodes = self.scheduler.filternodes_by_state(
+            available_nodes,
+            self.sched_flex_alloc_nodes.lower()
         )
         getlogger().debug(
             f'[F] Total available in state='
             f'{self.sched_flex_alloc_nodes.lower()}: {len(available_nodes)}'
         )
-        available_nodes = self.scheduler.filternodes(self, available_nodes)
-        getlogger().debug(
-            f'[F] Total available after scheduler filter: '
-            f'{len(available_nodes)}'
-        )
         return len(available_nodes) * num_tasks_per_node
 
     def submit(self):
@@ -694,17 +702,6 @@ def in_state(self, state):
                      :class:`False` otherwise.
         '''
 
-    @abc.abstractmethod
-    def is_avail(self):
-        '''Check whether the node is available for scheduling jobs.'''
-
-    def is_down(self):
-        '''Check whether node is down.
-
-        This is the inverse of :func:`is_avail`.
-        '''
-        return not self.is_avail()
-
 
 class AlwaysIdleNode(Node):
     def __init__(self, name):
@@ -715,9 +712,6 @@ def __init__(self, name):
     def name(self):
         return self._name
 
-    def is_avail(self):
-        return True
-
     def in_statex(self, state):
         return state.lower() == self._state
 

reframe/core/schedulers/slurm.py

@@ -147,6 +147,13 @@ def __init__(self):
         self._sched_access_in_submit = self.get_option(
             'sched_access_in_submit'
         )
+        self._available_states = {
+            'ALLOCATED',
+            'COMPLETING',
+            'IDLE',
+            'PLANNED',
+            'RESERVED'
+        }
 
     def make_job(self, *args, **kwargs):
         return _SlurmJob(*args, **kwargs)
@@ -427,21 +434,30 @@ def filternodes(self, job, nodes):
 
         return nodes
 
-    def _get_reservation_nodes(self, reservation):
-        completed = _run_strict('scontrol -a show res %s' % reservation)
-        node_match = re.search(r'(Nodes=\S+)', completed.stdout)
+    def _get_reservation_nodes(self, resv):
+        completed = _run_strict(f'scontrol -a show -o reservations {resv}')
+        self.log(f'reservation info:\n{completed.stdout}')
+
+        node_match = re.search(r'Nodes=(\S+)', completed.stdout)
         if node_match:
             reservation_nodes = node_match[1]
         else:
-            raise JobSchedulerError("could not extract the node names for "
-                                    "reservation '%s'" % reservation)
+            raise JobSchedulerError('could not extract the node names for '
+                                    f'reservation {resv!r}')
+
+        flags_match = re.search(r'Flags=(\S+)', completed.stdout)
+        if flags_match:
+            if 'MAINT' in flags_match.group(1).split(','):
+                self._available_states.add('MAINTENANCE')
 
-        completed = _run_strict('scontrol -a show -o %s' % reservation_nodes)
+        completed = _run_strict(
+            f'scontrol -a show -o nodes {reservation_nodes}'
+        )
         node_descriptions = completed.stdout.splitlines()
         return _create_nodes(node_descriptions)
 
     def _get_nodes_by_name(self, nodespec):
-        completed = osext.run_command('scontrol -a show -o node %s' %
+        completed = osext.run_command('scontrol -a show -o nodes %s' %
                                       nodespec)
         node_descriptions = completed.stdout.splitlines()
         return _create_nodes(node_descriptions)
@@ -594,7 +610,7 @@ def _do_cancel_if_blocked(self, job, reason_descr):
                         self.log(f'Checking if nodes {node_names!r} '
                                  f'are indeed unavailable')
                         nodes = self._get_nodes_by_name(node_names)
-                        if not any(n.is_down() for n in nodes):
+                        if not any(self.is_node_down(n) for n in nodes):
                             return
 
                         self.cancel(job)
@@ -630,6 +646,12 @@ def cancel(self, job):
     def finished(self, job):
         return slurm_state_completed(job.state)
 
+    def is_node_avail(self, node):
+        return node.states <= self._available_states
+
+    def is_node_down(self, node):
+        return not self.is_node_avail(node)
+
 
 @register_scheduler('squeue')
 class SqueueJobScheduler(SlurmJobScheduler):
@@ -727,26 +749,16 @@ def __eq__(self, other):
     def __hash__(self):
         return hash(self.name)
 
+    def __repr__(self):
+        return f'_SlurmNode({self.name!r})'
+
     def in_state(self, state):
         return all([self._states >= set(state.upper().split('+')),
                     self._partitions, self._active_features, self._states])
 
     def in_statex(self, state):
         return self._states == set(state.upper().split('+'))
 
-    def is_avail(self):
-        available_states = {
-            'ALLOCATED',
-            'COMPLETING',
-            'IDLE',
-            'PLANNED',
-            'RESERVED'
-        }
-        return self._states <= available_states
-
-    def is_down(self):
-        return not self.is_avail()
-
     def satisfies(self, slurm_constraint):
         # Convert the Slurm constraint to a Python expression and evaluate it,
         # but restrict our syntax to accept only AND or OR constraints and

reframe/frontend/testgenerators.py

@@ -14,7 +14,7 @@
 from reframe.core.fields import make_convertible
 from reframe.core.logging import getlogger, time_function
 from reframe.core.meta import make_test
-from reframe.core.schedulers import Job, filter_nodes_by_state
+from reframe.core.schedulers import Job
 from reframe.frontend.executors import generate_testcases
 
 
@@ -38,7 +38,9 @@ def getallnodes(state, jobs_cli_options=None):
             f'Total available nodes for {part.name}: {len(available_nodes)}'
         )
 
-        available_nodes = filter_nodes_by_state(available_nodes, state)
+        available_nodes = part.scheduler.filternodes_by_state(
+            available_nodes, state
+        )
         nodes[part.fullname] = [n.name for n in available_nodes]
 
     return nodes