clearml/trains/backend_interface/logger.py

import logging
import sys
import threading

from ..backend_interface.task.development.worker import DevWorker
from ..backend_interface.task.log import TaskHandler
from ..config import running_remotely


class StdStreamPatch(object):
    _stdout_proxy = None
    _stderr_proxy = None
    _stdout_original_write = None
    _stderr_original_write = None

    @staticmethod
    def patch_std_streams(logger):
        if DevWorker.report_stdout and not PrintPatchLogger.patched and not running_remotely():
            StdStreamPatch._stdout_proxy = PrintPatchLogger(sys.stdout, logger, level=logging.INFO)
            StdStreamPatch._stderr_proxy = PrintPatchLogger(sys.stderr, logger, level=logging.ERROR)
            logger._task_handler = TaskHandler(task=logger._task, capacity=100)
            # noinspection PyBroadException
            try:
                if StdStreamPatch._stdout_original_write is None:
                    StdStreamPatch._stdout_original_write = sys.stdout.write
                if StdStreamPatch._stderr_original_write is None:
                    StdStreamPatch._stderr_original_write = sys.stderr.write
                # this will only work in python 3, guard it with try/catch
                if not hasattr(sys.stdout, '_original_write'):
                    sys.stdout._original_write = sys.stdout.write
                sys.stdout.write = StdStreamPatch._stdout__patched__write__
                if not hasattr(sys.stderr, '_original_write'):
                    sys.stderr._original_write = sys.stderr.write
                sys.stderr.write = StdStreamPatch._stderr__patched__write__
            except Exception:
                pass
            sys.stdout = StdStreamPatch._stdout_proxy
            sys.stderr = StdStreamPatch._stderr_proxy
            # patch the base streams of sys (this way colorama will keep its ANSI colors)
            # noinspection PyBroadException
            try:
                sys.__stderr__ = sys.stderr
            except Exception:
                pass
            # noinspection PyBroadException
            try:
                sys.__stdout__ = sys.stdout
            except Exception:
                pass

            # now check if we have loguru and make it re-register the handlers
            # because it stores internally the stream.write function, which we cant patch
            # noinspection PyBroadException
            try:
                from loguru import logger
                register_stderr = None
                register_stdout = None
                for k, v in logger._handlers.items():
                    if v._name == '<stderr>':
                        register_stderr = k
                    elif v._name == '<stdout>':
                        register_stderr = k
                if register_stderr is not None:
                    logger.remove(register_stderr)
                    logger.add(sys.stderr)
                if register_stdout is not None:
                    logger.remove(register_stdout)
                    logger.add(sys.stdout)
            except Exception:
                pass

        elif DevWorker.report_stdout and not running_remotely():
            logger._task_handler = TaskHandler(task=logger._task, capacity=100)
            if StdStreamPatch._stdout_proxy:
                StdStreamPatch._stdout_proxy.connect(logger)
            if StdStreamPatch._stderr_proxy:
                StdStreamPatch._stderr_proxy.connect(logger)

    @staticmethod
    def remove_std_logger(logger=None):
        if isinstance(sys.stdout, PrintPatchLogger):
            # noinspection PyBroadException
            try:
                sys.stdout.disconnect(logger)
            except Exception:
                pass
        if isinstance(sys.stderr, PrintPatchLogger):
            # noinspection PyBroadException
            try:
                sys.stderr.disconnect(logger)
            except Exception:
                pass

    @staticmethod
    def stdout_original_write(*args, **kwargs):
        if StdStreamPatch._stdout_original_write:
            StdStreamPatch._stdout_original_write(*args, **kwargs)

    @staticmethod
    def stderr_original_write(*args, **kwargs):
        if StdStreamPatch._stderr_original_write:
            StdStreamPatch._stderr_original_write(*args, **kwargs)

    @staticmethod
    def _stdout__patched__write__(*args, **kwargs):
        if StdStreamPatch._stdout_proxy:
            return StdStreamPatch._stdout_proxy.write(*args, **kwargs)
        return sys.stdout._original_write(*args, **kwargs)

    @staticmethod
    def _stderr__patched__write__(*args, **kwargs):
        if StdStreamPatch._stderr_proxy:
            return StdStreamPatch._stderr_proxy.write(*args, **kwargs)
        return sys.stderr._original_write(*args, **kwargs)


class PrintPatchLogger(object):
    """
    Allowed patching a stream into the logger.
    Used for capturing and logging stdin and stderr when running in development mode pseudo worker.
    """
    patched = False
    lock = threading.Lock()
    recursion_protect_lock = threading.RLock()

    def __init__(self, stream, logger=None, level=logging.INFO):
        PrintPatchLogger.patched = True
        self._terminal = stream
        self._log = logger
        self._log_level = level
        self._cur_line = ''

    def write(self, message):
        # make sure that we do not end up in infinite loop (i.e. log.console ends up calling us)
        if self._log and not PrintPatchLogger.recursion_protect_lock._is_owned():
            try:
                self.lock.acquire()
                with PrintPatchLogger.recursion_protect_lock:
                    if hasattr(self._terminal, '_original_write'):
                        self._terminal._original_write(message)
                    else:
                        self._terminal.write(message)

                do_flush = '\n' in message
                do_cr = '\r' in message
                self._cur_line += message
                if (not do_flush and not do_cr) or not message:
                    return
                last_lf = self._cur_line.rindex('\n' if do_flush else '\r')
                next_line = self._cur_line[last_lf + 1:]
                cur_line = self._cur_line[:last_lf + 1].rstrip()
                self._cur_line = next_line
            finally:
                self.lock.release()

            if cur_line:
                with PrintPatchLogger.recursion_protect_lock:
                    # noinspection PyBroadException
                    try:
                        if self._log:
                            self._log._console(cur_line, level=self._log_level, omit_console=True)
                    except Exception:
                        # what can we do, nothing
                        pass
        else:
            if hasattr(self._terminal, '_original_write'):
                self._terminal._original_write(message)
            else:
                self._terminal.write(message)

    def connect(self, logger):
        self._cur_line = ''
        self._log = logger

    def disconnect(self, logger=None):
        # disconnect the logger only if it was registered
        if not logger or self._log == logger:
            self.connect(None)

    def __getattr__(self, attr):
        if attr in ['_log', '_terminal', '_log_level', '_cur_line']:
            return self.__dict__.get(attr)
        return getattr(self._terminal, attr)

    def __setattr__(self, key, value):
        if key in ['_log', '_terminal', '_log_level', '_cur_line']:
            self.__dict__[key] = value
        else:
            return setattr(self._terminal, key, value)


class LogFlusher(threading.Thread):
    def __init__(self, logger, period, **kwargs):
        super(LogFlusher, self).__init__(**kwargs)
        self.daemon = True

        self._period = period
        self._logger = logger
        self._exit_event = threading.Event()

    @property
    def period(self):
        return self._period

    def run(self):
        self._logger.flush()
        # store original wait period
        while True:
            period = self._period
            while not self._exit_event.wait(period or 1.0):
                self._logger.flush()
            # check if period is negative or None we should exit
            if self._period is None or self._period < 0:
                break
            # check if period was changed, we should restart
            self._exit_event.clear()

    def exit(self):
        self._period = None
        self._exit_event.set()

    def set_period(self, period):
        self._period = period
        # make sure we exit the previous wait
        self._exit_event.set()