EAMxx: allow non-contiguous set of cpus in test-all-scream

Author: bartgol

components/eamxx/scripts/test_all_scream.py

@@ -398,7 +398,7 @@ def generate_cmake_config(self, test, for_ctest=False):
         return result
 
     ###############################################################################
-    def get_taskset_range(self, test, for_compile=True):
+    def get_taskset_resources(self, test, for_compile=True):
     ###############################################################################
         res_name = "compile_res_count" if for_compile else "testing_res_count"
 
@@ -419,13 +419,11 @@ def get_taskset_range(self, test, for_compile=True):
 
         expect(offset < len(affinity_cp),
                f"Offset {offset} out of bounds (max={len(affinity_cp)}) for test {test}\naffinity_cp: {affinity_cp}")
-        start = affinity_cp[offset]
-        end = start
-        for i in range(1, getattr(test, res_name)):
-            expect(affinity_cp[offset+i] == start+i, f"Could not get contiguous range for test {test}")
-            end = affinity_cp[offset+i]
+        resources = []
+        for i in range(0, getattr(test, res_name)):
+            resources.append(affinity_cp[offset+i])
 
-        return start, end
+        return resources
 
     ###############################################################################
     def create_ctest_resource_file(self, test, build_dir):
@@ -438,7 +436,7 @@ def create_ctest_resource_file(self, test, build_dir):
         # res group is where we usually bind an individual MPI rank.
         # The id of the res groups on is offset so that it is unique across all builds
 
-        start, end = self.get_taskset_range(test, for_compile=False)
+        resources = self.get_taskset_resources(test, for_compile=False)
 
         data = {}
 
@@ -448,7 +446,7 @@ def create_ctest_resource_file(self, test, build_dir):
         # We add leading zeroes to ensure that ids will sort correctly
         # both alphabetically and numerically
         devices = []
-        for res_id in range(start,end+1):
+        for res_id in resources:
             devices.append({"id":f"{res_id:05d}"})
 
         # Add resource groups
@@ -457,7 +455,7 @@ def create_ctest_resource_file(self, test, build_dir):
         with (build_dir/"ctest_resource_file.json").open("w", encoding="utf-8") as outfile:
             json.dump(data,outfile,indent=2)
 
-        return (end-start)+1
+        return len(resources)
 
     ###############################################################################
     def generate_ctest_config(self, cmake_config, extra_configs, test):
@@ -498,11 +496,11 @@ def generate_ctest_config(self, cmake_config, extra_configs, test):
 
         # Ctest can only competently manage test pinning across a single instance of ctest. For
         # multiple concurrent instances of ctest, we have to help it. It's OK to use the compile_res_count
-        # taskset range even though the ctest script is also running the tests
+        # taskset resources even though the ctest script is also running the tests
         if self._parallel:
-            start, end = self.get_taskset_range(test)
+            resources = self.get_taskset_resources(test)
             result = result.replace("'", r"'\''") # handle nested quoting
-            result = f"taskset -c {start}-{end} sh -c '{result}'"
+            result = f"taskset -c {','.join([str(r) for r in resources])} sh -c '{result}'"
 
         return result
 
@@ -539,8 +537,8 @@ def generate_baselines(self, test):
 
         cmd = f"make -j{test.compile_res_count}"
         if self._parallel:
-            start, end = self.get_taskset_range(test)
-            cmd = f"taskset -c {start}-{end} sh -c '{cmd}'"
+            resources = self.get_taskset_resources(test)
+            cmd = f"taskset -c {','.join([str(r) for r in resources])} sh -c '{cmd}'"
 
         stat, _, err = run_cmd(cmd, from_dir=test_dir, verbose=True)