Unify the logger usage (#1581)

yiliu30 · web-flow · commit f50baf2e9107 · 2024-01-29T12:49:50.000+08:00
Signed-off-by: yiliu30 &lt;yi4.liu@intel.com&gt;
Signed-off-by: chensuyue &lt;suyue.chen@intel.com&gt;
diff --git a/.azure-pipelines/scripts/ut/3x/collect_log_3x.sh b/.azure-pipelines/scripts/ut/3x/collect_log_3x.sh
@@ -1,6 +1,6 @@
 source /neural-compressor/.azure-pipelines/scripts/change_color.sh
 
-set -xe
+set -e
 pip install coverage
 export COVERAGE_RCFILE=/neural-compressor/.azure-pipelines/scripts/ut/3x/coverage.${1}
 coverage_log="/neural-compressor/log_dir/coverage_log"
diff --git a/neural_compressor/common/__init__.py b/neural_compressor/common/__init__.py
@@ -13,38 +13,34 @@
 # limitations under the License.
 
 from neural_compressor.common.utils import (
-    level,
     log,
     info,
-    DEBUG,
     debug,
-    warn,
     warning,
     error,
     fatal,
+    level,
+    logger,
+    Logger,
     set_random_seed,
     set_workspace,
     set_resume_from,
     set_tensorboard,
-    Logger,
-    logger,
 )
 from neural_compressor.common.base_config import options
 
 
 __all__ = [
-    "level",
-    "log",
-    "info",
-    "DEBUG",
     "debug",
-    "warn",
-    "warning",
     "error",
     "fatal",
-    "options",
-    "Logger",
+    "info",
+    "level",
     "logger",
+    "log",
+    "warning",
+    "Logger",
+    "options",
     "set_workspace",
     "set_random_seed",
     "set_resume_from",
diff --git a/neural_compressor/common/base_config.py b/neural_compressor/common/base_config.py
@@ -24,7 +24,7 @@
 from itertools import product
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.utils import (
     BASE_CONFIG,
     COMPOSABLE_CONFIG,
@@ -36,8 +36,6 @@
     OP_NAME_OR_MODULE_TYPE,
 )
 
-logger = Logger().get_logger()
-
 __all__ = [
     "options",
     "register_config",
diff --git a/neural_compressor/common/base_tuning.py b/neural_compressor/common/base_tuning.py
@@ -18,11 +18,9 @@
 import uuid
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.base_config import BaseConfig, ComposableConfig
 
-logger = Logger().get_logger()
-
 __all__ = [
     "Evaluator",
     "TuningConfig",
diff --git a/neural_compressor/common/utils/logger.py b/neural_compressor/common/utils/logger.py
@@ -19,6 +19,18 @@
 import logging
 import os
 
+__all__ = [
+    "debug",
+    "error",
+    "fatal",
+    "info",
+    "level",
+    "logger",
+    "log",
+    "warning",
+    "Logger",
+]
+
 
 class Logger(object):
     """Logger class."""
@@ -67,68 +79,60 @@ def _pretty_dict(value, indent=0):
         return repr(value)
 
 
-level = Logger().get_logger().level
-DEBUG = logging.DEBUG
-
-
 def log(level, msg, *args, **kwargs):
     """Output log with the level as a parameter."""
     if isinstance(msg, dict):
         for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().log(level, line, *args, **kwargs)
+            Logger().get_logger().log(level, line, *args, **kwargs, stacklevel=2)
     else:
-        Logger().get_logger().log(level, msg, *args, **kwargs)
+        Logger().get_logger().log(level, msg, *args, **kwargs, stacklevel=2)
 
 
 def debug(msg, *args, **kwargs):
     """Output log with the debug level."""
     if isinstance(msg, dict):
         for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().debug(line, *args, **kwargs)
+            Logger().get_logger().debug(line, *args, **kwargs, stacklevel=2)
     else:
-        Logger().get_logger().debug(msg, *args, **kwargs)
+        Logger().get_logger().debug(msg, *args, **kwargs, stacklevel=2)
 
 
 def error(msg, *args, **kwargs):
     """Output log with the error level."""
     if isinstance(msg, dict):
         for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().error(line, *args, **kwargs)
+            Logger().get_logger().error(line, *args, **kwargs, stacklevel=2)
     else:
-        Logger().get_logger().error(msg, *args, **kwargs)
+        Logger().get_logger().error(msg, *args, **kwargs, stacklevel=2)
 
 
 def fatal(msg, *args, **kwargs):
     """Output log with the fatal level."""
     if isinstance(msg, dict):
         for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().fatal(line, *args, **kwargs)
+            Logger().get_logger().fatal(line, *args, **kwargs, stacklevel=2)
     else:
-        Logger().get_logger().fatal(msg, *args, **kwargs)
+        Logger().get_logger().fatal(msg, *args, **kwargs, stacklevel=2)
 
 
 def info(msg, *args, **kwargs):
     """Output log with the info level."""
     if isinstance(msg, dict):
         for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().info(line, *args, **kwargs)
+            Logger().get_logger().info(line, *args, **kwargs, stacklevel=2)
     else:
-        Logger().get_logger().info(msg, *args, **kwargs)
-
-
-def warn(msg, *args, **kwargs):
-    """Output log with the warning level."""
-    if isinstance(msg, dict):
-        for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().warning(line, *args, **kwargs)
-    else:
-        Logger().get_logger().warning(msg, *args, **kwargs)
+        Logger().get_logger().info(msg, *args, **kwargs, stacklevel=2)
 
 
 def warning(msg, *args, **kwargs):
     """Output log with the warning level (Alias of the method warn)."""
     if isinstance(msg, dict):
         for _, line in enumerate(_pretty_dict(msg).split("\n")):
-            Logger().get_logger().warning(line, *args, **kwargs)
+            Logger().get_logger().warning(line, *args, **kwargs, stacklevel=2)
     else:
-        Logger().get_logger().warning(msg, *args, **kwargs)
+        Logger().get_logger().warning(msg, *args, **kwargs, stacklevel=2)
+
+
+level = Logger().get_logger().level
+
+logger = Logger().get_logger()
diff --git a/neural_compressor/onnxrt/algorithms/weight_only/algo_entry.py b/neural_compressor/onnxrt/algorithms/weight_only/algo_entry.py
@@ -18,13 +18,11 @@
 
 import onnx
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.utils import RTN
 from neural_compressor.onnxrt.quantization.config import RTNConfig
 from neural_compressor.onnxrt.utils.utility import register_algo
 
-logger = Logger().get_logger()
-
 
 ###################### RTN Algo Entry ##################################
 @register_algo(name=RTN)
diff --git a/neural_compressor/onnxrt/quantization/config.py b/neural_compressor/onnxrt/quantization/config.py
@@ -23,12 +23,10 @@
 
 import onnx
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.base_config import BaseConfig, register_config, register_supported_configs_for_fwk
 from neural_compressor.common.utils import DEFAULT_WHITE_LIST, OP_NAME_OR_MODULE_TYPE, RTN
 
-logger = Logger().get_logger()
-
 FRAMEWORK_NAME = "onnxrt"
 
 
diff --git a/neural_compressor/onnxrt/quantization/quantize.py b/neural_compressor/onnxrt/quantization/quantize.py
@@ -18,13 +18,11 @@
 import onnx
 from onnxruntime.quantization import CalibrationDataReader
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.base_config import BaseConfig, ComposableConfig, config_registry
 from neural_compressor.onnxrt.quantization.config import FRAMEWORK_NAME
 from neural_compressor.onnxrt.utils.utility import algos_mapping
 
-logger = Logger().get_logger()
-
 
 def need_apply(quant_config: BaseConfig, algo_name):
     return quant_config.name == algo_name if hasattr(quant_config, "name") else False
diff --git a/neural_compressor/onnxrt/utils/onnx_model.py b/neural_compressor/onnxrt/utils/onnx_model.py
@@ -22,11 +22,9 @@
 
 import onnx
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.onnxrt.utils.utility import MAXIMUM_PROTOBUF, find_by_name
 
-logger = Logger().get_logger()
-
 
 class ONNXModel:
     """Build ONNX model."""
diff --git a/neural_compressor/onnxrt/utils/utility.py b/neural_compressor/onnxrt/utils/utility.py
@@ -18,9 +18,7 @@
 import onnx
 from packaging.version import Version
 
-from neural_compressor.common import Logger
-
-logger = Logger().get_logger()
+from neural_compressor.common import logger
 
 ONNXRT116_VERSION = Version("1.16.0")
 ONNXRT1161_VERSION = Version("1.16.1")
diff --git a/neural_compressor/tensorflow/algorithms/static_quantize/keras.py b/neural_compressor/tensorflow/algorithms/static_quantize/keras.py
@@ -26,11 +26,9 @@
 import tensorflow as tf
 import yaml
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.tensorflow.utils import deep_get, dump_elapsed_time
 
-logger = Logger().get_logger()
-
 
 def _add_supported_quantized_objects(custom_objects):
     """Map all the quantized objects."""
diff --git a/neural_compressor/tensorflow/quantization/quantize.py b/neural_compressor/tensorflow/quantization/quantize.py
@@ -16,14 +16,12 @@
 
 import tensorflow as tf
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.base_config import BaseConfig
 from neural_compressor.common.utils import STATIC_QUANT
 from neural_compressor.tensorflow.quantization.config import parse_config_from_dict
 from neural_compressor.tensorflow.utils import algos_mapping
 
-logger = Logger().get_logger()
-
 
 def quantize_model(
     model: tf.keras.Model, quant_config: BaseConfig, calib_dataloader: Callable = None, calib_iteration: int = 100
diff --git a/neural_compressor/torch/algorithms/weight_only_algos.py b/neural_compressor/torch/algorithms/weight_only_algos.py
@@ -17,13 +17,11 @@
 
 import torch
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.utils import GPTQ, RTN
 from neural_compressor.torch.quantization.config import GPTQConfig, RTNConfig
 from neural_compressor.torch.utils.utility import fetch_module, register_algo, set_module
 
-logger = Logger().get_logger()
-
 
 ###################### RTN Algo Entry ##################################
 @register_algo(name=RTN)
diff --git a/neural_compressor/torch/quantization/autotune.py b/neural_compressor/torch/quantization/autotune.py
@@ -17,15 +17,12 @@
 
 import torch
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.base_config import BaseConfig, get_all_config_set_from_config_registry
 from neural_compressor.common.base_tuning import TuningConfig, evaluator, init_tuning
 from neural_compressor.torch import quantize
 from neural_compressor.torch.quantization.config import FRAMEWORK_NAME
 
-logger = Logger().get_logger()
-
-
 __all__ = [
     "autotune",
     "get_all_config_set",
diff --git a/neural_compressor/torch/quantization/layers.py b/neural_compressor/torch/quantization/layers.py
@@ -25,7 +25,7 @@
 from torch.autograd import Function
 from torch.nn import functional as F
 
-from neural_compressor.common import DEBUG, level, logger
+from neural_compressor.common import logger
 from neural_compressor.torch.algorithms.weight_only import quant_tensor
 
 
diff --git a/neural_compressor/torch/quantization/quantize.py b/neural_compressor/torch/quantization/quantize.py
@@ -17,13 +17,11 @@
 
 import torch
 
-from neural_compressor.common import Logger
+from neural_compressor.common import logger
 from neural_compressor.common.base_config import BaseConfig, ComposableConfig, config_registry
 from neural_compressor.torch.quantization.config import FRAMEWORK_NAME
 from neural_compressor.torch.utils.utility import WHITE_MODULE_LIST, algos_mapping, get_model_info
 
-logger = Logger().get_logger()
-
 
 def need_apply(configs_mapping: Dict[Tuple[str, callable], BaseConfig], algo_name):
     return any(config.name == algo_name for config in configs_mapping.values())
diff --git a/neural_compressor/torch/utils/utility.py b/neural_compressor/torch/utils/utility.py
@@ -15,9 +15,7 @@
 
 from typing import Callable, Dict, List, Tuple
 
-from neural_compressor.common import Logger
-
-logger = Logger().get_logger()
+from neural_compressor.common import logger
 
 # Dictionary to store a mapping between algorithm names and corresponding algo implementation(function)
 algos_mapping: Dict[str, Callable] = {}
diff --git a/test/3x/onnxrt/quantization/weight_only/test_rtn.py b/test/3x/onnxrt/quantization/weight_only/test_rtn.py
@@ -4,9 +4,7 @@
 
 from optimum.exporters.onnx import main_export
 
-from neural_compressor.common import Logger
-
-logger = Logger().get_logger()
+from neural_compressor.common import logger
 
 
 def find_onnx_file(folder_path):
diff --git a/test/3x/onnxrt/test_config.py b/test/3x/onnxrt/test_config.py
@@ -7,9 +7,7 @@
 import onnx
 from optimum.exporters.onnx import main_export
 
-from neural_compressor.common import Logger
-
-logger = Logger().get_logger()
+from neural_compressor.common import logger
 
 
 def find_onnx_file(folder_path):
diff --git a/test/3x/tensorflow/test_config.py b/test/3x/tensorflow/test_config.py
@@ -25,9 +25,7 @@
 import tensorflow as tf
 from tensorflow import keras
 
-from neural_compressor.common import Logger
-
-logger = Logger().get_logger()
+from neural_compressor.common import logger
 
 
 def build_model():
diff --git a/test/3x/torch/quantization/weight_only/test_gptq_algo.py b/test/3x/torch/quantization/weight_only/test_gptq_algo.py
@@ -3,9 +3,7 @@
 
 import torch
 
-from neural_compressor.common import Logger
-
-logger = Logger().get_logger()
+from neural_compressor.common import logger
 
 
 def get_gpt_j():
diff --git a/test/3x/torch/quantization/weight_only/test_rtn.py b/test/3x/torch/quantization/weight_only/test_rtn.py
diff --git a/test/3x/torch/test_autotune.py b/test/3x/torch/test_autotune.py
diff --git a/test/3x/torch/test_config.py b/test/3x/torch/test_config.py
diff --git a/test/3x/torch/test_logger.py b/test/3x/torch/test_logger.py
diff --git a/test/3x/torch/test_utils.py b/test/3x/torch/test_utils.py
