Include RESERVED nodes for Slurm is_avail()

reframe-hpc · Oct 22, 2024 · d56bf83 · d56bf83
1 parent 4c7fef0
commit d56bf83
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 8 deletions.
diff --git a/reframe/core/schedulers/__init__.py b/reframe/core/schedulers/__init__.py
@@ -160,7 +160,7 @@ def filter_nodes_by_state(nodelist, state):
     :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 interpretes as a state string.
+        All other values are interpreted as a state string.
         State match is exclusive unless the ``*`` is added at the end of the
         state string.
     :returns: the filtered node list
@@ -169,7 +169,7 @@ def filter_nodes_by_state(nodelist, state):
         nodelist = {n for n in nodelist if n.is_avail()}
     elif state != 'all':
         if state.endswith('*'):
-            # non-exclusive stat match
+            # non-exclusive state match
             state = state[:-1]
             nodelist = {
                 n for n in nodelist if n.in_state(state)
@@ -180,8 +180,6 @@ def filter_nodes_by_state(nodelist, state):
             }
 
     return nodelist
-    nodes[part.fullname] = [n.name for n in nodelist]
-
 
 
 class Job(jsonext.JSONSerializable, metaclass=JobMeta):
@@ -603,7 +601,13 @@ def guess_num_tasks(self):
         available_nodes = filter_nodes_by_state(
             available_nodes, self.sched_flex_alloc_nodes.lower()
         )
+        getlogger().debug(
+            f'[F] Total available in state={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: {len(available_nodes)}'
+        )
         return len(available_nodes) * num_tasks_per_node
 
     def submit(self):

diff --git a/reframe/core/schedulers/slurm.py b/reframe/core/schedulers/slurm.py
@@ -358,8 +358,12 @@ def filternodes(self, job, nodes):
         if reservation:
             reservation = reservation.strip()
             nodes &= self._get_reservation_nodes(reservation)
-            self.log(f'[F] Filtering nodes by reservation {reservation}: '
-                     f'available nodes now: {len(nodes)}')
+        else:
+            nodes = {node for node in nodes if not node.in_state('RESERVED')}
+
+        self.log(f'[F] Filtering nodes by reservation={reservation}: '
+                 f'available nodes now: {len(nodes)}')
+
         if partitions:
             partitions = set(partitions.strip().split(','))
         else:
@@ -693,8 +697,13 @@ def in_statex(self, state):
         return self._states == set(state.upper().split('+'))
 
     def is_avail(self):
-        return any(self.in_statex(s)
-                   for s in ('ALLOCATED', 'COMPLETING', 'IDLE'))
+        available_states = {
+            'ALLOCATED',
+            'COMPLETING',
+            'IDLE',
+            'RESERVED',
+        }
+        return self._states <= available_states
 
     def is_down(self):
         return not self.is_avail()