File size: 5,873 Bytes
9aa90e0 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""Process-global wall-time ledger with cache-miss anomaly flags.
A run's wall time hides in a few cost centers (weight conversion/load, warmup
compile, generation, audio decode). When something that should be cached isn't —
e.g. TT_DIT_CACHE_DIR unset — the run is silently slow. This ledger collects
per-category time and, at the end of a run, prints where the time went, an
``untracked`` remainder that reconciles the table to true wall time (so the
breakdown can't mislead), and an ANOMALIES section so "should've been fast,
wasn't" is obvious instead of buried in thousands of log lines.
Fed from existing cost centers (no scattered timers): ``utils.cache.load_model``
(weight loads, with HIT/MISS), ``utils.progress.Watchdog`` (phases), and a few
dominant pipeline phases wrapped at their existing timing sites. Disable with
TT_WALLTIME=0.
NOT covered: JIT kernel-build-cache hit rate. That counter ("JIT cache stats:
H/T hits") is emitted by a C++ static destructor at process teardown — after
Python has exited — and ttnn exposes no pre-teardown Python accessor for it, so
it cannot be folded into this in-process ledger. Read it from the device log.
"""
from __future__ import annotations
import atexit
import os
import threading
import time
from collections import OrderedDict
from contextlib import contextmanager
from dataclasses import dataclass, field
_LOCK = threading.RLock()
# Wall baseline: render() falls back to (now - import) when no explicit wall is handed in.
_T_IMPORT = time.monotonic()
@dataclass
class _Cat:
seconds: float = 0.0
count: int = 0
hits: int = 0 # cached is True
misses: int = 0 # cached is False
@dataclass
class _Ledger:
cats: "OrderedDict[str, _Cat]" = field(default_factory=OrderedDict)
# (label, seconds, detail) for every cached=False record, surfaced as anomalies.
misses: list = field(default_factory=list)
# One process-global ledger, rendered once at teardown (see ``_atexit``).
_ledger = _Ledger()
# Read once at import: an instrumentation toggle should hold for a whole process, and record()
# consults it on every span.
_ENABLED = os.environ.get("TT_WALLTIME", "1") != "0"
def _enabled() -> bool:
return _ENABLED
def record(
category: str, label: str, seconds: float, *, cached: bool | None = None, count: int = 1, detail: str = ""
) -> None:
"""Accumulate ``seconds`` into ``category``. ``cached`` tallies HIT (True) / MISS (False); a MISS
is also retained for the anomaly list."""
if not _enabled():
return
with _LOCK:
cat = _ledger.cats.get(category)
if cat is None:
cat = _Cat()
_ledger.cats[category] = cat
cat.seconds += seconds
cat.count += count
if cached is True:
cat.hits += 1
elif cached is False:
cat.misses += 1
_ledger.misses.append((label, seconds, detail))
@contextmanager
def timed(category: str, label: str, *, cached: bool | None = None):
t0 = time.monotonic()
try:
yield
finally:
record(category, label, time.monotonic() - t0, cached=cached)
def render(title: str, ledger: _Ledger | None = None, wall: float | None = None) -> str:
"""Render a ledger as a table reconciled to wall time.
``wall`` is the true end-to-end wall for the scope (an explicit duration, or
process time since import when omitted). The table lists each category's
seconds and share of wall, then an ``untracked`` remainder (wall − tracked)
and ``TOTAL (wall)`` so the breakdown always reconciles. Anomalies list every
cached=False weight load. Categories are assumed disjoint; a phase span that
nests a weight_load would double-count and clamp ``untracked`` to 0.
"""
led = ledger if ledger is not None else _ledger
with _LOCK:
cats = list(led.cats.items())
misses = list(led.misses)
tracked = sum(c.seconds for _, c in cats)
if wall is None:
wall = time.monotonic() - _T_IMPORT
wall = max(wall, tracked) # never report a wall smaller than what we measured
untracked = wall - tracked
denom = wall or 1.0
width = 74
lines = ["", "=" * width, f"WALL-TIME LEDGER · {title}", "-" * width]
lines.append(f"{'category':<18}{'seconds':>10}{'%wall':>8}{'count':>8} notes")
for name, c in cats:
note = f"{c.hits} HIT / {c.misses} MISS" if (c.hits or c.misses) else ""
lines.append(f"{name:<18}{c.seconds:>10.1f}{100.0 * c.seconds / denom:>7.1f}%{c.count:>8} {note}")
lines.append(f"{'untracked':<18}{untracked:>10.1f}{100.0 * untracked / denom:>7.1f}%")
lines.append("-" * width)
lines.append(f"{'TOTAL (tracked)':<18}{tracked:>10.1f}{100.0 * tracked / denom:>7.1f}%")
lines.append(f"{'TOTAL (wall)':<18}{wall:>10.1f}{100.0:>7.1f}%")
anomalies = [
f"CACHE MISS: weight_load {label} ({seconds:.1f}s) — {detail or 'TT_DIT_CACHE_DIR unset or blocking key changed'}"
for label, seconds, detail in misses
]
lines.append("")
if anomalies:
lines.append("ANOMALIES:")
lines.extend(f" - {a}" for a in anomalies)
else:
lines.append("ANOMALIES: none")
lines.append("=" * width)
return "\n".join(lines)
def _atexit() -> None:
# The ledger surfaces once, at process teardown — under pytest (the tests wire no per-item hook) and
# the ``python -m ...`` entrypoints alike.
if not _enabled():
return
with _LOCK:
nonempty = bool(_ledger.cats)
if not nonempty:
return
try:
print(render("end of run", _ledger)) # noqa: T201
except (ValueError, OSError):
pass # stdout may already be closed during interpreter shutdown
atexit.register(_atexit)
|