File size: 3,047 Bytes
bb6d2aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """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()
|