mindstudio-probe 1.0.3__py3-none-any.whl → 1.0.4__py3-none-any.whl

msprobe/core/common/log.py

@@ -1,69 +1,76 @@
-import os
-import time
-import sys
-from functools import wraps
-from msprobe.core.common.const import MsgConst
-
-
-class BaseLogger:
-    def __init__(self):
-        self.warning_level = "WARNING"
-        self.error_level = "ERROR"
-        self.info_level = "INFO"
-        self.rank = None
-
-    @staticmethod
-    def _print_log(level, msg, end='\n'):
-        current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
-        pid = os.getpid()
-        full_msg = f"{current_time} ({pid}) [{level}] {msg}"
-        print(full_msg, end=end)
-        sys.stdout.flush()
-
-    def get_rank(self):
-        return self.rank
-
-    def filter_special_chars(func):
-        @wraps(func)
-        def func_level(self, msg):
-            for char in MsgConst.SPECIAL_CHAR:
-                msg = msg.replace(char, '_')
-            return func(self, msg)
-        return func_level
-
-    @filter_special_chars
-    def info(self, msg):
-        self._print_log(self.info_level, msg)
-
-    @filter_special_chars
-    def error(self, msg):
-        self._print_log(self.error_level, msg)
-
-    @filter_special_chars
-    def warning(self, msg):
-        self._print_log(self.warning_level, msg)
-
-    def on_rank_0(self, func):
-        def func_rank_0(*args, **kwargs):
-            current_rank = self.get_rank()
-            if current_rank is None or current_rank == 0:
-                return func(*args, **kwargs)
-            else:
-                return None
-        return func_rank_0
-
-    def info_on_rank_0(self, msg):
-        return self.on_rank_0(self.info)(msg)
-
-    def error_on_rank_0(self, msg):
-        return self.on_rank_0(self.error)(msg)
-
-    def warning_on_rank_0(self, msg):
-        return self.on_rank_0(self.warning)(msg)
-    
-    def error_log_with_exp(self, msg, exception):
-        self.error(msg)
-        raise exception
-
-
-logger = BaseLogger()
+import os
+import time
+import sys
+from functools import wraps
+from msprobe.core.common.const import MsgConst
+
+MSPROBE_LOG_LEVEL = os.environ.get(MsgConst.MSPROBE_LOG_LEVEL, "")
+
+
+class BaseLogger:
+    def __init__(self):
+        self.rank = None
+
+
+    def get_rank(self):
+        return self.rank
+
+    def filter_special_chars(func):
+        @wraps(func)
+        def func_level(self, msg, **kwargs):
+            for char in MsgConst.SPECIAL_CHAR:
+                msg = msg.replace(char, '_')
+            return func(self, msg, **kwargs)
+        return func_level
+
+    @filter_special_chars
+    def info(self, msg, **kwargs):
+        self._print_log(MsgConst.LEVEL[0], msg, **kwargs)
+
+    @filter_special_chars
+    def error(self, msg):
+        self._print_log(MsgConst.LEVEL[2], msg)
+
+    @filter_special_chars
+    def warning(self, msg):
+        self._print_log(MsgConst.LEVEL[1], msg)
+
+    @filter_special_chars
+    def debug(self, msg):
+        if MSPROBE_LOG_LEVEL == MsgConst.LEVEL[3]:
+            self._print_log(MsgConst.LEVEL[3], msg)
+
+    def on_rank_0(self, func):
+        def func_rank_0(*args, **kwargs):
+            current_rank = self.get_rank()
+            if current_rank is None or current_rank == 0:
+                return func(*args, **kwargs)
+            else:
+                return None
+        return func_rank_0
+
+    def info_on_rank_0(self, msg):
+        return self.on_rank_0(self.info)(msg)
+
+    def error_on_rank_0(self, msg):
+        return self.on_rank_0(self.error)(msg)
+
+    def warning_on_rank_0(self, msg):
+        return self.on_rank_0(self.warning)(msg)
+    
+    def error_log_with_exp(self, msg, exception):
+        self.error(msg)
+        raise exception
+
+    def _print_log(self, level, msg, end='\n'):
+        current_rank = self.get_rank()
+        current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+        pid = os.getpid()
+        if current_rank is not None:
+            full_msg = f"{current_time} ({pid}) [rank {current_rank}] [{level}] {msg}"
+        else:
+            full_msg = f"{current_time} ({pid}) [{level}] {msg}"
+        print(full_msg, end=end)
+        sys.stdout.flush()
+
+logger = BaseLogger()