Fix logger (#1583)

yiliu30 · web-flow · commit 83bc779a4e97 · 2024-01-31T17:39:38.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/run_3x_pt.sh b/.azure-pipelines/scripts/ut/3x/run_3x_pt.sh
@@ -51,7 +51,7 @@ fi
 
 # check pytest issue
 if [ -s run_pytest.sh ]; then
-    if [ $(grep -c '== FAILURES ==' ${ut_log_name}) != 0 ] || [ $(grep -c '== ERRORS ==' ${ut_log_name}) != 0 ] || [ $(grep -c 'passed,' ${ut_log_name}) == 0 ]; then
+    if [ $(grep -c '== FAILURES ==' ${ut_log_name}) != 0 ] || [ $(grep -c '== ERRORS ==' ${ut_log_name}) != 0 ] || [ $(grep -c ' passed ' ${ut_log_name}) == 0 ]; then
         echo "Find errors in pytest case, please check the output..."
         echo "Please search for '== FAILURES ==' or '== ERRORS =='"
         ut_status="failed"
diff --git a/docs/source/adaptor.md b/docs/source/adaptor.md
@@ -44,22 +44,17 @@ For example, a user can implement an `Abc` adaptor like below:
 ```python
 @adaptor_registry
 class AbcAdaptor(Adaptor):
-    def __init__(self, framework_specific_info):
-        ...
+    def __init__(self, framework_specific_info): ...
 
-    def quantize(self, tune_cfg, model, dataloader, q_func=None):
-        ...
+    def quantize(self, tune_cfg, model, dataloader, q_func=None): ...
 
     def evaluate(
         self, model, dataloader, postprocess=None, metric=None, measurer=None, iteration=-1, tensorboard=False
-    ):
-        ...
+    ): ...
 
-    def query_fw_capability(self, model):
-        ...
+    def query_fw_capability(self, model): ...
 
-    def query_fused_patterns(self, model):
-        ...
+    def query_fused_patterns(self, model): ...
 ```
 
 * `quantize` function is used to perform quantization for post-training quantization and quantization-aware training. Quantization processing includes calibration and conversion processing for post-training quantization, while for quantization-aware training, it includes training and conversion processing.
diff --git a/docs/source/migration.md b/docs/source/migration.md
@@ -41,8 +41,7 @@ val_dataloader = torch.utils.data.Dataloader(
 )
 
 
-def eval_func(model):
-    ...
+def eval_func(model): ...
 
 
 # Quantization code
@@ -115,8 +114,7 @@ val_dataloader = torch.utils.data.Dataloader(
 )
 
 
-def eval_func(model):
-    ...
+def eval_func(model): ...
 
 
 # Quantization code
@@ -147,12 +145,10 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path)
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
 
 
-def eval_func(model):
-    ...
+def eval_func(model): ...
 
 
-def train_func(model):
-    ...
+def train_func(model): ...
 
 
 trainer = Trainer(...)
@@ -213,12 +209,10 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path)
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
 
 
-def eval_func(model):
-    ...
+def eval_func(model): ...
 
 
-def train_func(model):
-    ...
+def train_func(model): ...
 
 
 trainer = Trainer(...)
diff --git a/docs/source/pruning.md b/docs/source/pruning.md
@@ -342,6 +342,7 @@ The following section exemplifies how to use hooks in user pass-in training func
   on_after_optimizer_step() # Update weights' criteria, mask weights
   on_train_end() # End of pruner, print sparse information
   """
+
   from neural_compressor.training import prepare_compression, WeightPruningConfig
 
   config = WeightPruningConfig(configs)
diff --git a/docs/source/quantization.md b/docs/source/quantization.md
@@ -319,8 +319,7 @@ criterion = ...
 
 
 # Quantization code
-def train_func(model):
-    ...
+def train_func(model): ...
 
 
 from neural_compressor import QuantizationAwareTrainingConfig
diff --git a/docs/source/tuning_strategies.md b/docs/source/tuning_strategies.md
@@ -507,8 +507,7 @@ For example, user can implement an `Abc` strategy like below:
 ```python
 @strategy_registry
 class AbcTuneStrategy(TuneStrategy):
-    def __init__(self, model, conf, q_dataloader, q_func=None, eval_dataloader=None, eval_func=None, dicts=None):
-        ...
+    def __init__(self, model, conf, q_dataloader, q_func=None, eval_dataloader=None, eval_func=None, dicts=None): ...
 
     def next_tune_cfg(self):
         # generate the next tuning config
diff --git a/neural_compressor/common/__init__.py b/neural_compressor/common/__init__.py
@@ -14,37 +14,21 @@
 
 from neural_compressor.common.utils import (
     level,
-    log,
-    info,
-    DEBUG,
-    debug,
-    warn,
-    warning,
-    error,
-    fatal,
+    logger,
+    Logger,
     set_random_seed,
-    set_workspace,
     set_resume_from,
+    set_workspace,
     set_tensorboard,
-    Logger,
-    logger,
 )
 from neural_compressor.common.base_config import options
 
 
 __all__ = [
-    "level",
-    "log",
-    "info",
-    "DEBUG",
-    "debug",
-    "warn",
-    "warning",
-    "error",
-    "fatal",
     "options",
-    "Logger",
+    "level",
     "logger",
+    "Logger",
     "set_workspace",
     "set_random_seed",
     "set_resume_from",
diff --git a/neural_compressor/common/utils/logger.py b/neural_compressor/common/utils/logger.py
@@ -19,6 +19,28 @@
 import logging
 import os
 
+__all__ = [
+    "level",
+    "Logger",  # TODO: not expose it
+    "logger",
+]
+
+
+def _pretty_dict(value, indent=0):
+    """Make the logger dict pretty."""
+    prefix = "\n" + " " * (indent + 4)
+    if isinstance(value, dict):
+        items = [prefix + repr(key) + ": " + _pretty_dict(value[key], indent + 4) for key in value]
+        return "{%s}" % (",".join(items) + "\n" + " " * indent)
+    elif isinstance(value, list):
+        items = [prefix + _pretty_dict(item, indent + 4) for item in value]
+        return "[%s]" % (",".join(items) + "\n" + " " * indent)
+    elif isinstance(value, tuple):
+        items = [prefix + _pretty_dict(item, indent + 4) for item in value]
+        return "(%s)" % (",".join(items) + "\n" + " " * indent)
+    else:
+        return repr(value)
+
 
 class Logger(object):
     """Logger class."""
@@ -50,85 +72,61 @@ def get_logger(self):
         """Get the logger."""
         return self._logger
 
-
-def _pretty_dict(value, indent=0):
-    """Make the logger dict pretty."""
-    prefix = "\n" + " " * (indent + 4)
-    if isinstance(value, dict):
-        items = [prefix + repr(key) + ": " + _pretty_dict(value[key], indent + 4) for key in value]
-        return "{%s}" % (",".join(items) + "\n" + " " * indent)
-    elif isinstance(value, list):
-        items = [prefix + _pretty_dict(item, indent + 4) for item in value]
-        return "[%s]" % (",".join(items) + "\n" + " " * indent)
-    elif isinstance(value, tuple):
-        items = [prefix + _pretty_dict(item, indent + 4) for item in value]
-        return "(%s)" % (",".join(items) + "\n" + " " * indent)
-    else:
-        return repr(value)
+    @staticmethod
+    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, stacklevel=2)
+        else:
+            Logger().get_logger().log(level, msg, *args, **kwargs, stacklevel=2)
+
+    @staticmethod
+    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, stacklevel=2)
+        else:
+            Logger().get_logger().debug(msg, *args, **kwargs, stacklevel=2)
+
+    @staticmethod
+    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, stacklevel=2)
+        else:
+            Logger().get_logger().error(msg, *args, **kwargs, stacklevel=2)
+
+    @staticmethod
+    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, stacklevel=2)
+        else:
+            Logger().get_logger().fatal(msg, *args, **kwargs, stacklevel=2)
+
+    @staticmethod
+    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, stacklevel=2)
+        else:
+            Logger().get_logger().info(msg, *args, **kwargs, stacklevel=2)
+
+    @staticmethod
+    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, stacklevel=2)
+        else:
+            Logger().get_logger().warning(msg, *args, **kwargs, stacklevel=2)
 
 
 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)
-    else:
-        Logger().get_logger().log(level, msg, *args, **kwargs)
-
-
-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)
-    else:
-        Logger().get_logger().debug(msg, *args, **kwargs)
-
 
-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)
-    else:
-        Logger().get_logger().error(msg, *args, **kwargs)
-
-
-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)
-    else:
-        Logger().get_logger().fatal(msg, *args, **kwargs)
-
-
-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)
-    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)
-
-
-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)
-    else:
-        Logger().get_logger().warning(msg, *args, **kwargs)
+logger = Logger
diff --git a/neural_compressor/compression/pruner/README.md b/neural_compressor/compression/pruner/README.md
@@ -343,6 +343,7 @@ The following section exemplifies how to use hooks in user pass-in training func
   on_after_optimizer_step() # Update weights' criteria, mask weights
   on_train_end() # End of pruner, print sparse information
   """
+
   from neural_compressor.training import prepare_compression, WeightPruningConfig
 
   config = WeightPruningConfig(configs)
diff --git a/neural_compressor/torch/quantization/modules.py b/neural_compressor/torch/quantization/modules.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/test/3x/common/test_logger.py b/test/3x/common/test_logger.py
diff --git a/test/3x/torch/test_logger.py b/test/3x/torch/test_logger.py
