Refactor fw_info handling in pytest unittests (#1488)

* Switch pytest unittests to fw_info patch instead of changing the global _current_framework_info state.
* Remove activation_quantizer_mapping from DummyFrameworkInfo in tests.

Author: elad-c

tests_pytest/common_tests/unit_tests/core/graph/test_base_node.py

@@ -15,11 +15,9 @@
 import numpy as np
 
 from tests_pytest._test_util.graph_builder_utils import build_node, build_nbits_qc
-from model_compression_toolkit.core.common.framework_info import set_fw_info
 
 
-def test_find_min_max_candidate_index(fw_info_mock):
-    set_fw_info(fw_info_mock)
+def test_find_min_max_candidate_index(patch_fw_info):
     qcs = []
     for ab in [4, 8, 16, 2]:
         for fb in [2, 8, 4]:

tests_pytest/common_tests/unit_tests/core/graph/test_node_quantization.py

@@ -21,7 +21,6 @@
 
 from model_compression_toolkit.quantization_preparation.load_fqc import set_quantization_configs_to_node
 from tests_pytest._test_util.graph_builder_utils import build_node, DummyLayer
-from model_compression_toolkit.core.common.framework_info import set_fw_info
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import QuantizationConfigOptions, \
     OpQuantizationConfig, AttributeQuantizationConfig, Signedness
 from mct_quantizers import QuantizationMethod
@@ -36,10 +35,6 @@ class NoActivationQuantNode:
 
 
 class TestSetNodeQuantizationConfig:
-    @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
-
     @staticmethod
     def _get_op_config(activation_n_bits,
                        supported_input_activation_n_bits,
@@ -55,10 +50,9 @@ def _get_op_config(activation_n_bits,
                                     quantization_preserving=quantization_preserving,
                                     signedness=Signedness.AUTO)
 
-    def test_activation_preserving_with_2_inputs(self, fw_info_mock):
+    def test_activation_preserving_with_2_inputs(self, patch_fw_info):
         """ Tests that . """
-        fw_info_mock.activation_quantizer_mapping = {QuantizationMethod.POWER_OF_TWO: lambda x: 0}
-        fw_info_mock.get_kernel_op_attribute = lambda x: None
+        patch_fw_info.get_kernel_op_attribute = lambda x: None
 
         n1 = build_node('in1_node')
         n2 = build_node('in2_node')
@@ -71,9 +65,9 @@ def test_activation_preserving_with_2_inputs(self, fw_info_mock):
                                  Edge(n3, n4, 0, 0),
                                  Edge(n1, qp3, 0, 0), Edge(qp3, qp4, 0, 0)])
         q_op_config_kwargs = {"activation_n_bits": 7, "supported_input_activation_n_bits": 7,
-                            "enable_activation_quantization": True, "quantization_preserving": False}
+                              "enable_activation_quantization": True, "quantization_preserving": False}
         qp_op_config_kwargs = {"activation_n_bits": 7, "supported_input_activation_n_bits": 7,
-                            "enable_activation_quantization": False, "quantization_preserving": True}
+                               "enable_activation_quantization": False, "quantization_preserving": True}
         _filters = {DummyLayer: QuantizationConfigOptions(quantization_configurations=[self._get_op_config(**q_op_config_kwargs)]),
                     PreservingNode: QuantizationConfigOptions(quantization_configurations=[self._get_op_config(**qp_op_config_kwargs)])}
         fqc = Mock(filterlayer2qco=_filters, layer2qco=_filters)
@@ -85,12 +79,11 @@ def test_activation_preserving_with_2_inputs(self, fw_info_mock):
         assert qp3.is_quantization_preserving()
         assert qp4.is_quantization_preserving()
 
-    def test_node_quantization_by_next_nodes(self, fw_info_mock):
+    def test_node_quantization_by_next_nodes(self, patch_fw_info):
         """
         Test that node quantization n_bits is unaffected by preserving next node and not-enabled quantization next node.
         """
-        fw_info_mock.activation_quantizer_mapping = {QuantizationMethod.POWER_OF_TWO: lambda x: 0}
-        fw_info_mock.get_kernel_op_attribute = lambda x: None
+        patch_fw_info.get_kernel_op_attribute = lambda x: None
 
         first_node = build_node('first_node')
         preserving_node = build_node('preserving_node', layer_class=PreservingNode)
@@ -105,14 +98,14 @@ def test_node_quantization_by_next_nodes(self, fw_info_mock):
                                     "quantization_preserving": False}
 
         preserving_node_config_kwargs = {"activation_n_bits": 8,
-                                        "supported_input_activation_n_bits": [8, 16],
-                                        "enable_activation_quantization": False,
-                                        "quantization_preserving": True}
+                                         "supported_input_activation_n_bits": [8, 16],
+                                         "enable_activation_quantization": False,
+                                         "quantization_preserving": True}
 
         no_quant_node_config_kwargs = {"activation_n_bits": 8,
-                                        "supported_input_activation_n_bits": [8],
-                                        "enable_activation_quantization": False,
-                                        "quantization_preserving": False}
+                                       "supported_input_activation_n_bits": [8],
+                                       "enable_activation_quantization": False,
+                                       "quantization_preserving": False}
         _filters = {
             DummyLayer: QuantizationConfigOptions(quantization_configurations=[self._get_op_config(**first_node_config_kwargs)]),
             PreservingNode: QuantizationConfigOptions(quantization_configurations=[self._get_op_config(**preserving_node_config_kwargs)]),

tests_pytest/common_tests/unit_tests/core/graph/test_quantization_preserving_node.py

@@ -16,17 +16,12 @@
 
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.graph.edge import Edge
-from model_compression_toolkit.core.common.framework_info import set_fw_info
 
 from tests_pytest._test_util.graph_builder_utils import build_node, build_nbits_qc
 
 
 class TestQuantizationPreservingNode:
-    @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
-
-    def test_activation_preserving_candidate(self):
+    def test_activation_preserving_candidate(self, patch_fw_info):
         """ Tests that the correct activation quantization candidate is selected. """
         n1 = build_node('qact_node', qcs=[build_nbits_qc()])
         n2 = build_node('qp1a_node', qcs=[build_nbits_qc(a_enable=False, q_preserving=True)])
@@ -42,7 +37,7 @@ def test_activation_preserving_candidate(self):
         assert graph.retrieve_preserved_quantization_node(n4) is n4
         assert graph.retrieve_preserved_quantization_node(n5) is n4
 
-    def test_activation_preserving_disable_for_multi_input_node(self):
+    def test_activation_preserving_disable_for_multi_input_node(self, patch_fw_info):
         """ Tests that the retrieve_preserved_quantization_node raises an assertion error if node has more than 1 input. """
         n1 = build_node('qact_node', qcs=[build_nbits_qc()])
         n2 = build_node('qp1a_node', qcs=[build_nbits_qc(a_enable=False, q_preserving=True)])

tests_pytest/common_tests/unit_tests/core/graph/test_virtual_activation_weights_node.py

@@ -15,7 +15,6 @@
 import numpy as np
 import pytest
 
-from model_compression_toolkit.core.common.framework_info import set_fw_info
 from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualActivationWeightsNode
 from tests_pytest._test_util.graph_builder_utils import build_node, DummyLayer, build_nbits_qc
 
@@ -25,17 +24,13 @@ class DummyLayerWKernel:
 
 
 class TestVirtualActivationWeightsNode:
-    @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
-
     # TODO tests only cover combining weights from activation and weight nodes and errors.
-    def test_activation_with_weights(self, fw_info_mock):
+    def test_activation_with_weights(self, patch_fw_info):
         """ Tests that weights from activation and weight node are combined correctly. """
         # Each node has a unique weight attr and a unique positional weights. In addition, both nodes have
         # an identical canonical attribute (but different full name), and an identical positional weight.
         # All weights have different quantization.
-        fw_info_mock.get_kernel_op_attribute = lambda nt: 'weight' if nt is DummyLayerWKernel else None
+        patch_fw_info.get_kernel_op_attribute = lambda nt: 'weight' if nt is DummyLayerWKernel else None
 
         a_node = build_node('a', final_weights={'aaweightaa': np.ones((3, 14)), 'foo': np.ones(15),
                                                 1: np.ones(15), 2: np.ones((5, 9))}, qcs=[build_nbits_qc(a_nbits=5,
@@ -101,8 +96,8 @@ def test_activation_with_weights(self, fw_info_mock):
             3: w_qc.weights_quantization_cfg.pos_attributes_config_mapping[3]
         }
 
-    def test_invalid_configurable_w_node_weight(self, fw_info_mock):
-        fw_info_mock.get_kernel_op_attribute = lambda nt: 'kernel' if nt is DummyLayerWKernel else None
+    def test_invalid_configurable_w_node_weight(self, patch_fw_info):
+        patch_fw_info.get_kernel_op_attribute = lambda nt: 'kernel' if nt is DummyLayerWKernel else None
 
         w_node = build_node('w', canonical_weights={'kernel': np.ones(3), 'foo': np.ones(14)}, qcs=[
             build_nbits_qc(w_attr={'kernel': (8, True), 'foo': (8, True)}),
@@ -113,8 +108,8 @@ def test_invalid_configurable_w_node_weight(self, fw_info_mock):
         with pytest.raises(NotImplementedError, match='Only kernel weight can be configurable. Got configurable .*foo'):
             VirtualActivationWeightsNode(a_node, w_node)
 
-    def test_invalid_a_node_configurable_weight(self, fw_info_mock):
-        fw_info_mock.get_kernel_op_attribute = lambda nt: 'kernel' if nt is DummyLayerWKernel else None
+    def test_invalid_a_node_configurable_weight(self, patch_fw_info):
+        patch_fw_info.get_kernel_op_attribute = lambda nt: 'kernel' if nt is DummyLayerWKernel else None
 
         w_node = build_node('w', canonical_weights={'kernel': np.ones(3), 'foo': np.ones(14)}, qcs=[
             build_nbits_qc(w_attr={'kernel': (8, True), 'foo': (8, True)}),
@@ -128,8 +123,8 @@ def test_invalid_a_node_configurable_weight(self, fw_info_mock):
                                                       'activation for VirtualActivationWeightsNode'):
             VirtualActivationWeightsNode(a_node, w_node)
 
-    def test_invalid_a_node_kernel(self, fw_info_mock):
-        fw_info_mock.get_kernel_op_attribute = lambda nt: 'weight' if nt is DummyLayerWKernel else 'kernel'
+    def test_invalid_a_node_kernel(self, patch_fw_info):
+        patch_fw_info.get_kernel_op_attribute = lambda nt: 'weight' if nt is DummyLayerWKernel else 'kernel'
         w_node = build_node('w', canonical_weights={'weight': np.ones(3)},
                             qcs=[build_nbits_qc(w_attr={'weight': (8, True)})], layer_class=DummyLayerWKernel)
         a_node = build_node('aaa', canonical_weights={'kernel': np.ones(3)},

tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py

@@ -22,7 +22,6 @@
 
 from model_compression_toolkit.constants import FLOAT_BITWIDTH, FUSED_LAYER_PATTERN, FUSED_OP_QUANT_CONFIG
 from model_compression_toolkit.core import ResourceUtilization
-from model_compression_toolkit.core.common.framework_info import set_fw_info
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.fusion.fusing_info import FusingInfo
 from model_compression_toolkit.core.common.graph.edge import Edge
@@ -69,9 +68,8 @@ class TestComputeResourceUtilization:
         compute_resource_utilization on a virtual graph is tested in TestBOPSAndVirtualGraph
     """
     @pytest.fixture(autouse=True)
-    def setup(self, graph_mock, fw_impl_mock, fw_info_mock):
-        fw_info_mock.get_kernel_op_attribute = Mock(return_value='foo')    # for bops
-        set_fw_info(fw_info_mock)
+    def setup(self, graph_mock, fw_impl_mock, patch_fw_info):
+        patch_fw_info.get_kernel_op_attribute = Mock(return_value='foo')    # for bops
         fw_impl_mock.get_node_mac_operations = lambda n: 42 if n == n2 else 0    # for bops
         n1 = build_node('n1', qcs=[build_qc()], output_shape=(None, 5, 10))
         n2 = build_node('n2', canonical_weights={'foo': np.zeros((3, 14))}, qcs=[build_qc(w_attr={'foo': (4, True)})],
@@ -226,8 +224,8 @@ def _validate(self, ret, detailed, exp_ru: ResourceUtilization):
 
 class TestActivationUtilizationMethods:
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     """ Tests for non-public activation utilization api. """
     def test_get_a_nbits_configurable(self, graph_mock, fw_impl_mock):
@@ -337,8 +335,8 @@ def test_get_target_activation_nodes(self, graph_mock, fw_impl_mock):
 
 class TestComputeActivationTensorsUtilization:
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     """ Tests for activation tensors utilization public apis. """
     def test_compute_node_activation_tensor_utilization(self, graph_mock, fw_impl_mock):
@@ -436,8 +434,8 @@ def test_compute_act_tensors_util_invalid_custom_qcs(self, graph_mock, fw_impl_m
 
 class TestActivationMaxCutUtilization:
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     """ Tests for activation max cut utilization. """
     def test_compute_cuts_integration(self, graph_mock, fw_impl_mock, mocker):
@@ -737,8 +735,8 @@ def test_compute_act_utilization_by_cut_invalid_custom_qcs(self, graph_mock, fw_
 class TestWeightUtilizationMethods:
     """ Tests for weights utilization non-public api. """
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     def test_get_w_nbits(self, graph_mock, fw_impl_mock):
         ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock)
@@ -850,8 +848,8 @@ def test_collect_target_nodes_w_attrs(self, graph_mock, fw_impl_mock):
 class TestComputeNodeWeightsUtilization:
     """ Tests for compute_node_weight_utilization public method. """
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     @pytest.fixture
     def setup_node_w_test(self, graph_mock, fw_impl_mock):
@@ -940,8 +938,8 @@ def test_compute_node_w_utilization_errors(self, graph_mock, fw_impl_mock, setup
 class TestComputeWeightUtilization:
     """ Tests for compute_weight_utilization public method. """
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     @pytest.fixture
     def prepare_compute_w_util(self, fw_impl_mock):
@@ -1050,8 +1048,8 @@ def test_compute_w_utilization_invalid_custom_qcs(self, graph_mock, fw_impl_mock
 class TestCalculatorMisc:
     """ Calculator tests that don't belong to other test classes """
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     def test_calculator_init(self, fw_impl_mock):
         n1 = build_node('n1', qcs=[build_qc(a_enable=False)], output_shape=(None, 5, 10))
@@ -1089,8 +1087,8 @@ class BOPNode:
 
 class TestBOPSAndVirtualGraph:
     @pytest.fixture(autouse=True)
-    def setup(self, fw_info_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, patch_fw_info):
+        pass
 
     def test_compute_regular_node_bops(self, fw_impl_mock, fw_info_mock):
         fw_info_mock.get_kernel_op_attribute = lambda node_type: 'foo' if node_type == BOPNode else None
@@ -1177,7 +1175,7 @@ def test_node_bops_invalid_criterion(self, graph_mock, fw_impl_mock, target_crit
         with pytest.raises(ValueError, match='BOPS computation is supported only for Any, AnyQuantized and AnyQuantizedNonFused targets.'):
             ru_calc.compute_node_bops(Mock(), target_criterion, BM.Float)
 
-    def test_compute_bops(self, fw_impl_mock, fw_info_mock,):
+    def test_compute_bops(self, fw_impl_mock, fw_info_mock):
         class BOPNode2:
             pass
 

tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_data.py

@@ -26,7 +26,7 @@
 
 class TestResourceUtilizationData:
     @pytest.mark.parametrize('error_method', [QuantizationErrorMethod.MSE, QuantizationErrorMethod.HMSE])
-    def test_resource_utilization_data(self, fw_info_mock, fw_impl_mock, error_method, mocker):
+    def test_resource_utilization_data(self, fw_impl_mock, error_method, mocker):
         core_cfg = CoreConfig()
         core_cfg.quantization_config.weights_error_method = error_method
         core_cfg.bit_width_config = BitWidthConfig([1, 2])
@@ -42,7 +42,6 @@ def test_resource_utilization_data(self, fw_info_mock, fw_impl_mock, error_metho
         prep_runner = mocker.patch('model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.'
                                    'resource_utilization_data.graph_preparation_runner')
 
-        _current_framework_info = fw_info_mock
         compute_resource_utilization_data(model_mock,
                                           data_gen_mock,
                                           core_cfg,

tests_pytest/common_tests/unit_tests/core/mixed_precision/sensitivity_eval/test_distance_calculator.py

@@ -17,7 +17,6 @@
 import numpy as np
 import pytest
 
-from model_compression_toolkit.core.common.framework_info import set_fw_info
 from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, MpDistanceWeighting
 from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.mixed_precision.sensitivity_eval.metric_calculators import \
@@ -34,8 +33,7 @@ class TestDistanceWeighting:
     out_pts = np.array([[1, 2], [3, 4], [5, 6]])
 
     @pytest.fixture
-    def setup(self, mocker, graph_mock, fw_info_mock, fw_impl_mock):
-        set_fw_info(fw_info_mock)
+    def setup(self, mocker, graph_mock, patch_fw_info, fw_impl_mock):
         mocker.patch.object(DistanceMetricCalculator, 'get_mp_interest_points', return_value=[None, None])
         mocker.patch.object(DistanceMetricCalculator, 'get_output_nodes_for_metric', return_value=[None])
         mocker.patch.object(DistanceMetricCalculator, '_init_metric_points_lists', return_value=(None, None))