From 03ce00ab9f13698901e98fcef2219ecb7efb5daa Mon Sep 17 00:00:00 2001
From: Vasileios Karakasis <vkarakasis@nvidia.com>
Date: Fri, 22 Nov 2024 20:08:14 +0100
Subject: [PATCH] Treat correctly Slurm constrains with hyphens when filtering
 nodes

---
 reframe/core/schedulers/slurm.py | 39 +++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/reframe/core/schedulers/slurm.py b/reframe/core/schedulers/slurm.py
index 7b4a4bff0..235759ab0 100644
--- a/reframe/core/schedulers/slurm.py
+++ b/reframe/core/schedulers/slurm.py
@@ -351,6 +351,7 @@ def filternodes(self, job, nodes):
         option_parser.add_argument('-w', '--nodelist')
         option_parser.add_argument('-C', '--constraint')
         option_parser.add_argument('-x', '--exclude')
+        self.log(f'Filtering by Slurm options: {" ".join(options)}')
         parsed_args, _ = option_parser.parse_known_args(options)
         reservation = parsed_args.reservation
         partitions = parsed_args.partition
@@ -363,7 +364,7 @@ def filternodes(self, job, nodes):
         else:
             nodes = {node for node in nodes if not node.in_state('RESERVED')}
 
-        self.log(f'[F] Filtering nodes by reservation={reservation}: '
+        self.log(f'Filtering nodes by reservation={reservation}: '
                  f'available nodes now: {len(nodes)}')
 
         if partitions:
@@ -377,27 +378,27 @@ def filternodes(self, job, nodes):
             )
             partitions = {default_partition} if default_partition else set()
             self.log(
-                f'[F] No partition specified; using {default_partition!r}'
+                f'No partition specified; using {default_partition!r}'
             )
 
         nodes = {n for n in nodes if n.partitions >= partitions}
-        self.log(f'[F] Filtering nodes by partition(s) {partitions}: '
+        self.log(f'Filtering nodes by partition(s) {partitions}: '
                  f'available nodes now: {len(nodes)}')
         if constraints:
             nodes = {n for n in nodes if n.satisfies(constraints)}
-            self.log(f'[F] Filtering nodes by constraint(s) {constraints}: '
+            self.log(f'Filtering nodes by constraint(s) {constraints}: '
                      f'available nodes now: {len(nodes)}')
 
         if nodelist:
             nodelist = nodelist.strip()
             nodes &= self._get_nodes_by_name(nodelist)
-            self.log(f'[F] Filtering nodes by nodelist: {nodelist}: '
+            self.log(f'Filtering nodes by nodelist: {nodelist}: '
                      f'available nodes now: {len(nodes)}')
 
         if exclude_nodes:
             exclude_nodes = exclude_nodes.strip()
             nodes -= self._get_nodes_by_name(exclude_nodes)
-            self.log(f'[F] Excluding node(s): {exclude_nodes}: '
+            self.log(f'Excluding node(s): {exclude_nodes}: '
                      f'available nodes now: {len(nodes)}')
 
         return nodes
@@ -711,17 +712,29 @@ def is_down(self):
         return not self.is_avail()
 
     def satisfies(self, slurm_constraint):
+        def _replacemany(s, replacements):
+            for src, dst in replacements:
+                s = s.replace(src, dst)
+
+            return s
+
         # Convert the Slurm constraint to a Python expression and evaluate it,
         # but restrict our syntax to accept only AND or OR constraints and
-        # their combinations
-        if not re.match(r'^[\w\d\(\)\|\&]*$', slurm_constraint):
+        # their combinations; to properly treat `-` in constraints we need to
+        # convert them to valid Python identifiers before evaluating the
+        # constraint.
+        if not re.match(r'^[\-\w\d\(\)\|\&]*$', slurm_constraint):
             return False
 
-        names = {grp[0]
-                 for grp in re.finditer(r'(\w(\w|\d)*)', slurm_constraint)}
-        expr = slurm_constraint.replace('|', ' or ').replace('&', ' and ')
-        vars = {n: True for n in self.active_features}
-        vars.update({n: False for n in names - self.active_features})
+        names = {
+            grp[0] for grp in re.finditer(r'[\-\w][\-\w\d]*', slurm_constraint)
+        }
+        expr = _replacemany(slurm_constraint,
+                            [('-', '_'), ('|', ' or '), ('&', ' and ')])
+        vars = {n.replace('-', '_'): True for n in self.active_features}
+        vars.update({
+            n.replace('-', '_'): False for n in names - self.active_features
+        })
         try:
             return eval(expr, {}, vars)
         except BaseException: