Changed serialization library to json for security

Author: Ashwin Ramesh

model_analyzer/perf_analyzer/perf_config.py

@@ -95,6 +95,10 @@ def update_config(self, params=None):
             for key in params:
                 self[key] = params[key]
 
+    def deserialize(self, perf_config_dict):
+        for key, val in perf_config_dict.items():
+            setattr(self, key, val)
+
     def representation(self):
         """
         Returns

model_analyzer/record/record.py

@@ -90,7 +90,6 @@ class Record(metaclass=RecordType):
     This class is used for representing
     records
     """
-
     def __init__(self, value, timestamp):
         """
         Parameters
@@ -144,6 +143,13 @@ def tag(self):
             the name tag of the record type.
         """
 
+    def serialize(self):
+        return (self.tag, self.__dict__)
+    
+    def deserialize(self, record_dict):
+        for key, val in record_dict.items():
+            setattr(self, key, val)
+
     def value(self):
         """
         This method returns the value of recorded metric

model_analyzer/reports/report_manager.py

@@ -604,7 +604,7 @@ def _get_gpu_stats(self, measurements):
         for measurement in measurements:
             for gpu_uuid, gpu_info in self._gpu_info.items():
                 if gpu_uuid in measurement.gpus_used():
-                    gpu_name = (gpu_info['name']).decode('ascii')
+                    gpu_name = gpu_info['name']
                     max_memory = round(gpu_info['total_memory'] / (2**30), 1)
                     if gpu_name not in gpu_dict:
                         gpu_dict[gpu_name] = max_memory

model_analyzer/result/measurement.py

@@ -14,7 +14,9 @@
 
 from functools import total_ordering
 import logging
+from model_analyzer.perf_analyzer.perf_config import PerfAnalyzerConfig
 
+from model_analyzer.record.record import RecordType
 from model_analyzer.model_analyzer_exceptions \
     import TritonModelAnalyzerException
 
@@ -39,10 +41,44 @@ def __init__(self, gpu_data, non_gpu_data, perf_config):
 
         # average values over all GPUs
         self._gpu_data = gpu_data
-        self._avg_gpu_data = self._average_list(list(self._gpu_data.values()))
         self._non_gpu_data = non_gpu_data
         self._perf_config = perf_config
 
+        self._avg_gpu_data = self._average_list(list(self._gpu_data.values()))
+        self._gpu_data_from_tag = {
+            type(metric).tag: metric
+            for metric in self._avg_gpu_data
+        }
+        self._non_gpu_data_from_tag = {
+            type(metric).tag: metric
+            for metric in self._non_gpu_data
+        }
+
+    def deserialize(self, measurement_dict):
+        # Deserialize gpu_data
+        for gpu_uuid, gpu_data_list in measurement_dict['_gpu_data'].items():
+            metric_list = []
+            for [tag, record_dict] in gpu_data_list:
+                record_type = RecordType.get(tag)
+                record = record_type(0)
+                record.deserialize(record_dict)
+                metric_list.append(record)
+            self._gpu_data[gpu_uuid] = metric_list
+
+        # non gpu data
+        self._non_gpu_data = []
+        for [tag, record_dict] in measurement_dict['_non_gpu_data']:
+            record_type = RecordType.get(tag)
+            record = record_type(0)
+            record.deserialize(record_dict)
+            self._non_gpu_data.append(record)
+
+        # perf config
+        self._perf_config = PerfAnalyzerConfig()
+        self._perf_config.deserialize(measurement_dict['_perf_config'])
+
+        # Compute contigent data structures
+        self._avg_gpu_data = self._average_list(list(self._gpu_data.values()))
         self._gpu_data_from_tag = {
             type(metric).tag: metric
             for metric in self._avg_gpu_data

model_analyzer/state/analyzer_state.py

@@ -12,19 +12,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from model_analyzer.model_analyzer_exceptions \
-    import TritonModelAnalyzerException
+from model_analyzer.triton.model.model_config import ModelConfig
+from model_analyzer.result.measurement import Measurement
+from model_analyzer.record.record import RecordType
 
 
 class AnalyzerState:
     """
     All the state information needed by 
     Model Analyzer in one place
     """
-
     def __init__(self):
         self._state_dict = {}
 
+    def serialize(self):
+        return self._state_dict
+
+    def deserialize(self, state_dict):
+        # Fill results
+        self._state_dict['ResultManager.results'] = {}
+        for model_name in state_dict['ResultManager.results']:
+            self._state_dict['ResultManager.results'][model_name] = {}
+            for model_config_name in state_dict['ResultManager.results'][
+                    model_name]:
+                model_config_dict, measurements = state_dict[
+                    'ResultManager.results'][model_name][model_config_name]
+
+                # Deserialize model config
+                model_config = ModelConfig(None)
+                model_config.deserialize(model_config_dict)
+
+                # Deserialize measurements
+                measurements_dict = {}
+                for measurement_key, measurement_dict in measurements.items():
+                    measurement = Measurement({}, [], None)
+                    measurement.deserialize(measurement_dict)
+                    measurements_dict[measurement_key] = measurement
+                self._state_dict['ResultManager.results'][model_name][
+                    model_config_name] = (model_config, measurements_dict)
+
+        # Server data
+        self._state_dict['ResultManager.server_only_data'] = {}
+        for gpu_uuid, gpu_data_list in state_dict[
+                'ResultManager.server_only_data'].items():
+            metric_list = []
+            for [tag, record_dict] in gpu_data_list:
+                record_type = RecordType.get(tag)
+                record = record_type(0)
+                record.deserialize(record_dict)
+                metric_list.append(record)
+            self._state_dict['ResultManager.server_only_data'][
+                gpu_uuid] = metric_list
+
+        # GPU data
+        self._state_dict['MetricsManager.gpus'] = state_dict[
+            'MetricsManager.gpus']
+
     def get(self, name):
         if name in self._state_dict:
             return self._state_dict[name]

model_analyzer/state/analyzer_state_manager.py

@@ -20,7 +20,7 @@
 
 import signal
 import logging
-import pickle
+import json
 import os
 import glob
 
@@ -105,19 +105,27 @@ def load_checkpoint(self):
         if os.path.exists(latest_checkpoint_file):
             logging.info(
                 f"Loaded checkpoint from file {latest_checkpoint_file}")
-            with open(latest_checkpoint_file, 'rb') as f:
+            with open(latest_checkpoint_file, 'r') as f:
                 try:
-                    state = pickle.load(f)
+
+                    self._current_state.deserialize(json.load(f))
                 except EOFError:
                     raise TritonModelAnalyzerException(
                         f'Checkpoint file {latest_checkpoint_file} is'
                         ' empty or corrupted. Remove it from checkpoint'
                         ' directory.')
-            self._current_state = state
             self._starting_fresh_run = False
         else:
             logging.info("No checkpoint file found, starting a fresh run.")
 
+    def default_encode(self, obj):
+        if isinstance(obj, bytes):
+            return obj.decode('utf-8')
+        elif hasattr(obj, 'serialize'):
+            return obj.serialize()
+        else:
+            return obj.__dict__
+
     def save_checkpoint(self):
         """
         Saves the state of the model analyzer to disk
@@ -132,8 +140,8 @@ def save_checkpoint(self):
         ckpt_filename = os.path.join(self._checkpoint_dir,
                                      f"{self._checkpoint_index}.ckpt")
         if self._state_changed:
-            with open(ckpt_filename, 'wb') as f:
-                pickle.dump(self._current_state, f)
+            with open(ckpt_filename, 'w') as f:
+                json.dump(self._current_state, f, default=self.default_encode)
             logging.info(f"Saved checkpoint to {ckpt_filename}.")
 
             self._checkpoint_index += 1

model_analyzer/triton/model/model_config.py

@@ -28,7 +28,6 @@ class ModelConfig:
     """
     A class that encapsulates all the metadata about a Triton model.
     """
-
     def __init__(self, model_config):
         """
         Parameters
@@ -39,6 +38,11 @@ def __init__(self, model_config):
         self._model_config = model_config
         self._cpu_only = False
 
+    def serialize(self):
+        model_config_dict = json_format.MessageToDict(self._model_config)
+        model_config_dict['cpu_only'] = self._cpu_only
+        return model_config_dict
+
     def __getstate__(self):
         """
         Allows serialization of
@@ -49,6 +53,13 @@ def __getstate__(self):
         model_config_dict['cpu_only'] = self._cpu_only
         return model_config_dict
 
+    def deserialize(self, model_config_dict):
+        self._cpu_only = model_config_dict['cpu_only']
+        del model_config_dict['cpu_only']
+        protobuf_message = json_format.ParseDict(
+            model_config_dict, model_config_pb2.ModelConfig())
+        self._model_config = protobuf_message
+
     def __setstate__(self, model_config_dict):
         """
         Allows deserialization of
@@ -57,8 +68,8 @@ def __setstate__(self, model_config_dict):
 
         self._cpu_only = model_config_dict['cpu_only']
         del model_config_dict['cpu_only']
-        protobuf_message = json_format.ParseDict(model_config_dict,
-                                                 model_config_pb2.ModelConfig())
+        protobuf_message = json_format.ParseDict(
+            model_config_dict, model_config_pb2.ModelConfig())
         self._model_config = protobuf_message
 
     @staticmethod
@@ -114,8 +125,8 @@ def create_from_dictionary(model_dict):
         ModelConfig
         """
 
-        protobuf_message = json_format.ParseDict(model_dict,
-                                                 model_config_pb2.ModelConfig())
+        protobuf_message = json_format.ParseDict(
+            model_dict, model_config_pb2.ModelConfig())
 
         return ModelConfig(protobuf_message)
 
@@ -159,7 +170,8 @@ def cpu_only(self):
 
         return self._cpu_only
 
-    def write_config_to_file(self, model_path, src_model_path, last_model_path):
+    def write_config_to_file(self, model_path, src_model_path,
+                             last_model_path):
         """
         Writes a protobuf config file.
 

qa/L0_server_launch_modes/test.sh

@@ -41,7 +41,7 @@ CLIENT_PROTOCOLS="http grpc"
 TRITON_DOCKER_IMAGE="nvcr.io/nvidia/tritonserver:21.05-py3"
 
 mkdir $CHECKPOINT_DIRECTORY
-cp $CHECKPOINT_REPOSITORY/server_launch_modes.ckpt $CHECKPOINT_DIRECTORY/0.ckpt
+# cp $CHECKPOINT_REPOSITORY/server_launch_modes.ckpt $CHECKPOINT_DIRECTORY/0.ckpt
 
 # Run the model-analyzer, both client protocols
 RET=0

tests/mocks/mock_pickle.py renamed to tests/mocks/mock_json.py

@@ -16,41 +16,41 @@
 from unittest.mock import Mock, MagicMock, patch
 
 
-class MockPickleMethods(MockBase):
+class MockJSONMethods(MockBase):
     """
     Mocks the methods for the os module
     """
     def __init__(self):
-        pickle_attrs = {'load': MagicMock(), 'dump': MagicMock()}
-        self.patcher_pickle = patch(
-            'model_analyzer.state.analyzer_state_manager.pickle',
-            Mock(**pickle_attrs))
+        json_attrs = {'load': MagicMock(), 'dump': MagicMock()}
+        self.patcher_json = patch(
+            'model_analyzer.state.analyzer_state_manager.json',
+            Mock(**json_attrs))
         super().__init__()
 
     def start(self):
         """
         start the patchers
         """
 
-        self.pickle_mock = self.patcher_pickle.start()
+        self.json_mock = self.patcher_json.start()
 
     def _fill_patchers(self):
         """
         Fills the patcher list for destruction
         """
 
-        self._patchers.append(self.patcher_pickle)
+        self._patchers.append(self.patcher_json)
 
-    def set_pickle_load_return_value(self, value):
+    def set_json_load_return_value(self, value):
         """
-        Sets the return value for pickle load
+        Sets the return value for json load
         """
 
-        self.pickle_mock.load.return_value = value
+        self.json_mock.load.return_value = value
 
-    def set_pickle_load_side_effect(self, effect):
+    def set_json_load_side_effect(self, effect):
         """
         Sets a side effet
         """
 
-        self.pickle_mock.load.side_effect = effect
+        self.json_mock.load.side_effect = effect

tests/test_report_manager.py

@@ -28,7 +28,7 @@
 from .mocks.mock_io import MockIOMethods
 from .mocks.mock_matplotlib import MockMatplotlibMethods
 from .mocks.mock_os import MockOSMethods
-from .mocks.mock_pickle import MockPickleMethods
+from .mocks.mock_json import MockJSONMethods
 
 from .common.test_utils import construct_measurement
 from .common import test_result_collector as trc
@@ -115,8 +115,8 @@ def setUp(self):
         self.io_mock.start()
         self.matplotlib_mock = MockMatplotlibMethods()
         self.matplotlib_mock.start()
-        self.pickle_mock = MockPickleMethods()
-        self.pickle_mock.start()
+        self.json_mock = MockJSONMethods()
+        self.json_mock.start()
 
     def test_add_results(self):
         for mode in ['online', 'offline']:
@@ -221,7 +221,7 @@ def tearDown(self):
         self.matplotlib_mock.stop()
         self.io_mock.stop()
         self.os_mock.stop()
-        self.pickle_mock.stop()
+        self.json_mock.stop()
 
 
 if __name__ == '__main__':