| """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( |
| 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) |
| |
| |
| 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) |
|
|
