Skip to content

Commit

Permalink
[stream] Set number of CPUs/task based on available CPUs (#346)
Browse files Browse the repository at this point in the history
* [stream] Set number of CPUs/task based on available CPUs

* [stream] Always run without MPI launcher
  • Loading branch information
giordano authored Oct 29, 2024
1 parent e290808 commit 939514f
Showing 1 changed file with 26 additions and 5 deletions.
31 changes: 26 additions & 5 deletions benchmarks/examples/stream/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
# Import modules from reframe and excalibur-tests
import reframe as rfm
import reframe.utility.sanity as sn
from reframe.core.backends import getlauncher
from benchmarks.modules.utils import SpackTest

@rfm.simple_test
Expand All @@ -22,11 +23,6 @@ class StreamBenchmark(SpackTest):
num_tasks = 1

time_limit = '5m'
num_cpus_per_task = 128
env_vars = {
'OMP_NUM_THREADS': f'{num_cpus_per_task}',
'OMP_PLACES': 'cores'
}
use_multithreading = False

## Reference performance values for Archer2
Expand All @@ -39,6 +35,31 @@ class StreamBenchmark(SpackTest):
}
}


# Automatically set default value of `num_cpus_per_task` based on number of
# CPUs on a node.
@run_after('setup')
def setup_num_tasks(self):
self.set_var_default(
'num_cpus_per_task',
(self.current_partition.processor.num_cpus or 1) //
min(1, (self.current_partition.processor.num_cpus_per_core or 1)))
self.env_vars['OMP_NUM_THREADS'] = f'{self.num_cpus_per_task}'
self.env_vars['OMP_PLACES'] = 'cores'


# Unlike many of the other benchmarks we support, this one doesn't use
# MPI. Since in principle an MPI launcher may not be available on the
# compute node out-of-the-box unless explicitly requested, to avoid issues
# in case `mpirun`/`mpiexec` aren't readily available we always force the
# local launcher:
# <https://reframe-hpc.readthedocs.io/en/v4.6.3/tutorial.html#replacing-the-launch-command>.
# This function is not needed for all other benchmarks which do need MPI.
@run_before('run')
def replace_launcher(self):
self.job.launcher = getlauncher('local')()


## Build configuration
## Comment/uncomment the appropriate one

Expand Down

0 comments on commit 939514f

Please sign in to comment.