test: Moved file deletion to autouse fixture, fixed torch.load warning, and changed qconfig_save warning to logger.info

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

fms_mo/utils/qconfig_utils.py

@@ -816,8 +816,7 @@ def qconfig_save(
 
     # Save config as json
     if os.path.isfile(fname):
-        message = f"{fname} already exist, will overwrite."
-        warnings.warn(message, UserWarning)
+        logger.info(f"{fname} already exist, will overwrite.")
     with open(fname, "w", encoding="utf-8") as outfile:
         json.dump(temp_qcfg, outfile, indent=4)
 

tests/models/test_model_utils.py

@@ -140,9 +140,9 @@ def qmodule_error(
 ###############################
 
 
-def delete_config(file_path: str = "qcfg.json"):
+def delete_file(file_path: str = "qcfg.json"):
     """
-    Delete a qconfig at the file path provided
+    Delete a file at the file path provided
 
     Args:
         file_path (str, optional): Qconfig file to delete. Defaults to "qcfg.json".
@@ -214,7 +214,7 @@ def load_state_dict(fname: str = "qmodel_for_aiu.pt") -> dict:
     Returns:
         dict: Model state dictionary
     """
-    return torch.load(fname)
+    return torch.load(fname, weights_only=True)
 
 
 def check_linear_dtypes(state_dict: dict, linear_names: list):

tests/models/test_qmodelprep.py

@@ -27,7 +27,16 @@
 from fms_mo import qconfig_init, qmodel_prep
 from fms_mo.prep import has_quantized_module
 from fms_mo.utils.utils import patch_torch_bmm
-from tests.models.test_model_utils import count_qmodules, delete_config, qmodule_error
+from tests.models.test_model_utils import count_qmodules, delete_file, qmodule_error
+
+
+@pytest.fixture(autouse=True)
+def delete_files():
+    """
+    Delete any known files lingering before starting test
+    """
+    delete_file("qcfg.json")
+
 
 ################
 # Qmodel tests #
@@ -49,7 +58,6 @@ def test_model_quantized(
             sample_input_fp32 (torch.FloatTensor): Sample fp32 input for calibration.
             config_fp32 (dict): Config w/ fp32 settings
         """
-        delete_config()
         with pytest.raises(RuntimeError):
             qmodel_prep(model_quantized, sample_input_fp32, config_fp32)
 
@@ -81,11 +89,9 @@ def test_double_qmodel_prep_assert(
         sample_input_fp32 (torch.FloatTensor): Sample fp32 input for calibration
         config_fp32 (dict): Config w/ fp32 settings
     """
-    delete_config()
-
     # Run qmodel_prep once
     qmodel_prep(model_fp32, sample_input_fp32, config_fp32)
-    delete_config()
+    delete_file()
 
     # If model now has a quantized node, ensure it has raises a RuntimeError
     if has_quantized_module(model_fp32):
@@ -107,7 +113,6 @@ def test_config_recipes_fp32(
         sample_input_fp32 (torch.FloatTensor): Sample fp32 input for calibration
         config (dict): Recipe Config w/ int8 settings
     """
-    delete_config()
     qmodel_prep(model_fp32, sample_input_fp32, config_int8)
 
 
@@ -124,7 +129,6 @@ def test_config_recipes_fp16(
         sample_input_fp16 (torch.FloatTensor): Sample fp16 input for calibration
         config (dict): Recipe Config w/ int8 settings
     """
-    delete_config()
     qmodel_prep(model_fp16, sample_input_fp16, config_int8)
 
 
@@ -145,7 +149,6 @@ def test_config_fp32_qmodes(
         qa_mode (str): Activation quantizer
         qw_mode (str): Weight quantizer
     """
-    delete_config()
     config_int8["qa_mode"] = qa_mode
     config_int8["qw_mode"] = qw_mode
     qmodel_prep(model_config_fp32, sample_input_fp32, config_int8)
@@ -170,7 +173,6 @@ def test_resnet50_torchscript(
         config (dict): Recipe Config w/ int8 settings
     """
     # Run qmodel_prep w/ default torchscript tracer
-    delete_config()
     qmodel_prep(model_resnet, batch_resnet, config_int8, use_dynamo=False)
     qmodule_error(model_resnet, 6, 48)
 
@@ -189,7 +191,6 @@ def test_resnet50_dynamo(
         config (dict): Recipe Config w/ int8 settings
     """
     # Run qmodel_prep w/ Dynamo tracer
-    delete_config()
     qmodel_prep(model_resnet, batch_resnet, config_int8, use_dynamo=True)
     qmodule_error(model_resnet, 6, 48)
 
@@ -209,7 +210,6 @@ def test_resnet50_dynamo_layers(
         config (dict): Recipe Config w/ int8 settings
     """
     # Run qmodel_prep w/ qlayer_name_pattern + Dynamo tracer
-    delete_config()
     config_int8["qlayer_name_pattern"] = ["layer[1,2,4]"]  # allow regex
     qmodel_prep(model_resnet, batch_resnet, config_int8, use_dynamo=True)
     qmodule_error(model_resnet, 21, 33)
@@ -230,7 +230,6 @@ def test_vit_torchscript(
         config (dict): Recipe Config w/ int8 settings
     """
     # Run qmodel_prep w/ default torchscript tracer
-    delete_config()
     qmodel_prep(model_vit, batch_vit, config_int8, use_dynamo=False)
     qmodule_error(model_vit, 2, 36)
 
@@ -249,7 +248,6 @@ def test_vit_dynamo(
         config (dict): Recipe Config w/ int8 settings
     """
     # Run qmodel_prep w/ Dynamo tracer
-    delete_config()
     qmodel_prep(model_vit, batch_vit, config_int8, use_dynamo=True)
     qmodule_error(model_vit, 2, 36)
 
@@ -268,7 +266,6 @@ def test_bert_dynamo(
         config (dict): Recipe Config w/ int8 settings
     """
     # Run qmodel_prep w/ Dynamo tracer
-    delete_config()
     qmodel_prep(model_bert, input_bert, config_int8, use_dynamo=True)
     qmodule_error(model_bert, 1, 72)
 
@@ -295,7 +292,6 @@ def test_bert_dynamo_wi_qbmm(
         input_bert (torch.FloatTensor): Tokenized input for BERT
         config (dict): Recipe Config w/ int8 settings
     """
-    delete_config()
     config_int8["nbits_bmm1"] = 8
     config_int8["nbits_bmm2"] = 8
     qmodel_prep(model_bert_eager, input_bert, config_int8, use_dynamo=True)

tests/models/test_save_aiu.py

@@ -3,7 +3,7 @@
 import pytest
 
 # Local
-from .test_model_utils import check_linear_dtypes, delete_config, load_state_dict
+from .test_model_utils import check_linear_dtypes, delete_file, load_state_dict
 from fms_mo import qmodel_prep
 from fms_mo.utils.aiu_utils import save_for_aiu
 
@@ -13,9 +13,9 @@ def delete_files():
     """
     Delete any known files lingering before starting test
     """
-    delete_config("qcfg.json")
-    delete_config("keys_to_save.json")
-    delete_config("qmodel_for_aiu.pt")
+    delete_file("qcfg.json")
+    delete_file("keys_to_save.json")
+    delete_file("qmodel_for_aiu.pt")
 
 
 def test_save_model_bert(

tests/models/test_saveconfig.py

@@ -22,12 +22,22 @@
 # Local
 from fms_mo.utils.qconfig_utils import qconfig_load, qconfig_save
 from tests.models.test_model_utils import (
-    delete_config,
+    delete_file,
     load_json,
     save_json,
     save_serialized_json,
 )
 
+
+@pytest.fixture(autouse=True)
+def delete_files():
+    """
+    Delete any known files lingering before starting test
+    """
+    delete_file("qcfg.json")
+    delete_file("keys_to_save.json")
+
+
 #########
 # Tests #
 #########
@@ -45,7 +55,6 @@ def test_save_config_warn_bad_pair(
         bad_pair (tuple): A pair that can't be serialized for qconfig_save
     """
     key, val = bad_pair
-    delete_config()
 
     # Add bad key,val pair and save ; should generate UserWarning(s) for removing bad pair
     config_fp32[key] = val
@@ -56,8 +65,6 @@ def test_save_config_warn_bad_pair(
     loaded_config = load_json("qcfg.json")  # load json as is - do not modify
     assert key not in loaded_config
 
-    delete_config()
-
 
 def test_save_config_wanted_pairs(
     config_fp32: dict,
@@ -71,7 +78,6 @@ def test_save_config_wanted_pairs(
         wanted_pair (tuple): A pair that needs to be re-init if not present in qconfig_save
     """
     key, default_val = wanted_pair
-    delete_config()
 
     # Delete wanted pair from config and save ; should be reset to default
     if key in config_fp32:
@@ -82,8 +88,6 @@ def test_save_config_wanted_pairs(
     loaded_config = load_json()
     assert loaded_config.get(key) == default_val
 
-    delete_config()
-
 
 def test_save_config_with_qcfg_save(
     config_fp32: dict,
@@ -96,7 +100,6 @@ def test_save_config_with_qcfg_save(
         config_fp32 (dict): Config for fp32 quantization
         save_list (list): List of variables to save in a quantized config.
     """
-    delete_config()
     config_fp32["keys_to_save"] = save_list
 
     qconfig_save(config_fp32, minimal=False)
@@ -114,7 +117,6 @@ def test_save_config_with_qcfg_save(
         assert key in loaded_config
         assert loaded_config.get(key) == config_fp32.get(key)
 
-    delete_config()
     del config_fp32["keys_to_save"]
 
 
@@ -129,10 +131,6 @@ def test_save_config_with_recipe_save(
         config_fp32 (dict): Config for fp32 quantization
         save_list (list): List of variables to save in a quantized config.
     """
-    # Delete both qcfg and the save.json before starting
-    delete_config()
-    delete_config("keys_to_save.json")
-
     # Save new "save.json"
     save_path = "keys_to_save.json"
     save_json(save_list, file_path=save_path)
@@ -153,9 +151,6 @@ def test_save_config_with_recipe_save(
         assert key in loaded_config
         assert loaded_config.get(key) == config_fp32.get(key)
 
-    delete_config()
-    delete_config("keys_to_save.json")
-
 
 def test_save_config_minimal(
     config_fp32: dict,
@@ -166,8 +161,6 @@ def test_save_config_minimal(
     Args:
         config_fp32 (dict): Config for fp32 quantization
     """
-    delete_config()
-
     qconfig_save(config_fp32, minimal=True)
 
     # Check that saved qcfg matches
@@ -180,8 +173,6 @@ def test_save_config_minimal(
     # No items should exist - default config should be completely removed
     assert len(loaded_config) == 0
 
-    delete_config()
-
 
 def test_double_qconfig_save(
     config_fp32: dict,
@@ -192,14 +183,8 @@ def test_double_qconfig_save(
     Args:
         config_fp32 (dict): Config for fp32 quantization
     """
-    delete_config()
-
-    # Creating a qcfg, then saving again will cause a warning -> ignore it
-    with pytest.warns(UserWarning, match="qcfg.json already exist, will overwrite."):
-        qconfig_save(config_fp32, minimal=False)
-        qconfig_save(config_fp32, minimal=False)
-
-    delete_config()
+    qconfig_save(config_fp32, minimal=False)
+    qconfig_save(config_fp32, minimal=False)
 
 
 def test_qconfig_save_list_as_dict(
@@ -211,7 +196,7 @@ def test_qconfig_save_list_as_dict(
     Args:
         config_fp32 (dict): Config for fp32 quantization
     """
-    delete_config()
+    delete_file()
 
     # Fill in keys_to_save as dict with nonsense val
     config_fp32["keys_to_save"] = {
@@ -226,8 +211,6 @@ def test_qconfig_save_list_as_dict(
     with pytest.raises(ValueError):
         qconfig_save(config_fp32, minimal=True)
 
-    delete_config()
-
 
 def test_qconfig_save_recipe_as_dict(
     config_fp32: dict,
@@ -238,8 +221,6 @@ def test_qconfig_save_recipe_as_dict(
     Args:
         config_fp32 (dict): Config for fp32 quantization
     """
-    delete_config()
-
     # Fill in keys_to_save as dict with nonsense val
     save_dict = {
         "qa_mode": None,
@@ -254,8 +235,6 @@ def test_qconfig_save_recipe_as_dict(
     with pytest.raises(ValueError):
         qconfig_save(config_fp32, recipe="keys_to_save.json", minimal=True)
 
-    delete_config()
-
 
 def test_qconfig_load_with_recipe_as_list(
     config_fp32: dict,
@@ -266,17 +245,13 @@ def test_qconfig_load_with_recipe_as_list(
     Args:
         config_fp32 (dict): Config for fp32 quantization
     """
-    delete_config()
-
     config_list = list(config_fp32.keys())
 
     save_json(config_list, file_path="qcfg.json")
 
     with pytest.raises(ValueError):
         _ = qconfig_load(fname="qcfg.json")
 
-    delete_config()
-
 
 def test_load_config_restored_pair(
     config_fp32: dict,
@@ -290,7 +265,6 @@ def test_load_config_restored_pair(
         wanted_pair (tuple): A pair that needs to be re-init if not present in qconfig_load
     """
     key, default_val = wanted_pair
-    delete_config()
 
     if key in config_fp32:
         del config_fp32[key]
@@ -302,8 +276,6 @@ def test_load_config_restored_pair(
     loaded_config = qconfig_load("qcfg.json")
     assert loaded_config.get(key) == default_val
 
-    delete_config()
-
 
 def test_load_config_required_pair(
     config_fp32: dict,
@@ -317,7 +289,6 @@ def test_load_config_required_pair(
         required_pair (tuple): A pair that needs to be re-init if not present in qconfig_load
     """
     key, default_val = required_pair
-    delete_config()
 
     if key in config_fp32:
         del config_fp32[key]