Move is_avail and is_down functions in scheduler

Author: ekouts

reframe/core/schedulers/__init__.py

@@ -153,7 +153,7 @@ def log(self, message, level=DEBUG2):
         getlogger().log(level, f'[S] {self.registered_name}: {message}')
 
 
-def filter_nodes_by_state(nodelist, state):
+def filter_nodes_by_state(nodelist, state, scheduler):
     '''Filter nodes by their state
 
     :arg nodelist: List of :class:`Node` instances to filter.
@@ -178,11 +178,13 @@ def filter_nodes_by_state(nodelist, state):
         allowed_states = state.split('|')
         final_nodelist = set()
         for s in allowed_states:
-            final_nodelist.update(filter_nodes_by_state(nodelist, s))
+            final_nodelist.update(
+                filter_nodes_by_state(nodelist, s, scheduler)
+            )
 
         nodelist = final_nodelist
     elif state == 'avail':
-        nodelist = {n for n in nodelist if n.is_avail()}
+        nodelist = {n for n in nodelist if scheduler.is_node_avail(n)}
     elif state != 'all':
         if state.endswith('*'):
             # non-exclusive state match
@@ -618,19 +620,22 @@ 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.sched_flex_alloc_nodes.lower(),
+            self.scheduler
         )
         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 +699,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 +709,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,7 +147,13 @@ def __init__(self):
         self._sched_access_in_submit = self.get_option(
             'sched_access_in_submit'
         )
-        self.addl_avail_states = set()
+        self.node_available_states = {
+            'ALLOCATED',
+            'COMPLETING',
+            'IDLE',
+            'PLANNED',
+            'RESERVED'
+        }
 
     def make_job(self, *args, **kwargs):
         return _SlurmJob(*args, **kwargs)
@@ -324,7 +330,7 @@ def allnodes(self):
                 'could not retrieve node information') from e
 
         node_descriptions = completed.stdout.splitlines()
-        return _create_nodes(node_descriptions, self.addl_avail_states)
+        return _create_nodes(node_descriptions)
 
     def _get_default_partition(self):
         completed = _run_strict('scontrol -a show -o partitions')
@@ -440,20 +446,20 @@ def _get_reservation_nodes(self, reservation):
         flags_match = re.search(r'Flags=(\S+)', completed.stdout)
         if flags_match:
             if 'MAINT' in flags_match[1].split(','):
-                self.addl_avail_states.add('MAINTENANCE')
+                self.node_available_states.add('MAINTENANCE')
         else:
             self.log(f"could not extract the reservation flags for "
                      f"reservation '{reservation}'")
 
         completed = _run_strict('scontrol -a show -o %s' % reservation_nodes)
         node_descriptions = completed.stdout.splitlines()
-        return _create_nodes(node_descriptions, self.addl_avail_states)
+        return _create_nodes(node_descriptions)
 
     def _get_nodes_by_name(self, nodespec):
         completed = osext.run_command('scontrol -a show -o node %s' %
                                       nodespec)
         node_descriptions = completed.stdout.splitlines()
-        return _create_nodes(node_descriptions, self.addl_avail_states)
+        return _create_nodes(node_descriptions)
 
     def _update_completion_time(self, job, timestamps):
         if job._completion_time is not None:
@@ -603,7 +609,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)
@@ -639,6 +645,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.node_available_states
+
+    def is_node_down(self, node):
+        return not self.is_node_avail(node)
+
 
 @register_scheduler('squeue')
 class SqueueJobScheduler(SlurmJobScheduler):
@@ -700,19 +712,19 @@ def poll(self, *jobs):
             self._cancel_if_pending_too_long(job)
 
 
-def _create_nodes(descriptions, addl_avail_states=None):
+def _create_nodes(descriptions):
     nodes = set()
     for descr in descriptions:
         with suppress(JobSchedulerError):
-            nodes.add(_SlurmNode(descr, addl_avail_states=addl_avail_states))
+            nodes.add(_SlurmNode(descr))
 
     return nodes
 
 
 class _SlurmNode(sched.Node):
     '''Class representing a Slurm node.'''
 
-    def __init__(self, node_descr, addl_avail_states=None):
+    def __init__(self, node_descr):
         self._name = self._extract_attribute('NodeName', node_descr)
         if not self._name:
             raise JobSchedulerError(
@@ -727,15 +739,6 @@ def __init__(self, node_descr, addl_avail_states=None):
             'State', node_descr, sep='+') or set()
         self._descr = node_descr
 
-        self.addl_avail_states = addl_avail_states or set()
-        self.available_states = {
-            'ALLOCATED',
-            'COMPLETING',
-            'IDLE',
-            'PLANNED',
-            'RESERVED'
-        } | self.addl_avail_states
-
     def __eq__(self, other):
         if not isinstance(other, type(self)):
             return NotImplemented
@@ -752,12 +755,6 @@ def in_state(self, state):
     def in_statex(self, state):
         return self._states == set(state.upper().split('+'))
 
-    def is_avail(self):
-        return self._states <= self.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

@@ -38,7 +38,11 @@ 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 = filter_nodes_by_state(
+            available_nodes,
+            state,
+            part.scheduler
+        )
         nodes[part.fullname] = [n.name for n in available_nodes]
 
     return nodes

unittests/test_schedulers.py

@@ -41,6 +41,8 @@ def slurm_only(scheduler):
     if scheduler.registered_name not in ('slurm', 'squeue'):
         pytest.skip('test is relevant only for Slurm backends')
 
+    return scheduler
+
 
 @pytest.fixture
 def local_only(scheduler):
@@ -1455,7 +1457,8 @@ def test_slurm_node_in_state(slurm_node_allocated,
 
 def test_slurm_node_is_down(slurm_node_allocated,
                             slurm_node_idle,
-                            slurm_node_nopart):
-    assert not slurm_node_allocated.is_down()
-    assert not slurm_node_idle.is_down()
-    assert slurm_node_nopart.is_down()
+                            slurm_node_nopart,
+                            slurm_only):
+    assert not slurm_only().is_node_down(slurm_node_allocated)
+    assert not slurm_only().is_node_down(slurm_node_idle)
+    assert slurm_only().is_node_down(slurm_node_nopart)