diff --git a/reframe/core/schedulers/slurm.py b/reframe/core/schedulers/slurm.py index 773ef893d9..a0c9f5620a 100644 --- a/reframe/core/schedulers/slurm.py +++ b/reframe/core/schedulers/slurm.py @@ -161,6 +161,15 @@ def _get_all_nodes(self): node_descriptions = completed.stdout.splitlines() return {SlurmNode(descr) for descr in node_descriptions} + def _get_default_partition(self): + completed = _run_strict('scontrol -a show -o partitions') + partition_match = re.search(r'PartitionName=(?P\S+)\s+' + r'.*Default=YES.*', completed.stdout) + if partition_match: + return partition_match.group('partition') + + return None + def get_partition_nodes(self): nodes = self._get_all_nodes() return self.filter_nodes(nodes, self.sched_access) @@ -187,10 +196,14 @@ def filter_nodes(self, nodes, options): if partitions: partitions = set(partitions.strip().split(',')) - nodes = {n for n in nodes if n.partitions >= partitions} - getlogger().debug( - 'flex_alloc_tasks: filtering nodes by partition(s) %s: ' - 'available nodes now: %s' % (partitions, len(nodes))) + else: + default_partition = self._get_default_partition() + partitions = {default_partition} if default_partition else set() + + nodes = {n for n in nodes if n.partitions >= partitions} + getlogger().debug( + 'flex_alloc_tasks: filtering nodes by partition(s) %s: ' + 'available nodes now: %s' % (partitions, len(nodes))) if constraints: constraints = set(constraints.strip().split(',')) diff --git a/unittests/test_schedulers.py b/unittests/test_schedulers.py index d8bdc982fb..cbab7176e6 100644 --- a/unittests/test_schedulers.py +++ b/unittests/test_schedulers.py @@ -396,6 +396,9 @@ def test_guess_num_tasks(self): # monkey patch `get_partition_nodes()` to simulate extraction of # slurm nodes through the use of `scontrol show` self.testjob.get_partition_nodes = lambda: set() + # monkey patch `_get_default_partition()` to simulate extraction + # of the default partition through the use of `scontrol show` + self.testjob._get_default_partition = lambda: 'pdef' self.assertEqual(self.testjob.guess_num_tasks(), 0) @@ -502,7 +505,7 @@ def create_dummy_nodes(obj): 'RealMemory=32220 AllocMem=0 FreeMem=10000 ' 'Sockets=1 Boards=1 State=MAINT+DRAIN ' 'ThreadsPerCore=2 TmpDisk=0 Weight=1 Owner=N/A ' - 'MCS_label=N/A Partitions=p1,p2 ' + 'MCS_label=N/A Partitions=p1,p2,pdef ' 'BootTime=01 Jan 2018 ' 'SlurmdStartTime=01 Jan 2018 ' 'CfgTRES=cpu=24,mem=32220M ' @@ -520,7 +523,7 @@ def create_dummy_nodes(obj): 'RealMemory=32220 AllocMem=0 FreeMem=10000 ' 'Sockets=1 Boards=1 State=MAINT+DRAIN ' 'ThreadsPerCore=2 TmpDisk=0 Weight=1 Owner=N/A ' - 'MCS_label=N/A Partitions=p2,p3' + 'MCS_label=N/A Partitions=p2,p3,pdef ' 'BootTime=01 Jan 2018 ' 'SlurmdStartTime=01 Jan 2018 ' 'CfgTRES=cpu=24,mem=32220M ' @@ -538,7 +541,7 @@ def create_dummy_nodes(obj): 'RealMemory=32220 AllocMem=0 FreeMem=10000 ' 'Sockets=1 Boards=1 State=IDLE ' 'ThreadsPerCore=2 TmpDisk=0 Weight=1 Owner=N/A ' - 'MCS_label=N/A Partitions=p1,p3 ' + 'MCS_label=N/A Partitions=p1,p3,pdef ' 'BootTime=01 Jan 2018 ' 'SlurmdStartTime=01 Jan 2018 ' 'CfgTRES=cpu=24,mem=32220M ' @@ -556,7 +559,7 @@ def create_dummy_nodes(obj): 'RealMemory=32220 AllocMem=0 FreeMem=10000 ' 'Sockets=1 Boards=1 State=IDLE ' 'ThreadsPerCore=2 TmpDisk=0 Weight=1 Owner=N/A ' - 'MCS_label=N/A Partitions=p1,p3 ' + 'MCS_label=N/A Partitions=p1,p3,pdef ' 'BootTime=01 Jan 2018 ' 'SlurmdStartTime=01 Jan 2018 ' 'CfgTRES=cpu=24,mem=32220M ' @@ -589,6 +592,9 @@ def setUp(self): # monkey patch `_get_all_nodes` to simulate extraction of # slurm nodes through the use of `scontrol show` self.testjob._get_all_nodes = self.create_dummy_nodes + # monkey patch `_get_default_partition` to simulate extraction + # of the default partition + self.testjob._get_default_partition = lambda: 'pdef' self.testjob._sched_flex_alloc_tasks = 'all' self.testjob._num_tasks_per_node = 4 self.testjob._num_tasks = 0