102 lines
4.1 KiB
Python
102 lines
4.1 KiB
Python
import sys
|
|
import os
|
|
import time
|
|
from toolkit.accelerator import get_accelerator
|
|
|
|
|
|
def print_acc(*args, **kwargs):
|
|
if get_accelerator().is_local_main_process:
|
|
print(*args, **kwargs)
|
|
|
|
|
|
# Progress bars (tqdm etc.) refresh many times a second with \r / cursor-up
|
|
# rewrites. The terminal gets every refresh untouched, but writing them all to
|
|
# the log file makes it enormous, so transient refreshes are buffered and
|
|
# written at most once per interval — each new refresh of the same line(s)
|
|
# replaces the buffered one. Real content (anything with actual text and a
|
|
# newline) always writes through immediately, preceded by any buffered refresh
|
|
# to preserve stream order.
|
|
TRANSIENT_WRITE_INTERVAL = 1.0
|
|
# Safety valve: an unterminated refresh stream writes through past this size.
|
|
MAX_TRANSIENT_BUFFER = 65536
|
|
|
|
|
|
class Logger:
|
|
def __init__(self, terminal, log_file):
|
|
self.terminal = terminal
|
|
self.log = log_file
|
|
# Last completed refresh cycle (a new cycle replaces the previous one).
|
|
self._cycle = ''
|
|
# In-progress refresh: a lone \r line rewrite, or a multi-write
|
|
# cursor-movement cycle (nested bars) still being assembled.
|
|
self._tail = ''
|
|
self._tail_is_lone_refresh = False
|
|
self._last_transient_write = 0.0
|
|
|
|
def write(self, message):
|
|
self.terminal.write(message)
|
|
self._write_log(message)
|
|
|
|
def _write_log(self, message):
|
|
has_move = '\r' in message or '\x1b[A' in message
|
|
has_newline = '\n' in message
|
|
if has_move and not has_newline:
|
|
if message.startswith('\r') and '\x1b[A' not in message and (not self._tail or self._tail_is_lone_refresh):
|
|
# Single-line rewrite (plain tqdm bar) — replaces the previous one.
|
|
self._tail = message
|
|
self._tail_is_lone_refresh = True
|
|
else:
|
|
# Part of a multi-write cursor-movement cycle (nested bars).
|
|
self._tail += message
|
|
self._tail_is_lone_refresh = False
|
|
elif self._tail and has_newline and message.strip('\r\n') == '':
|
|
# Pure newline movement closes a multi-write cycle; the completed
|
|
# cycle replaces the previously buffered one.
|
|
self._cycle = self._tail + message
|
|
self._tail = ''
|
|
self._tail_is_lone_refresh = False
|
|
else:
|
|
# Real content — write any buffered refresh first to keep order.
|
|
self._flush_transient()
|
|
self.log.write(message)
|
|
self.log.flush()
|
|
self._last_transient_write = time.monotonic()
|
|
return
|
|
now = time.monotonic()
|
|
# Only flush between cycles (or on a lone \r rewrite, which leaves the
|
|
# cursor on the same row) — flushing mid-cycle would leave the file's
|
|
# cursor moved up and misalign everything written after.
|
|
mid_cycle = self._tail and not self._tail_is_lone_refresh
|
|
if (now - self._last_transient_write >= TRANSIENT_WRITE_INTERVAL and not mid_cycle) or len(
|
|
self._tail) > MAX_TRANSIENT_BUFFER:
|
|
self._flush_transient()
|
|
self._last_transient_write = now
|
|
|
|
def _flush_transient(self):
|
|
if self._cycle or self._tail:
|
|
self.log.write(self._cycle + self._tail)
|
|
self.log.flush()
|
|
self._cycle = ''
|
|
self._tail = ''
|
|
self._tail_is_lone_refresh = False
|
|
|
|
def flush(self):
|
|
self.terminal.flush()
|
|
self.log.flush()
|
|
|
|
def isatty(self):
|
|
return self.terminal.isatty()
|
|
|
|
|
|
def setup_log_to_file(filename):
|
|
if get_accelerator().is_local_main_process:
|
|
if not os.path.exists(os.path.dirname(filename)):
|
|
os.makedirs(os.path.dirname(filename))
|
|
# Capture the real streams before replacing them — wrapping the
|
|
# already-replaced sys.stdout as the stderr Logger's "terminal" would
|
|
# double-write every stderr message to the file. Both wrappers share a
|
|
# single file handle.
|
|
log_file = open(filename, 'a')
|
|
sys.stdout = Logger(sys.stdout, log_file)
|
|
sys.stderr = Logger(sys.stderr, log_file)
|