File size: 2,186 Bytes
629b123 | 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 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
from pathlib import Path
import torch
from loguru import logger
# Environment variable to specify the golden data directory.
# If not set, golden data will be computed inline and not cached.
TT_DIT_TEST_GOLDEN_CACHE_DIR = "TT_DIT_TEST_GOLDEN_CACHE_DIR"
def _golden_dir() -> Path | None:
"""Return the golden data directory from the environment, or None if not set."""
val = os.environ.get(TT_DIT_TEST_GOLDEN_CACHE_DIR)
if val is None:
return None
return Path(val)
def golden_path(commit_hash: str, test_name: str, param_id: str) -> Path | None:
"""Return the full path for a golden data file, or None if the directory is not configured.
Golden data is namespaced by the model ``commit_hash`` so that golden values produced
against different model revisions do not collide.
"""
root = _golden_dir()
if root is None:
return None
return root / test_name / param_id / f"{commit_hash}.pt"
def load_golden(commit_hash: str, test_name: str, param_id: str) -> dict[str, torch.Tensor] | None:
"""Load golden data from disk if it exists.
Returns a dict of tensors if the file exists, None otherwise.
"""
path = golden_path(commit_hash, test_name, param_id)
if path is None or not path.exists():
logger.warning(
f"No golden data found for commit={commit_hash} test={test_name} param={param_id} (looked in {path})"
)
return None
logger.info(f"Loading golden data from {path}")
return torch.load(path, weights_only=True)
def save_golden(commit_hash: str, test_name: str, param_id: str, data: dict[str, torch.Tensor]) -> None:
"""Save golden data to disk. Only saves if the directory is configured."""
path = golden_path(commit_hash, test_name, param_id)
if path is None:
logger.info(f"No golden path for {_golden_dir()} {commit_hash} {test_name} {param_id}")
return
path.parent.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving golden data to {path}")
torch.save(data, path)
|