From c86c7827f125b9af17bf5ddff97e8a52db52550d Mon Sep 17 00:00:00 2001
From: Aiden Grossman <agrossman154@yahoo.com>
Date: Mon, 30 Dec 2024 11:35:56 +0000
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20change?=
 =?UTF-8?q?s=20introduced=20through=20rebase?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4

[skip ci]
---
 .../pipelines/benchmark_cpu_scheduler.py      | 58 +++++++++++++++++--
 1 file changed, 52 insertions(+), 6 deletions(-)

diff --git a/gematria/datasets/pipelines/benchmark_cpu_scheduler.py b/gematria/datasets/pipelines/benchmark_cpu_scheduler.py
index 0bd370e6..665224fd 100644
--- a/gematria/datasets/pipelines/benchmark_cpu_scheduler.py
+++ b/gematria/datasets/pipelines/benchmark_cpu_scheduler.py
@@ -12,34 +12,78 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractmethod
+import abc
+import typing_extensions
 from collections.abc import Iterable
 import os
 import re
 from enum import Enum
 
 
-class BenchmarkScheduler(ABC):
+class BenchmarkScheduler(metaclass=abc.ABCMeta):
+  """Schedules a benchmark and the parent process to reduce noise.
 
-  @abstractmethod
+  BenchmarkScheduler is an abstraction that provides two main pieces of
+  functionality. Firstly, it provides a function
+  (setup_and_get_benchmark_core) that allows an implementation to perform any
+  necessary setup in the parent process and provide a core ID that should be
+  used to perform any benchmarking. Additionally, implementations are
+  intended to hold state to verify that the expected state is maintained and
+  not changed by external software.
+  """
+
+  @abc.abstractmethod
   def setup_and_get_benchmark_core(self) -> int | None:
-    pass
+    """Sets up the parent process and chooses a benchmark core.
+
+    This function will perform any relevant setup in the parent process,
+    and return a core ID that can be used to run benchmarks on.
+
+    Returns:
+      Returns an integer core ID to specify a core that should be used for
+      running benchmarks, or None to indicate that any core can be used.
+    """
 
-  @abstractmethod
+  @abc.abstractmethod
   def verify(self):
-    pass
+    """Verifies that conditions match what is expected.
+
+    This function allows for implementations to verify that the original
+    setup created in setup_and_get_benchmark_core is maintained for every
+    benchmark that is run.
+    """
 
 
 class NoSchedulingBenchmarkScheduler(BenchmarkScheduler):
+  """A basic BenchmarkScheduler implementation that does nothing.
+
+  This BenchmarkScheduler implementation does nothing. It leaves scheduling
+  of the benchmarking process to the operating system by specifying any core
+  can be used for benchmarking, performs no setup in the parent process, and
+  performs no verification.
+  """
 
+  @typing_extensions.override
   def setup_and_get_benchmark_core(self) -> int | None:
     return None
 
+  @typing_extensions.override
   def verify(self):
     pass
 
 
 class DefaultBenchmarkScheduler(BenchmarkScheduler):
+  """A BenchmarkScheduler that schedules processes separately.
+
+  DefaultBenchmarkScheduler schedules the main process and the benchmark
+  subprocess on separate cores, making sure to reserve the second hyperthread
+  on the benchmarking core to prevent interference. It expects that the main
+  process is initially given a CPU Mask with three active threads, additionally
+  assuming that two of the threads are neighboring (part of the same core).
+  Errors are raised if these conditions are not met. The benchmarking core
+  returned is one of the two neighboring threads. The main process has its
+  COU mask limited to the thread that neighbors neither of the other threads.
+  """
 
   def __init__(self):
     self._cpu_mask = []
@@ -78,6 +122,7 @@ def _get_aux_core_and_hyperthread_pair(
         'Expected a pair of neighboring hyperthreads in the CPU mask.'
     )
 
+  @typing_extensions.override
   def setup_and_get_benchmark_core(self) -> int | None:
     cpu_mask = os.sched_getaffinity(0)
 
@@ -92,6 +137,7 @@ def setup_and_get_benchmark_core(self) -> int | None:
 
     return hyperthread_pair[0]
 
+  @typing_extensions.override
   def verify(self):
     cpu_mask = list(os.sched_getaffinity(0))
     if self._cpu_mask != cpu_mask: