File size: 2,654 Bytes
916755e | 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 63 64 65 66 67 68 69 70 | """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()
|