Backends now provide their own record_property

mir-ci/mir_ci/benchmarker.py

@@ -12,47 +12,6 @@
 logger = logging.getLogger(__name__)
 
 
-class AggregatedProcessInfo:
-    name: str
-    start_time_seconds: float
-    cpu_time_microseconds_total: float
-    mem_bytes_total: int
-    mem_bytes_max: int
-    num_data_points: int
-
-    def __init__(self, name: str) -> None:
-        self.name = name
-        self.start_time_seconds = time.time()
-        self.cpu_time_microseconds_total = 0
-        self.mem_bytes_total = 0
-        self.mem_bytes_max = 0
-        self.num_data_points = 0
-
-
-class ProcessInfoFrame:
-    program: Benchmarkable
-    current_memory_bytes: int
-    cpu_time_microseconds_total: float
-
-    peak_memory_bytes: Optional[int]
-    """
-    A frame may contain the peak memory in bytes if it is available.
-    This may be more reliable than finding the largest in a sample of
-    current memory bytes data points
-    """
-
-    def __init__(
-            self,
-            program: Benchmarkable,
-            current_memory_bytes: int,
-            cpu_time_microseconds_total: float,
-            peak_memory_bytes: Optional[int] = None) -> None:
-        self.program = program
-        self.current_memory_bytes = current_memory_bytes
-        self.cpu_time_microseconds_total = cpu_time_microseconds_total
-        self.peak_memory_bytes = peak_memory_bytes
-
-
 class BenchmarkBackend(ABC):
     """
     Abstract class that aggregates programs together and emits process stats as it is requested
@@ -65,34 +24,26 @@ def add(self, program: Benchmarkable) -> None:
         pass
 
     @abstractmethod
-    def poll(self, cb: Callable[[ProcessInfoFrame], None]) -> None:
+    def poll(self) -> None:
+        pass
+
+    @abstractmethod
+    def generate_report(self, record_property: Callable[[str, object], None]) -> None:
         pass
 
 
 class Benchmarker:
     def __init__(self, *programs: Benchmarkable, poll_time_seconds: float = 1.0):
         self.programs = programs
-        self.data_records: Dict[Benchmarkable, AggregatedProcessInfo] = {}
         self.running = False
         self.backend: BenchmarkBackend = CgroupsBackend()
         self.task: Optional[asyncio.Task[None]] = None
         self.poll_time_seconds = poll_time_seconds
 
-    def _on_packet(self, packet: ProcessInfoFrame) -> None:        
-        self.data_records[packet.program].cpu_time_microseconds_total = packet.cpu_time_microseconds_total
-        self.data_records[packet.program].mem_bytes_total += packet.current_memory_bytes
-
-        if packet.peak_memory_bytes is not None:
-            self.data_records[packet.program].mem_bytes_max = packet.peak_memory_bytes
-        elif self.data_records[packet.program].mem_bytes_max < packet.current_memory_bytes:
-            self.data_records[packet.program].mem_bytes_max = packet.current_memory_bytes
-
-        self.data_records[packet.program].num_data_points += 1
-
     async def _run(self) -> None:
         while self.running:
             try:
-                self.backend.poll(self._on_packet)
+                self.backend.poll()
             except:
                 pass
             await asyncio.sleep(self.poll_time_seconds)
@@ -104,7 +55,6 @@ async def start(self) -> None:
         self.running = True
         for program in self.programs:
             await program.__aenter__()
-            self.data_records[program] = AggregatedProcessInfo(program.get_name())
             self.backend.add(program)
 
         self.task = asyncio.ensure_future(self._run())
@@ -128,32 +78,57 @@ async def __aenter__(self):
     async def __aexit__(self, *args):
         await self.stop()
 
-    def generate_report(self, record_property: Callable[[str, object], None]):
-        for program, info in self.data_records.items():
-            cpu_time_microseconds_total = info.cpu_time_microseconds_total
-            max_mem_bytes = info.mem_bytes_max
-            avg_mem_bytes = 0 if info.num_data_points == 0 else info.mem_bytes_total / info.num_data_points
-            record_property(f"{info.name}_cpu_time_microseconds", cpu_time_microseconds_total)
-            record_property(f"{info.name}_max_mem_bytes", max_mem_bytes)
-            record_property(f"{info.name}_avg_mem_bytes", avg_mem_bytes)
+    def generate_report(self, record_property: Callable[[str, object], None]) -> None:
+        self.backend.generate_report(record_property)
 
 
 class CgroupsBackend(BenchmarkBackend):
+    class ProcessInfo:
+        name: str
+        start_time_seconds: float
+        cpu_time_microseconds_total: float
+        mem_bytes_total: int
+        mem_bytes_max: int
+        num_data_points: int
+
+        def __init__(self, name: str) -> None:
+            self.name = name
+            self.start_time_seconds = time.time()
+            self.cpu_time_microseconds_total = 0
+            self.mem_bytes_total = 0
+            self.mem_bytes_max = 0
+            self.num_data_points = 0
+
     def __init__(self) -> None:
-        self._program_list: List[Benchmarkable] = []
+        self.data_records: Dict[Benchmarkable, CgroupsBackend.ProcessInfo] = {}
 
     def add(self, program: Benchmarkable) -> None:
-        self._program_list.append(program)
+        self.data_records[program] = CgroupsBackend.ProcessInfo(program.get_name())
 
-    def poll(self, cb: Callable[[ProcessInfoFrame], None]) -> None:
-        for program in self._program_list:
+    def poll(self) -> None:
+        for program in self.data_records:
             try:
                 cgroup = program.get_cgroup()
-                cb(ProcessInfoFrame(
-                    program,
-                    cgroup.get_current_memory(),
-                    cgroup.get_cpu_time_microseconds(),
-                    cgroup.get_peak_memory()
-                ))
+                current_memory_bytes = cgroup.get_current_memory()
+                cpu_time_microseconds_total = cgroup.get_cpu_time_microseconds()
+                peak_memory_bytes = cgroup.get_peak_memory()
+                self.data_records[program].cpu_time_microseconds_total = cpu_time_microseconds_total
+                self.data_records[program].mem_bytes_total += current_memory_bytes
+
+                if peak_memory_bytes is not None:
+                    self.data_records[program].mem_bytes_max = peak_memory_bytes
+                elif self.data_records[program].mem_bytes_max < current_memory_bytes:
+                    self.data_records[program].mem_bytes_max = current_memory_bytes
+
+                self.data_records[program].num_data_points += 1
             except (RuntimeError, FileNotFoundError, PermissionError, OSError) as ex:
                 logger.error(ex)
+
+    def generate_report(self, record_property: Callable[[str, object], None]) -> None:
+        for program, info in self.data_records.items():
+            cpu_time_microseconds_total = info.cpu_time_microseconds_total
+            max_mem_bytes = info.mem_bytes_max
+            avg_mem_bytes = 0 if info.num_data_points == 0 else info.mem_bytes_total / info.num_data_points
+            record_property(f"{info.name}_cpu_time_microseconds", cpu_time_microseconds_total)
+            record_property(f"{info.name}_max_mem_bytes", max_mem_bytes)
+            record_property(f"{info.name}_avg_mem_bytes", avg_mem_bytes)