Merge pull request #3569 from teojgo/feat/support_or_in_node_state

[feat] Support OR-ing node states in flexible node allocation

Author: vkarak

docs/manpage.rst

@@ -1059,6 +1059,8 @@ The way the tests are generated and how they interact with the test filtering op
       ``--distribute=NODESTATE`` now matches nodes that are exclusively in state ``NODESTATE``, so that the default ``--distribute=idle`` will match only the Slurm nodes that are in the ``IDLE`` state exclusively.
       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.
 
 .. option:: -P, --parameterize=[TEST.]VAR=VAL0,VAL1,...
 

reframe/core/schedulers/__init__.py

@@ -161,11 +161,27 @@ def filter_nodes_by_state(nodelist, state):
         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 state == 'avail':
+    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('*'):

unittests/test_schedulers.py

@@ -1247,6 +1247,31 @@ def test_flex_alloc_not_enough_nodes_constraint_expr(make_flexible_job):
         prepare_job(job)
 
 
+def test_flex_alloc_alloc_state_OR(make_flexible_job):
+    job = make_flexible_job('allocated|idle')
+    job.options = ['--partition=p3']
+    prepare_job(job)
+    assert job.num_tasks == 12
+
+    job = make_flexible_job('maint*|idle')
+    prepare_job(job)
+    assert job.num_tasks == 16
+
+    job = make_flexible_job('maint|avail')
+    job.options = ['--partition=p1']
+    prepare_job(job)
+    assert job.num_tasks == 12
+
+    job = make_flexible_job('all|idle')
+    prepare_job(job)
+    assert job.num_tasks == 16
+
+    job = make_flexible_job('allocated|idle|maint')
+    job.options = ['--partition=p1']
+    prepare_job(job)
+    assert job.num_tasks == 12
+
+
 @pytest.fixture
 def slurm_node_allocated():
     return _SlurmNode(