"""Process-shared, authority-bound cache for immutable NoNE page weights. This module is an external storage boundary. It may cache a reconstruction only after the caller has authorized the accepted session/generation/manifest and verified the immutable source object. It never stores optimizer state, candidate state, routing state, or trainable tensor aliases. """ from __future__ import annotations import fcntl import hashlib import json import os import threading import time from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path from typing import Any, Final import orjson import torch from safetensors import safe_open from safetensors.torch import save_file SHARED_PAGE_WEIGHTS_CACHE_SCHEMA: Final[str] = ( "nnf.resynthesis.none_shared_page_weights_cache.v1" ) SHARED_PAGE_WEIGHTS_CACHE_LOCATOR_SCHEMA: Final[str] = ( "nnf.resynthesis.none_shared_page_weights_cache_locator.v1" ) SHARED_PAGE_WEIGHTS_CACHE_CONTRACT_REVISION: Final[int] = 1 DEFAULT_SHARED_PAGE_WEIGHTS_CACHE_ROOT: Final[Path] = Path( "/dev/shm/nnf-resynthesis-none-page-weights-v1" ) _WEIGHT_NAMES: Final[tuple[str, ...]] = ( "ffn_mode_t", "gate_t", "up_t", "down_t", "glyph_down_t", "glyph_up_t", "translation_gate_t", "outcome_memory_t", "repair_memory_t", "transfer_memory_t", ) _ARTIFACT_NAMES: Final[frozenset[str]] = frozenset( ("page_ids_t", *_WEIGHT_NAMES) ) _FileIdentity = tuple[int, int, int, int, int] @dataclass(frozen=True) class SharedNoNEPageWeights: """One immutable, single-page, CPU weights reconstruction.""" page_ids_t: torch.Tensor ffn_mode_t: torch.Tensor gate_t: torch.Tensor up_t: torch.Tensor down_t: torch.Tensor glyph_down_t: torch.Tensor glyph_up_t: torch.Tensor translation_gate_t: torch.Tensor outcome_memory_t: torch.Tensor repair_memory_t: torch.Tensor transfer_memory_t: torch.Tensor @dataclass(frozen=True) class SharedPageCacheAuthority: """Immutable source authority supplied after accepted-branch validation.""" page_id: int source_object_sha256: str source_object_bytes: int source_file_identity: _FileIdentity source_format_revision: int materialized_dtype: torch.dtype def external_record_boundary(self) -> dict[str, Any]: """Serialize cache authority only at this explicit I/O boundary.""" _validate_sha256(self.source_object_sha256) if ( self.page_id < 0 or self.source_object_bytes < 1 or self.source_format_revision < 1 or self.source_file_identity[2] != self.source_object_bytes ): raise RuntimeError("shared page cache source authority differs") return { "contractRevision": SHARED_PAGE_WEIGHTS_CACHE_CONTRACT_REVISION, "pageId": self.page_id, "sourceObjectSha256": self.source_object_sha256, "sourceObjectBytes": self.source_object_bytes, "sourceFileIdentity": list(self.source_file_identity), "sourceFormatRevision": self.source_format_revision, "materializedDtype": _dtype_name(self.materialized_dtype), } def _validate_sha256(value: object) -> None: if not ( isinstance(value, str) and len(value) == 64 and all(character in "0123456789abcdef" for character in value) ): raise RuntimeError("shared page cache SHA-256 authority differs") def _dtype_name(dtype: torch.dtype) -> str: value = str(dtype) if not value.startswith("torch."): raise RuntimeError("shared page cache dtype authority differs") return value.removeprefix("torch.") def _canonical_json_bytes(payload: Mapping[str, Any]) -> bytes: return json.dumps( payload, sort_keys=True, separators=(",", ":"), ).encode("utf-8") def _payload_sha256(payload: Mapping[str, Any]) -> str: return hashlib.sha256(_canonical_json_bytes(payload)).hexdigest() def _file_identity(path: Path) -> _FileIdentity: status = path.stat() return ( status.st_dev, status.st_ino, status.st_size, status.st_mtime_ns, status.st_ctime_ns, ) def _stable_file_sha256(path: Path) -> str: identity = _file_identity(path) digest = hashlib.sha256() with path.open("rb") as handle: while block := handle.read(8 * 1024 * 1024): digest.update(block) if _file_identity(path) != identity: raise RuntimeError("shared page cache file changed during hashing") return digest.hexdigest() def _fsync_directory(path: Path) -> None: descriptor = os.open(path, os.O_RDONLY) try: os.fsync(descriptor) finally: os.close(descriptor) def _protected_root(root: Path) -> Path: root.mkdir(mode=0o700, parents=True, exist_ok=True) root.chmod(0o700) resolved = root.resolve() status = resolved.stat() if ( not resolved.is_dir() or status.st_uid != os.geteuid() or status.st_mode & 0o077 ): raise RuntimeError("shared page cache root is not protected") return resolved def _protected_regular_file(path: Path) -> None: status = path.lstat() if ( not path.is_file() or path.is_symlink() or status.st_uid != os.geteuid() or status.st_mode & 0o077 ): raise RuntimeError("shared page cache artifact is not protected") def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None: temporary = path.with_name( f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp" ) temporary.unlink(missing_ok=True) with temporary.open("xb") as handle: os.fchmod(handle.fileno(), 0o600) handle.write( orjson.dumps( payload, option=orjson.OPT_APPEND_NEWLINE | orjson.OPT_SORT_KEYS, ) ) handle.flush() os.fsync(handle.fileno()) os.replace(temporary, path) path.chmod(0o600) _fsync_directory(path.parent) def _record_with_payload_sha256(payload: Mapping[str, Any]) -> dict[str, Any]: record = dict(payload) record["recordPayloadSha256"] = _payload_sha256(record) return record def _validated_record(path: Path, *, schema: str) -> dict[str, Any]: _protected_regular_file(path) try: raw = json.loads(path.read_text(encoding="utf-8")) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: raise RuntimeError("shared page cache record is malformed") from error if not isinstance(raw, dict): raise RuntimeError("shared page cache record is malformed") recorded_sha256 = raw.get("recordPayloadSha256") unsigned = dict(raw) unsigned.pop("recordPayloadSha256", None) if ( raw.get("schema") != schema or not isinstance(recorded_sha256, str) or _payload_sha256(unsigned) != recorded_sha256 ): raise RuntimeError("shared page cache record authority differs") return raw def _weights_tensors( weights: SharedNoNEPageWeights, ) -> tuple[tuple[str, torch.Tensor], ...]: return ( ("page_ids_t", weights.page_ids_t), *((name, getattr(weights, name)) for name in _WEIGHT_NAMES), ) def _validate_weights( weights: SharedNoNEPageWeights, *, authority: SharedPageCacheAuthority, ) -> None: tensors = _weights_tensors(weights) if ( weights.page_ids_t.shape != (1,) or weights.page_ids_t.dtype != torch.long or int(weights.page_ids_t.reshape(())) != authority.page_id or any(tensor.device.type != "cpu" for _name, tensor in tensors) or any(tensor.requires_grad for _name, tensor in tensors) or any(tensor.shape[0] != 1 for _name, tensor in tensors) or any( tensor.dtype != authority.materialized_dtype for name, tensor in tensors if name != "page_ids_t" ) ): raise RuntimeError("shared page cache tensor authority differs") if any( not torch.isfinite(tensor).all() for name, tensor in tensors if name != "page_ids_t" ): raise RuntimeError("shared page cache contains nonfinite weights") def _geometry_record( weights: SharedNoNEPageWeights, ) -> dict[str, dict[str, Any]]: return { name: { "shape": list(tensor.shape), "dtype": _dtype_name(tensor.dtype), } for name, tensor in _weights_tensors(weights) } def _private_weights(weights: SharedNoNEPageWeights) -> SharedNoNEPageWeights: values = { name: tensor.detach().to(device="cpu", copy=True).contiguous() for name, tensor in _weights_tensors(weights) } return SharedNoNEPageWeights(**values) def _artifact_payload( weights: SharedNoNEPageWeights, ) -> dict[str, torch.Tensor]: return { name: tensor.detach().to(device="cpu", copy=True).contiguous() for name, tensor in _weights_tensors(weights) } class SharedPageWeightsCache: """Bounded process-shared cache with per-authority single-flight.""" def __init__( self, root: Path = DEFAULT_SHARED_PAGE_WEIGHTS_CACHE_ROOT, *, budget_bytes: int | None = None, ) -> None: self.root = _protected_root(root) physical_memory_bytes = ( os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") ) filesystem_bytes = ( os.statvfs(self.root).f_blocks * os.statvfs(self.root).f_frsize ) computed_budget = max( os.sysconf("SC_PAGE_SIZE"), min(physical_memory_bytes // 16, filesystem_bytes // 4), ) self.budget_bytes = ( computed_budget if budget_bytes is None else budget_bytes ) if self.budget_bytes < 1: raise ValueError("shared page cache budget must be positive") @staticmethod def authority_prefix(authority: SharedPageCacheAuthority) -> str: return _payload_sha256(authority.external_record_boundary()) def _lock_path(self, authority_prefix: str) -> Path: _validate_sha256(authority_prefix) return self.root / f".{authority_prefix}.lock" def _locator_path(self, authority_prefix: str) -> Path: return self.root / f"{authority_prefix}.locator.json" def _artifact_path(self, authority_prefix: str, cache_key: str) -> Path: _validate_sha256(cache_key) return self.root / f"{authority_prefix}.{cache_key}.safetensors" def _receipt_path(self, authority_prefix: str, cache_key: str) -> Path: _validate_sha256(cache_key) return self.root / f"{authority_prefix}.{cache_key}.receipt.json" def _open_lock(self, authority_prefix: str) -> Any: lock_path = self._lock_path(authority_prefix) descriptor = os.open( lock_path, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600, ) os.fchmod(descriptor, 0o600) return os.fdopen(descriptor, "a+b") def _load_locator( self, authority: SharedPageCacheAuthority, authority_prefix: str, ) -> dict[str, Any] | None: locator_path = self._locator_path(authority_prefix) if not locator_path.exists(): return None locator = _validated_record( locator_path, schema=SHARED_PAGE_WEIGHTS_CACHE_LOCATOR_SCHEMA, ) cache_key = locator.get("cacheKey") receipt_sha256 = locator.get("receiptSha256") if ( locator.get("authority") != authority.external_record_boundary() or locator.get("authorityPrefix") != authority_prefix ): raise RuntimeError("shared page cache locator authority differs") _validate_sha256(cache_key) _validate_sha256(receipt_sha256) return locator def _load_hit( self, *, authority: SharedPageCacheAuthority, authority_prefix: str, locator: Mapping[str, Any], ) -> SharedNoNEPageWeights | None: cache_key = str(locator["cacheKey"]) artifact_path = self._artifact_path(authority_prefix, cache_key) receipt_path = self._receipt_path(authority_prefix, cache_key) if not artifact_path.exists() or not receipt_path.exists(): return None if _stable_file_sha256(receipt_path) != locator["receiptSha256"]: raise RuntimeError("shared page cache receipt hash differs") receipt = _validated_record( receipt_path, schema=SHARED_PAGE_WEIGHTS_CACHE_SCHEMA, ) artifact = receipt.get("artifact") geometry = receipt.get("geometry") if ( receipt.get("authority") != authority.external_record_boundary() or receipt.get("authorityPrefix") != authority_prefix or receipt.get("cacheKey") != cache_key or receipt.get("weightsOnly") is not True or receipt.get("optimizerStateStored") is not False or receipt.get("candidateStateStored") is not False or not isinstance(artifact, dict) or artifact.get("path") != artifact_path.name or not isinstance(artifact.get("bytes"), int) or not isinstance(geometry, dict) ): raise RuntimeError("shared page cache receipt authority differs") artifact_sha256 = artifact.get("sha256") geometry_sha256 = receipt.get("geometrySha256") _validate_sha256(artifact_sha256) _validate_sha256(geometry_sha256) _protected_regular_file(artifact_path) if ( artifact_path.stat().st_size != artifact["bytes"] or _stable_file_sha256(artifact_path) != artifact_sha256 ): raise RuntimeError("shared page cache artifact hash differs") identity = _file_identity(artifact_path) with safe_open( # type: ignore[no-untyped-call] str(artifact_path), framework="pt", device="cpu", ) as handle: if set(handle.keys()) != _ARTIFACT_NAMES: raise RuntimeError("shared page cache tensor set differs") loaded = SharedNoNEPageWeights( **{ name: handle.get_tensor(name) for name in ("page_ids_t", *_WEIGHT_NAMES) } ) if _file_identity(artifact_path) != identity: raise RuntimeError( "shared page cache artifact changed during materialization" ) _validate_weights(loaded, authority=authority) rebuilt_geometry = _geometry_record(loaded) if ( rebuilt_geometry != geometry or _payload_sha256(rebuilt_geometry) != geometry_sha256 ): raise RuntimeError("shared page cache geometry differs") key_payload = { "authority": authority.external_record_boundary(), "artifactSha256": artifact_sha256, "geometrySha256": geometry_sha256, } if _payload_sha256(key_payload) != cache_key: raise RuntimeError("shared page cache content key differs") os.utime(self._locator_path(authority_prefix), None) return _private_weights(loaded) def _publish( self, *, authority: SharedPageCacheAuthority, authority_prefix: str, weights: SharedNoNEPageWeights, ) -> SharedNoNEPageWeights: _validate_weights(weights, authority=authority) geometry = _geometry_record(weights) geometry_sha256 = _payload_sha256(geometry) temporary = self.root / ( f".{authority_prefix}.{os.getpid()}." f"{threading.get_ident()}.{time.time_ns()}.tmp" ) save_file(_artifact_payload(weights), str(temporary)) temporary.chmod(0o600) with temporary.open("rb") as handle: os.fsync(handle.fileno()) artifact_sha256 = _stable_file_sha256(temporary) key_payload = { "authority": authority.external_record_boundary(), "artifactSha256": artifact_sha256, "geometrySha256": geometry_sha256, } cache_key = _payload_sha256(key_payload) artifact_path = self._artifact_path(authority_prefix, cache_key) receipt_path = self._receipt_path(authority_prefix, cache_key) locator_path = self._locator_path(authority_prefix) pair_exists = artifact_path.exists() and receipt_path.exists() if (artifact_path.exists() or receipt_path.exists()) and not pair_exists: artifact_path.unlink(missing_ok=True) receipt_path.unlink(missing_ok=True) _fsync_directory(self.root) if pair_exists: temporary.unlink(missing_ok=True) else: os.replace(temporary, artifact_path) artifact_path.chmod(0o400) _fsync_directory(self.root) receipt = _record_with_payload_sha256( { "schema": SHARED_PAGE_WEIGHTS_CACHE_SCHEMA, "authority": authority.external_record_boundary(), "authorityPrefix": authority_prefix, "cacheKey": cache_key, "geometry": geometry, "geometrySha256": geometry_sha256, "artifact": { "path": artifact_path.name, "sha256": artifact_sha256, "bytes": artifact_path.stat().st_size, }, "weightsOnly": True, "optimizerStateStored": False, "candidateStateStored": False, "acceptedSourceAuthorityRequiredOnEveryAccess": True, } ) _atomic_json(receipt_path, receipt) receipt_sha256 = _stable_file_sha256(receipt_path) locator = _record_with_payload_sha256( { "schema": SHARED_PAGE_WEIGHTS_CACHE_LOCATOR_SCHEMA, "authority": authority.external_record_boundary(), "authorityPrefix": authority_prefix, "cacheKey": cache_key, "receiptSha256": receipt_sha256, } ) _atomic_json(locator_path, locator) loaded = self._load_hit( authority=authority, authority_prefix=authority_prefix, locator=locator, ) if loaded is None: raise RuntimeError("shared page cache publication is incomplete") return loaded def load_or_build( self, *, authority: SharedPageCacheAuthority, builder: Callable[[], SharedNoNEPageWeights], ) -> SharedNoNEPageWeights: """Load or publish one exact weight row under a per-key process lock.""" authority_prefix = self.authority_prefix(authority) with self._open_lock(authority_prefix) as handle: fcntl.flock(handle.fileno(), fcntl.LOCK_EX) try: locator = self._load_locator(authority, authority_prefix) loaded = ( self._load_hit( authority=authority, authority_prefix=authority_prefix, locator=locator, ) if locator is not None else None ) if loaded is None: if locator is not None: stale_cache_key = locator.get("cacheKey") _validate_sha256(stale_cache_key) assert isinstance(stale_cache_key, str) self._locator_path(authority_prefix).unlink( missing_ok=True ) self._receipt_path( authority_prefix, stale_cache_key, ).unlink(missing_ok=True) self._artifact_path( authority_prefix, stale_cache_key, ).unlink(missing_ok=True) _fsync_directory(self.root) loaded = self._publish( authority=authority, authority_prefix=authority_prefix, weights=builder(), ) finally: fcntl.flock(handle.fileno(), fcntl.LOCK_UN) # Eviction owns a separate global lock and begins only after this key's # single-flight lock is released. Returned tensors are private copies. self._evict_to_budget() return loaded def _evict_to_budget(self) -> None: eviction_path = self.root / ".eviction.lock" descriptor = os.open( eviction_path, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600, ) os.fchmod(descriptor, 0o600) with os.fdopen(descriptor, "a+b") as eviction_handle: fcntl.flock(eviction_handle.fileno(), fcntl.LOCK_EX) try: artifacts = tuple(self.root.glob("*.safetensors")) resident_bytes = sum(path.stat().st_size for path in artifacts) if resident_bytes <= self.budget_bytes: return locators = sorted( self.root.glob("*.locator.json"), key=lambda path: path.stat().st_mtime_ns, ) for locator_path in locators: if resident_bytes <= self.budget_bytes: break authority_prefix = locator_path.name.removesuffix( ".locator.json" ) try: lock_handle = self._open_lock(authority_prefix) except (OSError, RuntimeError): continue with lock_handle: try: fcntl.flock( lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB, ) except BlockingIOError: continue try: locator = _validated_record( locator_path, schema=( SHARED_PAGE_WEIGHTS_CACHE_LOCATOR_SCHEMA ), ) cache_key = locator.get("cacheKey") _validate_sha256(cache_key) assert isinstance(cache_key, str) artifact_path = self._artifact_path( authority_prefix, cache_key, ) receipt_path = self._receipt_path( authority_prefix, cache_key, ) artifact_bytes = ( artifact_path.stat().st_size if artifact_path.is_file() else 0 ) locator_path.unlink(missing_ok=True) receipt_path.unlink(missing_ok=True) artifact_path.unlink(missing_ok=True) resident_bytes -= artifact_bytes _fsync_directory(self.root) finally: fcntl.flock( lock_handle.fileno(), fcntl.LOCK_UN, ) finally: fcntl.flock(eviction_handle.fileno(), fcntl.LOCK_UN)