| |
| """C06 goldens: golden tensor files, a tap registry, gates that are never loosened, JSON compare reports. |
| |
| - **Golden files** are ``.npz`` (loaded with ``allow_pickle=False``) holding named tensors ("taps") and an optional |
| JSON metadata record under the key ``__meta__`` (source file hashes, reference script, sample, date). |
| - **Taps** are named intermediates recorded by :class:`TapRegistry` from the CPU reference or from an *eager* |
| device run (never inside a trace capture: a traced intermediate is freed at the end of the capture). |
| - **Gates** (PLAN.md section 0.3): thresholds are declared in the test (from the spec targets) and frozen in a JSON |
| file next to the test at the first green run (``test_pcc_device.py`` -> ``test_pcc_device.gates.json``). From then |
| on a declaration looser than the frozen threshold raises :class:`GateLoosenedError`; tightening is allowed and |
| re-frozen. Changing a frozen gate means editing the JSON by hand and disclosing it (VERIFICATION_<date>.md). |
| - **Reports**: :class:`CompareReport` collects per-tap statistics and gate verdicts and writes JSON. |
| |
| numpy only (device tensors are converted through :func:`.tensors.to_numpy`, which imports ttnn lazily). |
| """ |
| from __future__ import annotations |
|
|
| import contextlib |
| import datetime as _dt |
| import fnmatch |
| import json |
| import math |
| import os |
| import sys |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Union |
|
|
| import numpy as np |
|
|
| from . import metrics as _m |
| from .weights import file_sha256 |
|
|
| __all__ = [ |
| "save_goldens", |
| "load_goldens", |
| "Goldens", |
| "TapRegistry", |
| "Gate", |
| "GateResult", |
| "GateLoosenedError", |
| "GateRegistry", |
| "METRICS", |
| "DIRECTIONS", |
| "TapComparison", |
| "CompareReport", |
| "compare_taps", |
| ] |
|
|
| PathLike = Union[str, os.PathLike] |
| META_KEY = "__meta__" |
|
|
|
|
| def _now() -> str: |
| return _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="seconds") |
|
|
|
|
| def _host(value: Any) -> np.ndarray: |
| from .tensors import to_numpy |
|
|
| return np.asarray(to_numpy(value)) |
|
|
|
|
| def _atomic_write_text(path: Path, text: str) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp") |
| tmp.write_text(text) |
| os.replace(tmp, path) |
|
|
|
|
| |
|
|
| def save_goldens(path: PathLike, tensors: Mapping[str, Any], meta: Optional[Mapping[str, Any]] = None, *, |
| compress: bool = False) -> Path: |
| """Write ``tensors`` (numpy / torch / ttnn) and ``meta`` (JSON-able) to an ``.npz`` file.""" |
| path = Path(path) |
| arrays = {str(k): _host(v) for k, v in tensors.items()} |
| if META_KEY in arrays: |
| raise ValueError(f"{META_KEY!r} is reserved for metadata") |
| arrays[META_KEY] = np.array(json.dumps(dict(meta or {}), sort_keys=True, default=str)) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_name(f".{path.stem}.{os.getpid()}.tmp.npz") |
| (np.savez_compressed if compress else np.savez)(tmp, **arrays) |
| os.replace(tmp, path) |
| return path |
|
|
|
|
| class Goldens(Mapping): |
| """Read-only, lazily loaded golden tensors of one ``.npz`` file (``with load_goldens(p) as g: g["neck"]``).""" |
|
|
| def __init__(self, path: PathLike): |
| self.path = Path(path) |
| self._npz = np.load(self.path, allow_pickle=False) |
| self._cache: Dict[str, np.ndarray] = {} |
| self.meta: Dict[str, Any] = (json.loads(str(self._npz[META_KEY])) if META_KEY in self._npz.files else {}) |
| self._sha256: Optional[str] = None |
|
|
| @property |
| def sha256(self) -> str: |
| if self._sha256 is None: |
| self._sha256 = file_sha256(self.path) |
| return self._sha256 |
|
|
| def __getitem__(self, name: str) -> np.ndarray: |
| if name == META_KEY or name not in self._npz.files: |
| raise KeyError(f"{name!r} not in {self.path.name}; have {sorted(self)[:20]}") |
| if name not in self._cache: |
| self._cache[name] = self._npz[name] |
| return self._cache[name] |
|
|
| def __iter__(self) -> Iterator[str]: |
| return (k for k in self._npz.files if k != META_KEY) |
|
|
| def __len__(self) -> int: |
| return sum(1 for _ in self) |
|
|
| def close(self) -> None: |
| self._npz.close() |
|
|
| def __enter__(self) -> "Goldens": |
| return self |
|
|
| def __exit__(self, *exc) -> None: |
| self.close() |
|
|
|
|
| def load_goldens(path: PathLike) -> Goldens: |
| return Goldens(path) |
|
|
|
|
| |
|
|
| class TapRegistry: |
| """Records named intermediates. ``tap(name, value)`` returns ``value`` unchanged, so it can be inlined:: |
| |
| taps = TapRegistry(include=["backbone.*", "neck"]) |
| x = taps.tap("backbone.layer1", layer1(x)) # stored as numpy (host read for device tensors) |
| with taps.scope("head"): |
| y = taps.tap("cls", cls_head(x)) # recorded as "head.cls" |
| taps.save("golden/sample0_taps.npz", meta={...}) |
| |
| A disabled registry (``TapRegistry(enabled=False)``, or :data:`NULL_TAPS`) does nothing, so production code can |
| keep its ``tap`` calls. Device tensors are read immediately (they may be deallocated later by the model); |
| calling ``tap`` on a device tensor while a trace capture is active raises.""" |
|
|
| def __init__(self, enabled: bool = True, *, include: Sequence[str] = ("*",), exclude: Sequence[str] = ()): |
| self.enabled = enabled |
| self.include = tuple(include) |
| self.exclude = tuple(exclude) |
| self._prefix: List[str] = [] |
| self._values: Dict[str, np.ndarray] = {} |
|
|
| def _full(self, name: str) -> str: |
| return ".".join(self._prefix + [name]) |
|
|
| def wants(self, name: str) -> bool: |
| """Whether ``name`` (full dotted name) would be recorded.""" |
| return (self.enabled and any(fnmatch.fnmatchcase(name, p) for p in self.include) |
| and not any(fnmatch.fnmatchcase(name, p) for p in self.exclude)) |
|
|
| def tap(self, name: str, value: Any) -> Any: |
| if not self.enabled: |
| return value |
| full = self._full(name) |
| if not self.wants(full): |
| return value |
| ttnn = sys.modules.get("ttnn") |
| if ttnn is not None and isinstance(value, getattr(ttnn, "Tensor", ())): |
| if value.storage_type() == ttnn.StorageType.DEVICE and ttnn.is_trace_capture_active(value.device()): |
| raise RuntimeError(f"tap {full!r} inside a trace capture: return it as a trace output instead") |
| self._values[full] = _host(value) |
| return value |
|
|
| @contextlib.contextmanager |
| def scope(self, name: str) -> Iterator["TapRegistry"]: |
| self._prefix.append(name) |
| try: |
| yield self |
| finally: |
| self._prefix.pop() |
|
|
| def names(self) -> List[str]: |
| return list(self._values) |
|
|
| def __contains__(self, name: str) -> bool: |
| return name in self._values |
|
|
| def __getitem__(self, name: str) -> np.ndarray: |
| return self._values[name] |
|
|
| def items(self): |
| return self._values.items() |
|
|
| def to_dict(self) -> Dict[str, np.ndarray]: |
| return dict(self._values) |
|
|
| def clear(self) -> None: |
| self._values.clear() |
|
|
| def save(self, path: PathLike, meta: Optional[Mapping[str, Any]] = None, *, compress: bool = False) -> Path: |
| return save_goldens(path, self._values, meta, compress=compress) |
|
|
|
|
| NULL_TAPS = TapRegistry(enabled=False) |
|
|
| |
|
|
| |
| METRICS: Dict[str, Any] = { |
| "pcc": (_m.pcc, "min"), |
| "masked_pcc": (_m.masked_pcc, "min"), |
| "argmax_agreement": (_m.argmax_agreement, "min"), |
| "label_agreement": (_m.label_agreement, "min"), |
| "mask_iou": (_m.mask_iou, "min"), |
| "max_abs": (lambda t, r: _m.error_stats(t, r)["max_abs"], "max"), |
| "mean_abs": (lambda t, r: _m.error_stats(t, r)["mean_abs"], "max"), |
| "rel_l2": (lambda t, r: _m.error_stats(t, r)["rel_l2"], "max"), |
| } |
|
|
| |
| |
| |
| DIRECTIONS: Dict[str, str] = { |
| **{name: direction for name, (_, direction) in METRICS.items()}, |
| **dict.fromkeys(("valid_row_pcc", "agreement", "iou", "miou", "mean_iou", "recall", "precision", "topk_overlap", |
| "topk_set_overlap"), "min"), |
| **dict.fromkeys(("max_abs_err", "ade", "fde", "center_err", "center_err_max", "center_err_mean", "score_err", |
| "score_err_max", "score_err_mean"), "max"), |
| } |
|
|
|
|
| @dataclass(frozen=True) |
| class Gate: |
| """``direction="min"``: pass when ``value >= threshold`` (PCC, agreement, IoU, recall); ``"max"``: pass when |
| ``value <= threshold`` (errors, ADE/FDE). NaN never passes.""" |
|
|
| threshold: float |
| direction: str = "min" |
| metric: str = "pcc" |
|
|
| def __post_init__(self) -> None: |
| if self.direction not in ("min", "max"): |
| raise ValueError(f"gate direction {self.direction!r}: expected 'min' or 'max'") |
| object.__setattr__(self, "threshold", float(self.threshold)) |
|
|
| @classmethod |
| def of(cls, spec: Union["Gate", float, int, Sequence, Mapping], metric: Optional[str] = None) -> "Gate": |
| """``0.999`` | ``(0.999, "min")`` | ``(0.05, "max", "max_abs")`` | ``{"threshold": ..., ...}`` | Gate. |
| A bare number takes the direction of ``metric`` (default ``pcc``) from :data:`DIRECTIONS` and raises |
| ``ValueError`` for a metric whose direction is unknown; a direction that contradicts a known metric's |
| (e.g. a ``"max"`` PCC gate) raises too.""" |
| if isinstance(spec, Gate): |
| gate = spec |
| elif isinstance(spec, Mapping): |
| gate = cls(**spec) |
| elif isinstance(spec, (int, float)): |
| name = metric or "pcc" |
| if name not in DIRECTIONS: |
| raise ValueError(f"gate {spec!r} on metric {name!r}: its direction is unknown; give " |
| f"(threshold, 'min' | 'max', {name!r})") |
| gate = cls(float(spec), DIRECTIONS[name], name) |
| else: |
| parts = list(spec) |
| gate = cls(*parts) if len(parts) == 3 else cls(parts[0], parts[1], metric or "pcc") |
| known = DIRECTIONS.get(gate.metric) |
| if known is not None and gate.direction != known: |
| better = "higher" if known == "min" else "lower" |
| raise ValueError(f"gate {gate}: {gate.metric!r} is a {known!r} metric ({better} is better)") |
| return gate |
|
|
| def passes(self, value: float) -> bool: |
| if value is None or (isinstance(value, float) and math.isnan(value)): |
| return False |
| return value >= self.threshold if self.direction == "min" else value <= self.threshold |
|
|
| def looser_than(self, other: "Gate") -> bool: |
| """True if this gate accepts something ``other`` rejects (same direction), or the direction differs.""" |
| if self.direction != other.direction or self.metric != other.metric: |
| return True |
| return self.threshold < other.threshold if self.direction == "min" else self.threshold > other.threshold |
|
|
|
|
| @dataclass |
| class GateResult: |
| name: str |
| value: Optional[float] |
| gate: Gate |
| passed: bool |
| newly_frozen: bool = False |
|
|
| def __str__(self) -> str: |
| op = ">=" if self.gate.direction == "min" else "<=" |
| v = "nan" if self.value is None else f"{self.value:.6g}" |
| return f"{'PASS' if self.passed else 'FAIL'} {self.name}: {self.gate.metric} {v} {op} {self.gate.threshold:g}" |
|
|
|
|
| class GateLoosenedError(AssertionError): |
| """A declared gate is looser than the threshold frozen at the first green run.""" |
|
|
|
|
| class GateRegistry: |
| """Declared gates + the JSON record of frozen gates next to the test file. |
| |
| Usage in a device test:: |
| |
| GATES = GateRegistry.for_test(__file__, {"backbone": 0.999, "head.reg": 0.99, "dets.recall": (0.95, "min")}) |
| |
| @pytest.mark.parametrize("stage", GATES.names()) |
| def test_stage(stage, stage_outputs): |
| GATES.require(stage, pcc(*stage_outputs[stage])) # AssertionError on failure; frozen on first pass |
| |
| ``write=False`` (or ``TTAW_GATES_READONLY=1``) never writes the JSON (CI / verification runs).""" |
|
|
| VERSION = 1 |
|
|
| def __init__(self, path: PathLike, gates: Optional[Mapping[str, Any]] = None, *, write: Optional[bool] = None): |
| self.path = Path(path) |
| if write is None: |
| write = os.environ.get("TTAW_GATES_READONLY", "0").strip() in ("", "0", "false", "no") |
| self.write = bool(write) |
| self.declared: Dict[str, Gate] = {k: Gate.of(v) for k, v in (gates or {}).items()} |
| self.frozen: Dict[str, Dict[str, Any]] = {} |
| if self.path.is_file(): |
| data = json.loads(self.path.read_text()) |
| self.frozen = dict(data.get("gates", {})) |
| self.results: List[GateResult] = [] |
| loosened = [] |
| for name, g in self.declared.items(): |
| if name in self.frozen and g.looser_than(self._frozen_gate(name)): |
| f = self._frozen_gate(name) |
| loosened.append(f"{name}: declared {g.direction} {g.threshold:g} ({g.metric}) vs frozen " |
| f"{f.direction} {f.threshold:g} ({f.metric})") |
| if loosened: |
| raise GateLoosenedError(f"gates in {self.path.name} may never be loosened once recorded: " |
| + "; ".join(loosened)) |
|
|
| @classmethod |
| def for_test(cls, test_file: PathLike, gates: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> "GateRegistry": |
| """Registry stored as ``<test file stem>.gates.json`` next to ``test_file`` (pass ``__file__``).""" |
| p = Path(test_file) |
| return cls(p.with_name(p.stem + ".gates.json"), gates, **kwargs) |
|
|
| def _frozen_gate(self, name: str) -> Gate: |
| rec = self.frozen[name] |
| return Gate(rec["threshold"], rec["direction"], rec["metric"]) |
|
|
| def names(self) -> List[str]: |
| return sorted(set(self.declared) | set(self.frozen)) |
|
|
| def gate(self, name: str) -> Gate: |
| """The effective gate: the declared one (never looser than the frozen one), else the frozen one.""" |
| if name in self.declared: |
| return self.declared[name] |
| if name in self.frozen: |
| return self._frozen_gate(name) |
| raise KeyError(f"no gate {name!r} declared or frozen in {self.path.name}") |
|
|
| def check(self, name: str, value: Optional[float], gate: Optional[Any] = None) -> GateResult: |
| """Evaluate ``value``; on a pass, freeze the gate if it is new or tighter than the frozen one.""" |
| g = Gate.of(gate) if gate is not None else self.gate(name) |
| if name in self.frozen and g.looser_than(self._frozen_gate(name)): |
| raise GateLoosenedError(f"gate {name!r} {g} is looser than the frozen {self.frozen[name]}") |
| v = None if value is None else float(value) |
| passed = g.passes(v) |
| newly = False |
| if passed and (name not in self.frozen or self._frozen_gate(name) != g): |
| self.frozen[name] = {**asdict(g), "value_at_freeze": v, "frozen_at": _now()} |
| newly = True |
| if self.write: |
| self.save() |
| result = GateResult(name, v, g, passed, newly) |
| self.results.append(result) |
| return result |
|
|
| def require(self, name: str, value: Optional[float], gate: Optional[Any] = None) -> GateResult: |
| """:meth:`check`, raising ``AssertionError`` when the gate fails.""" |
| result = self.check(name, value, gate) |
| if not result.passed: |
| raise AssertionError(str(result)) |
| return result |
|
|
| def save(self) -> None: |
| body = {"version": self.VERSION, "note": "Frozen at the first green run; never loosen (PLAN.md 0.3).", |
| "gates": dict(sorted(self.frozen.items()))} |
| _atomic_write_text(self.path, json.dumps(body, indent=1, sort_keys=False) + "\n") |
|
|
|
|
| |
|
|
| @dataclass |
| class TapComparison: |
| name: str |
| metric: str |
| value: Optional[float] |
| gate: Optional[Dict[str, Any]] |
| passed: Optional[bool] |
| stats: Dict[str, Any] = field(default_factory=dict) |
| note: str = "" |
|
|
|
|
| class CompareReport: |
| """Per-tap comparison records -> JSON (``logs/<model>/compare_<tag>.json``) and a text table.""" |
|
|
| def __init__(self, title: str = "", meta: Optional[Mapping[str, Any]] = None, |
| gates: Optional[GateRegistry] = None): |
| self.title = title |
| self.meta = dict(meta or {}) |
| self.gates = gates |
| self.entries: List[TapComparison] = [] |
|
|
| def _gate_for(self, name: str, gate: Any, metric: str) -> Optional[Gate]: |
| if gate is not None: |
| return Gate.of(gate, metric) |
| if self.gates is not None and name in self.gates.names(): |
| return self.gates.gate(name) |
| return None |
|
|
| def add(self, name: str, test: Any, ref: Any, *, metric: str = "pcc", gate: Any = None, |
| **metric_kwargs: Any) -> TapComparison: |
| """Compare two tensors with ``metric`` (a :data:`METRICS` name); always records the error statistics.""" |
| t, r = _host(test), _host(ref) |
| stats: Dict[str, Any] = {"shape_test": list(t.shape), "shape_ref": list(r.shape), |
| "dtype_test": str(t.dtype), "dtype_ref": str(r.dtype)} |
| note = "" |
| if t.shape == r.shape: |
| stats.update(_m.error_stats(t, r)) |
| fn = METRICS[metric][0] |
| value: Optional[float] = float(fn(t, r, **metric_kwargs)) |
| else: |
| value, note = None, "shape mismatch" |
| return self._record(name, metric, value, gate, stats, note) |
|
|
| def add_value(self, name: str, value: Optional[float], *, metric: str, gate: Any = None, |
| stats: Optional[Mapping[str, Any]] = None) -> TapComparison: |
| """Record an already computed scalar (detection recall, ADE, argmax agreement ...).""" |
| return self._record(name, metric, None if value is None else float(value), gate, dict(stats or {}), "") |
|
|
| def _record(self, name: str, metric: str, value: Optional[float], gate: Any, stats: Dict[str, Any], |
| note: str) -> TapComparison: |
| g = self._gate_for(name, gate, metric) |
| passed: Optional[bool] = None |
| if g is not None: |
| passed = (self.gates.check(name, value, g).passed if self.gates is not None else g.passes(value)) |
| entry = TapComparison(name, metric, value, None if g is None else asdict(g), passed, stats, note) |
| self.entries.append(entry) |
| return entry |
|
|
| @property |
| def passed(self) -> bool: |
| """All gated entries passed (ungated entries are informational).""" |
| return all(e.passed for e in self.entries if e.passed is not None) |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| return {"title": self.title, "created": _now(), "meta": self.meta, "passed": self.passed, |
| "entries": [asdict(e) for e in self.entries]} |
|
|
| def save_json(self, path: PathLike) -> Path: |
| path = Path(path) |
| _atomic_write_text(path, json.dumps(self.to_dict(), indent=1, default=str) + "\n") |
| return path |
|
|
| def table(self) -> str: |
| rows = [f"{self.title}".strip(), f"{'name':40s} {'metric':16s} {'value':>12s} {'gate':>14s} verdict"] |
| for e in self.entries: |
| gate = "" |
| if e.gate is not None: |
| gate = (">=" if e.gate["direction"] == "min" else "<=") + f"{e.gate['threshold']:g}" |
| val = "nan" if e.value is None else f"{e.value:.6g}" |
| verdict = {True: "PASS", False: "FAIL", None: "-"}[e.passed] |
| rows.append(f"{e.name:40s} {e.metric:16s} {val:>12s} {gate:>14s} {verdict} {e.note}".rstrip()) |
| return "\n".join(r for r in rows if r) |
|
|
|
|
| def compare_taps(test: Mapping[str, Any], golden: Mapping[str, Any], *, gates: Optional[Any] = None, |
| metric: str = "pcc", names: Optional[Sequence[str]] = None, title: str = "", |
| meta: Optional[Mapping[str, Any]] = None) -> CompareReport: |
| """Compare every tap present in both mappings (or ``names``). ``gates``: a :class:`GateRegistry` or a |
| ``{name: gate spec}`` mapping. Taps missing on either side are recorded as failures when gated.""" |
| registry = gates if isinstance(gates, GateRegistry) else None |
| plain = {} if registry is not None or gates is None else {k: Gate.of(v, metric) for k, v in gates.items()} |
| report = CompareReport(title, meta, registry) |
| for name in (names if names is not None else [n for n in golden if n in test]): |
| if name not in test or name not in golden: |
| report.add_value(name, None, metric=metric, gate=plain.get(name), |
| stats={"missing": "test" if name not in test else "golden"}) |
| continue |
| report.add(name, test[name], golden[name], metric=metric, gate=plain.get(name)) |
| return report |
|
|
