Logging: Added _primary_loggers

bittensor/btlogging/format.py

@@ -69,6 +69,18 @@ def success(self, message, *args, **kws):
     for level, color in log_level_color_prefix.items()
 }
 
+DEFAULT_LOG_FORMAT = (
+    f"{Fore.BLUE}%(asctime)s{Fore.RESET} | "
+    f"{Style.BRIGHT}{Fore.WHITE}%(levelname)s{Style.RESET_ALL} | "
+    f"%(name)s:%(filename)s:%(lineno)s | %(message)s"
+)
+
+DEFAULT_TRACE_FORMAT = (
+    f"{Fore.BLUE}%(asctime)s{Fore.RESET} | "
+    f"{Style.BRIGHT}{Fore.WHITE}%(levelname)s{Style.RESET_ALL} | "
+    f"%(name)s:%(filename)s:%(lineno)s | %(message)s"
+)
+
 
 class BtStreamFormatter(logging.Formatter):
     def __init__(self, *args, **kwargs):
@@ -93,10 +105,15 @@ def format(self, record):
         format_orig = self._style._fmt
         record.levelname = f"{record.levelname:^16}"
 
-        if self.trace is True:
-            self._style._fmt = LOG_TRACE_FORMATS[record.levelno]
+        if record.levelno not in LOG_FORMATS:
+            self._style._fmt = (
+                DEFAULT_TRACE_FORMAT if self.trace else DEFAULT_LOG_FORMAT
+            )
         else:
-            self._style._fmt = LOG_FORMATS[record.levelno]
+            if self.trace is True:
+                self._style._fmt = LOG_TRACE_FORMATS[record.levelno]
+            else:
+                self._style._fmt = LOG_FORMATS[record.levelno]
 
         for text, emoji in emoji_map.items():
             record.msg = record.msg.replace(text, emoji)

bittensor/btlogging/loggingmachine.py

@@ -69,7 +69,7 @@ def __init__(self, config: bittensor.config, name: str = BITTENSOR_LOGGER_NAME):
         # basics
         super(LoggingMachine, self).__init__()
         self._queue = mp.Queue(-1)
-        self._name = name
+        self._primary_loggers = {name}
         self._config = config
 
         # Formatters
@@ -89,7 +89,7 @@ def __init__(self, config: bittensor.config, name: str = BITTENSOR_LOGGER_NAME):
         self._listener = self._create_and_start_listener(self._handlers)
 
         # set up all the loggers
-        self._logger = self._initialize_bt_logger(name, config)
+        self._logger = self._initialize_bt_logger(name)
         self.disable_third_party_loggers()
 
     def _configure_handlers(self, config) -> list[stdlogging.Handler]:
@@ -118,9 +118,8 @@ def set_config(self, config):
         """
         self._config = config
         if config.logging_dir and config.record_log:
-            logfile = os.path.abspath(
-                os.path.join(config.logging_dir, DEFAULT_LOG_FILE_NAME)
-            )
+            expanded_dir = os.path.expanduser(config.logging_dir)
+            logfile = os.path.abspath(os.path.join(expanded_dir, DEFAULT_LOG_FILE_NAME))
             self._enable_file_logging(logfile)
         if config.trace:
             self.enable_trace()
@@ -149,7 +148,7 @@ def get_queue(self):
         """
         return self._queue
 
-    def _initialize_bt_logger(self, name, config):
+    def _initialize_bt_logger(self, name):
         """
         Initialize logging for bittensor.
 
@@ -162,6 +161,17 @@ def _initialize_bt_logger(self, name, config):
         logger.addHandler(queue_handler)
         return logger
 
+    def _deinitialize_bt_logger(self, name):
+        """
+        Find the logger by name and remove the
+        queue handler associated with it.
+        """
+        logger = stdlogging.getLogger(name)
+        for handler in list(logger.handlers):
+            if isinstance(handler, QueueHandler):
+                logger.removeHandler(handler)
+        return logger
+
     def _create_file_handler(self, logfile: str):
         file_handler = RotatingFileHandler(
             logfile,
@@ -172,9 +182,30 @@ def _create_file_handler(self, logfile: str):
         file_handler.setLevel(stdlogging.TRACE)
         return file_handler
 
+    def register_primary_logger(self, name: str):
+        """
+        Register a logger as primary logger
+
+        This adds a logger to the _primary_loggers set to ensure
+        it doesn't get disabled when disabling third-party loggers.
+        A queue handler is also associated with it.
+        """
+        self._primary_loggers.add(name)
+        self._initialize_bt_logger(name)
+
+    def deregister_primary_logger(self, name: str):
+        """
+        De-registers a primary logger
+
+        This function removes the logger from the _primary_loggers
+        set and deinitializes its queue handler
+        """
+        self._primary_loggers.remove(name)
+        self._deinitialize_bt_logger(name)
+
     def enable_third_party_loggers(self):
         for logger in all_loggers():
-            if logger.name == self._name:
+            if logger.name in self._primary_loggers:
                 continue
             queue_handler = QueueHandler(self._queue)
             logger.addHandler(queue_handler)
@@ -183,7 +214,7 @@ def enable_third_party_loggers(self):
     def disable_third_party_loggers(self):
         # remove all handlers
         for logger in all_loggers():
-            if logger.name == self._name:
+            if logger.name in self._primary_loggers:
                 continue
             for handler in logger.handlers:
                 logger.removeHandler(handler)
@@ -211,7 +242,7 @@ def before_enable_default(self):
         self._logger.info(f"Enabling default logging.")
         self._logger.setLevel(stdlogging.INFO)
         for logger in all_loggers():
-            if logger.name == self._name:
+            if logger.name in self._primary_loggers:
                 continue
             logger.setLevel(stdlogging.CRITICAL)