| """Versioned cache for processed graph samples.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
|
|
| from .graph_dataset import GraphSample |
| from .metadata import CACHE_SCHEMA_VERSION, FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION |
|
|
|
|
| @dataclass(frozen=True) |
| class CacheMetadata: |
| feature_schema_version: int = FEATURE_SCHEMA_VERSION |
| graph_schema_version: int = GRAPH_SCHEMA_VERSION |
| cache_schema_version: int = CACHE_SCHEMA_VERSION |
| source: str | None = None |
| preprocessing: dict[str, Any] = field(default_factory=dict) |
|
|
|
|
| class GraphSampleCache: |
| """Save/load graph samples with schema metadata checked before use.""" |
|
|
| def __init__(self, path: str | Path, metadata: CacheMetadata | None = None) -> None: |
| self.path = Path(path) |
| self.metadata = metadata or CacheMetadata() |
|
|
| def save(self, samples: list[GraphSample] | tuple[GraphSample, ...]) -> None: |
| payload = { |
| "cache_metadata": asdict(self.metadata), |
| "samples": tuple(samples), |
| } |
| self.path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(payload, self.path) |
|
|
| def load(self) -> tuple[GraphSample, ...]: |
| payload = torch.load(self.path, map_location="cpu", weights_only=False) |
| if not isinstance(payload, dict): |
| raise ValueError("cache payload must be a mapping") |
| actual = payload.get("cache_metadata", {}) |
| if not isinstance(actual, dict): |
| raise ValueError("cache metadata is missing or malformed") |
| expected = asdict(self.metadata) |
| for key in ( |
| "feature_schema_version", |
| "graph_schema_version", |
| "cache_schema_version", |
| ): |
| if actual.get(key) != expected[key]: |
| raise ValueError( |
| f"cache schema mismatch for {key}: " |
| f"found {actual.get(key)!r}, expected {expected[key]!r}" |
| ) |
| if actual.get("preprocessing", {}) != expected["preprocessing"]: |
| raise ValueError("cache preprocessing fingerprint mismatch") |
| if "samples" not in payload: |
| raise ValueError("cache samples are missing") |
| samples = payload["samples"] |
| if not isinstance(samples, (list, tuple)): |
| raise ValueError("cache samples are malformed") |
| if any(not isinstance(sample, GraphSample) for sample in samples): |
| raise ValueError("cache contains invalid graph samples") |
| return tuple(samples) |
|
|
| def exists(self) -> bool: |
| return self.path.is_file() |
|
|
