Added YAML security fixes

Author: Ashwin Ramesh

model_analyzer/config/input/config_command.py

@@ -21,7 +21,6 @@ class ConfigCommand:
     """
     Model Analyzer config object.
     """
-
     def __init__(self):
         """
         Create a new config.
@@ -69,7 +68,7 @@ def _load_config_file(self, file_path):
         """
 
         with open(file_path, 'r') as config_file:
-            config = yaml.load(config_file, Loader=yaml.FullLoader)
+            config = yaml.safe_load(config_file, Loader=yaml.FullLoader)
             return config
 
     def set_config_values(self, args):

qa/L0_config_range/config_generator.py

@@ -115,7 +115,7 @@ def get_all_configurations():
         with open(f'./config-{i}.yml', 'w') as file:
             yaml.dump(configuration, file)
         with open(f'./config-{i}.yml', 'r') as file:
-            config = yaml.load(file, Loader=yaml.FullLoader)
+            config = yaml.safe_load(file, Loader=yaml.FullLoader)
         with open(f'./config-{i}.txt', 'w') as file:
             file.write(str(total_param))
 

qa/L0_custom_flags/check_results.py

@@ -52,7 +52,8 @@ def check_perf_global(self):
         True if test passes else False
         """
 
-        if 'perf_analyzer_flags' in self._config and 'percentile' in self._config['perf_analyzer_flags']:
+        if 'perf_analyzer_flags' in self._config and 'percentile' in self._config[
+                'perf_analyzer_flags']:
             with open(self._analyzer_log, 'r') as f:
                 contents = f.read()
 
@@ -80,13 +81,15 @@ def check_perf_mode_time_window(self):
         True if test passes else False
         """
 
-        if 'perf_analyzer_flags' in self._config and 'measurement-mode' in self._config['perf_analyzer_flags']:
+        if 'perf_analyzer_flags' in self._config and 'measurement-mode' in self._config[
+                'perf_analyzer_flags']:
             with open(self._analyzer_log, 'r') as f:
                 contents = f.read()
 
             # In contents, search for "stabilizing with px latency"
-            measurement_mode = self._config['perf_analyzer_flags']['measurement-mode']
-            assert(measurement_mode == 'time_windows')
+            measurement_mode = self._config['perf_analyzer_flags'][
+                'measurement-mode']
+            assert (measurement_mode == 'time_windows')
             token = "time_windows"
 
             # Ensure the token appears in the text
@@ -229,7 +232,7 @@ def check_triton_per_model(self):
     args = parser.parse_args()
 
     with open(args.config_file, 'r') as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+        config = yaml.safe_load(f, Loader=yaml.FullLoader)
 
     TestOutputValidator(config, args.profile_models, args.analyzer_log_file,
                         args.triton_log_file)

qa/L0_perf_analyzer/check_results.py

@@ -22,7 +22,6 @@ class TestOutputValidator:
     Functions that validate the output
     of the test
     """
-
     def __init__(self, config, config_file, analyzer_log, test_name):
         self._config = config
         self._config_file = config_file
@@ -72,7 +71,7 @@ def check_time_window_50(self):
     args = parser.parse_args()
 
     with open(args.config_file, 'r') as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+        config = yaml.safe_load(f, Loader=yaml.FullLoader)
 
     TestOutputValidator(config, args.config_file, args.analyzer_log,
                         args.test_name)

qa/L0_profile/check_results.py

@@ -24,7 +24,6 @@ class TestOutputValidator:
     Functions that validate the output
     of the test
     """
-
     def __init__(self, config, test_name, analyzer_log):
         self._config = config
         self._models = config['profile_models']
@@ -87,6 +86,6 @@ def check_profile_logs(self):
     args = parser.parse_args()
 
     with open(args.config_file, 'r') as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+        config = yaml.safe_load(f, Loader=yaml.FullLoader)
 
     TestOutputValidator(config, args.test_name, args.analyzer_log_file)

qa/L0_results/check_results.py

@@ -106,6 +106,6 @@ def check_detailed_reports(self):
     args = parser.parse_args()
 
     with open(args.config_file, 'r') as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+        config = yaml.safe_load(f, Loader=yaml.FullLoader)
 
     TestOutputValidator(config, args.test_name, args.export_path)

qa/L0_state_management/check_results.py

@@ -180,7 +180,7 @@ def check_measurements_consistent_with_config(self):
     args = parser.parse_args()
 
     with open(args.config_file, 'r') as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+        config = yaml.safe_load(f, Loader=yaml.FullLoader)
 
     TestOutputValidator(config, args.test_name, args.checkpoint_dir,
                         args.analyzer_log_file)