From 939514fcddcae1cdce1aa3b47cd02034e1499585 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mos=C3=A8=20Giordano?=
 <765740+giordano@users.noreply.github.com>
Date: Tue, 29 Oct 2024 11:28:26 +0000
Subject: [PATCH] [stream] Set number of CPUs/task based on available CPUs
 (#346)

* [stream] Set number of CPUs/task based on available CPUs

* [stream] Always run without MPI launcher
---
 benchmarks/examples/stream/stream.py | 31 +++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/benchmarks/examples/stream/stream.py b/benchmarks/examples/stream/stream.py
index 4786e5e7..3ccbd179 100644
--- a/benchmarks/examples/stream/stream.py
+++ b/benchmarks/examples/stream/stream.py
@@ -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
@@ -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
@@ -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