postprocess/upscale/profiling.py

"""Optional profiling for `upscale_frames()`, gated by `UPSCALE_PROFILE=1` (see `upscale.PROFILE`).

Buckets each tiled `model()` call into "first call at this (tile-shape, batch-size)" vs. "steady
call" (a repeat of a shape/batch-size combo already seen this run): the first call at a new
shape is where one-off costs land (CUDA/cuDNN lazy init today; per-shape recompiles back when the
model was torch.compile'd — dropped per issue #50). Also tracks peak CUDA memory for the run. See
issue #48.

The memory line also reports what was already allocated when the run started (e.g. the diffusion
pipeline's weights, resident in the same process — `reset_peak_memory_stats` resets the peak, not
the baseline), the upscale run's own delta above that, and the device's total/free VRAM, so
`TILE_SIZE`/`FRAME_BATCH_SIZE`/`MAX_TILE_BATCH` can be tuned against real headroom. See issue #51.

Also buckets the non-model spans identified in issue #61 (frame/tensor prep, tile-shape grouping,
tile merge, the post-tile host transfer, and per-frame PIL conversion) via `span()`, so a dev-Space
run localizes exactly where the ~74%-of-wall-time "non_model" cost goes instead of leaving it as
one opaque number.

Off by default — the `synchronize()` calls this needs for accurate per-call timing would
otherwise skew real request latency by killing async kernel overlap.
"""

from __future__ import annotations

import time as _time
from contextlib import contextmanager
from typing import Iterator

import torch


class UpscaleProfiler:
    def __init__(self, enabled: bool, device: torch.device) -> None:
        self.enabled = enabled
        self._device = device
        self._seen_keys: set[tuple[tuple[int, int], int]] = set()
        self._first_call_count = 0
        self._first_call_time = 0.0
        self._steady_call_count = 0
        self._steady_call_time = 0.0
        self._span_times: dict[str, float] = {}
        self._span_counts: dict[str, int] = {}
        self._baseline_allocated = 0
        self._baseline_reserved = 0
        self._wall_t0 = _time.perf_counter()
        if self.enabled and device.type == "cuda":
            torch.cuda.reset_peak_memory_stats(device)
            self._baseline_allocated = torch.cuda.memory_allocated(device)
            self._baseline_reserved = torch.cuda.memory_reserved(device)

    @contextmanager
    def timed(self, shape: tuple[int, int], batch_size: int) -> Iterator[None]:
        """Wrap a single tiled model() call, bucketing its elapsed time by (shape, batch_size)."""
        if not self.enabled:
            yield
            return
        if self._device.type == "cuda":
            torch.cuda.synchronize()
        t0 = _time.perf_counter()
        yield
        if self._device.type == "cuda":
            torch.cuda.synchronize()
        elapsed = _time.perf_counter() - t0

        key = (shape, batch_size)
        if key in self._seen_keys:
            self._steady_call_count += 1
            self._steady_call_time += elapsed
        else:
            self._seen_keys.add(key)
            self._first_call_count += 1
            self._first_call_time += elapsed

    @contextmanager
    def span(self, name: str) -> Iterator[None]:
        """Wrap an arbitrary non-model span (e.g. tile merge, host transfer, PIL conversion),
        accumulating elapsed time under `name` for the non_model breakdown in `report()`."""
        if not self.enabled:
            yield
            return
        if self._device.type == "cuda":
            torch.cuda.synchronize()
        t0 = _time.perf_counter()
        yield
        if self._device.type == "cuda":
            torch.cuda.synchronize()
        elapsed = _time.perf_counter() - t0
        self._span_times[name] = self._span_times.get(name, 0.0) + elapsed
        self._span_counts[name] = self._span_counts.get(name, 0) + 1

    def report(self, frame_count: int, tile_size: int, frame_batch_size: int, max_tile_batch: int) -> None:
        if not self.enabled:
            return
        wall_elapsed = _time.perf_counter() - self._wall_t0
        model_time = self._first_call_time + self._steady_call_time
        steady_avg = (
            self._steady_call_time / self._steady_call_count if self._steady_call_count else float("nan")
        )
        print(
            f"[upscale] profile: frames={frame_count} wall={wall_elapsed:.2f}s model_time={model_time:.2f}s "
            f"(non_model={wall_elapsed - model_time:.2f}s) | first-at-shape calls={self._first_call_count} "
            f"total={self._first_call_time:.2f}s avg={self._first_call_time / max(self._first_call_count, 1):.3f}s/call "
            f"| steady calls={self._steady_call_count} total={self._steady_call_time:.2f}s avg={steady_avg:.3f}s/call",
            flush=True,
        )
        if self._span_times:
            non_model = wall_elapsed - model_time
            span_total = sum(self._span_times.values())
            breakdown = " ".join(
                f"{name}={t:.2f}s(n={self._span_counts[name]})"
                for name, t in sorted(self._span_times.items(), key=lambda kv: -kv[1])
            )
            print(
                f"[upscale] profile: non_model breakdown: {breakdown} | "
                f"tracked_total={span_total:.2f}s unaccounted={non_model - span_total:.2f}s",
                flush=True,
            )
        if self._device.type == "cuda":
            gib = 1024**3
            peak_allocated = torch.cuda.max_memory_allocated(self._device)
            peak_reserved = torch.cuda.max_memory_reserved(self._device)
            # mem_get_info is driver-level: free/total across every process on the device, not just
            # this caching allocator — "free" is what's actually left right now, after the run.
            free, total = torch.cuda.mem_get_info(self._device)
            print(
                f"[upscale] profile: peak CUDA memory allocated={peak_allocated / gib:.2f} GiB "
                f"reserved={peak_reserved / gib:.2f} GiB (TILE_SIZE={tile_size}, "
                f"FRAME_BATCH_SIZE={frame_batch_size}, MAX_TILE_BATCH={max_tile_batch})",
                flush=True,
            )
            print(
                f"[upscale] profile: baseline at start allocated={self._baseline_allocated / gib:.2f} GiB "
                f"reserved={self._baseline_reserved / gib:.2f} GiB | upscale delta "
                f"allocated=+{(peak_allocated - self._baseline_allocated) / gib:.2f} GiB "
                f"reserved=+{(peak_reserved - self._baseline_reserved) / gib:.2f} GiB | device "
                f"total={total / gib:.2f} GiB free_now={free / gib:.2f} GiB "
                f"headroom_vs_peak_reserved={(total - peak_reserved) / gib:.2f} GiB",
                flush=True,
            )