clearml/trains/backend_interface/task/log.py

import os
import sys
import time
from logging import LogRecord, getLogger, basicConfig
from logging.handlers import BufferingHandler
from threading import Thread, Event
from six.moves.queue import Queue

from ...backend_api.services import events
from ...config import config

buffer_capacity = config.get('log.task_log_buffer_capacity', 100)


class TaskHandler(BufferingHandler):
    __flush_max_history_seconds = 30.
    __wait_for_flush_timeout = 10.
    __once = False

    @property
    def task_id(self):
        return self._task_id

    @task_id.setter
    def task_id(self, value):
        self._task_id = value

    def __init__(self, session, task_id, capacity=buffer_capacity):
        super(TaskHandler, self).__init__(capacity)
        self.task_id = task_id
        self.session = session
        self.last_timestamp = 0
        self.counter = 1
        self._last_event = None
        self._exit_event = None
        self._queue = None
        self._thread = None
        self._pending = 0

    def shouldFlush(self, record):
        """
        Should the handler flush its buffer?

        Returns true if the buffer is up to capacity. This method can be
        overridden to implement custom flushing strategies.
        """
        if self._task_id is None:
            return False
        # Notice! protect against infinite loops, i.e. flush while sending previous records
        # if self.lock._is_owned():
        #     return False

        # if we need to add handlers to the base_logger,
        # it will not automatically create stream one when first used, so we must manually configure it.
        if not TaskHandler.__once:
            base_logger = getLogger()
            if len(base_logger.handlers) == 1 and isinstance(base_logger.handlers[0], TaskHandler):
                if record.name != 'console' and not record.name.startswith('trains.'):
                    base_logger.removeHandler(self)
                    basicConfig()
                    base_logger.addHandler(self)
                    TaskHandler.__once = True
            else:
                TaskHandler.__once = True

        # if we passed the max buffer
        if len(self.buffer) >= self.capacity:
            return True

        # if the first entry in the log was too long ago.
        try:
            if len(self.buffer) and (time.time() - self.buffer[0].created) > self.__flush_max_history_seconds:
                return True
        except:
            pass

        return False

    def _record_to_event(self, record):
        # type: (LogRecord) -> events.TaskLogEvent
        if self._task_id is None:
            return None
        timestamp = int(record.created * 1000)
        if timestamp == self.last_timestamp:
            timestamp += self.counter
            self.counter += 1
        else:
            self.last_timestamp = timestamp
            self.counter = 1

        # ignore backspaces (they are often used)
        msg = record.getMessage().replace('\x08', '')

        # unite all records in a single second
        if self._last_event and timestamp - self._last_event.timestamp < 1000 and \
                record.levelname.lower() == str(self._last_event.level):
            # ignore backspaces (they are often used)
            self._last_event.msg += '\n' + msg
            return None

        # if we have a previous event and it timed out, return it.
        new_event = events.TaskLogEvent(
            task=self.task_id,
            timestamp=timestamp,
            level=record.levelname.lower(),
            worker=self.session.worker,
            msg=msg
        )
        if self._last_event:
            event = self._last_event
            self._last_event = new_event
            return event

        self._last_event = new_event
        return None

    def flush(self):
        if self._task_id is None:
            return

        if not self.buffer:
            return

        self.acquire()
        if not self.buffer:
            self.release()
            return
        buffer = self.buffer
        self.buffer = []
        try:
            record_events = [self._record_to_event(record) for record in buffer] + [self._last_event]
            self._last_event = None
            batch_requests = events.AddBatchRequest(requests=[events.AddRequest(e) for e in record_events if e])
        except Exception:
            self.__log_stderr("WARNING: trains.log - Failed logging task to backend ({:d} lines)".format(len(buffer)))
            batch_requests = None

        if batch_requests:
            self._pending += 1
            self._add_to_queue(batch_requests)

        self.release()

    def _add_to_queue(self, request):
        if not self._queue:
            self._queue = Queue()
            self._exit_event = Event()
            self._exit_event.clear()
            # multiple workers could be supported as well
            self._thread = Thread(target=self._daemon)
            # self._thread.daemon = True
            self._thread.start()
        self._queue.put(request)

    def close(self, wait=False):
        self.__log_stderr('Closing {} wait={}'.format(os.getpid(), wait))
        # flush pending logs
        if not self._task_id:
            return
        self.flush()
        # shut down the TaskHandler, from this point onwards. No events will be logged
        self.acquire()
        _thread = self._thread
        self._thread = None
        if self._queue:
            self._exit_event.set()
            self._queue.put(None)
        self._task_id = None
        self.release()
        if wait and _thread:
            try:
                _thread.join(timeout=self.__wait_for_flush_timeout)
                self.__log_stderr('Closing {} wait done'.format(os.getpid()))
            except:
                pass
        # call super and remove the handler
        super(TaskHandler, self).close()

    def _send_events(self, a_request):
        try:
            if self._thread is None:
                self.__log_stderr('INFO: trains.log - '
                                  'Task.close() flushing remaining logs ({})'.format(self._pending))
            self._pending -= 1
            res = self.session.send(a_request)
            if not res.ok():
                self.__log_stderr("WARNING: trains.log._send_events: failed logging task to backend "
                                  "({:d} lines, {})".format(len(a_request.requests), str(res.meta)))
        except Exception as ex:
            self.__log_stderr("WARNING: trains.log._send_events: Retrying, "
                              "failed logging task to backend ({:d} lines): {}".format(len(a_request.requests), ex))
            # we should push ourselves back into the thread pool
            if self._queue:
                self._pending += 1
                self._queue.put(a_request)

    def _daemon(self):
        # multiple daemons are supported
        leave = self._exit_event.wait(0)
        request = True
        while not leave or request:
            # pull from queue
            request = None
            if self._queue:
                try:
                    request = self._queue.get(block=not leave)
                except:
                    pass
            if request:
                self._send_events(request)
            leave = self._exit_event.wait(0)
        self.__log_stderr('INFO: trains.log - leaving {}'.format(os.getpid()))

    @staticmethod
    def __log_stderr(t):
        write = sys.stderr._original_write if hasattr(sys.stderr, '_original_write') else sys.stderr.write
        write(t + '\n')