| """Lightweight per-stage timing accumulator. |
| |
| A single global ``TimingAccumulator`` collects wall-clock time and call counts |
| for named stages (e.g. ``reference_logits_time``, ``ESM2_prior_time``, |
| ``PeptiVerse_time``, ``backward_time``, ...). Callers wrap their code in |
| ``with STAGE_TIMER.section("stage_name"):`` and periodically call |
| ``STAGE_TIMER.report_and_reset()`` to print/log a table. |
| |
| Cache hit/miss counters (``bump("esm2_cache_hit")``) live on the same object. |
| |
| The whole module is process-local and thread-safe enough for our single-process |
| training loop; no cross-process aggregation is attempted. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import contextlib |
| import threading |
| import time |
| from collections import defaultdict |
|
|
|
|
| class TimingAccumulator: |
| def __init__(self) -> None: |
| self._lock = threading.Lock() |
| self._time: dict[str, float] = defaultdict(float) |
| self._calls: dict[str, int] = defaultdict(int) |
| self._counters: dict[str, int] = defaultdict(int) |
|
|
| @contextlib.contextmanager |
| def section(self, name: str): |
| t0 = time.perf_counter() |
| try: |
| yield |
| finally: |
| dt = time.perf_counter() - t0 |
| with self._lock: |
| self._time[name] += dt |
| self._calls[name] += 1 |
|
|
| def add(self, name: str, seconds: float) -> None: |
| with self._lock: |
| self._time[name] += float(seconds) |
| self._calls[name] += 1 |
|
|
| def bump(self, name: str, amount: int = 1) -> None: |
| with self._lock: |
| self._counters[name] += int(amount) |
|
|
| def snapshot(self) -> dict[str, float]: |
| with self._lock: |
| snap: dict[str, float] = {} |
| for k, v in self._time.items(): |
| snap[k] = float(v) |
| snap[f"{k}_calls"] = int(self._calls.get(k, 0)) |
| for k, v in self._counters.items(): |
| snap[k] = int(v) |
| return snap |
|
|
| def reset(self) -> None: |
| with self._lock: |
| self._time.clear() |
| self._calls.clear() |
| self._counters.clear() |
|
|
| def format_table(self, title: str = "timings") -> str: |
| with self._lock: |
| rows: list[tuple[str, float, int]] = [] |
| for k in sorted(self._time.keys()): |
| rows.append((k, float(self._time[k]), int(self._calls.get(k, 0)))) |
| counters = dict(self._counters) |
| lines = [f"[{title}]"] |
| for name, secs, calls in rows: |
| per = (secs / calls) if calls else 0.0 |
| lines.append( |
| f" {name:<28s} total={secs:>8.3f}s calls={calls:>8d} avg={per*1000:>8.3f}ms" |
| ) |
| if counters: |
| lines.append(" -- counters --") |
| for k in sorted(counters.keys()): |
| lines.append(f" {k:<28s} {counters[k]}") |
| return "\n".join(lines) |
|
|
| def report_and_reset(self, title: str = "timings") -> str: |
| s = self.format_table(title) |
| self.reset() |
| return s |
|
|
|
|
| STAGE_TIMER = TimingAccumulator() |
|
|
