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

This commit is contained in:
allegroai 2022-12-18 23:01:41 +02:00
parent e53c4008a1
commit 8ef25a0030
2 changed files with 58 additions and 41 deletions

View File

@ -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

View File

@ -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: