test: Added tests for recipe save feature in qconfig_save

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Author: BrandonGroth

tests/models/conftest.py

@@ -798,6 +798,36 @@ def config_fp16(request):
     qconfig["nbits_w"] = 16
     return qconfig
 
+save_list_params = [
+    ["qa_mode", "qw_mode", "nbits_a", "nbits_w", "qskip_layer_name"],
+]
+@pytest.fixture(scope="session", params=save_list_params)
+def save_list(request):
+    """
+    Generate a save list for testing user-requested save config.
+
+    Args:
+        request (list): List of variables to save in a quantized config.
+
+    Returns:
+        list: List of variables to save in a quantized config.
+    """
+    return request.param
+
+bad_recipe_params = ["qat_int7", "pzq_int8"]
+
+@pytest.fixture(scope="session", params=bad_recipe_params)
+def bad_recipe(request):
+    """
+    Get a bad recipe json file name in fms_mo/recipes
+
+    Args:
+        request (str): Bad recipe name in fms_mo/recipes
+
+    Returns:
+        str: Bad recipe name
+    """
+    return request.param
 
 # Create QAT/PTQ int8 config fixture.
 config_params = ["qat_int8", "ptq_int8"]

tests/models/test_model_utils.py

@@ -173,6 +173,16 @@ def load_json(file_path: str = "qcfg.json"):
     assert json_file is not None, f"JSON at {file_path} was not found"
     return json_file
 
+def save_json(data, file_path: str = "qcfg.json"):
+    """
+    Save data object to json file
+
+    Args:
+        data (_type_): _description_
+        file_path (str, optional): _description_. Defaults to "qcfg.json".
+    """
+    with open(file_path, "w", encoding="utf-8") as outfile:
+        json.dump(data, outfile, indent=4)
 
 def save_serialized_json(config: dict, file_path: str = "qcfg.json"):
     """
@@ -189,5 +199,4 @@ def save_serialized_json(config: dict, file_path: str = "qcfg.json"):
             del config[key]
 
     serialize_config(config)  # Only remove stuff necessary to dump
-    with open(file_path, "w", encoding="utf-8") as outfile:
-        json.dump(config, outfile, indent=4)
+    save_json(config, file_path)

tests/models/test_qmodelprep.py

@@ -24,7 +24,7 @@
 
 # Local
 # fms_mo imports
-from fms_mo import qmodel_prep
+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
@@ -53,6 +53,18 @@ def test_model_quantized(
         with pytest.raises(RuntimeError):
             qmodel_prep(model_quantized, sample_input_fp32, config_fp32)
 
+def test_bad_recipe(
+    bad_recipe: str,
+):
+    """
+    Test if giving a bad recipe .json file name results in a ValueError.
+
+    Args:
+        bad_recipe (str): Bad .json file name
+    """
+    with pytest.raises(ValueError):
+        qconfig_init(recipe=bad_recipe)
+
 
 def test_double_qmodel_prep_assert(
     model_fp32: torch.nn.Module,

tests/models/test_saveconfig.py

@@ -21,7 +21,12 @@
 
 # Local
 from fms_mo.utils.qconfig_utils import qconfig_load, qconfig_save
-from tests.models.test_model_utils import delete_config, load_json, save_serialized_json
+from tests.models.test_model_utils import (
+    delete_config,
+    load_json,
+    save_json,
+    save_serialized_json,
+)
 
 #########
 # Tests #
@@ -45,7 +50,7 @@ def test_save_config_warn_bad_pair(
     # Add bad key,val pair and save ; should generate UserWarning(s) for removing bad pair
     config_fp32[key] = val
     with pytest.warns(UserWarning):
-        qconfig_save(config_fp32)
+        qconfig_save(config_fp32, minimal=False)
 
     # Load saved config and assert the key is not saved
     loaded_config = load_json("qcfg.json")  # load json as is - do not modify
@@ -71,14 +76,102 @@ def test_save_config_wanted_pairs(
     # Delete wanted pair from config and save ; should be reset to default
     if key in config_fp32:
         del config_fp32[key]
-    qconfig_save(config_fp32)
+    qconfig_save(config_fp32, minimal=False)
 
     # Load saved config and check the wanted pair was reset to default
     loaded_config = load_json()
     assert loaded_config.get(key) == default_val
 
     delete_config()
 
+def test_save_config_with_qcfg_save(
+    config_fp32: dict,
+    save_list: list,
+):
+    """
+    Test for checking that the "save_list" functionality works from within a quantized config
+
+    Args:
+        config_fp32 (dict): Config for fp32 quantization
+        save_list (list): List of variables to save in a quantized config.
+    """
+    delete_config()
+    config_fp32["save"] = save_list
+
+    qconfig_save(config_fp32, minimal=False)
+    
+    loaded_config = load_json()
+
+    # Remove pkg_versions and date before processing
+    del loaded_config["pkg_versions"]
+    del loaded_config["date"]
+
+    assert len(loaded_config) == len(save_list)
+
+    # Now ensure that every value in save_list was properly saved
+    for key in save_list:
+        assert key in loaded_config
+        assert loaded_config.get(key) == config_fp32.get(key)
+
+    delete_config()
+    del config_fp32["save"]
+
+def test_save_config_with_recipe_save(
+    config_fp32: dict,
+    save_list: list,
+):
+    """
+    Test for checking that the "save_list" functionality works from a saved json file
+
+    Args:
+        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("save.json")
+
+    # Save new "save.json"
+    save_path = "save_list.json"
+    save_json(save_list, file_path=save_path)
+
+    qconfig_save(config_fp32, recipe="save_list")
+    
+    # Check that saved qcfg matches
+    loaded_config = load_json()
+    
+    # Remove pkg_versions and date before processing
+    del loaded_config["pkg_versions"]
+    del loaded_config["date"]
+
+    assert len(loaded_config) == len(save_list)
+
+    # Now ensure that every value in save_list was properly saved
+    for key in save_list:
+        assert key in loaded_config
+        assert loaded_config.get(key) == config_fp32.get(key)
+
+    delete_config()
+    delete_config("save.json")
+
+def test_save_config_minimal(
+    config_fp32: dict,
+):
+    delete_config()
+
+    qconfig_save(config_fp32, minimal=True)
+
+    # Check that saved qcfg matches
+    loaded_config = load_json()
+
+    # Remove pkg_versions and date before processing
+    del loaded_config["pkg_versions"]
+    del loaded_config["date"]
+
+    # No items should exist - default config should be completely removed
+    assert len(loaded_config) == 0
+
+    delete_config()
 
 def test_load_config_restored_pair(
     config_fp32: dict,
@@ -96,13 +189,16 @@ def test_load_config_restored_pair(
 
     if key in config_fp32:
         del config_fp32[key]
+
     save_serialized_json(
         config_fp32
     )  # Save config as is, no other edits other than to serialize
 
     loaded_config = qconfig_load("qcfg.json")
     assert loaded_config.get(key) == default_val
 
+    delete_config()
+
 
 def test_load_config_required_pair(
     config_fp32: dict,