Create checkpoint class (#134)

* initial changes

* Adding GAP option to optimize

* Fixing codeQL issues

* Results class initial coding w/ testing

* Fixing issue in RCM testing

* Fixing codeQL issues

* Checkpoint class creation

* Fixing codeQL issues

* fixing codeql issue

* Removing checkpoint file

* Removing checkpoint file

* Fixing json to properly format checkpoint file

* Minor typing cleanup

* Adding records for ISL/OSL and testing this in checkpoint creation

* Changing method name

* Changing read/write checkpoint method names

genai-perf/genai_perf/checkpoint/checkpoint.py

@@ -0,0 +1,98 @@
+# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+
+from genai_perf.config.input.config_command import ConfigCommand
+from genai_perf.config.run.results import Results
+from genai_perf.exceptions import GenAIPerfException
+from genai_perf.types import CheckpointObject
+
+
+@dataclass(frozen=True)
+class CheckpointDefaults:
+    FILENAME = "checkpoint.json"
+
+
+@dataclass
+class Checkpoint:
+    """
+    Contains the methods necessary for reading and writing GenAI-Perf
+    state to a file so that stateful subcommands (such as Optimize or
+    Analyze) can resume or continue (ex: running Analyze then Visualize)
+    """
+
+    config: ConfigCommand
+
+    # Every top-level class that needs to store state is passed in
+    results: Results
+
+    def __post_init__(self):
+        self._create_class_from_checkpoint()
+
+    ###########################################################################
+    # Read/Write Methods
+    ###########################################################################
+    def create_checkpoint_object(self) -> None:
+        state_dict = {"Results": self.results.create_checkpoint_object()}
+
+        checkpoint_file_path = self._create_checkpoint_file_path()
+        with open(checkpoint_file_path, "w") as checkpoint_file:
+            json.dump(state_dict, checkpoint_file, default=checkpoint_encoder)
+
+    def _create_class_from_checkpoint(self) -> None:
+        checkpoint_file_path = self._create_checkpoint_file_path()
+
+        if os.path.isfile(checkpoint_file_path):
+            self.checkpoint_exists = True
+            try:
+                with open(checkpoint_file_path, "r") as checkpoint_file:
+                    checkpoint_json = json.load(checkpoint_file)
+                    self._state: CheckpointObject = {
+                        "Results": Results.create_class_from_checkpoint(
+                            checkpoint_json["Results"]
+                        )
+                    }
+
+            except EOFError:
+                raise (
+                    GenAIPerfException(
+                        f"Checkpoint file {checkpoint_file} is"
+                        " empty or corrupted. Delete it and rerun GAP"
+                    )
+                )
+        else:
+            self.checkpoint_exists = False
+            self._state = {}
+
+    def _create_checkpoint_file_path(self) -> str:
+        checkpoint_file_path = os.path.join(
+            self.config.checkpoint_directory, CheckpointDefaults.FILENAME
+        )
+
+        return checkpoint_file_path
+
+
+###########################################################################
+# Encoder
+###########################################################################
+def checkpoint_encoder(obj):
+    if isinstance(obj, bytes):
+        return obj.decode("utf-8")
+    elif hasattr(obj, "create_checkpoint_object"):
+        return obj.create_checkpoint_object()
+    else:
+        return obj.__dict__

genai-perf/genai_perf/config/generate/genai_perf_config.py

@@ -14,7 +14,6 @@
 
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, Dict
 
 from genai_perf.config.generate.search_parameter import SearchUsage
 from genai_perf.config.input.config_command import (
@@ -23,7 +22,7 @@
     ConfigOutputTokens,
     ConfigSyntheticTokens,
 )
-from genai_perf.types import ModelObjectiveParameters
+from genai_perf.types import CheckpointObject, ModelObjectiveParameters
 
 
 @dataclass
@@ -63,7 +62,7 @@ def _set_options_based_on_objective(
     ###########################################################################
     # Checkpoint Methods
     ###########################################################################
-    def write_to_checkpoint(self) -> Dict[str, Any]:
+    def create_checkpoint_object(self) -> CheckpointObject:
         """
         Converts the class data into a dictionary that can be written to
         the checkpoint file
@@ -73,8 +72,8 @@ def write_to_checkpoint(self) -> Dict[str, Any]:
         return genai_perf_config_dict
 
     @classmethod
-    def read_from_checkpoint(
-        cls, genai_perf_config_dict: Dict[str, Any]
+    def create_class_from_checkpoint(
+        cls, genai_perf_config_dict: CheckpointObject
     ) -> "GenAIPerfConfig":
         """
         Takes the checkpoint's representation of the class and creates (and populates)

genai-perf/genai_perf/config/generate/perf_analyzer_config.py

@@ -19,7 +19,7 @@
 from genai_perf.config.generate.search_parameter import SearchUsage
 from genai_perf.config.input.config_command import ConfigCommand, ConfigPerfAnalyzer
 from genai_perf.exceptions import GenAIPerfException
-from genai_perf.types import ModelName, ModelObjectiveParameters
+from genai_perf.types import CheckpointObject, ModelName, ModelObjectiveParameters
 
 
 @dataclass
@@ -54,6 +54,15 @@ def _set_options_based_on_objective(
                 if parameter.usage == SearchUsage.RUNTIME_PA:
                     self._parameters[name] = parameter.get_value_based_on_category()
 
+    ###########################################################################
+    # Get Accessor Methods
+    ###########################################################################
+    def get_parameters(self) -> Dict[str, Any]:
+        """
+        Returns a dictionary of parameters and their values
+        """
+        return self._parameters
+
     ###########################################################################
     # CLI String Creation Methods
     ###########################################################################
@@ -154,7 +163,7 @@ def _remove_option_from_cli_string(
     ###########################################################################
     # Checkpoint Methods
     ###########################################################################
-    def write_to_checkpoint(self) -> Dict[str, Any]:
+    def create_checkpoint_object(self) -> CheckpointObject:
         """
         Converts the class data into a dictionary that can be written to
         the checkpoint file
@@ -164,8 +173,8 @@ def write_to_checkpoint(self) -> Dict[str, Any]:
         return pa_config_dict
 
     @classmethod
-    def read_from_checkpoint(
-        cls, perf_analyzer_config_dict: Dict[str, Any]
+    def create_class_from_checkpoint(
+        cls, perf_analyzer_config_dict: CheckpointObject
     ) -> "PerfAnalyzerConfig":
         """
         Takes the checkpoint's representation of the class and creates (and populates)

genai-perf/genai_perf/config/input/config_command.py

@@ -43,6 +43,9 @@ class Range:
 # These will be moved to RunConfig once it's created
 @dataclass(frozen=True)
 class RunConfigDefaults:
+    # Top-level Defaults
+    CHECKPOINT_DIRECTORY = "./"
+
     # Optimize: Top-level Defaults
     OBJECTIVE = "throughput"
     CONSTRAINT = None
@@ -212,6 +215,7 @@ class ConfigOutputTokens:
 @dataclass
 class ConfigCommand:
     model_names: List[ModelName]
+    checkpoint_directory: str = default_field(RunConfigDefaults.CHECKPOINT_DIRECTORY)
     optimize: ConfigOptimize = ConfigOptimize()
     analyze: ConfigAnalyze = ConfigAnalyze()
     perf_analyzer: ConfigPerfAnalyzer = ConfigPerfAnalyzer()

genai-perf/genai_perf/config/run/results.py

@@ -11,11 +11,16 @@
 # limitations under the License.
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, List
+from typing import List
 
 from genai_perf.config.run.run_config import RunConfig
 from genai_perf.measurements.run_constraints import RunConstraints
-from genai_perf.types import GpuMetricObjectives, ModelWeights, PerfMetricObjectives
+from genai_perf.types import (
+    CheckpointObject,
+    GpuMetricObjectives,
+    ModelWeights,
+    PerfMetricObjectives,
+)
 
 
 @dataclass
@@ -30,28 +35,28 @@ class Results:
     ###########################################################################
     # Checkpoint Methods
     ###########################################################################
-    def write_to_checkpoint(self) -> Dict[str, Any]:
+    def create_checkpoint_object(self) -> CheckpointObject:
         """
         Converts the class data into a dictionary that can be written to
         the checkpoint file
         """
         run_config_dicts = [
-            run_config.write_to_checkpoint() for run_config in self.run_configs
+            run_config.create_checkpoint_object() for run_config in self.run_configs
         ]
 
         results_dict = {"run_configs": run_config_dicts}
 
         return results_dict
 
     @classmethod
-    def read_from_checkpoint(cls, results_dict: Dict[str, Any]) -> "Results":
+    def create_class_from_checkpoint(cls, results_dict: CheckpointObject) -> "Results":
         """
         Takes the checkpoint's representation of the class and creates (and populates)
         a new instance of Results
         """
         results = Results()
         for run_config_dict in results_dict["run_configs"]:
-            run_config = RunConfig.read_from_checkpoint(run_config_dict)
+            run_config = RunConfig.create_class_from_checkpoint(run_config_dict)
             results.add_run_config(run_config)
 
         return results

genai-perf/genai_perf/config/run/run_config.py

@@ -22,6 +22,7 @@
 from genai_perf.measurements.run_constraints import RunConstraints
 from genai_perf.record.record import Record
 from genai_perf.types import (
+    CheckpointObject,
     GpuRecords,
     MetricObjectives,
     ModelName,
@@ -52,34 +53,36 @@ class RunConfig:
     ###########################################################################
     # Checkpoint Methods
     ###########################################################################
-    def write_to_checkpoint(self) -> Dict[str, Any]:
+    def create_checkpoint_object(self) -> CheckpointObject:
         """
         Converts the class data into a dictionary that can be written to
         the checkpoint file
         """
         run_config_dict = {
             "name": self.name,
-            "genai_perf_config": self.genai_perf_config.write_to_checkpoint(),
-            "perf_analyzer_config": self.perf_analyzer_config.write_to_checkpoint(),
-            "measurement": self.measurement.write_to_checkpoint(),
+            "genai_perf_config": self.genai_perf_config.create_checkpoint_object(),
+            "perf_analyzer_config": self.perf_analyzer_config.create_checkpoint_object(),
+            "measurement": self.measurement.create_checkpoint_object(),
         }
 
         return run_config_dict
 
     @classmethod
-    def read_from_checkpoint(cls, run_config_dict: Dict[str, Any]) -> "RunConfig":
+    def create_class_from_checkpoint(
+        cls, run_config_dict: CheckpointObject
+    ) -> "RunConfig":
         """
         Takes the checkpoint's representation of the class and creates (and populates)
         a new instance of a RCM
         """
         name = run_config_dict["name"]
-        genai_perf_config = GenAIPerfConfig.read_from_checkpoint(
+        genai_perf_config = GenAIPerfConfig.create_class_from_checkpoint(
             run_config_dict["genai_perf_config"]
         )
-        perf_analyzer_config = PerfAnalyzerConfig.read_from_checkpoint(
+        perf_analyzer_config = PerfAnalyzerConfig.create_class_from_checkpoint(
             run_config_dict["perf_analyzer_config"]
         )
-        measurement = RunConfigMeasurement.read_from_checkpoint(
+        measurement = RunConfigMeasurement.create_class_from_checkpoint(
             run_config_dict["measurement"]
         )
 
@@ -92,6 +95,9 @@ def read_from_checkpoint(cls, run_config_dict: Dict[str, Any]) -> "RunConfig":
     ###########################################################################
     # Get Accessor Methods
     ###########################################################################
+    def get_perf_analyzer_parameters(self) -> Dict[str, Any]:
+        return self.perf_analyzer_config.get_parameters()
+
     def get_all_gpu_metrics(self) -> GpuRecords:
         return self.measurement.get_all_gpu_metrics()
 

genai-perf/genai_perf/measurements/model_config_measurement.py

@@ -20,7 +20,13 @@
 from typing import Any, Dict, Optional, TypeAlias
 
 from genai_perf.record.record import Record
-from genai_perf.types import MetricObjectives, ModelName, PerfMetricName, PerfRecords
+from genai_perf.types import (
+    CheckpointObject,
+    MetricObjectives,
+    ModelName,
+    PerfMetricName,
+    PerfRecords,
+)
 
 ###########################################################################
 # Typing
@@ -96,7 +102,7 @@ def set_metric_objectives(self, metric_objectives: MetricObjectives) -> None:
     ###########################################################################
     # Checkpoint Methods
     ###########################################################################
-    def write_to_checkpoint(self) -> Dict[str, Any]:
+    def create_checkpoint_object(self) -> CheckpointObject:
         """
         Converts the class data into a dictionary that can be written to
         the checkpoint file
@@ -110,7 +116,9 @@ def write_to_checkpoint(self) -> Dict[str, Any]:
         return mcm_dict
 
     @classmethod
-    def read_from_checkpoint(cls, mcm_dict: Dict[str, Any]) -> "ModelConfigMeasurement":
+    def create_class_from_checkpoint(
+        cls, mcm_dict: CheckpointObject
+    ) -> "ModelConfigMeasurement":
         """
         Takes the checkpoint's representation of the class and creates (and populates)
         a new instance of a MCM
@@ -129,7 +137,7 @@ def _read_perf_metrics_from_checkpoint(
 
         for [tag, record_dict] in perf_metrics_dict.values():
             record = Record.get(tag)
-            record = record.read_from_checkpoint(record_dict)  # type: ignore
+            record = record.create_class_from_checkpoint(record_dict)  # type: ignore
             perf_metrics[tag] = record  # type: ignore
 
         return perf_metrics