Fix task.close() removes all log handlers (should only remove ClearML handlers)

clearml/debugging/log.py

@@ -94,29 +94,6 @@ class _LevelRangeFilter(logging.Filter):
 class LoggerRoot(object):
     __base_logger = None
 
-    @classmethod
-    def _make_stream_handler(cls, level=None, stream=sys.stdout, colored=False):
-        ch = logging.StreamHandler(stream=stream)
-        ch.setLevel(level)
-        formatter = None
-
-        # if colored, try to import colorama & coloredlogs (by default, not in the requirements)
-        if colored:
-            try:
-                import colorama
-                from coloredlogs import ColoredFormatter
-                colorama.init()
-                formatter = ColoredFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-            except ImportError:
-                colored = False
-
-        # if we don't need or failed getting colored formatter
-        if not colored or not formatter:
-            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-
-        ch.setFormatter(formatter)
-        return ch
-
     @classmethod
     def get_base_logger(cls, level=None, stream=sys.stdout, colored=False):
         if LoggerRoot.__base_logger:
@@ -153,17 +130,17 @@ class LoggerRoot(object):
             # and redirect is set for ERROR, in which case we redirect from CRITICAL)
             redirect_level = max(level, redirect_level)
             LoggerRoot.__base_logger.addHandler(
-                cls._make_stream_handler(redirect_level, sys.stderr, colored)
+                ClearmlStreamHandler(redirect_level, sys.stderr, colored)
             )
 
             if level < redirect_level:
                 # Not all levels were redirected, remaining should be sent to requested stream
-                handler = cls._make_stream_handler(level, stream, colored)
+                handler = ClearmlStreamHandler(level, stream, colored)
                 handler.addFilter(_LevelRangeFilter(min_level=level, max_level=redirect_level - 1))
                 LoggerRoot.__base_logger.addHandler(handler)
         else:
             LoggerRoot.__base_logger.addHandler(
-                cls._make_stream_handler(level, stream, colored)
+                ClearmlStreamHandler(level, stream, colored)
             )
 
         LoggerRoot.__base_logger.propagate = False
@@ -175,6 +152,16 @@ class LoggerRoot(object):
             for h in LoggerRoot.__base_logger.handlers:
                 h.flush()
 
+    @staticmethod
+    def clear_logger_handlers():
+        # https://github.com/pytest-dev/pytest/issues/5502#issuecomment-647157873
+        loggers = [logging.getLogger()] + list(logging.Logger.manager.loggerDict.values())
+        for logger in loggers:
+            handlers = getattr(logger, 'handlers', [])
+            for handler in handlers:
+                if isinstance(handler, ClearmlLoggerHandler):
+                    logger.removeHandler(handler)
+
 
 def add_options(parser):
     """ Add logging options to an argparse.ArgumentParser object """
@@ -204,7 +191,7 @@ def get_logger(path=None, level=None, stream=None, colored=False):
     if level is not None:
         log.setLevel(level)
     if stream:
-        ch = logging.StreamHandler(stream=stream)
+        ch = ClearmlStreamHandler(stream=stream, dont_set_formater=True)
         if level is not None:
             ch.setLevel(level)
         log.addHandler(ch)
@@ -226,7 +213,7 @@ def _add_file_handler(logger, log_dir, fh, formatter=None):
 def add_rotating_file_handler(logger, log_dir, log_file_prefix, max_bytes=10 * 1024 * 1024, backup_count=20,
                               formatter=None):
     """ Create and add a rotating file handler to a logger """
-    fh = logging.handlers.RotatingFileHandler(
+    fh = ClearmlRotatingFileHandler(
         str(Path(log_dir) / ('%s.log' % log_file_prefix)), maxBytes=max_bytes, backupCount=backup_count)
     _add_file_handler(logger, log_dir, fh, formatter)
 
@@ -237,7 +224,7 @@ def add_time_rotating_file_handler(logger, log_dir, log_file_prefix, when='midni
         Possible values for when are 'midnight', weekdays ('w0'-'W6', when 0 is Monday), and 's', 'm', 'h' amd 'd' for
             seconds, minutes, hours and days respectively (case-insensitive)
     """
-    fh = logging.handlers.TimedRotatingFileHandler(
+    fh = ClearmlTimedRotatingFileHandler(
         str(Path(log_dir) / ('%s.log' % log_file_prefix)), when=when)
     _add_file_handler(logger, log_dir, fh, formatter)
 
@@ -249,7 +236,7 @@ def get_null_logger(name=None):
         # avoid nested imports
         from ..config import config
 
-        log.addHandler(logging.NullHandler())
+        log.addHandler(ClearmlNullHandler())
         log.propagate = config.get("log.null_log_propagate", False)
     return PickledLogger.wrapper(log, func=get_null_logger, name=name)
 
@@ -287,3 +274,43 @@ class TqdmLog(object):
 
     def close(self):
         self._tqdm.close()
+
+
+class ClearmlLoggerHandler:
+    pass
+
+
+class ClearmlStreamHandler(logging.StreamHandler, ClearmlLoggerHandler):
+    def __init__(self, level=None, stream=sys.stdout, colored=False, dont_set_formater=False):
+        super(ClearmlStreamHandler, self).__init__(stream=stream)
+        self.setLevel(level)
+        if dont_set_formater:
+            return
+
+        formatter = None
+        # if colored, try to import colorama & coloredlogs (by default, not in the requirements)
+        if colored:
+            try:
+                import colorama
+                from coloredlogs import ColoredFormatter
+
+                colorama.init()
+                formatter = ColoredFormatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+            except ImportError:
+                colored = False
+        # if we don't need or failed getting colored formatter
+        if not colored or not formatter:
+            formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        self.setFormatter(formatter)
+
+
+class ClearmlRotatingFileHandler(logging.handlers.RotatingFileHandler, ClearmlLoggerHandler):
+    pass
+
+
+class ClearmlTimedRotatingFileHandler(logging.handlers.TimedRotatingFileHandler, ClearmlLoggerHandler):
+    pass
+
+
+class ClearmlNullHandler(logging.NullHandler, ClearmlLoggerHandler):
+    pass

clearml/task.py

@@ -3552,16 +3552,6 @@ class Task(_Task):
         #     # we have to forcefully shutdown if we have forked processes, sometimes they will get stuck
         #     os._exit(self.__exit_hook.exit_code if self.__exit_hook and self.__exit_hook.exit_code else 0)
 
-    @staticmethod
-    def _clear_loggers():
-        # https://github.com/pytest-dev/pytest/issues/5502#issuecomment-647157873
-        import logging
-        loggers = [logging.getLogger()] + list(logging.Logger.manager.loggerDict.values())
-        for logger in loggers:
-            handlers = getattr(logger, 'handlers', [])
-            for handler in handlers:
-                logger.removeHandler(handler)
-
     def __shutdown(self):
         """
         Will happen automatically once we exit code, i.e. atexit
@@ -3588,7 +3578,7 @@ class Task(_Task):
             # from here only a single thread can re-enter
             self._at_exit_called = get_current_thread_id()
 
-        Task._clear_loggers()
+        LoggerRoot.clear_logger_handlers()
 
         # disable lock on signal callbacks, to avoid deadlocks.
         if self.__exit_hook and self.__exit_hook.signal is not None: