"""Model-independent canonical Planner Cache protocol and storage.""" from __future__ import annotations from dataclasses import dataclass import hashlib import json from pathlib import Path import torch from torch import Tensor from safetensors import safe_open from safetensors.torch import load_file, save_file from pcm.planner.cache import PlannerCache, PlannerCacheConfig, SlotType CANONICAL_VALUE_LABELS = tuple( "Alice Bob Clara David Elena Frank Grace Henry Irene James Karen Louis Maria Nancy " "Oscar Peter Queen Robert Sarah Thomas Victor Wendy Xavier London Paris Berlin Rome " "Cairo Tokyo Sydney garden kitchen cellar library forest castle".split() ) def tensor_state_checksum(state: dict[str, Tensor]) -> str: """Deterministic checksum for portable tensor-only state.""" digest = hashlib.sha256() for name, tensor in sorted(state.items()): value = tensor.detach().cpu().contiguous() digest.update(name.encode()) digest.update(str(value.dtype).encode()) digest.update(json.dumps(list(value.shape)).encode()) digest.update(value.reshape(-1).view(torch.uint8).numpy().tobytes()) return digest.hexdigest() def canonical_snapshot_checksum( tensors: dict[str, Tensor], metadata: dict[str, str] ) -> str: digest = hashlib.sha256(tensor_state_checksum(tensors).encode()) digest.update(json.dumps(metadata, sort_keys=True, separators=(",", ":")).encode()) return digest.hexdigest() @dataclass(frozen=True) class CanonicalPConfig: slots: int = 128 width: int = 512 dtype: torch.dtype = torch.float16 device: str | torch.device = "cpu" merge_similarity: float = 0.92 class CanonicalPStore: """Fixed P slots whose serialized state has no model-hidden representation.""" FORMAT = "pcm-canonical-p-v1" def __init__(self, config: CanonicalPConfig) -> None: self.config = config self.cache = PlannerCache(PlannerCacheConfig( slots=config.slots, width=config.width, dtype=config.dtype, device=config.device, merge_similarity=config.merge_similarity, )) device = self.cache.device self.entity_id = torch.full((config.slots,), -1, dtype=torch.int64, device=device) self.relation_id = torch.full((config.slots,), -1, dtype=torch.int64, device=device) self.value_id = torch.full((config.slots,), -1, dtype=torch.int64, device=device) self.canonical_metadata_id = torch.full( (config.slots,), -1, dtype=torch.int64, device=device ) @property def canonical_values(self) -> Tensor: return self.cache.values @property def valid(self) -> Tensor: return self.cache.valid def create( self, canonical_value: Tensor, *, entity_id: int, relation_id: int, value_id: int, metadata_id: int, slot_type: SlotType = SlotType.FACT, **cache_metadata, ) -> tuple[int, object]: merge_mask = ( (self.entity_id == entity_id) & (self.relation_id == relation_id) & self.cache.valid ) slot, operation = self.cache.create( canonical_value, slot_type=slot_type, merge_mask=merge_mask, **cache_metadata ) if slot >= 0: self.entity_id[slot] = entity_id self.relation_id[slot] = relation_id self.value_id[slot] = value_id self.canonical_metadata_id[slot] = metadata_id return slot, operation def allocation_signature(self): fields = (self.entity_id, self.relation_id, self.value_id, self.canonical_metadata_id) return self.cache.allocation_signature() + tuple( (field.data_ptr(), tuple(field.shape)) for field in fields ) def modify( self, slot: int, canonical_value: Tensor, *, entity_id: int, relation_id: int, value_id: int, metadata_id: int, **cache_metadata, ) -> int: result = self.cache.modify(slot, canonical_value, **cache_metadata) self.entity_id[slot] = entity_id self.relation_id[slot] = relation_id self.value_id[slot] = value_id self.canonical_metadata_id[slot] = metadata_id return result def invalidate(self, slot: int) -> int: result = self.cache.invalidate(slot) self.entity_id[slot] = -1 self.relation_id[slot] = -1 self.value_id[slot] = -1 self.canonical_metadata_id[slot] = -1 return result def save(self, path: str | Path) -> None: tensors = { "canonical_values": self.cache.values.detach().cpu(), "valid": self.cache.valid.detach().cpu(), "slot_type": self.cache.slot_type.detach().cpu(), "confidence": self.cache.confidence.detach().cpu(), "importance": self.cache.importance.detach().cpu(), "freshness": self.cache.freshness.detach().cpu(), "persistence": self.cache.persistence.detach().cpu(), "last_updated": self.cache.last_updated.detach().cpu(), "source": self.cache.source.detach().cpu(), "entity_id": self.entity_id.detach().cpu(), "relation_id": self.relation_id.detach().cpu(), "value_id": self.value_id.detach().cpu(), "canonical_metadata_id": self.canonical_metadata_id.detach().cpu(), } metadata = { "format": self.FORMAT, "config": json.dumps({ "slots": self.config.slots, "width": self.config.width, "merge_similarity": self.config.merge_similarity, }), "labels": json.dumps(self.cache.labels), } metadata["content_sha256"] = canonical_snapshot_checksum(tensors, metadata) save_file(tensors, str(Path(path)), metadata=metadata) @classmethod def load( cls, path: str | Path, *, device: str | torch.device = "cpu", dtype: torch.dtype = torch.float16, ) -> "CanonicalPStore": path = Path(path) with safe_open(str(path), framework="pt", device="cpu") as handle: metadata = handle.metadata() if metadata.get("format") != cls.FORMAT: raise ValueError("unsupported canonical P snapshot") config = json.loads(metadata["config"]) tensors = load_file(str(path), device=str(device)) checksum_metadata = { key: value for key, value in metadata.items() if key != "content_sha256" } if canonical_snapshot_checksum(tensors, checksum_metadata) != metadata.get( "content_sha256" ): raise ValueError("canonical P snapshot checksum does not match") result = cls(CanonicalPConfig( slots=config["slots"], width=config["width"], dtype=dtype, device=device, merge_similarity=config.get("merge_similarity", 0.92), )) result.cache.values.copy_(tensors["canonical_values"].to(dtype=dtype)) for name in ( "valid", "slot_type", "confidence", "importance", "freshness", "persistence", "last_updated", "source", ): getattr(result.cache, name).copy_(tensors[name]) for name in ("entity_id", "relation_id", "value_id", "canonical_metadata_id"): getattr(result, name).copy_(tensors[name]) result.cache.labels = json.loads(metadata["labels"]) result.cache._clock = int(result.cache.last_updated.max()) return result CANONICAL_P_PROTOCOL = CanonicalPStore.FORMAT def model_config_checksum(config: object) -> str: payload = config.to_dict() if hasattr(config, "to_dict") else config if isinstance(payload, dict): payload = dict(payload) configured_path = payload.get("_name_or_path") if configured_path and Path(str(configured_path)).is_absolute(): payload["_name_or_path"] = Path(str(configured_path)).name encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") return hashlib.sha256(encoded).hexdigest()