| | #!/usr/bin/env -S uv run --script |
| | # SPDX-License-Identifier: MIT |
| | # /// script | | | # dependencies = ["rich", "plotext<6"] # plotext 6 dropped clf/plotsize/build | | | # /// | | | """ | | | loratop - a tiny TUI monitor for Ostris AI Toolkit LoRA training. | | | It reads a run's loss_log.db (ai-toolkit writes it when logging.use_ui_logger | | | is true) or, failing that, parses loss/step/lr from the tqdm output in a | | | training log. It renders loss and lr graphs under two lines of progress. | | | Designed to be readable on tiny screens, so it adapts to narrow terminals | | | and degrades gracefully when optional deps are missing. | | | Linux only: finding the trainer reads /proc. | | | Usage | | | ----- | | | loratop # the run in this folder, else the one running trainer | | | loratop output/my_lora # a run folder: reads its loss_log.db, live or finished | | | loratop train.log # a training log (python run.py cfg > train.log 2>&1) | | | loratop 12345 # a running trainer's PID: its loss_log.db, else its log | | | Run it with uv (uv run loratop.py), or chmod +x it and put it on your PATH | | | as loratop. | | | Deps (plotext optional, falls back to a sparkline): | | | rich, plotext. Declared inline above, so uv installs them on first run. | | | """ | | | import argparse | | | import os | | | import re | | | import select | | | import shutil | | | import signal | | | import sqlite3 | | | import sys | | | import termios | | | import time | | | import tty | | | from collections import deque | | | # ---- optional dependencies ------------------------------------------------- | | | try: | | | import plotext as plt | | | if not hasattr(plt, "clf"): # plotext 6+: fall back to the sparkline | | | plt = None | | | except ImportError: | | | plt = None | | | try: | | | from rich.console import Console, Group | | | from rich.live import Live | | | from rich.panel import Panel | | | from rich.table import Table | | | from rich.text import Text | | | except ImportError: | | | sys.exit("loratop needs 'rich': run it with uv, or pip install rich") | | | # ---- parsing patterns (tweak these if your log format differs) ------------- | | | # matches "loss: 0.0823", "loss=1.2e-01", "avg_loss 0.05", etc. | | | RE_LOSS = re.compile(r"loss[=:\s]+([0-9]*.?[0-9]+(?:[eE][-+]?[0-9]+)?)", re.I) | | | # matches the tqdm " 500/10000" current/total fragment |
| | RE_STEP = re.compile(r"(\d+)\s*/\s*(\d+)") |
| | # matches "lr: 1.0e-04" |
| | RE_LR = re.compile(r"lr[=:\s]+([0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)", re.I) |
| | # tqdm's own "[elapsed<remaining, rate]" fragment: "[01:10<2:30:00, 1.2it/s" |
| | RE_TQDM_ELAPSED = re.compile(r"\[\s*((?:\d+:)?\d{1,2}:\d{2})\s*<") |
| | RE_TQDM_ETA = re.compile(r"<\s*((?:\d+:)?\d{1,2}:\d{2})") |
| | RE_TQDM_RATE = re.compile(r"([\d.]+)\s*(it/s\|s/it)") |
| | SPARK = "▁▂▃▄▅▆▇█" | | | # ai-toolkit's UILogger database, in the run's output folder | | | DB_NAME = "loss_log.db" | | | # train.steps in the config.yaml ai-toolkit saves beside it; the database | | | # itself doesn't record the total | | | RE_CONFIG_STEPS = re.compile(r"^\ssteps:\s(\d+)\s*$", re.M) | | | # Graph layout, after nvtop's: side by side by default, stacked once each | | | # graph would get at least STACK_MIN_ROWS rows, loss alone when too narrow. | | | STACK_MIN_ROWS = 11 | | | MIN_GRAPH_COLS = 30 |
| | # --------------------------------------------------------------------------- |
| | def _pids(): |
| | """Every PID but ours, or none where there's no /proc.""" | | | try: | | | names = os.listdir("/proc") | | | except OSError: |
| | return [] |
| | me = os.getpid() |
| | return [int(n) for n in names if n.isdigit() and int(n) != me] |
| | def discover_log_from_pid(pid): |
| | """Return the file stdout/stderr of pid is redirected to, if any.""" |
| | for fd in (1, 2): |
| | link = f"/proc/{pid}/fd/{fd}" |
| | try: | | | target = os.readlink(link) | | | except OSError: | | | continue | | | if os.path.isfile(target): | | | return target | | | return None | | | def _comm(pid): | | | try: |
| | with open(f"/proc/{pid}/comm") as fh: |
| | return fh.read().strip() |
| | except OSError: |
| | return "" |
| | def _writes_to(pid, target): |
| | """Yield (fd, link) for fds of pid open for writing on target path |
| | (or any fd if target is None).""" |
| | fd_dir = f"/proc/{pid}/fd" |
| | try: |
| | fds = os.listdir(fd_dir) |
| | except OSError: |
| | return |
| | for fd in fds: |
| | try: |
| | link = os.readlink(f"{fd_dir}/{fd}") |
| | if target is not None and os.path.realpath(link) != target: |
| | continue |
| | with open(f"/proc/{pid}/fdinfo/{fd}") as fh: |
| | info = fh.read() |
| | m = re.search(r"flags:\s*(\d+)", info) |
| | if m and int(m.group(1), 8) & 0o3: # O_WRONLY or O_RDWR |
| | yield fd, link | | | except OSError: | | | continue | | | def discover_pid_from_log(log_path): | | | """Find the training PID given only the log file: look for processes | | | writing to it, and if the writer is tee, hop across tee's stdin pipe to | | | the process feeding it.""" |
| | real = os.path.realpath(log_path) |
| | writers = [pid for pid in _pids() if any(True for _ in _writes_to(pid, real))] |
| | # direct redirect (python run.py > log 2>&1): the trainer holds the file |
| | for pid in writers: | | | if _comm(pid) != "tee": | | | return pid | | | # piped through tee: follow tee's stdin pipe back to its feeder | | | for tee_pid in writers: | | | try: | | | pipe = os.readlink(f"/proc/{tee_pid}/fd/0") | | | except OSError: | | | continue | | | if not pipe.startswith("pipe:"): | | | continue |
| | for pid in _pids(): |
| | if pid == tee_pid: |
| | continue |
| | for fd in (1, 2): |
| | try: |
| | if os.readlink(f"/proc/{pid}/fd/{fd}") == pipe: |
| | return pid |
| | except OSError: |
| | continue |
| | return None |
| | def db_of_pid(pid): |
| | """The loss_log.db pid has open, if any.""" |
| | for _, link in _writes_to(pid, None): |
| | if os.path.basename(link) == DB_NAME and os.path.isfile(link): |
| | return link | | | return None | | | def _argv(pid): | | | try: |
| | with open(f"/proc/{pid}/cmdline", "rb") as fh: |
| | return fh.read().decode(errors="replace").split("\0") |
| | except OSError: |
| | return [] |
| | def find_trainers(): |
| | """{path: pid} for running trainers: processes holding a loss_log.db open, | | | or ai-toolkit's run.py with its output going to a file.""" |
| | found = {} |
| | for pid in _pids(): |
| | db = db_of_pid(pid) |
| | if db: |
| | found.setdefault(db, pid) |
| | elif any(os.path.basename(a) == "run.py" for a in _argv(pid)): |
| | log = discover_log_from_pid(pid) |
| | if log: | | | found.setdefault(log, pid) | | | # a trainer with a database also writes a log; show the database only |
| | pids_with_db = {pid for path, pid in found.items() if os.path.basename(path) == DB_NAME} |
| | return {path: pid for path, pid in found.items() |
| | if os.path.basename(path) == DB_NAME or pid not in pids_with_db} |
| | def resolve_source(target): |
| | """Turn the command-line target (or nothing) into (path, pid or None). |
| | A folder means its loss_log.db, a file a database or log, a number a PID.""" | | | if target is None: |
| | db = os.path.join(os.getcwd(), DB_NAME) |
| | if os.path.isfile(db): |
| | return db, None |
| | found = find_trainers() |
| | if len(found) == 1: |
| | (path, pid), = found.items() |
| | print(f"found a running trainer: pid {pid}, {path}", file=sys.stderr) |
| | return path, pid | | | if found: | | | sys.exit("more than one trainer is running. pick one:\n" + "\n".join( | | | f" loratop {pid} # {path}" for path, pid in sorted(found.items()))) | | | sys.exit("nothing to watch: no loss_log.db here and no running trainer.\n" | | | "pass a run folder, a training log, or a PID") |
| | if os.path.isdir(target): |
| | db = os.path.join(target, DB_NAME) |
| | if os.path.isfile(db): |
| | return db, None | | | sys.exit(f"{target} has no {DB_NAME}. ai-toolkit writes one when " | | | "logging.use_ui_logger is true; otherwise pass the training log") | | | if os.path.isfile(target): | | | return target, None |
| | if target.isdigit(): |
| | pid = int(target) |
| | path = db_of_pid(pid) or discover_log_from_pid(pid) |
| | if path: | | | print(f"pid {pid}: {path}", file=sys.stderr) | | | return path, pid | | | sys.exit(f"pid {pid} has no {DB_NAME} open and its output isn't going " | | | "to a file. pass the run folder or log instead") |
| | sys.exit(f"not a folder, file or PID: {target}") |
| | def fmt_secs(s): |
| | if s is None or s != s or s < 0: |
| | return "--:--" |
| | s = int(s) |
| | h, rem = divmod(s, 3600) |
| | m, sec = divmod(rem, 60) |
| | return f"{h}:{m:02d}:{sec:02d}" if h else f"{m:02d}:{sec:02d}" |
| | def term_size(): |
| | """Current terminal size, straight from the tty ioctl. Ignores the | | | COLUMNS/LINES env vars, which multiplexers can leave stale.""" | | | for fh in (sys.stdout, sys.stderr): | | | try: |
| | return os.get_terminal_size(fh.fileno()) |
| | except (OSError, ValueError, AttributeError): |
| | continue |
| | return shutil.get_terminal_size((80, 24)) |
| | def sparkline(vals, width): |
| | if not vals: | | | return "" |
| | vals = list(vals)[-width:] |
| | lo, hi = min(vals), max(vals) |
| | rng = (hi - lo) or 1.0 |
| | return "".join(SPARK[min(7, int((v - lo) / rng * 7))] for v in vals) |
| | # --------------------------------------------------------------------------- |
| | class Tailer: | | | """Incrementally reads a growing file, splitting on \r and \n so that | | | tqdm's in-place progress bar updates become a stream of parseable lines.""" | | | def init(self, path, backfill=8 * 1024 * 1024): | | | self.path = path |
| | self.fh = open(path, "r", errors="replace") |
| | self.inode = os.fstat(self.fh.fileno()).st_ino |
| | # seed history from the tail of the existing file | | | try: |
| | size = os.path.getsize(path) |
| | self.fh.seek(max(0, size - backfill)) |
| | if size > backfill: |
| | self.fh.readline() # discard partial line | | | except OSError: | | | pass | | | self.buf = "" | | | def lines(self): | | | # handle log rotation / truncation | | | try: |
| | if os.stat(self.path).st_ino != self.inode or \ |
| | os.path.getsize(self.path) < self.fh.tell(): |
| | self.fh.close() |
| | self.fh = open(self.path, "r", errors="replace") |
| | self.inode = os.fstat(self.fh.fileno()).st_ino |
| | self.buf = "" |
| | except OSError: |
| | pass |
| | chunk = self.fh.read() |
| | if not chunk: |
| | return |
| | self.buf += chunk |
| | parts = re.split(r"[\r\n]", self.buf) |
| | self.buf = parts.pop() # keep trailing partial fragment |
| | for line in parts: |
| | if line.strip(): |
| | yield line |
| | # --------------------------------------------------------------------------- |
| | class AggSeries: |
| | """Fixed-capacity series covering the whole run: (step, value) points are |
| | averaged into buckets of stride samples, and when the buffer fills, |
| | adjacent buckets merge and the stride doubles. Memory stays bounded while |
| | the x-axis always spans everything seen so far, in training steps.""" |
| | def __init__(self, cap): |
| | self.cap = max(4, cap - cap % 2) # even, so compaction pairs cleanly |
| | self.stride = 1 |
| | self.pts = [] |
| | self._sx = self._sv = 0.0 |
| | self._n = 0 | | | def append(self, x, v): | | | self._sx += x | | | self._sv += v | | | self._n += 1 |
| | if self._n >= self.stride: |
| | self.pts.append((self._sx / self._n, self._sv / self._n)) |
| | self._sx = self._sv = 0.0 |
| | self._n = 0 |
| | if len(self.pts) >= self.cap: |
| | self.pts = [((a[0] + b[0]) / 2, (a[1] + b[1]) / 2) |
| | for a, b in zip(self.pts[::2], self.pts[1::2])] |
| | self.stride *= 2 | | | def view(self): | | | """(xs, values), including the in-progress bucket so the tail stays live.""" | | | out = list(self.pts) | | | if self._n: |
| | out.append((self._sx / self._n, self._sv / self._n)) |
| | return [p[0] for p in out], [p[1] for p in out] |
| | class State: |
| | def __init__(self, history, smooth, raw_smooth=0.8): |
| | self._args = (history, smooth, raw_smooth) |
| | self.raw = AggSeries(history) |
| | self.ema = AggSeries(history) |
| | self.lr_series = AggSeries(history) |
| | self.smooth = smooth | | | # Ostris-style smoothing on the "raw" trace: keep raw_smooth of the | | | # previous value, so spikes don't blow out the y-axis | | | self.raw_alpha = 1.0 - raw_smooth | | | self._raw_ema = None | | | self._ema = None | | | self.step = self.total = None | | | self.lr = None | | | self.loss = None | | | self.min_loss = None | | | self.rate_pts = deque(maxlen=60) # (t, step), distinct steps only | | | self.eta_str = None # tqdm's own remaining-time, straight from the log | | | self.elapsed_str = None # tqdm's own elapsed, so it covers the whole run | | | self._loss_text = None | | | self.rate_str = None # tqdm's own it/s | | | self._t0 = None # first wall time, for database sources | | | self.finished = False # database source with no trainer writing it |
| | def reset(self): |
| | self.__init__(*self._args) |
| | def _add_point(self, x, v, lr): |
| | self.loss = v | | | self._raw_ema = v if self._raw_ema is None else \ | | | self.raw_alpha * v + (1 - self.raw_alpha) * self._raw_ema | | | self.raw.append(x, self._raw_ema) | | | self._ema = v if self._ema is None else \ |
| | self.smooth * v + (1 - self.smooth) * self._ema |
| | self.ema.append(x, self._ema) |
| | self.min_loss = v if self.min_loss is None else min(self.min_loss, v) |
| | if lr and lr > 0: |
| | self.lr_series.append(x, lr) |
| | def record(self, step, wall_time, loss, lr): |
| | """One step from loss_log.db.""" | | | self.step = step | | | if lr is not None: |
| | self.lr = f"{lr:.1e}" |
| | if loss is not None and loss == loss: # not NaN |
| | self._add_point(step, loss, lr) |
| | if wall_time is not None: | | | if self._t0 is None: | | | self._t0 = wall_time |
| | self.rate_pts.append((wall_time, step)) |
| | self.elapsed_str = fmt_secs(wall_time - self._t0) |
| | def feed(self, line): |
| | ms = RE_STEP.search(line) |
| | if ms: | | | cur, tot = int(ms.group(1)), int(ms.group(2)) | | | # the training bar is the one carrying a loss; until one shows up, | | | # take any long bar. Sampling and caching bars are ignored after. |
| | if tot >= 1 and (tot == self.total or RE_LOSS.search(line) |
| | or (self.total is None and tot > 50)): |
| | if self.step is not None and cur < self.step: | | | self.rate_pts.clear() # step went backwards: new run | | | # tqdm redraws the same step several times (set_postfix etc); | | | # only distinct steps carry rate information |
| | if cur != self.step: |
| | self.rate_pts.append((time.time(), cur)) |
| | self.step, self.total = cur, tot | | | else: | | | ms = None # another bar (sampling, caching): ignore its timing | | | if ms: | | | # tqdm computes its own smoothed remaining/rate; prefer that | | | mel = RE_TQDM_ELAPSED.search(line) | | | if mel: |
| | self.elapsed_str = mel.group(1) |
| | me = RE_TQDM_ETA.search(line) |
| | if me: |
| | self.eta_str = me.group(1) |
| | mr = RE_TQDM_RATE.search(line) |
| | if mr: | | | try: |
| | v, unit = float(mr.group(1)), mr.group(2) |
| | self.rate_str = f"{1 / v:.2f}" if unit == "s/it" and v > 0 \ |
| | else f"{v:.2f}" |
| | except ValueError: | | | pass | | | ml = RE_LR.search(line) | | | if ml: |
| | self.lr = ml.group(1) |
| | m = RE_LOSS.search(line) |
| | # tqdm prints each loss twice: once when it's posted, and again on the | | | # next redraw when the step counter ticks. Count it once. |
| | if m and m.group(1) != self._loss_text: |
| | self._loss_text = m.group(1) |
| | try: | | | v = float(m.group(1)) | | | except ValueError: | | | v = None | | | if v is not None and v == v: # not NaN | | | try: | | | lr = float(self.lr) if self.lr else None | | | except ValueError: | | | lr = None | | | self._add_point(self.step or 0, v, lr) |
| | def rate(self): |
| | if len(self.rate_pts) < 2: |
| | return None |
| | (t0, s0), (t1, s1) = self.rate_pts[0], self.rate_pts[-1] |
| | dt = t1 - t0 |
| | return (s1 - s0) / dt if dt >= 1.0 else None |
| | def eta(self): |
| | r = self.rate() |
| | if r and self.total and self.step is not None and r > 0: | | | return (self.total - self.step) / r | | | return None | | | class LogSource: |
| | def __init__(self, path): |
| | self.tail = Tailer(path) |
| | def poll(self, state): |
| | for line in self.tail.lines(): |
| | state.feed(line) |
| | class DbSource: | | | """Reads new steps from ai-toolkit's loss_log.db, read-only.""" | | | def init(self, path): | | | self.path = path | | | self.con = None | | | self.last = 0 | | | self.total = None | | | try: |
| | with open(os.path.join(os.path.dirname(path), "config.yaml")) as fh: |
| | m = RE_CONFIG_STEPS.search(fh.read()) |
| | self.total = int(m.group(1)) if m else None |
| | except OSError: | | | pass | | | def poll(self, state): | | | try: | | | if self.con is None: |
| | self.con = sqlite3.connect(f"file:{self.path}?mode=ro", uri=True, timeout=1) |
| | newest = self.con.execute("SELECT MAX(step) FROM steps").fetchone()[0] or 0 |
| | if newest < self.last: | | | # a resumed run deletes the steps after its checkpoint | | | state.reset() | | | self.last = 0 | | | rows = self.con.execute(""" | | | SELECT s.step, s.wall_time, | | | MAX(CASE WHEN m.key = 'loss/loss' THEN m.value_real END), | | | MAX(CASE WHEN m.key = 'learning_rate' THEN m.value_real END) | | | FROM steps s LEFT JOIN metrics m ON m.step = s.step | | | WHERE s.step > ? GROUP BY s.step ORDER BY s.step""", (self.last,)).fetchall() | | | except sqlite3.Error: | | | self.con = None # busy or mid-write; try again next tick | | | return | | | state.total = self.total | | | for step, wall_time, loss, lr in rows: | | | state.record(step, wall_time, loss, lr) | | | self.last = step |
| | # --------------------------------------------------------------------------- |
| | def _plot_start(w, h): |
| | plt.clf() |
| | # plotext clamps plotsize to its own terminal detection, which can |
| | # lag reality (and reads 80x24 off-tty); we size it ourselves |
| | for fn in ("limit_size", "limitsize"): |
| | if hasattr(plt, fn): |
| | getattr(plt, fn)(False, False) |
| | break |
| | plt.plotsize(max(20, w), max(7, h)) |
| | plt.theme("dark") |
| | def build_loss_graph(state, w, h): |
| | raw_x, raw = state.raw.view() |
| | ema_x, ema = state.ema.view() |
| | if plt and len(ema) >= 2: |
| | _plot_start(w, h) |
| | if len(raw) >= 2: |
| | plt.plot(raw_x, raw, color="red", label="raw", marker="braille") |
| | plt.plot(ema_x, ema, color="cyan", label="ema", marker="braille") |
| | plt.xlabel("step") |
| | stride = state.raw.stride |
| | plt.title(f"loss (avg of {stride} steps)" if stride > 1 else "loss") |
| | return Text.from_ansi(plt.build().rstrip("\n")) |
| | # fallback: unicode sparkline | | | if ema: |
| | return Text(sparkline(ema, w - 4), style="cyan") |
| | return Text("waiting for loss data…", style="dim") |
| | def build_lr_graph(state, w, h): |
| | xs, lrs = state.lr_series.view() |
| | if plt and len(lrs) >= 2: |
| | _plot_start(w, h) |
| | # log scale: automagic-style optimizers climb orders of magnitude |
| | plt.yscale("log") |
| | plt.plot(xs, lrs, color="orange", marker="braille") |
| | # plotext prints these as 0.00000600; label them in e-notation instead |
| | lo, hi = min(lrs), max(lrs) |
| | n = max(2, min(5, h // 3)) |
| | ticks = [lo * (hi / lo) ** (i / (n - 1)) for i in range(n)] if hi > lo else [lo] |
| | plt.yticks(ticks, [f"{t:.1e}" for t in ticks]) |
| | plt.xlabel("step") |
| | plt.title("lr (log)") |
| | return Text.from_ansi(plt.build().rstrip("\n")) |
| | if lrs: |
| | return Text(sparkline(lrs, w - 4), style="yellow") |
| | return Text("waiting for lr data…", style="dim") |
| | def header(state, pid_status): |
| | """Two lines: progress on top, loss and lr below.""" | | | pct = (state.step / state.total * 100) if state.step and state.total else None |
| | r = state.rate() |
| | def line(*pairs): |
| | t = Text() |
| | for i, (k, v) in enumerate(pairs): |
| | if i: |
| | t.append(" ") |
| | t.append(f"{k} ", style="dim") |
| | t.append(v, style="bold") |
| | return t |
| | top = line( |
| | ("step", f"{'--' if state.step is None else state.step}/{state.total}" if state.total |
| | else f"{'--' if state.step is None else state.step}"), |
| | ("progress", f"{pct:.1f}%" if pct is not None else "--"), |
| | ("it/s", state.rate_str or (f"{r:.2f}" if r else "--")), |
| | ("eta", "--" if state.finished else state.eta_str or fmt_secs(state.eta())), |
| | ("elapsed", state.elapsed_str or "--:--"), |
| | ) |
| | bottom = line( |
| | ("loss", f"{state.loss:.4f}" if state.loss is not None else "--"), |
| | ("lr", state.lr or "--"), |
| | ) | | | if pid_status: |
| | bottom.append(f" {pid_status}", style="dim") |
| | return [top, bottom] |
| | def pid_status(pid): |
| | if not pid: | | | return "" | | | try: |
| | os.kill(pid, 0) |
| | return f"pid {pid}" |
| | except OSError: |
| | return f"pid {pid}: not running" |
| | def render(state, pid, console): |
| | size = term_size() |
| | # cols inside the outer panel: border (2) + padding (2) |
| | cols = size.columns - 4 |
| | head = header(state, pid_status(pid)) |
| | # rows for graph panels: total minus outer borders (2) and the header, | | | # which wraps on narrow terminals | | | rows = size.lines - 2 - sum(len(t.wrap(console, cols)) for t in head) | | | # each graph panel spends 2 rows and 4 cols on its border and padding |
| | if (rows - 4) // 2 >= STACK_MIN_ROWS: |
| | h = (rows - 4) // 2 |
| | graphs = Group( |
| | Panel(build_loss_graph(state, cols - 4, h), border_style="white", padding=(0, 1)), |
| | Panel(build_lr_graph(state, cols - 4, rows - 4 - h), border_style="white", padding=(0, 1)), |
| | ) | | | elif cols // 2 - 4 >= MIN_GRAPH_COLS: | | | half = cols // 2 |
| | graphs = Table.grid(expand=True) |
| | graphs.add_column(ratio=1) |
| | graphs.add_column(ratio=1) |
| | graphs.add_row( |
| | Panel(build_loss_graph(state, half - 4, rows - 2), border_style="white", padding=(0, 1)), |
| | Panel(build_lr_graph(state, cols - half - 4, rows - 2), border_style="white", padding=(0, 1)), |
| | ) | | | else: |
| | graphs = Panel(build_loss_graph(state, cols - 4, rows - 2), |
| | border_style="white", padding=(0, 1)) |
| | return Panel(Group(*head, graphs), |
| | title="loratop", subtitle="q to quit", border_style="white") |
| | # --------------------------------------------------------------------------- |
| | def main(): |
| | ap = argparse.ArgumentParser(description="TUI monitor for Ostris AI Toolkit training") | | | ap.add_argument("target", nargs="?", | | | help="run folder, training log, or PID (default: the run in " | | | "this folder, else the one running trainer)") |
| | ap.add_argument("--pid", type=int, help=argparse.SUPPRESS) # old spelling of a PID target |
| | ap.add_argument("--interval", type=float, default=1.0, help="refresh seconds") |
| | ap.add_argument("--history", type=int, default=600, |
| | help="graph resolution: points kept before the plot " | | | "stride doubles to cover the whole run") |
| | ap.add_argument("--smooth", type=float, default=0.1, help="EMA factor 0-1") |
| | ap.add_argument("--raw-smooth", type=float, default=0.8, |
| | help="Ostris-style smoothing 0-1 on the raw trace " |
| | "(0 = truly raw)") |
| | ap.add_argument("--once", action="store_true", help="render one frame and exit (debug)") |
| | args = ap.parse_args() |
| | # stale COLUMNS/LINES from a parent shell would freeze rich at the old | | | # size; drop them so every layer measures the tty directly |
| | os.environ.pop("COLUMNS", None) |
| | os.environ.pop("LINES", None) |
| | target = args.target if args.target is not None else \ |
| | (str(args.pid) if args.pid else None) |
| | path, pid = resolve_source(target) |
| | pinned = pid is not None | | | pid = pid or discover_pid_from_log(path) | | | state = State(args.history, args.smooth, args.raw_smooth) |
| | is_db = os.path.basename(path) == DB_NAME |
| | source = DbSource(path) if is_db else LogSource(path) |
| | source.poll(state) # backfill history | | | if not is_db: | | | # backfilled log lines all carry "now" as their timestamp, useless | | | # (and wildly wrong) for rate; start measuring fresh from here | | | state.rate_pts.clear() | | | state.finished = is_db and not pid | | | console = Console() | | | if args.once: | | | console.print(render(state, pid, console)) | | | return | | | def pid_alive(p): | | | try: | | | os.kill(p, 0) | | | return True | | | except OSError: | | | return False | | | resized = [False] | | | try: |
| | signal.signal(signal.SIGWINCH, lambda *_: resized.__setitem__(0, True)) |
| | except (ValueError, AttributeError): |
| | pass # non-main thread or platform without SIGWINCH | | | # cbreak mode so a bare 'q' quits cleanly (exit 0, no SIGINT to the | | | # process group, plays nice with TUI wrappers that suspend around us) |
| | kb_fd = old_attrs = None |
| | if sys.stdin.isatty(): |
| | kb_fd = sys.stdin.fileno() |
| | old_attrs = termios.tcgetattr(kb_fd) |
| | tty.setcbreak(kb_fd) |
| | def key_pressed(timeout): |
| | if kb_fd is None: | | | time.sleep(timeout) | | | return None | | | r, _, _ = select.select([kb_fd], [], [], timeout) | | | return os.read(kb_fd, 1) if r else None | | | last_disc = 0.0 | | | try: |
| | with Live(render(state, pid, console), console=console, |
| | refresh_per_second=4, screen=True) as live: |
| | while True: | | | # re-discover if the trainer died/restarted (unless pinned) |
| | if not pinned and (pid is None or not pid_alive(pid)) \ |
| | and time.monotonic() - last_disc > 5: |
| | last_disc = time.monotonic() |
| | pid = discover_pid_from_log(path) |
| | source.poll(state) |
| | state.finished = is_db and not (pid and pid_alive(pid)) |
| | live.update(render(state, pid, console), refresh=True) |
| | deadline = time.monotonic() + args.interval |
| | while time.monotonic() < deadline: |
| | if resized[0]: |
| | resized[0] = False |
| | live.update(render(state, pid, console), |
| | refresh=True) |
| | if key_pressed(0.05) in (b"q", b"Q"): |
| | return | | | except KeyboardInterrupt: | | | pass | | | finally: | | | if old_attrs is not None: | | | termios.tcsetattr(kb_fd, termios.TCSADRAIN, old_attrs) |
| | if __name__ == "__main__": |
| | main() |