fix(mge/logger): fix set_mgb_log_level

GitOrigin-RevId: 6c94c51483

imperative/python/megengine/logger.py

@@ -200,12 +200,6 @@ try:
 
     _megbrain_logger = get_logger("megbrain", MegBrainLogFormatter)
     _imperative_rt_logger.set_log_handler(_megbrain_logger)
-    if _default_level == logging.getLevelName("ERROR"):
-        _imperative_rt_logger.set_log_level(_imperative_rt_logger.LogLevel.Error)
-    elif _default_level == logging.getLevelName("INFO"):
-        _imperative_rt_logger.set_log_level(_imperative_rt_logger.LogLevel.Info)
-    else:
-        _imperative_rt_logger.set_log_level(_imperative_rt_logger.LogLevel.Debug)
 
     def set_mgb_log_level(level):
         r"""
@@ -215,11 +209,23 @@ try:
         :param level: new log level
         :return: original log level
         """
-        logger = _megbrain_logger
-        rst = logger.getEffectiveLevel()
-        logger.setLevel(level)
+        _megbrain_logger.setLevel(level)
+        if level == logging.getLevelName("ERROR"):
+            rst = _imperative_rt_logger.set_log_level(
+                _imperative_rt_logger.LogLevel.Error
+            )
+        elif level == logging.getLevelName("INFO"):
+            rst = _imperative_rt_logger.set_log_level(
+                _imperative_rt_logger.LogLevel.Info
+            )
+        else:
+            rst = _imperative_rt_logger.set_log_level(
+                _imperative_rt_logger.LogLevel.Debug
+            )
         return rst
 
+    set_mgb_log_level(_default_level)
+
 
 except ImportError as exc:
 

imperative/python/test/unit/core/test_util.py

@@ -6,10 +6,19 @@
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+import logging
+
 from megengine.core._imperative_rt import Logger
+from megengine.logger import _imperative_rt_logger, set_mgb_log_level
 
 
 def test_logger():
-    orig_level = Logger().set_log_level(Logger.LogLevel.Info)
-    assert Logger().set_log_level(Logger.LogLevel.Info) == Logger.LogLevel.Info
+    orig_level = Logger().set_log_level(Logger.LogLevel.Debug)
+    assert Logger().set_log_level(Logger.LogLevel.Debug) == Logger.LogLevel.Debug
     Logger().set_log_level(orig_level)
+    orig_level = set_mgb_log_level(logging.DEBUG)
+    assert (
+        _imperative_rt_logger.set_log_level(Logger.LogLevel.Debug)
+        == Logger.LogLevel.Debug
+    )
+    _imperative_rt_logger.set_log_level(orig_level)