fix: prevent terminal ui opening log file (#130)

When running directly as a part of the worker with `enable_terminal_ui: true` we no longer open the log file, instead we hook Loguru directly. This prevents errors when the log file is rotated.

worker/ui.py

@@ -6,7 +6,6 @@
 import locale
 import os
 import re
-import sys
 import textwrap
 import threading
 import time
@@ -16,6 +15,7 @@
 import psutil
 import requests
 import yaml
+from loguru import logger
 from pynvml.smi import nvidia_smi
 
 
@@ -34,6 +34,13 @@ def set_size(self, size):
     def flush(self):
         pass
 
+    def isatty(self):
+        # No, we are not a TTY
+        return False
+
+    def close(self):
+        pass
+
 
 class GPUInfo:
     def __init__(self):
@@ -128,6 +135,9 @@ def get_info(self):
 class TerminalUI:
 
     REGEX = re.compile(r"(INIT|DEBUG|INFO|WARNING|ERROR).*(\d\d\d\d-\d\d-\d\d \d\d:\d\d:\d\d).*\| (.*) - (.*)$")
+    LOGURU_REGEX = re.compile(
+        r"(\d\d\d\d-\d\d-\d\d \d\d:\d\d:\d\d).*\| (INIT|INIT_OK|DEBUG|INFO|WARNING|ERROR).*\| (.*) - (.*)$"
+    )
     KUDOS_REGEX = re.compile(r".*average kudos per hour: (\d+)")
     JOBDONE_REGEX = re.compile(r".*(Generation for id.*finished successfully|Finished interrogation.*)")
 
@@ -175,8 +185,10 @@ def __init__(self, worker_name=None, apikey=None, url="https://stablehorde.net")
         self.show_debug = False
         self.last_key = None
         self.pause_log = False
+        self.use_log_file = False
+        self.log_file = None
+        self.input = DequeOutputCollector()
         self.output = DequeOutputCollector()
-        self.stdout = DequeOutputCollector()
         self.worker_name = worker_name
         self.apikey = apikey
         self.worker_id = self.load_worker_id()
@@ -210,38 +222,66 @@ def __init__(self, worker_name=None, apikey=None, url="https://stablehorde.net")
         self.last_audio_alert = 0
 
     def initialise(self):
+        if self.use_log_file:
+            self.open_log()
+        else:
+            # Hook loguru output
+            logger.add(self.input, level="DEBUG")
         locale.setlocale(locale.LC_ALL, "")
         self.initialise_main_window()
         self.resize()
-        self.open_log()
         self.get_remote_worker_info()
 
     def open_log(self):
         # We try a couple of times, log rotiation, etc
         for _ in range(2):
             try:
-                self.input = open("logs/bridge.log", "rt", encoding="utf-8", errors="ignore")
-                self.input.seek(0, os.SEEK_END)
+                self.log_file = open("logs/bridge.log", "rt", encoding="utf-8", errors="ignore")
+                self.log_file.seek(0, os.SEEK_END)
                 break
             except OSError:
                 time.sleep(1)
 
     def load_log(self):
-        while line := self.input.readline():
-            ignore = False
-            for skip in TerminalUI.JUNK:
-                if skip.lower() in line.lower():
-                    ignore = True
-            if ignore:
-                continue
+        if self.use_log_file:
+            while line := self.log_file.readline():
+                self.input.write(line)
+        self.load_log_queue()
+
+    def parse_log_line(self, line):
+        if self.use_log_file:
             if regex := TerminalUI.REGEX.match(line):
                 if not self.show_debug and regex.group(1) == "DEBUG":
-                    continue
+                    return
                 if regex.group(1) == "ERROR":
                     self.error_count += 1
                 elif regex.group(1) == "WARNING":
                     self.warning_count += 1
-                self.output.write(f"{regex.group(1)}::::{regex.group(2)}::::{regex.group(3)}::::{regex.group(4)}")
+                return f"{regex.group(1)}::::{regex.group(2)}::::{regex.group(3)}::::{regex.group(4)}"
+        else:
+            if regex := TerminalUI.LOGURU_REGEX.match(line):
+                if not self.show_debug and regex.group(2) == "DEBUG":
+                    return
+                if regex.group(2) == "ERROR":
+                    self.error_count += 1
+                elif regex.group(2) == "WARNING":
+                    self.warning_count += 1
+                return f"{regex.group(2)}::::{regex.group(1)}::::{regex.group(3)}::::{regex.group(4)}"
+
+    def load_log_queue(self):
+        lines = list(self.input.deque)
+        self.input.deque.clear()
+        for line in lines:
+            ignore = False
+            for skip in TerminalUI.JUNK:
+                if skip.lower() in line.lower():
+                    ignore = True
+            if ignore:
+                continue
+            log_line = self.parse_log_line(line)
+            if not log_line:
+                continue
+            self.output.write(log_line)
             if regex := TerminalUI.KUDOS_REGEX.match(line):
                 self.kudos_per_hour = int(regex.group(1))
             if regex := TerminalUI.JOBDONE_REGEX.match(line):
@@ -264,8 +304,6 @@ def initialise_main_window(self):
         curses.init_pair(5, curses.COLOR_MAGENTA, curses.COLOR_BLACK)
         curses.init_pair(6, curses.COLOR_CYAN, curses.COLOR_BLACK)
         curses.init_pair(7, curses.COLOR_WHITE, curses.COLOR_BLACK)
-        # Suppress all output to stdout
-        sys.stdout = self.stdout
 
     def resize(self):
         # Determine terminal size
@@ -541,7 +579,6 @@ def print_log(self):
 
         if not self.pause_log:
             self.load_log()
-
         output = list(self.output.deque)
         if not output:
             return
@@ -560,7 +597,7 @@ def print_log(self):
                 colour = TerminalUI.COLOUR_WHITE
             elif cat == "ERROR":
                 colour = TerminalUI.COLOUR_RED
-            elif cat == "INIT":
+            elif cat == "INIT" or cat == "INIT_OK":
                 colour = TerminalUI.COLOUR_MAGENTA
             elif cat == "WARNING":
                 colour = TerminalUI.COLOUR_YELLOW
@@ -744,4 +781,6 @@ def run(self):
             apikey = config.get("api_key", "")
 
     term = TerminalUI(workername, apikey)
+    # Standalone UI we need to inspect the log file
+    term.use_log_file = True
     term.run()