Fixes + all unit testing completed

genai-perf/genai_perf/measurements/model_config_measurement.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass
 from functools import total_ordering
 from statistics import mean
-from typing import Any, Dict, List, Optional, TypeAlias
+from typing import Any, Dict, Optional, TypeAlias
 
 from genai_perf.record.record import Record
 
@@ -121,7 +121,7 @@ def _read_perf_metrics_from_checkpoint(
     ) -> Records:
         perf_metrics: Records = {}
 
-        for [tag, record_dict] in perf_metrics_dict.items():
+        for [tag, record_dict] in perf_metrics_dict.values():
             record = Record.get(tag)
             record = record.read_from_checkpoint(record_dict)  # type: ignore
             perf_metrics[tag] = record  # type: ignore
@@ -149,13 +149,11 @@ def __gt__(self, other: "ModelConfigMeasurement") -> bool:
             == ModelConfigMeasurementDefaults.SELF_IS_BETTER
         )
 
-    # TODO: OPTIMIZE
-    # Why is mypy complaining about this?
-    # def __eq__(self, other: "ModelConfigMeasurement") -> bool:
-    #     return (
-    #         self._compare_measurements(other)
-    #         == ModelConfigMeasurementDefaults.EQUALIVILENT
-    #     )
+    def __eq__(self, other: "ModelConfigMeasurement") -> bool:  # type: ignore
+        return (
+            self._compare_measurements(other)
+            == ModelConfigMeasurementDefaults.EQUALIVILENT
+        )
 
     def _compare_measurements(self, other: "ModelConfigMeasurement") -> int:
         """

genai-perf/genai_perf/record/record.py

@@ -159,7 +159,7 @@ def tag(self) -> str:
             the name tag of the record type.
         """
 
-    def to_dict(self):
+    def write_to_checkpoint(self):
         return (self.tag, self.__dict__)
 
     @classmethod

genai-perf/genai_perf/utils.py

@@ -108,3 +108,14 @@ def get_enum_entry(name: str, enum: Type[Enum]) -> Optional[Enum]:
 
 def scale(value, factor):
     return value * factor
+
+
+# FIXME: OPTIMIZE
+# This will move to the checkpoint class when it's created
+def checkpoint_encoder(obj):
+    if isinstance(obj, bytes):
+        return obj.decode("utf-8")
+    elif hasattr(obj, "write_to_checkpoint"):
+        return obj.write_to_checkpoint()
+    else:
+        return obj.__dict__