[mypy] nncf/config (#3175)

### Changes

Enable mypy check for `nncf/config`

Author: AlexanderDokuchaev

nncf/common/accuracy_aware_training/training_loop.py

@@ -277,7 +277,7 @@ def __init__(
         super().__init__(compression_controller)
         accuracy_aware_training_params = extract_accuracy_aware_training_params(nncf_config)
         runner_factory = EarlyExitTrainingRunnerCreator(
-            accuracy_aware_training_params,  # type: ignore
+            accuracy_aware_training_params,
             compression_controller,
             uncompressed_model_accuracy,
             verbose,
@@ -330,7 +330,7 @@ def __init__(
 
         accuracy_aware_training_params = extract_accuracy_aware_training_params(nncf_config)
         runner_factory = AdaptiveCompressionLevelTrainingRunnerCreator(
-            accuracy_aware_training_params,  # type: ignore
+            accuracy_aware_training_params,
             self.adaptive_controller,
             uncompressed_model_accuracy,
             verbose,

nncf/config/config.py

@@ -11,10 +11,10 @@
 
 from copy import deepcopy
 from pathlib import Path
-from typing import Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
 import jsonschema
-import jstyleson as json
+import jstyleson as json  # type: ignore
 
 import nncf
 from nncf.common.logging import nncf_logger
@@ -28,27 +28,27 @@
 
 
 @api(canonical_alias="nncf.NNCFConfig")
-class NNCFConfig(dict):
+class NNCFConfig(dict[str, Any]):
     """Contains the configuration parameters required for NNCF to apply the selected algorithms.
 
     This is a regular dictionary object extended with some utility functions, such as the ability to attach well-defined
     structures to pass non-serializable objects as parameters. It is primarily built from a .json file, or from a
     Python JSON-like dictionary - both data types will be checked against a JSONSchema. See the definition of the
     schema at https://openvinotoolkit.github.io/nncf/schema/, or by calling NNCFConfig.schema()."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.__nncf_extra_structs: Dict[str, NNCFExtraConfigStruct] = {}
 
     @classmethod
-    def from_dict(cls, nncf_dict: Dict) -> "NNCFConfig":
+    def from_dict(cls, nncf_dict: Dict[str, Any]) -> "NNCFConfig":
         """
         Load NNCF config from a Python dictionary. The dict must contain only JSON-supported primitives.
 
         :param nncf_dict: A Python dict with the JSON-style configuration for NNCF.
         """
 
-        NNCFConfig.validate(nncf_dict)
+        cls.validate(nncf_dict)
         return cls(deepcopy(nncf_dict))
 
     @classmethod
@@ -63,7 +63,7 @@ def from_json(cls, path: str) -> "NNCFConfig":
             loaded_json = json.load(f)
         return cls.from_dict(loaded_json)
 
-    def register_extra_structs(self, struct_list: List[NNCFExtraConfigStruct]):
+    def register_extra_structs(self, struct_list: List[NNCFExtraConfigStruct]) -> None:
         """
         Attach the supplied list of extra configuration structures to this configuration object.
 
@@ -78,7 +78,7 @@ def register_extra_structs(self, struct_list: List[NNCFExtraConfigStruct]):
     def get_extra_struct(self, struct_cls: Type[NNCFExtraConfigStruct]) -> NNCFExtraConfigStruct:
         return self.__nncf_extra_structs[struct_cls.get_id()]
 
-    def has_extra_struct(self, struct_cls: Type[NNCFExtraConfigStruct]) -> NNCFExtraConfigStruct:
+    def has_extra_struct(self, struct_cls: Type[NNCFExtraConfigStruct]) -> bool:
         return struct_cls.get_id() in self.__nncf_extra_structs
 
     def get_all_extra_structs(self) -> List[NNCFExtraConfigStruct]:
@@ -108,7 +108,7 @@ def get_redefinable_global_param_value_for_algo(self, param_name: str, algo_name
         return param
 
     @staticmethod
-    def schema() -> Dict:
+    def schema() -> Dict[str, Any]:
         """
         Returns the JSONSchema against which the input data formats (.json or Python dict) are validated.
         """
@@ -124,15 +124,15 @@ def _is_path_to_algorithm_name(path_parts: List[str]) -> bool:
         )
 
     @staticmethod
-    def validate(loaded_json):
+    def validate(loaded_json: Dict[str, Any]) -> None:
         try:
             jsonschema.validate(loaded_json, NNCFConfig.schema())
         except jsonschema.ValidationError as e:
             nncf_logger.error("Invalid NNCF config supplied!")
             absolute_path_parts = [str(x) for x in e.absolute_path]
             if not NNCFConfig._is_path_to_algorithm_name(absolute_path_parts):
                 e.message += f"\nRefer to the NNCF config schema documentation at {SCHEMA_VISUALIZATION_URL}"
-                e.schema = "*schema too long for stdout display*"
+                e.schema = "*schema too long for stdout display*"  # type: ignore[assignment]
                 raise e
 
             # Need to make the error more algo-specific in case the config was so bad that no

nncf/config/extractors.py

@@ -31,7 +31,7 @@ def extract_algorithm_names(config: NNCFConfig) -> List[str]:
     return retval
 
 
-def extract_algo_specific_config(config: NNCFConfig, algo_name_to_match: str) -> Dict:
+def extract_algo_specific_config(config: NNCFConfig, algo_name_to_match: str) -> Dict[str, Any]:
     """
     Extracts a .json sub-dictionary for a given compression algorithm from the
     common NNCFConfig.
@@ -77,7 +77,7 @@ def extract_algo_specific_config(config: NNCFConfig, algo_name_to_match: str) ->
     return next(iter(matches))
 
 
-def extract_range_init_params(config: NNCFConfig, algorithm_name: str = "quantization") -> Optional[Dict[str, object]]:
+def extract_range_init_params(config: NNCFConfig, algorithm_name: str = "quantization") -> Optional[Dict[str, Any]]:
     """
     Extracts parameters of the quantization range initialization algorithm from the
     compression algorithm NNCFconfig.
@@ -90,7 +90,6 @@ def extract_range_init_params(config: NNCFConfig, algorithm_name: str = "quantiz
     algo_config = extract_algo_specific_config(config, algorithm_name)
     init_range_config_dict_or_list = algo_config.get("initializer", {}).get("range", {})
 
-    range_init_args = None
     try:
         range_init_args = config.get_extra_struct(QuantizationRangeInitArgs)
     except KeyError:
@@ -120,7 +119,7 @@ def extract_range_init_params(config: NNCFConfig, algorithm_name: str = "quantiz
 
     if max_num_init_samples == 0:
         return None
-    if range_init_args is None:
+    if not isinstance(range_init_args, QuantizationRangeInitArgs):
         raise ValueError(
             "Should run range initialization as specified via config,"
             "but the initializing data loader is not provided as an extra struct. "
@@ -161,7 +160,7 @@ class BNAdaptDataLoaderNotFoundError(RuntimeError):
     pass
 
 
-def get_bn_adapt_algo_kwargs(nncf_config: NNCFConfig, params: Dict[str, Any]) -> Dict[str, Any]:
+def get_bn_adapt_algo_kwargs(nncf_config: NNCFConfig, params: Dict[str, Any]) -> Optional[Dict[str, Any]]:
     num_bn_adaptation_samples = params.get("num_bn_adaptation_samples", NUM_BN_ADAPTATION_SAMPLES)
 
     if num_bn_adaptation_samples == 0:
@@ -175,6 +174,11 @@ def get_bn_adapt_algo_kwargs(nncf_config: NNCFConfig, params: Dict[str, Any]) ->
             "because the data loader is not provided as an extra struct. Refer to the "
             "`NNCFConfig.register_extra_structs` method and the `BNAdaptationInitArgs` class."
         ) from None
+
+    if not isinstance(args, BNAdaptationInitArgs):
+        raise BNAdaptDataLoaderNotFoundError(
+            "The extra struct for batch-norm adaptation must be an instance of the BNAdaptationInitArgs class."
+        )
     params = {
         "num_bn_adaptation_samples": num_bn_adaptation_samples,
         "data_loader": args.data_loader,
@@ -183,7 +187,7 @@ def get_bn_adapt_algo_kwargs(nncf_config: NNCFConfig, params: Dict[str, Any]) ->
     return params
 
 
-def extract_accuracy_aware_training_params(config: NNCFConfig) -> Dict[str, object]:
+def extract_accuracy_aware_training_params(config: NNCFConfig) -> Dict[str, Any]:
     """
     Extracts accuracy aware training parameters from NNCFConfig.
 
@@ -196,7 +200,7 @@ class NNCFAlgorithmNames:
         FILTER_PRUNING = "filter_pruning"
         SPARSITY = ["rb_sparsity", "magnitude_sparsity", "const_sparsity"]
 
-    def validate_accuracy_aware_schema(config: NNCFConfig, params: Dict[str, object]):
+    def validate_accuracy_aware_schema(config: NNCFConfig, params: Dict[str, Any]) -> None:
         from nncf.common.accuracy_aware_training import AccuracyAwareTrainingMode
 
         if params["mode"] == AccuracyAwareTrainingMode.EARLY_EXIT:

nncf/config/schema.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 import logging
-from typing import Dict
+from typing import Any, Dict
 
 import jsonschema
 
@@ -148,7 +148,9 @@
 }
 
 
-def validate_single_compression_algo_schema(single_compression_algo_dict: Dict, ref_vs_algo_schema: Dict):
+def validate_single_compression_algo_schema(
+    single_compression_algo_dict: Dict[str, Any], ref_vs_algo_schema: Dict[str, Any]
+) -> None:
     """single_compression_algo_dict must conform to BASIC_COMPRESSION_ALGO_SCHEMA (and possibly has other
     algo-specific properties"""
     algo_name = single_compression_algo_dict["algorithm"]
@@ -172,7 +174,7 @@ def validate_single_compression_algo_schema(single_compression_algo_dict: Dict,
         raise e
 
 
-def validate_accuracy_aware_training_schema(single_compression_algo_dict: Dict):
+def validate_accuracy_aware_training_schema(single_compression_algo_dict: Dict[str, Any]) -> None:
     """
     Checks accuracy_aware_training section.
     """

nncf/config/schemata/algo/quantization.py

@@ -160,7 +160,7 @@
 PER_LAYER_RANGE_INIT_CONFIG_PROPERTIES = {
     "type": "object",
     "properties": {
-        **BASIC_RANGE_INIT_CONFIG_PROPERTIES["properties"],
+        **BASIC_RANGE_INIT_CONFIG_PROPERTIES["properties"],  # type: ignore[dict-item]
         **SCOPING_PROPERTIES,
         "target_quantizer_group": with_attributes(
             STRING,
@@ -338,7 +338,7 @@
     "parameters have a better chance to get fine-tuned to values that result in good accuracy.",
     "properties": {
         "batchnorm_adaptation": BATCHNORM_ADAPTATION_SCHEMA,
-        **RANGE_INIT_CONFIG_PROPERTIES["initializer"]["properties"],
+        **RANGE_INIT_CONFIG_PROPERTIES["initializer"]["properties"],  # type: ignore[dict-item]
         "precision": PRECISION_INITIALIZER_SCHEMA,
     },
     "additionalProperties": False,

nncf/config/schemata/basic.py

@@ -8,7 +8,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Dict
+from typing import Any, Dict, List, Optional
 
 NUMBER = {"type": "number"}
 STRING = {"type": "string"}
@@ -17,22 +17,22 @@
 ARRAY_OF_STRINGS = {"type": "array", "items": STRING}
 
 
-def annotated_enum(names_vs_description: Dict[str, str]) -> Dict:
+def annotated_enum(names_vs_description: Dict[str, str]) -> Dict[str, List[Dict[str, str]]]:
     retval_list = []
     for name, descr in names_vs_description.items():
         retval_list.append({"const": name, "title": name, "description": descr})
     return {"oneOf": retval_list}
 
 
-def make_string_or_array_of_strings_schema(addtl_dict_entries: Dict = None) -> Dict:
+def make_string_or_array_of_strings_schema(addtl_dict_entries: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
     if addtl_dict_entries is None:
         addtl_dict_entries = {}
     retval = {"type": ["array", "string"], "items": {"type": "string"}}
     retval.update(addtl_dict_entries)
     return retval
 
 
-def make_object_or_array_of_objects_schema(single_object_schema: Dict = None) -> Dict:
+def make_object_or_array_of_objects_schema(single_object_schema: Dict[str, Any]) -> Dict[str, Any]:
     retval = {
         "oneOf": [
             {
@@ -45,6 +45,6 @@ def make_object_or_array_of_objects_schema(single_object_schema: Dict = None) ->
     return retval
 
 
-def with_attributes(schema: Dict, **kwargs) -> Dict:
+def with_attributes(schema: Dict[str, Any], **kwargs: Any) -> Dict[str, Any]:
     retval = {**schema, **kwargs}
     return retval

nncf/config/schemata/experimental_schema.py

@@ -35,7 +35,7 @@
 # Experimental Quantization
 ########################################################################################################################
 EXPERIMENTAL_QUANTIZATION_SCHEMA = copy.deepcopy(QUANTIZATION_SCHEMA)
-EXPERIMENTAL_QUANTIZATION_SCHEMA["properties"]["algorithm"]["const"] = EXPERIMENTAL_QUANTIZATION_ALGO_NAME_IN_CONFIG
+EXPERIMENTAL_QUANTIZATION_SCHEMA["properties"]["algorithm"]["const"] = EXPERIMENTAL_QUANTIZATION_ALGO_NAME_IN_CONFIG  # type: ignore[index]
 
 ########################################################################################################################
 # BootstrapNAS

nncf/config/structures.py

@@ -11,7 +11,7 @@
 """
 Structures for passing live Python objects into NNCF algorithms.
 """
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 
 from nncf.common.initialization.dataloader import NNCFDataLoader
 from nncf.common.utils.api_marker import api
@@ -47,7 +47,7 @@ def data_loader(self) -> NNCFDataLoader:
         return self._data_loader
 
     @property
-    def device(self) -> str:
+    def device(self) -> Optional[str]:
         return self._device
 
     @classmethod
@@ -74,7 +74,7 @@ def data_loader(self) -> NNCFDataLoader:
         return self._data_loader
 
     @property
-    def device(self) -> str:
+    def device(self) -> Optional[str]:
         return self._device
 
     @classmethod
@@ -91,7 +91,7 @@ class ModelEvaluationArgs(NNCFExtraConfigStruct):
         the evaluation split of the dataset corresponding to the model.
     """
 
-    def __init__(self, eval_fn: Callable):
+    def __init__(self, eval_fn: Callable[..., Any]):
         self.eval_fn = eval_fn
 
     @classmethod

pyproject.toml

@@ -93,6 +93,7 @@ strict = true
 # https://github.com/hauntsaninja/no_implicit_optional
 implicit_optional = true
 files = [
+    "nncf/config",
     "nncf/common/sparsity",
     "nncf/common/graph",
     "nncf/common/accuracy_aware_training/",