utils/timing.py

"""Phase timing for the Space container logs.

Every line is prefixed `[nv-generate]` so the app's own output can be separated
from Gradio, MONAI, and tqdm noise with a single grep. Phases nest, so a run
reads as a tree and it is obvious which stage a slow generation spent its time
in — which matters because everything inside an @spaces.GPU function is billed
against the ZeroGPU allocation.

    [nv-generate] -> ct.generate
    [nv-generate]   -> ct.inference
    [nv-generate]   <- ct.inference            24.61s
    [nv-generate]   -> ct.label-scan
    [nv-generate]   <- ct.label-scan            2.09s
    [nv-generate] <- ct.generate                38.42s
"""
from __future__ import annotations

import contextlib
import threading
import time

PREFIX = "[nv-generate]"

_state = threading.local()


def _depth() -> int:
    return getattr(_state, "depth", 0)


def log(message: str, indent: int | None = None) -> None:
    """Emit one prefixed, flushed line. Spaces buffers aggressively otherwise."""
    pad = "  " * (_depth() if indent is None else indent)
    print(f"{PREFIX} {pad}{message}", flush=True)


def split(name: str, seconds: float, extra: str = "") -> None:
    """Log an already-measured duration in the same column layout as phase()."""
    log(_fmt(name, seconds, "<-") + (f"  {extra}" if extra else ""))


def _fmt(name: str, seconds: float, marker: str) -> str:
    # Pad the name so the durations line up in a column and can be eyeballed.
    return f"{marker} {name:<34} {seconds:7.2f}s"


@contextlib.contextmanager
def phase(name: str, quiet_start: bool = False):
    """Time a named block, logging entry, exit, and duration.

    Re-raises on failure after logging, so a phase marker never swallows an
    error — it only annotates where it happened.
    """
    if not quiet_start:
        log(f"-> {name}")
    _state.depth = _depth() + 1
    t0 = time.perf_counter()
    try:
        yield
    except BaseException as e:
        elapsed = time.perf_counter() - t0
        _state.depth = _depth() - 1
        log(_fmt(name, elapsed, "xx") + f"  FAILED: {type(e).__name__}: {e}")
        raise
    else:
        elapsed = time.perf_counter() - t0
        _state.depth = _depth() - 1
        log(_fmt(name, elapsed, "<-"))


class Stopwatch:
    """Collects named splits, then prints them as a block with a total.

    Used for startup, where the phases are sequential and the useful output is
    one summary a reader can scan rather than interleaved progress.
    """

    def __init__(self, title: str):
        self.title = title
        self.splits: list[tuple[str, float]] = []
        self._t0 = time.perf_counter()

    @contextlib.contextmanager
    def split(self, name: str):
        t0 = time.perf_counter()
        try:
            yield
        finally:
            self.splits.append((name, time.perf_counter() - t0))

    def record(self, name: str, seconds: float) -> None:
        self.splits.append((name, seconds))

    @property
    def total(self) -> float:
        return time.perf_counter() - self._t0

    def banner(self, headline: str) -> None:
        rule = "=" * 58
        log(rule, indent=0)
        log(headline, indent=0)
        for name, seconds in self.splits:
            log(f"  {name:<38} {seconds:7.2f}s", indent=0)
        log(f"  {'TOTAL':<38} {self.total:7.2f}s", indent=0)
        log(rule, indent=0)