| |
|
|
| |
|
|
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
|
|
| import torch |
| from loguru import logger |
|
|
| |
| |
| 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) |
|
|
