| """Strict safetensors cache for matched-manifest live-teacher features.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import tempfile |
| from typing import Any, Mapping |
|
|
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| import torch |
|
|
| from utils.matched_video_manifest import MatchedVideoManifestEntry |
|
|
|
|
| CACHE_FORMAT = "sf_matched_manifest_video_vjepa2_teacher_feature" |
| CACHE_FORMAT_VERSION = 1 |
| CACHE_METADATA_KEY = "sf_matched_video_teacher_feature" |
| CACHE_TENSOR_KEY = "features" |
|
|
| |
| |
| MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT: dict[str, Any] = { |
| "adapter_contract_version": 1, |
| "contract_version": 1, |
| "crop_size": 256, |
| "expected_source_frame_count": 81, |
| "feature_dim": 1024, |
| "feature_type": "vjepa2_dense_latent", |
| "frame_indices": [17, 80], |
| "frame_sampling": "last_contiguous", |
| "interpolation": "bilinear_antialias", |
| "layer_index": 15, |
| "model_family": "vjepa2-vitl-fpc64-256", |
| "model_config_sha256": ( |
| "3dec96fe962e94e569182d3a7b9ef0dd74b6b8c89c337a428e43e10d593e70c9" |
| ), |
| "normalization": "imagenet_default", |
| "num_frames": 64, |
| "num_tokens": 1024, |
| "output_dtype": "float32", |
| "output_shape": [1024, 1024], |
| "processor_config_sha256": ( |
| "d2fab4418fc0390b62c4cd72ade56908a7929f80c62288adbe10dd8d23421227" |
| ), |
| "resize_short_side": 292, |
| "sample_axis": "video", |
| "token_alignment": "same_deterministic_indices", |
| "token_indices_sha256": ( |
| "a77ee1f9d20cba6b2f2a9a971c3d36610ca7f8aa2a55c616192582665be69d10" |
| ), |
| "token_sampling": "uniform_linspace_round", |
| "transformers_version": "4.53.3", |
| } |
|
|
| MATCHED_PYAV_LAST64_DECODER_CONTRACT: dict[str, Any] = { |
| "backend": "pyav", |
| "backend_version": "13.1.0", |
| "contract_version": 1, |
| "expected_frame_indices": [17, 80], |
| "expected_source_frame_count": 81, |
| "frame_count_source": "container_decode_order_with_stream_frames_hint", |
| "num_frames": 64, |
| "output_dtype": "float32", |
| "output_layout": "TCHW", |
| "output_range": [0.0, 1.0], |
| "pixel_format": "rgb24", |
| "selection": "last_contiguous", |
| } |
|
|
|
|
| class MatchedVideoFeatureCacheValidationError(ValueError): |
| """An existing cache file is stale, corrupt, or identity-incompatible.""" |
|
|
|
|
| def canonical_json(value: Any) -> str: |
| return json.dumps( |
| value, |
| ensure_ascii=True, |
| sort_keys=True, |
| separators=(",", ":"), |
| allow_nan=False, |
| ) |
|
|
|
|
| def canonical_sha256(value: Any) -> str: |
| return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest() |
|
|
|
|
| MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT_SHA256 = canonical_sha256( |
| MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT |
| ) |
| MATCHED_PYAV_LAST64_DECODER_CONTRACT_SHA256 = canonical_sha256( |
| MATCHED_PYAV_LAST64_DECODER_CONTRACT |
| ) |
|
|
|
|
| def _is_sha256(value: str) -> bool: |
| return len(value) == 64 and all(c in "0123456789abcdef" for c in value) |
|
|
|
|
| def _entry_context( |
| entry: MatchedVideoManifestEntry | Mapping[str, Any], |
| ) -> dict[str, Any]: |
| if isinstance(entry, MatchedVideoManifestEntry): |
| result = { |
| "sample_id": entry.sample_id, |
| "video_mtime_ns": int(entry.video_mtime_ns), |
| "video_name": entry.video_name, |
| "video_size_bytes": int(entry.video_size_bytes), |
| } |
| elif isinstance(entry, Mapping): |
| required = { |
| "sample_id", "video_mtime_ns", "video_name", "video_size_bytes", |
| } |
| missing = required - set(entry) |
| if missing: |
| raise ValueError(f"cache sample context is missing {sorted(missing)}") |
| result = { |
| "sample_id": str(entry["sample_id"]), |
| "video_mtime_ns": int(entry["video_mtime_ns"]), |
| "video_name": str(entry["video_name"]), |
| "video_size_bytes": int(entry["video_size_bytes"]), |
| } |
| else: |
| raise TypeError("entry must be MatchedVideoManifestEntry or mapping") |
| if not result["sample_id"] or result["sample_id"] != result["video_name"]: |
| raise ValueError("cache sample_id must equal video_name") |
| if result["video_size_bytes"] <= 0 or result["video_mtime_ns"] <= 0: |
| raise ValueError("cache video stat must be positive") |
| return result |
|
|
|
|
| def _tensor_sha256(value: torch.Tensor) -> str: |
| value = value.detach().cpu().contiguous() |
| return hashlib.sha256(value.numpy().tobytes(order="C")).hexdigest() |
|
|
|
|
| class MatchedVideoFeatureCache: |
| """Lazy one-file-per-video FP32 cache with immutable identity checks.""" |
|
|
| def __init__( |
| self, |
| cache_root: str | os.PathLike, |
| *, |
| expected_model_sha256: str, |
| expected_feature_contract_sha256: str, |
| feature_contract: Mapping[str, Any] | None = None, |
| decoder_contract: Mapping[str, Any] | None = None, |
| expected_shape: tuple[int, int] | None = None, |
| ) -> None: |
| self.cache_root = Path(cache_root).expanduser().resolve() |
| self.model_sha256 = str(expected_model_sha256).lower() |
| if not _is_sha256(self.model_sha256): |
| raise ValueError("expected_model_sha256 must be lowercase SHA-256") |
| self.feature_contract = json.loads(canonical_json( |
| feature_contract or MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT |
| )) |
| self.feature_contract_sha256 = canonical_sha256(self.feature_contract) |
| expected_hash = str(expected_feature_contract_sha256).lower() |
| if not _is_sha256(expected_hash): |
| raise ValueError( |
| "expected_feature_contract_sha256 must be lowercase SHA-256" |
| ) |
| if self.feature_contract_sha256 != expected_hash: |
| raise ValueError( |
| "feature contract SHA-256 mismatch: " |
| f"expected={expected_hash}, actual={self.feature_contract_sha256}" |
| ) |
| self.decoder_contract = json.loads(canonical_json( |
| decoder_contract or MATCHED_PYAV_LAST64_DECODER_CONTRACT |
| )) |
| self.decoder_contract_sha256 = canonical_sha256(self.decoder_contract) |
| contract_shape = tuple( |
| int(value) for value in self.feature_contract.get("output_shape", ()) |
| ) |
| self.expected_shape = ( |
| tuple(int(value) for value in expected_shape) |
| if expected_shape is not None else contract_shape |
| ) |
| if len(self.expected_shape) != 2 or min(self.expected_shape) <= 0: |
| raise ValueError("expected feature shape must be positive [N,D]") |
| if contract_shape and self.expected_shape != contract_shape: |
| raise ValueError( |
| "expected_shape must match feature_contract.output_shape" |
| ) |
|
|
| def _key_payload( |
| self, entry: MatchedVideoManifestEntry | Mapping[str, Any], |
| ) -> dict[str, Any]: |
| return { |
| "decoder_contract": self.decoder_contract, |
| "decoder_contract_sha256": self.decoder_contract_sha256, |
| "feature_contract": self.feature_contract, |
| "feature_contract_sha256": self.feature_contract_sha256, |
| "format": CACHE_FORMAT, |
| "format_version": CACHE_FORMAT_VERSION, |
| "model_sha256": self.model_sha256, |
| "sample": _entry_context(entry), |
| } |
|
|
| def key_for( |
| self, entry: MatchedVideoManifestEntry | Mapping[str, Any], |
| ) -> str: |
| digest = hashlib.sha256() |
| digest.update(b"sf-matched-video-teacher-feature-cache\0") |
| digest.update(canonical_json(self._key_payload(entry)).encode("utf-8")) |
| return digest.hexdigest() |
|
|
| def path_for_key(self, cache_key: str) -> Path: |
| key = str(cache_key).lower() |
| if not _is_sha256(key): |
| raise ValueError("cache_key must be lowercase SHA-256") |
| return self.cache_root / key[:2] / key[2:4] / f"{key}.safetensors" |
|
|
| def _expected_metadata( |
| self, |
| entry: MatchedVideoManifestEntry | Mapping[str, Any], |
| cache_key: str, |
| tensor_sha256: str, |
| ) -> dict[str, Any]: |
| result = self._key_payload(entry) |
| result.update({ |
| "cache_key": cache_key, |
| "tensor": { |
| "dtype": "float32", "key": CACHE_TENSOR_KEY, |
| "sha256": tensor_sha256, "shape": list(self.expected_shape), |
| }, |
| }) |
| return result |
|
|
| def _validate_tensor(self, feature: torch.Tensor) -> torch.Tensor: |
| if not isinstance(feature, torch.Tensor): |
| raise TypeError("teacher feature must be a torch.Tensor") |
| if feature.ndim == 3 and int(feature.shape[0]) == 1: |
| feature = feature[0] |
| if feature.ndim != 2 or tuple(feature.shape) != self.expected_shape: |
| raise ValueError( |
| f"teacher feature shape mismatch: {tuple(feature.shape)} " |
| f"!= {self.expected_shape}" |
| ) |
| if not feature.is_floating_point(): |
| raise TypeError("teacher feature must be floating point") |
| result = feature.detach().to(device="cpu", dtype=torch.float32).contiguous() |
| if not bool(torch.isfinite(result).all().item()): |
| raise FloatingPointError("teacher feature contains non-finite values") |
| return result |
|
|
| def load( |
| self, |
| entry: MatchedVideoManifestEntry | Mapping[str, Any], |
| *, |
| cache_key: str | None = None, |
| ) -> torch.Tensor | None: |
| key = self.key_for(entry) |
| if cache_key is not None and str(cache_key) != key: |
| raise MatchedVideoFeatureCacheValidationError( |
| "caller cache_key does not match sample/model/contracts" |
| ) |
| path = self.path_for_key(key) |
| if not path.exists(): |
| return None |
| if not path.is_file(): |
| raise MatchedVideoFeatureCacheValidationError( |
| f"cache path is not a regular file: {path}" |
| ) |
| try: |
| with safe_open(path, framework="pt", device="cpu") as handle: |
| if list(handle.keys()) != [CACHE_TENSOR_KEY]: |
| raise MatchedVideoFeatureCacheValidationError( |
| f"cache tensor keys mismatch in {path}" |
| ) |
| serialized = (handle.metadata() or {}).get(CACHE_METADATA_KEY) |
| if serialized is None: |
| raise MatchedVideoFeatureCacheValidationError( |
| f"cache metadata is missing in {path}" |
| ) |
| metadata = json.loads(serialized) |
| feature = handle.get_tensor(CACHE_TENSOR_KEY).clone() |
| except MatchedVideoFeatureCacheValidationError: |
| raise |
| except Exception as exc: |
| raise MatchedVideoFeatureCacheValidationError( |
| f"failed to read safetensors cache {path}: {exc}" |
| ) from exc |
| feature = self._validate_tensor(feature) |
| expected = self._expected_metadata(entry, key, _tensor_sha256(feature)) |
| if metadata != expected: |
| raise MatchedVideoFeatureCacheValidationError( |
| f"cache metadata/identity mismatch (stale or corrupt) in {path}" |
| ) |
| return feature |
|
|
| def write( |
| self, |
| entry: MatchedVideoManifestEntry | Mapping[str, Any], |
| feature: torch.Tensor, |
| *, |
| cache_key: str | None = None, |
| overwrite: bool = False, |
| ) -> Path: |
| key = self.key_for(entry) |
| if cache_key is not None and str(cache_key) != key: |
| raise MatchedVideoFeatureCacheValidationError( |
| "caller cache_key does not match sample/model/contracts" |
| ) |
| feature = self._validate_tensor(feature) |
| tensor_hash = _tensor_sha256(feature) |
| destination = self.path_for_key(key) |
| if destination.exists() and not overwrite: |
| existing = self.load(entry, cache_key=key) |
| if existing is None or _tensor_sha256(existing) != tensor_hash: |
| raise MatchedVideoFeatureCacheValidationError( |
| "refusing to replace a valid cache entry with different features: " |
| f"{destination}" |
| ) |
| return destination |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| metadata = self._expected_metadata(entry, key, tensor_hash) |
| descriptor, temporary_name = tempfile.mkstemp( |
| dir=destination.parent, |
| prefix=f".{destination.name}.", suffix=".tmp", |
| ) |
| os.close(descriptor) |
| temporary = Path(temporary_name) |
| try: |
| temporary.unlink() |
| save_file( |
| {CACHE_TENSOR_KEY: feature}, str(temporary), |
| metadata={CACHE_METADATA_KEY: canonical_json(metadata)}, |
| ) |
| with temporary.open("rb") as handle: |
| os.fsync(handle.fileno()) |
| if overwrite: |
| os.replace(temporary, destination) |
| else: |
| try: |
| os.link(temporary, destination) |
| except FileExistsError: |
| existing = self.load(entry, cache_key=key) |
| if existing is None or _tensor_sha256(existing) != tensor_hash: |
| raise MatchedVideoFeatureCacheValidationError( |
| "concurrent cache writer produced different features " |
| f"for {key}" |
| ) |
| temporary.unlink(missing_ok=True) |
| finally: |
| temporary.unlink(missing_ok=True) |
| self.load(entry, cache_key=key) |
| return destination |
|
|
|
|
| __all__ = [ |
| "CACHE_FORMAT", "CACHE_FORMAT_VERSION", "CACHE_METADATA_KEY", |
| "CACHE_TENSOR_KEY", "MATCHED_PYAV_LAST64_DECODER_CONTRACT", |
| "MATCHED_PYAV_LAST64_DECODER_CONTRACT_SHA256", |
| "MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT", |
| "MATCHED_VJEPA2_L15_LAST64_FEATURE_CONTRACT_SHA256", |
| "MatchedVideoFeatureCache", "MatchedVideoFeatureCacheValidationError", |
| "canonical_json", "canonical_sha256", |
| ] |
|
|
