Add default benchmark scheduler

gematria/datasets/pipelines/benchmark_cpu_scheduler.py

@@ -14,7 +14,7 @@
 
 import abc
 from typing_extensions import override
-from collections.abc import Iterable
+from collections.abc import Collection
 import os
 import re
 
@@ -71,6 +71,19 @@ def verify(self):
     pass
 
 
+def _get_neighboring_threads(cpu_index: int) -> list[int]:
+  with open(
+      f'/sys/devices/system/cpu/cpu{cpu_index}/topology/thread_siblings_list'
+  ) as thread_sibling_list_handle:
+    neighboring_threads_strings = re.split(
+        '[-,]+', thread_sibling_list_handle.read().strip()
+    )
+    neighboring_threads = [
+        int(cpu_index_str) for cpu_index_str in neighboring_threads_strings
+    ]
+  return neighboring_threads
+
+
 class DefaultBenchmarkScheduler(BenchmarkScheduler):
   """A BenchmarkScheduler that schedules processes separately.
 
@@ -87,38 +100,32 @@ class DefaultBenchmarkScheduler(BenchmarkScheduler):
   def __init__(self):
     self._cpu_mask = []
 
-  @staticmethod
-  def _get_neighboring_threads(cpu_index: int) -> list[int]:
-    with open(
-        f'/sys/devices/system/cpu/cpu{cpu_index}/topology/thread_siblings_list'
-    ) as thread_sibling_list_handle:
-      neighboring_threads_strings = re.split(
-          r'[-,]+', thread_sibling_list_handle.read().strip()
-      )
-      neighboring_threads = [
-          int(cpu_index_str) for cpu_index_str in neighboring_threads_strings
-      ]
-    return neighboring_threads
-
   def _get_aux_core_and_hyperthread_pair(
       self,
-      cpu_mask: Iterable[int],
+      cpu_mask: Collection[int],
   ) -> tuple[int, list[int]]:
     for cpu_index in cpu_mask:
-      neighboring_threads = self._get_neighboring_threads(cpu_index)
+      neighboring_threads = _get_neighboring_threads(cpu_index)
       if len(neighboring_threads) != 2:
         raise ValueError('Expected two hyperthreads per CPU.')
 
       if (
-          neighboring_threads[0] in cpu_mask
-          and neighboring_threads[1] in cpu_mask
+          neighboring_threads[0] not in cpu_mask
+          or neighboring_threads[1] not in cpu_mask
       ):
-        cpus = list(cpu_mask)
-        cpus.remove(neighboring_threads[0])
-        cpus.remove(neighboring_threads[1])
-        return (cpus[0], [neighboring_threads[0], neighboring_threads[1]])
+        continue
+
+      cpus = list(cpu_mask)
+      cpus.remove(neighboring_threads[0])
+      cpus.remove(neighboring_threads[1])
+
+      if not cpus:
+        continue
+
+      return (cpus[0], [neighboring_threads[0], neighboring_threads[1]])
     raise ValueError(
-        'Expected a pair of neighboring hyperthreads in the CPU mask.'
+        'Expected a pair of neighboring hyperthreads and an aux thread in the'
+        ' CPU mask.'
     )
 
   @override

gematria/datasets/pipelines/benchmark_cpu_scheduler_test.py

@@ -22,73 +22,83 @@
 
 class BenchmarkSchedulerTests(absltest.TestCase):
 
+  def _set_normal_affinity(self):
+    cpu_mask = os.sched_getaffinity(0)
+    cpu_mask_list = list(cpu_mask)
+
+    self.addCleanup(lambda: os.sched_setaffinity(0, cpu_mask_list))
+
+    # We first select a CPU to use for the hyperthread pair and then remove
+    # both cores in the pair from the mask set to ensure that we do not
+    # accidentally select an auxiliary CPU that is part of the pair.
+    hyperthread_pair_part = next(iter(cpu_mask))
+    hyperthread_pair = benchmark_cpu_scheduler._get_neighboring_threads(
+        hyperthread_pair_part
+    )
+    cpu_mask.remove(hyperthread_pair[0])
+    cpu_mask.remove(hyperthread_pair[1])
+    aux_cpu = cpu_mask.pop()
+    new_cpu_mask = [aux_cpu, *hyperthread_pair]
+
+    os.sched_setaffinity(0, new_cpu_mask)
+    return (aux_cpu, hyperthread_pair, cpu_mask_list)
+
   def test_no_scheduling(self):
     scheduler = benchmark_cpu_scheduler.NoSchedulingBenchmarkScheduler()
     self.assertIsNone(scheduler.setup_and_get_benchmark_core())
     scheduler.verify()
 
   def test_default_scheduler_get_neighboring_threads(self):
-    scheduler = benchmark_cpu_scheduler.DefaultBenchmarkScheduler()
-    neighboring_threads = scheduler._get_neighboring_threads(0)
+    neighboring_threads = benchmark_cpu_scheduler._get_neighboring_threads(0)
 
     # Just check that we get two CPU ids back that are not the same. We cannot
     # do much more without knowing more about the system topology, and this
     # should be a reasonable enough test.
     self.assertLen(neighboring_threads, 2)
     self.assertNotEqual(neighboring_threads[0], neighboring_threads[1])
 
-  @staticmethod
-  def _set_normal_affinity():
-    cpu_mask = os.sched_getaffinity(0)
-    cpu_mask_list = list(cpu_mask)
-    aux_cpu = cpu_mask.pop()
-    hyperthread_pair_part = cpu_mask.pop()
-    hyperthread_pair = benchmark_cpu_scheduler.DefaultBenchmarkScheduler._get_neighboring_threads(
-        hyperthread_pair_part
-    )
-    new_cpu_mask = [aux_cpu, *hyperthread_pair]
-
-    os.sched_setaffinity(0, new_cpu_mask)
-    return (aux_cpu, hyperthread_pair, cpu_mask_list)
-
-  @staticmethod
-  def _reset_cpu_affinity(cpu_mask: Iterable[int]):
-    os.sched_setaffinity(0, cpu_mask)
-
   def test_default_scheduler_get_cores(self):
-    expected_aux_cpu, expected_hyperthread_pair, old_cpu_mask = (
-        self._set_normal_affinity()
-    )
+    expected_aux_cpu, expected_hyperthread_pair, _ = self._set_normal_affinity()
     scheduler = benchmark_cpu_scheduler.DefaultBenchmarkScheduler()
     cpu_mask = os.sched_getaffinity(0)
     aux_cpu, hyperthread_pair = scheduler._get_aux_core_and_hyperthread_pair(
         cpu_mask
     )
     self.assertEqual(aux_cpu, expected_aux_cpu)
     self.assertContainsSubsequence(hyperthread_pair, expected_hyperthread_pair)
-    self._reset_cpu_affinity(old_cpu_mask)
 
   def test_default_scheduler_get_cores_no_neighboring_threads(self):
     cpu_mask = os.sched_getaffinity(0)
+
+    # If we have less than five CPUs, that means we are guaranteed to get
+    # a pair when selecting three of them. Skip this test in those
+    # instances.
+    if len(cpu_mask) < 5:
+      self.skipTest('Not enough cores to complete setup properly.')
+
     three_cores = [cpu_mask.pop(), cpu_mask.pop(), cpu_mask.pop()]
 
     scheduler = benchmark_cpu_scheduler.DefaultBenchmarkScheduler()
-    with self.assertRaises(ValueError):
+    with self.assertRaisesRegex(
+        ValueError,
+        'Expected a pair of neighboring hyperthreads and an aux thread in the'
+        ' CPU mask.',
+    ):
       scheduler._get_aux_core_and_hyperthread_pair(three_cores)
 
   def test_default_scheduler_setup(self):
-    expected_aux_cpu, expected_hyperthread_pair, old_cpu_mask = (
-        self._set_normal_affinity()
-    )
+    expected_aux_cpu, expected_hyperthread_pair, _ = self._set_normal_affinity()
 
     scheduler = benchmark_cpu_scheduler.DefaultBenchmarkScheduler()
     benchmark_core = scheduler.setup_and_get_benchmark_core()
     self.assertIn(benchmark_core, expected_hyperthread_pair)
     set_cpu_mask = os.sched_getaffinity(0)
-    self.assertLen(set_cpu_mask, 1)
-    self.assertEqual(set_cpu_mask.pop(), expected_aux_cpu)
-
-    self._reset_cpu_affinity(old_cpu_mask)
+    self.assertSequenceEqual(
+        set_cpu_mask,
+        {
+            expected_aux_cpu,
+        },
+    )
 
   def test_default_scheduler_not_three_cpus(self):
     old_cpu_mask = os.sched_getaffinity(0)
@@ -102,14 +112,12 @@ def test_default_scheduler_not_three_cpus(self):
     os.sched_setaffinity(0, old_cpu_mask)
 
   def test_default_scheduler_verify(self):
-    _, _, old_cpu_mask = self._set_normal_affinity()
+    self._set_normal_affinity()
 
     scheduler = benchmark_cpu_scheduler.DefaultBenchmarkScheduler()
     scheduler.setup_and_get_benchmark_core()
     scheduler.verify()
 
-    self._reset_cpu_affinity(old_cpu_mask)
-
   def test_default_scheduler_verify_mask_changed(self):
     _, _, old_cpu_mask = self._set_normal_affinity()
 
@@ -121,8 +129,6 @@ def test_default_scheduler_verify_mask_changed(self):
     with self.assertRaises(ValueError):
       scheduler.verify()
 
-    self._reset_cpu_affinity(old_cpu_mask)
-
 
 if __name__ == '__main__':
   absltest.main()