rpx_benchmark.cli_ux

  1"""Industry-grade colored logging + progress UI for the RPX CLI scripts.
  2
  3Built on `rich`. One canonical helper module so every runnable script
  4(``run_depth.py``, ``run_relative_pose.py``, ``run_nvs.py``,
  5``sync_results_to_box.py``, …) produces a consistent, scannable
  6terminal experience: sectioned headers, spinners for long ops,
  7progress bars with rate + ETA, and a final summary panel.
  8
  9Public surface
 10--------------
 11* :func:`setup(name)` — return a configured ``logging.Logger`` with a
 12  ``RichHandler`` attached. Idempotent; call once at the top of ``main()``.
 13* :func:`banner(title, subtitle)` — splash header at the top of a run.
 14* :func:`section(title, subtitle)` — colored section divider.
 15* :func:`step(msg)`, :func:`note(msg)`, :func:`warn(msg)`, :func:`error(msg)`,
 16  :func:`success(msg)`, :func:`bullet(msg)` — short status lines.
 17* :func:`kv(key, value)` — inline key-value (renders as ``key: value``).
 18* :func:`config(rows)` — multi-row config echo (used right after argparse).
 19* :func:`working(msg)` — context manager: spinner while a long op runs,
 20  ✓ + elapsed time on exit.
 21* :func:`progress(description, total)` — wrapped ``rich.progress.Progress``.
 22* :func:`summary(rows, title)` — final 2-column table panel.
 23* :func:`fmt_duration(s)`, :func:`fmt_rate(n, s)`, :func:`fmt_bytes(n)` —
 24  human-friendly formatters.
 25
 26Environment
 27-----------
 28- ``NO_COLOR=1`` — strip all colour escapes (respects the universal
 29  no-color convention).
 30- ``CI=true`` (set by GitHub Actions, GitLab CI, etc.) — auto-degrade
 31  spinners + live progress to plain log lines, so captured run output
 32  stays readable in CI consoles.
 33- ``RPX_QUIET=1`` — suppress all decoration; only ERROR-level logging
 34  lines reach stdout. For piping into scripts.
 35- ``RPX_NO_ANIMATION=1`` — render the splash logo statically (no
 36  letter-by-letter reveal). Useful when piping into tools that don't
 37  handle the live-redraw escape sequences gracefully.
 38- ``RPX_ANIM=fast|normal|slow|off`` — fine-grained reveal speed knob.
 39  ``off`` is equivalent to ``RPX_NO_ANIMATION=1``. ``fast`` ≈ 550ms
 40  total, ``slow`` ≈ 1.4s total. Default ``normal`` ≈ 900ms. Honoured
 41  only when the terminal supports the animation in the first place
 42  (i.e. not in CI or with ``NO_COLOR``).
 43
 44Brand palette
 45-------------
 46Matches the LaTeX draft (see ``paper-submission/overleaf/root.tex``):
 47``rpxR = #4F46E5`` (indigo-600), ``rpxP = #DB2777`` (pink-600),
 48``rpxX = #C2410C`` (orange-700).
 49"""
 50
 51from __future__ import annotations
 52
 53import logging
 54import os
 55import time
 56from contextlib import contextmanager
 57from typing import Iterator, Mapping, Optional
 58
 59from rich.console import Console
 60from rich.logging import RichHandler
 61from rich.panel import Panel
 62from rich.progress import (
 63    BarColumn,
 64    MofNCompleteColumn,
 65    Progress,
 66    SpinnerColumn,
 67    TextColumn,
 68    TimeElapsedColumn,
 69    TimeRemainingColumn,
 70)
 71from rich.table import Table
 72from rich.text import Text
 73
 74__all__ = [
 75    "setup",
 76    "banner",
 77    "section",
 78    "step",
 79    "note",
 80    "warn",
 81    "error",
 82    "success",
 83    "bullet",
 84    "kv",
 85    "config",
 86    "working",
 87    "progress",
 88    "summary",
 89    "fmt_duration",
 90    "fmt_rate",
 91    "fmt_bytes",
 92]
 93
 94
 95# ─────────────────────────────────────────────────────────────────────────────
 96# Brand palette (kept in sync with paper-submission/overleaf/root.tex)
 97# ─────────────────────────────────────────────────────────────────────────────
 98
 99BRAND_R = "#4F46E5"   # rpxR — indigo-600
100BRAND_P = "#DB2777"   # rpxP — pink-600
101BRAND_X = "#C2410C"   # rpxX — orange-700
102
103
104# ─────────────────────────────────────────────────────────────────────────────
105# Console singleton + environment detection
106# ─────────────────────────────────────────────────────────────────────────────
107
108
109_console: Optional[Console] = None
110
111
112def _is_ci() -> bool:
113    """Detect common CI environments. Suppresses spinners (which become
114    ANSI cursor-junk in captured CI logs)."""
115    return any(os.environ.get(k) for k in ("CI", "GITHUB_ACTIONS", "GITLAB_CI", "BUILDKITE"))
116
117
118def _is_quiet() -> bool:
119    return bool(os.environ.get("RPX_QUIET"))
120
121
122def _get_console() -> Console:
123    global _console
124    if _console is None:
125        # Force colors when stdout is captured but the user wants them
126        # (helpful when running `python script.py | tee out.log`); NO_COLOR
127        # still takes precedence and strips everything.
128        no_color = bool(os.environ.get("NO_COLOR"))
129        force = not no_color and not _is_ci()
130        _console = Console(force_terminal=force, no_color=no_color, highlight=False)
131    return _console
132
133
134# ─────────────────────────────────────────────────────────────────────────────
135# Logger setup
136# ─────────────────────────────────────────────────────────────────────────────
137
138
139def setup(name: str = "rpx", level: int = logging.INFO) -> logging.Logger:
140    """Return a logger with a single ``RichHandler``. Idempotent — safe to
141    call from every script's ``main()``.
142
143    In CI environments and under ``RPX_QUIET=1`` the handler still works
144    but spinner/progress helpers degrade to plain log lines.
145    """
146    logger = logging.getLogger(name)
147    if getattr(logger, "_rich_configured", False):
148        return logger
149    if _is_quiet():
150        level = logging.ERROR
151    logger.setLevel(level)
152    logger.propagate = False
153    handler = RichHandler(
154        console=_get_console(),
155        show_time=True,
156        show_path=False,
157        markup=True,
158        rich_tracebacks=True,
159        log_time_format="[%H:%M:%S]",
160    )
161    handler.setLevel(level)
162    # Replace, don't append, so calling setup twice doesn't double-log.
163    logger.handlers = [handler]
164    logger._rich_configured = True  # type: ignore[attr-defined]
165    return logger
166
167
168# ─────────────────────────────────────────────────────────────────────────────
169# Splash logo (ANSI Shadow R / P / X coloured to match the LaTeX brand palette)
170# ─────────────────────────────────────────────────────────────────────────────
171
172# Each letter is 6 rows tall. Rows are constant-width within a letter so the
173# three blocks line up cleanly when concatenated column-wise.
174_RPX_GLYPHS: tuple[tuple[str, ...], ...] = (
175    # R — indigo
176    (
177        " ██████╗ ",
178        " ██╔══██╗",
179        " ██████╔╝",
180        " ██╔══██╗",
181        " ██║  ██║",
182        " ╚═╝  ╚═╝",
183    ),
184    # P — pink
185    (
186        " ██████╗ ",
187        " ██╔══██╗",
188        " ██████╔╝",
189        " ██╔═══╝ ",
190        " ██║     ",
191        " ╚═╝     ",
192    ),
193    # X — orange (leading space added so the inter-letter gap matches R→P)
194    (
195        " ██╗  ██╗",
196        " ╚██╗██╔╝",
197        "  ╚███╔╝ ",
198        "  ██╔██╗ ",
199        " ██╔╝ ██╗",
200        " ╚═╝  ╚═╝",
201    ),
202)
203_RPX_COLORS = (BRAND_R, BRAND_P, BRAND_X)
204_RPX_ROWS = 6
205
206
207_RPX_LETTER_WIDTHS: tuple[int, ...] = tuple(len(g[0]) for g in _RPX_GLYPHS)
208_RPX_TOTAL_COLS: int = sum(_RPX_LETTER_WIDTHS)
209
210
211def _logo_text(n_cols: int = _RPX_TOTAL_COLS) -> Text:
212    """Compose the RPX logo revealed up to ``n_cols`` columns from the left.
213
214    The block height stays at ``_RPX_ROWS`` rows regardless of how much is
215    revealed — unfilled columns are padded with blanks so the surrounding
216    layout never reflows mid-animation.
217    """
218    n_cols = max(0, min(n_cols, _RPX_TOTAL_COLS))
219    out = Text()
220    for row in range(_RPX_ROWS):
221        consumed = 0
222        for li in range(3):
223            w = _RPX_LETTER_WIDTHS[li]
224            if consumed + w <= n_cols:
225                # Letter fully revealed on this row.
226                out.append(_RPX_GLYPHS[li][row], style=f"bold {_RPX_COLORS[li]}")
227            elif consumed < n_cols:
228                # Letter partially revealed — show prefix, pad the rest.
229                visible = n_cols - consumed
230                out.append(_RPX_GLYPHS[li][row][:visible], style=f"bold {_RPX_COLORS[li]}")
231                out.append(" " * (w - visible))
232            else:
233                # Letter still hidden — full-width blank pad.
234                out.append(" " * w)
235            consumed += w
236        if row != _RPX_ROWS - 1:
237            out.append("\n")
238    return out
239
240
241_ANIM_SPEEDS: dict[str, float] = {
242    "fast":   0.015,
243    "normal": 0.025,
244    "slow":   0.045,
245}
246
247
248def _splash_already_seen_this_shell() -> bool:
249    """Per-shell first-run gate. First call within a shell session returns
250    False (and records a marker); subsequent calls within the same parent
251    PID return True so the animation is skipped — first-time delight, never
252    annoying on the 5th invocation.
253
254    Disabled under CI / RPX_QUIET (where it's already disabled by other
255    means) and under ``RPX_FORCE_SPLASH=1`` for screencast capture.
256    """
257    if _is_ci() or _is_quiet() or os.environ.get("RPX_FORCE_SPLASH"):
258        return False
259    from pathlib import Path
260
261    try:
262        marker = Path(f"/tmp/rpx-splash-{os.getppid()}")  # noqa: S108
263    except Exception:
264        return False
265    fresh = marker.exists() and time.time() - marker.stat().st_mtime < 3600
266    if not fresh:
267        try:
268            marker.touch()
269        except OSError:
270            return False  # can't write marker → animate every time
271    return fresh
272
273
274def _animate_logo() -> None:
275    """Smooth column-wipe reveal of the RPX logo using ``rich.live``.
276
277    Each frame extends the reveal by one column, so the logo appears to be
278    "drawn" left-to-right at ~40 fps (≈900ms total budget by default).
279    Falls back to a single static print under CI / NO_COLOR /
280    RPX_NO_ANIMATION / RPX_ANIM=off / second-or-later invocation in the
281    same shell session.
282    """
283    console = _get_console()
284    speed = (os.environ.get("RPX_ANIM") or "normal").lower()
285    static = (
286        _is_ci()
287        or _is_quiet()
288        or bool(os.environ.get("NO_COLOR"))
289        or bool(os.environ.get("RPX_NO_ANIMATION"))
290        or speed == "off"
291        or _splash_already_seen_this_shell()
292    )
293    if static:
294        if not _is_quiet():
295            console.print(_logo_text())
296        return
297
298    from rich.live import Live  # local import: rich.live is a heavier submodule
299
300    frame_delay = _ANIM_SPEEDS.get(speed, _ANIM_SPEEDS["normal"])
301    with Live(
302        _logo_text(0),
303        console=console,
304        refresh_per_second=60,
305        transient=False,
306    ) as live:
307        for cols in range(1, _RPX_TOTAL_COLS + 1):
308            live.update(_logo_text(cols))
309            time.sleep(frame_delay)
310        time.sleep(0.18)  # brief settle on full logo before handing back
311
312
313# ─────────────────────────────────────────────────────────────────────────────
314# Headers / sectioning
315# ─────────────────────────────────────────────────────────────────────────────
316
317
318def banner(title: str, subtitle: Optional[str] = None) -> None:
319    """Top-of-run splash: animated RPX logo + a tightly-set tagline row.
320
321    Designed to be the single visual element at the top of every script —
322    no second bordered panel competing with the logo. The script title
323    appears in bold directly under the logo, followed by an optional dim
324    subtitle, separated by a thin indigo rule that visually grounds the
325    splash.
326
327    Honours ``RPX_QUIET=1`` (suppress entirely), ``RPX_NO_ANIMATION=1``
328    (render the logo statically), and ``RPX_ANIM=off`` (alias).
329    """
330    if _is_quiet():
331        return
332    console = _get_console()
333    console.print()
334    _animate_logo()
335    # Tagline row — bold title + dim subtitle, sitting directly under the
336    # logo. No bordered panel: the logo itself carries the visual weight.
337    console.print()
338    console.print(Text(f"  {title}", style=f"bold {BRAND_R}"))
339    if subtitle:
340        console.print(Text(f"  {subtitle}", style="dim"))
341    console.rule(style=BRAND_R)
342
343
344def section(title: str, subtitle: Optional[str] = None) -> None:
345    """Mid-run colored section divider."""
346    if _is_quiet():
347        return
348    console = _get_console()
349    text = Text(title, style=f"bold {BRAND_R}")
350    if subtitle:
351        text.append("  ", style="")
352        text.append(subtitle, style="dim")
353    console.print()
354    console.rule(text, style=BRAND_R)
355
356
357# ─────────────────────────────────────────────────────────────────────────────
358# Inline status lines
359# ─────────────────────────────────────────────────────────────────────────────
360
361
362def step(msg: str) -> None:
363    """Green ✓ — a completed step."""
364    if _is_quiet():
365        return
366    _get_console().print(f"[green]✓[/] {msg}")
367
368
369def note(msg: str) -> None:
370    """Dim two-space-indented info line (file paths, config echoes)."""
371    if _is_quiet():
372        return
373    _get_console().print(f"  [dim]{msg}[/]")
374
375
376def bullet(msg: str) -> None:
377    """List bullet, slightly stronger than ``note``."""
378    if _is_quiet():
379        return
380    _get_console().print(f"  [{BRAND_R}]•[/] {msg}")
381
382
383def warn(msg: str) -> None:
384    """Yellow ! — non-fatal warning."""
385    _get_console().print(f"[yellow]![/] {msg}")
386
387
388def error(msg: str) -> None:
389    """Red ✗ — fatal-class error (caller still decides whether to raise)."""
390    _get_console().print(f"[red]✗[/] {msg}", style="red")
391
392
393def success(msg: str) -> None:
394    """Bold green ✓ — a milestone, louder than ``step``."""
395    if _is_quiet():
396        return
397    _get_console().print(f"[bold green]✓[/] {msg}")
398
399
400def kv(key: str, value: object) -> None:
401    """Inline key-value pair, right-aligned key column."""
402    if _is_quiet():
403        return
404    _get_console().print(f"  [bold]{key:>18}[/]  [dim]:[/]  {value}")
405
406
407def config(rows: Mapping[str, object], title: str = "Configuration") -> None:
408    """Echo a config dict (typically the parsed argparse Namespace)
409    as a labelled section followed by aligned ``key: value`` rows."""
410    if _is_quiet():
411        return
412    section(title)
413    for k, v in rows.items():
414        kv(k, v)
415
416
417# ─────────────────────────────────────────────────────────────────────────────
418# Long-running operations
419# ─────────────────────────────────────────────────────────────────────────────
420
421
422@contextmanager
423def working(msg: str, success_msg: Optional[str] = None) -> Iterator[object]:
424    """Spinner + live message while a long op runs; on exit replaces the
425    spinner with a green ✓ line that includes elapsed time. In CI / quiet
426    mode the spinner is omitted and the message is logged plainly.
427
428    Usage::
429
430        with cli_ux.working("Loading checkpoint"):
431            model.load_weights(...)
432
433        with cli_ux.working("Evaluating") as status:
434            for i in range(N):
435                status.update(f"[cyan]Evaluating ({i}/{N})…[/]")
436    """
437    console = _get_console()
438    t0 = time.time()
439    if _is_ci() or _is_quiet():
440        # Plain mode — no live spinner.
441        if not _is_quiet():
442            console.print(f"… {msg}")
443
444        class _NullStatus:
445            def update(self, _msg: str) -> None:
446                pass
447
448        try:
449            yield _NullStatus()
450        except Exception:
451            console.print(f"[red]✗[/] {msg} (failed after {fmt_duration(time.time() - t0)})")
452            raise
453        else:
454            if not _is_quiet():
455                console.print(f"[green]✓[/] {success_msg or msg} ({fmt_duration(time.time() - t0)})")
456        return
457
458    status = console.status(f"[{BRAND_R}]{msg}…[/]", spinner="dots")
459    status.start()
460    try:
461        yield status
462    except Exception:
463        status.stop()
464        console.print(f"[red]✗[/] {msg} (failed after {fmt_duration(time.time() - t0)})")
465        raise
466    else:
467        status.stop()
468        dt = time.time() - t0
469        console.print(f"[green]✓[/] {success_msg or msg} ({fmt_duration(dt)})")
470
471
472@contextmanager
473def progress(description: str = "working", total: Optional[int] = None) -> Iterator[tuple[object, int]]:
474    """Context manager yielding ``(Progress, task_id)`` for a tracked loop.
475
476    Usage::
477
478        with cli_ux.progress("Frames", total=N) as (p, task):
479            for i in range(N):
480                ...
481                p.update(task, advance=1)
482
483    In CI / quiet mode the live bar is replaced with periodic plain log
484    lines (every 10% of progress).
485    """
486    if _is_ci() or _is_quiet():
487        # Plain mode — emit a log line every ~10% of progress.
488        console = _get_console()
489        state = {"done": 0, "last_pct": -10}
490
491        class _PlainProgress:
492            def update(self, _task: int, advance: int = 1) -> None:
493                state["done"] += advance
494                if total:
495                    pct = int(100 * state["done"] / total)
496                    if pct >= state["last_pct"] + 10:
497                        state["last_pct"] = pct
498                        if not _is_quiet():
499                            console.print(f"  [{description}] {state['done']}/{total} ({pct}%)")
500                elif state["done"] % 50 == 0 and not _is_quiet():
501                    console.print(f"  [{description}] {state['done']}")
502
503        yield _PlainProgress(), 0
504        return
505
506    p = Progress(
507        SpinnerColumn(),
508        TextColumn("[progress.description]{task.description}"),
509        BarColumn(bar_width=None),
510        MofNCompleteColumn(),
511        TimeElapsedColumn(),
512        TimeRemainingColumn(),
513        console=_get_console(),
514        transient=False,
515    )
516    p.start()
517    try:
518        task = p.add_task(description, total=total)
519        yield p, task
520    finally:
521        p.stop()
522
523
524# ─────────────────────────────────────────────────────────────────────────────
525# Summaries
526# ─────────────────────────────────────────────────────────────────────────────
527
528
529def summary(rows: Mapping[str, object], title: str = "Summary") -> None:
530    """Final 2-column table panel — last thing a CLI run prints."""
531    if _is_quiet():
532        return
533    table = Table(show_header=False, box=None, padding=(0, 2))
534    table.add_column("k", style=f"bold {BRAND_R}")
535    table.add_column("v")
536    for k, v in rows.items():
537        table.add_row(str(k), str(v))
538    panel = Panel(table, title=f"[bold]{title}[/]", border_style=BRAND_R, expand=False)
539    console = _get_console()
540    console.print()
541    console.print(panel)
542
543
544# ─────────────────────────────────────────────────────────────────────────────
545# Formatters
546# ─────────────────────────────────────────────────────────────────────────────
547
548
549def fmt_duration(seconds: float) -> str:
550    """``23 ms`` / ``1.4 s`` / ``2 m 13 s`` / ``1 h 02 m``."""
551    if seconds < 1.0:
552        return f"{seconds * 1000:.0f} ms"
553    if seconds < 60:
554        return f"{seconds:.2f} s"
555    if seconds < 3600:
556        m, s = divmod(int(seconds), 60)
557        return f"{m} m {s:02d} s"
558    h, rem = divmod(int(seconds), 3600)
559    m, _ = divmod(rem, 60)
560    return f"{h} h {m:02d} m"
561
562
563def fmt_rate(count: int, seconds: float) -> str:
564    """``42.3/s``, with em-dash for zero/negative durations."""
565    if seconds <= 0:
566        return "—"
567    return f"{count / seconds:.1f}/s"
568
569
570def fmt_bytes(n: float) -> str:
571    """``1.4 KB`` / ``128.0 MB`` / ``3.2 GB``."""
572    for unit in ("B", "KB", "MB", "GB", "TB"):
573        if n < 1024:
574            return f"{n:.1f} {unit}"
575        n /= 1024
576    return f"{n:.1f} PB"
577
578
579# ─────────────────────────────────────────────────────────────────────────────
580# Side effects when run as a smoke test
581# ─────────────────────────────────────────────────────────────────────────────
582
583
584if __name__ == "__main__":  # pragma: no cover
585    # Force the animation on every smoke-test invocation, regardless of
586    # whether this shell already saw the splash in a real script today.
587    os.environ["RPX_FORCE_SPLASH"] = "1"
588    log = setup("demo")
589    banner("rpx_benchmark.cli_ux — smoke test", "verify every helper renders")
590    section("Imports", "verify rich + helper functions wired correctly")
591    step("module imported")
592    note("import path: rpx_benchmark.cli_ux")
593    bullet("setup, banner, section, step, note, kv, working, progress, summary")
594
595    config({"split": "easy", "batch_size": 4, "device": "cuda"})
596
597    section("Live demo")
598    with working("Pretend-loading model"):
599        time.sleep(0.4)
600
601    with progress("Pretend frames", total=20) as (p, task):
602        for _ in range(20):
603            time.sleep(0.02)
604            p.update(task, advance=1)  # type: ignore[attr-defined]
605
606    warn("Pretend warning: token expires in 5 min")
607    success("Demo complete")
608
609    summary(
610        {
611            "frames": "20",
612            "elapsed": fmt_duration(0.42),
613            "rate": fmt_rate(20, 0.42),
614        },
615        title="Demo summary",
616    )