| """Content-addressed frozen-parent prefill caches. |
| |
| The frozen Resynthesis parent is identical across optimizer steps. When a decode arm |
| resets KV state and re-runs the same prompt prefill, this cache returns the |
| prior ``ResynthesisParentForward`` tensors from CPU-pinned storage instead of repeating |
| the 4M tiled parent forward. |
| |
| Training uses the narrower ``FrozenBackbonePrefillPacket`` contract below. |
| Only the frozen decoder's final hidden state and tiled summaries are durable; |
| the additive head, historical RBO, and both parent/current Fabric graphs remain |
| live on every replay. Its disk lifecycle intentionally matches the packed |
| token sidecars: content-addressed object roots, one nonblocking ownership lock, |
| partial artifacts, a durable progress frontier, fsync, atomic promotion, and a |
| manifest published last. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import fcntl |
| import hashlib |
| import json |
| import os |
| import stat |
| import threading |
| from collections import OrderedDict |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Final |
|
|
| import torch |
|
|
|
|
| FROZEN_BACKBONE_AUTHORITY_SCHEMA_V1: Final = ( |
| "nnf.resynthesis.frozen_backbone_prefill_authority.v1" |
| ) |
| FROZEN_BACKBONE_AUTHORITY_SCHEMA: Final = ( |
| "nnf.resynthesis.frozen_backbone_prefill_authority.v2" |
| ) |
| FROZEN_BACKBONE_PACKET_SCHEMA: Final = ( |
| "nnf.resynthesis.frozen_backbone_prefill_packet.v1" |
| ) |
| FROZEN_BACKBONE_SIDECAR_SCHEMA: Final = ( |
| "nnf.resynthesis.frozen_backbone_prefill_sidecar.v1" |
| ) |
| FROZEN_BACKBONE_SIDECAR_PROGRESS_SCHEMA: Final = ( |
| "nnf.resynthesis.frozen_backbone_prefill_sidecar.progress.v1" |
| ) |
| _FROZEN_BACKBONE_DTYPES: Final = { |
| "float16": torch.float16, |
| "bfloat16": torch.bfloat16, |
| "float32": torch.float32, |
| } |
| _FROZEN_BACKBONE_SOURCE_DIAGNOSTIC_FIELDS: Final = frozenset( |
| { |
| "parentSourceBundleSha256", |
| "sourceHashDiagnosticOnly", |
| "sourceHashAffectsExecution", |
| } |
| ) |
| _FROZEN_BACKBONE_VALIDATION_CACHE_ENTRIES: Final = 8_192 |
|
|
|
|
| @dataclass(frozen=True) |
| class _FrozenBackboneFileStamp: |
| """Kernel-owned identity observed at one immutable-file boundary.""" |
|
|
| device: int |
| inode: int |
| mode: int |
| size: int |
| mtime_ns: int |
| ctime_ns: int |
|
|
|
|
| @dataclass(frozen=True) |
| class _FrozenBackboneManifestValidation: |
| """Exact content proof reusable while every immutable file stamp is stable.""" |
|
|
| manifest: dict[str, object] |
| manifest_sha256: str |
| artifact_sha256s: tuple[str, str, str] |
| file_stamps: tuple[ |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| ] |
|
|
|
|
| _LEGACY_FROZEN_BACKBONE_INDEX_LOCK: Final = threading.Lock() |
| _LEGACY_FROZEN_BACKBONE_INDEXES: OrderedDict[ |
| Path, |
| dict[str, tuple[Path, ...]], |
| ] = OrderedDict() |
| _FROZEN_BACKBONE_VALIDATION_CACHE_LOCK: Final = threading.Lock() |
| _FROZEN_BACKBONE_VALIDATION_CACHE: OrderedDict[ |
| tuple[int, Path, str], |
| _FrozenBackboneManifestValidation, |
| ] = OrderedDict() |
|
|
|
|
| def _is_sha256_hex(value: object) -> bool: |
| return ( |
| isinstance(value, str) |
| and len(value) == 64 |
| and all(character in "0123456789abcdef" for character in value) |
| ) |
|
|
|
|
| def _canonical_json_bytes(value: object) -> bytes: |
| return json.dumps( |
| value, |
| sort_keys=True, |
| separators=(",", ":"), |
| ensure_ascii=True, |
| ).encode("utf-8") |
|
|
|
|
| def _string_object_mapping_boundary( |
| value: object, |
| ) -> dict[str, object] | None: |
| """Return a typed JSON-object mapping without coercing malformed keys.""" |
|
|
| if not isinstance(value, dict): |
| return None |
| result: dict[str, object] = {} |
| for key, item in value.items(): |
| if not isinstance(key, str): |
| return None |
| result[key] = item |
| return result |
|
|
|
|
| def _source_independent_authority_record_boundary( |
| value: object, |
| ) -> dict[str, object] | None: |
| """Normalize mutable source observations and the parent artifact alias.""" |
|
|
| record = _string_object_mapping_boundary(value) |
| if record is None: |
| return None |
| for diagnostic_field in _FROZEN_BACKBONE_SOURCE_DIAGNOSTIC_FIELDS: |
| record.pop(diagnostic_field, None) |
| parent_artifact_sha256 = record.get("parentModelArtifactSha256") |
| if _is_sha256_hex(parent_artifact_sha256): |
| record["parentCheckpointId"] = ( |
| f"resynthesis-native-parent:{parent_artifact_sha256}" |
| ) |
| if record.get("schema") == FROZEN_BACKBONE_AUTHORITY_SCHEMA_V1: |
| record["schema"] = FROZEN_BACKBONE_AUTHORITY_SCHEMA |
| return record |
|
|
|
|
| def _atomic_json_boundary(path: Path, payload: dict[str, object]) -> None: |
| """Publish one external-I/O record atomically and fsync its directory.""" |
|
|
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") |
| with temporary.open("wb") as handle: |
| handle.write(_canonical_json_bytes(payload)) |
| handle.write(b"\n") |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temporary, path) |
| directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) |
| try: |
| os.fsync(directory_fd) |
| finally: |
| os.close(directory_fd) |
|
|
|
|
| def _file_sha256_boundary(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| while chunk := handle.read(8 * 1024 * 1024): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def _tensor_storage_bytes_boundary(value: torch.Tensor) -> bytes: |
| """Serialize one contiguous CPU tensor without dtype conversion.""" |
|
|
| cpu = value.detach().to(device="cpu").contiguous() |
| return cpu.view(torch.uint8).numpy().tobytes() |
|
|
|
|
| def _write_all_fd_boundary(descriptor: int, payload: bytes) -> None: |
| """Write one complete staged artifact at the external-I/O boundary.""" |
|
|
| remaining = memoryview(payload) |
| while remaining: |
| written = os.write(descriptor, remaining) |
| if written < 1: |
| raise OSError("frozen-backbone staged artifact write made no progress") |
| remaining = remaining[written:] |
|
|
|
|
| def _synchronize_open_files_boundary(descriptors: tuple[int, ...]) -> None: |
| """Flush one already-written outer wave as a filesystem transaction group.""" |
|
|
| for descriptor in descriptors: |
| os.fdatasync(descriptor) |
|
|
|
|
| def _synchronize_directories_boundary(paths: tuple[Path, ...]) -> None: |
| """Persist manifest-last renames from deepest object root to cache root.""" |
|
|
| unique_paths = sorted( |
| {path.expanduser().resolve() for path in paths}, |
| key=lambda path: (len(path.parts), str(path)), |
| reverse=True, |
| ) |
| for path in unique_paths: |
| descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) |
| try: |
| os.fsync(descriptor) |
| finally: |
| os.close(descriptor) |
|
|
|
|
| def frozen_backbone_prompt_mask_sha256_boundary( |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| ) -> str: |
| """Hash exact prompt/mask geometry and bytes at the CPU staging boundary.""" |
|
|
| if ( |
| input_ids.device.type != "cpu" |
| or attention_mask.device.type != "cpu" |
| or input_ids.ndim != 2 |
| or input_ids.shape[1] < 1 |
| or attention_mask.shape != input_ids.shape |
| or input_ids.dtype not in {torch.int32, torch.int64} |
| or attention_mask.dtype not in { |
| torch.bool, |
| torch.int32, |
| torch.int64, |
| } |
| ): |
| raise ValueError( |
| "frozen-backbone prompt/mask identity requires CPU [batch, sequence]" |
| ) |
| digest = hashlib.sha256() |
| digest.update(b"nnf.resynthesis.frozen_backbone_prompt_mask.v1\x00") |
| for value in (input_ids, attention_mask): |
| digest.update(str(value.dtype).encode("ascii")) |
| digest.update(b"\x00") |
| digest.update(_canonical_json_bytes(tuple(value.shape))) |
| digest.update(b"\x00") |
| digest.update(_tensor_storage_bytes_boundary(value)) |
| return digest.hexdigest() |
|
|
|
|
| def frozen_backbone_position_policy_sha256_boundary( |
| *, |
| final_hidden_adapter_id: str, |
| dual_chunk_pretrain_length: int, |
| dual_chunk_local_size: int, |
| pretrained_rope_band_tokens: int, |
| prefill_tile_tokens: int, |
| prefill_summaries_per_tile: int, |
| ) -> str: |
| """Bind every position/tile rule that can change frozen decoder features.""" |
|
|
| policy = { |
| "schema": "nnf.resynthesis.frozen_backbone_position_policy.v1", |
| "finalHiddenAdapterId": final_hidden_adapter_id, |
| "dualChunkPretrainLength": dual_chunk_pretrain_length, |
| "dualChunkLocalSize": dual_chunk_local_size, |
| "pretrainedRopeBandTokens": pretrained_rope_band_tokens, |
| "prefillTileTokens": prefill_tile_tokens, |
| "prefillSummariesPerTile": prefill_summaries_per_tile, |
| } |
| if ( |
| not final_hidden_adapter_id |
| or any( |
| isinstance(value, bool) or not isinstance(value, int) or value < 1 |
| for value in ( |
| dual_chunk_pretrain_length, |
| dual_chunk_local_size, |
| pretrained_rope_band_tokens, |
| prefill_tile_tokens, |
| prefill_summaries_per_tile, |
| ) |
| ) |
| ): |
| raise ValueError("frozen-backbone position policy is malformed") |
| return hashlib.sha256(_canonical_json_bytes(policy)).hexdigest() |
|
|
|
|
| @dataclass(frozen=True) |
| class FrozenBackboneCacheAuthority: |
| """Complete immutable identity for one prompt-only decoder feature object.""" |
|
|
| qualified_work_id: str |
| prompt_sha256: str |
| prompt_mask_sha256: str |
| parent_checkpoint_id: str |
| parent_manifest_payload_sha256: str |
| parent_model_artifact_sha256: str |
| parent_source_bundle_sha256: str = field(compare=False) |
| position_policy_sha256: str |
| token_dtype: str |
| feature_dtype: str |
| hidden_size: int |
| summary_positions: int |
|
|
| def _record_boundary( |
| self, |
| *, |
| schema: str, |
| parent_checkpoint_id: str, |
| ) -> dict[str, object]: |
| """Serialize one validated cache-authority representation. |
| |
| Source-tree hashes intentionally do not participate. Training follows |
| a moving source tree; launch/status receipts retain that observation, |
| while reusable frozen-parent features remain bound to the immutable |
| model artifact, checkpoint, prompt geometry, and decoder policy. |
| """ |
|
|
| hash_fields = ( |
| self.prompt_sha256, |
| self.prompt_mask_sha256, |
| self.parent_manifest_payload_sha256, |
| self.parent_model_artifact_sha256, |
| self.position_policy_sha256, |
| ) |
| if ( |
| not self.qualified_work_id |
| or not parent_checkpoint_id |
| or not all(_is_sha256_hex(value) for value in hash_fields) |
| or self.token_dtype not in {"int32", "int64"} |
| or self.feature_dtype not in _FROZEN_BACKBONE_DTYPES |
| or isinstance(self.hidden_size, bool) |
| or self.hidden_size < 1 |
| or isinstance(self.summary_positions, bool) |
| or self.summary_positions < 1 |
| ): |
| raise ValueError("frozen-backbone cache authority is malformed") |
| return { |
| "schema": schema, |
| "qualifiedWorkId": self.qualified_work_id, |
| "qualifiedWorkIdSha256": hashlib.sha256( |
| self.qualified_work_id.encode("utf-8") |
| ).hexdigest(), |
| "promptSha256": self.prompt_sha256, |
| "promptMaskSha256": self.prompt_mask_sha256, |
| "parentCheckpointId": parent_checkpoint_id, |
| "parentManifestPayloadSha256": ( |
| self.parent_manifest_payload_sha256 |
| ), |
| "parentModelArtifactSha256": self.parent_model_artifact_sha256, |
| "positionPolicySha256": self.position_policy_sha256, |
| "tokenDtype": self.token_dtype, |
| "featureDtype": self.feature_dtype, |
| "hiddenSize": self.hidden_size, |
| "summaryPositions": self.summary_positions, |
| "backboneOnly": True, |
| "additiveHeadCached": False, |
| "historicalRboCached": False, |
| "fabricCached": False, |
| "targetEnteredForward": False, |
| } |
|
|
| def record_boundary(self) -> dict[str, object]: |
| """Return the canonical Resynthesis authority for new objects.""" |
|
|
| return self._record_boundary( |
| schema=FROZEN_BACKBONE_AUTHORITY_SCHEMA, |
| parent_checkpoint_id=( |
| "resynthesis-native-parent:" |
| f"{self.parent_model_artifact_sha256}" |
| ), |
| ) |
|
|
| def legacy_v1_record_boundary(self) -> dict[str, object] | None: |
| """Return the historical source-bound record for a direct legacy probe. |
| |
| A malformed or changed source observation cannot block current cache |
| identity. When the reported digest still matches the historical |
| object, this record gives lookup a constant-time read-only fast path. |
| A per-root normalized manifest index handles all other legacy sources. |
| """ |
|
|
| if not _is_sha256_hex(self.parent_source_bundle_sha256): |
| return None |
| record = self._record_boundary( |
| schema=FROZEN_BACKBONE_AUTHORITY_SCHEMA_V1, |
| parent_checkpoint_id=self.parent_checkpoint_id, |
| ) |
| record["parentSourceBundleSha256"] = ( |
| self.parent_source_bundle_sha256 |
| ) |
| return record |
|
|
| def legacy_v1_sha256_boundary(self) -> str | None: |
| record = self.legacy_v1_record_boundary() |
| if record is None: |
| return None |
| return hashlib.sha256(_canonical_json_bytes(record)).hexdigest() |
|
|
| def sha256_boundary(self) -> str: |
| return hashlib.sha256( |
| _canonical_json_bytes(self.record_boundary()) |
| ).hexdigest() |
|
|
| def digest_tensors_boundary( |
| self, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Return typed row identity tensors for the active packet.""" |
|
|
| cache_key_t = torch.tensor( |
| tuple(bytes.fromhex(self.sha256_boundary())), |
| dtype=torch.uint8, |
| ).reshape(1, 32) |
| work_id_t = torch.tensor( |
| tuple( |
| hashlib.sha256( |
| self.qualified_work_id.encode("utf-8") |
| ).digest() |
| ), |
| dtype=torch.uint8, |
| ).reshape(1, 32) |
| prompt_mask_t = torch.tensor( |
| tuple(bytes.fromhex(self.prompt_mask_sha256)), |
| dtype=torch.uint8, |
| ).reshape(1, 32) |
| return cache_key_t, work_id_t, prompt_mask_t |
|
|
|
|
| @dataclass(frozen=True) |
| class FrozenBackbonePrefillPacket: |
| """Raw frozen-decoder features; all additive/RBO/Fabric work stays live.""" |
|
|
| final_hidden_t: torch.Tensor |
| summary_hidden_t: torch.Tensor |
| summary_mask_t: torch.Tensor |
| input_positions_t: torch.Tensor |
| cache_key_sha256_t: torch.Tensor |
| work_id_sha256_t: torch.Tensor |
| prompt_mask_sha256_t: torch.Tensor |
|
|
| def validate_boundary(self) -> None: |
| batch_size = self.final_hidden_t.shape[0] |
| if ( |
| self.final_hidden_t.ndim != 3 |
| or self.final_hidden_t.shape[1] != 1 |
| or self.summary_hidden_t.ndim != 3 |
| or self.summary_hidden_t.shape[0] != batch_size |
| or self.summary_hidden_t.shape[1] < 1 |
| or self.summary_hidden_t.shape[2] |
| != self.final_hidden_t.shape[2] |
| or self.summary_hidden_t.dtype != self.final_hidden_t.dtype |
| or self.summary_hidden_t.device != self.final_hidden_t.device |
| or self.summary_mask_t.shape |
| != self.summary_hidden_t.shape[:2] |
| or self.summary_mask_t.dtype != torch.bool |
| or self.summary_mask_t.device != self.final_hidden_t.device |
| or self.input_positions_t.shape != (batch_size,) |
| or self.input_positions_t.dtype != torch.long |
| or self.input_positions_t.device != self.final_hidden_t.device |
| or any( |
| value.shape != (batch_size, 32) |
| or value.dtype != torch.uint8 |
| or value.device != self.final_hidden_t.device |
| for value in ( |
| self.cache_key_sha256_t, |
| self.work_id_sha256_t, |
| self.prompt_mask_sha256_t, |
| ) |
| ) |
| ): |
| raise ValueError("frozen-backbone prefill packet geometry differs") |
| torch._assert_async( |
| self.summary_mask_t.all(), |
| "frozen-backbone replay requires exact unpadded summary geometry", |
| ) |
| torch._assert_async( |
| self.input_positions_t.gt(0).all(), |
| "frozen-backbone input positions must be positive", |
| ) |
|
|
| @classmethod |
| def from_features_boundary( |
| cls, |
| *, |
| authority: FrozenBackboneCacheAuthority, |
| final_hidden_t: torch.Tensor, |
| summary_hidden_t: torch.Tensor, |
| input_positions_t: torch.Tensor, |
| ) -> FrozenBackbonePrefillPacket: |
| expected_feature_dtype = _FROZEN_BACKBONE_DTYPES.get( |
| authority.feature_dtype |
| ) |
| if ( |
| final_hidden_t.shape |
| != (1, 1, authority.hidden_size) |
| or summary_hidden_t.shape |
| != ( |
| 1, |
| authority.summary_positions, |
| authority.hidden_size, |
| ) |
| or final_hidden_t.dtype != expected_feature_dtype |
| or summary_hidden_t.dtype != expected_feature_dtype |
| or input_positions_t.numel() != 1 |
| ): |
| raise ValueError( |
| "one frozen-backbone sidecar object geometry differs" |
| ) |
| cache_key_t, work_id_t, prompt_mask_t = ( |
| authority.digest_tensors_boundary() |
| ) |
| device = final_hidden_t.device |
| packet = cls( |
| final_hidden_t=final_hidden_t, |
| summary_hidden_t=summary_hidden_t, |
| summary_mask_t=torch.ones( |
| summary_hidden_t.shape[:2], |
| dtype=torch.bool, |
| device=device, |
| ), |
| input_positions_t=input_positions_t.reshape(1).to( |
| device=device, |
| dtype=torch.long, |
| ), |
| cache_key_sha256_t=cache_key_t.to(device=device), |
| work_id_sha256_t=work_id_t.to(device=device), |
| prompt_mask_sha256_t=prompt_mask_t.to(device=device), |
| ) |
| packet.validate_boundary() |
| return packet |
|
|
| @classmethod |
| def stack_boundary( |
| cls, |
| packets: tuple[FrozenBackbonePrefillPacket, ...], |
| ) -> FrozenBackbonePrefillPacket: |
| if not packets: |
| raise ValueError("frozen-backbone packet stack is empty") |
| for packet in packets: |
| packet.validate_boundary() |
| summary_positions = {packet.summary_hidden_t.shape[1] for packet in packets} |
| hidden_sizes = {packet.final_hidden_t.shape[2] for packet in packets} |
| dtypes = {packet.final_hidden_t.dtype for packet in packets} |
| devices = {packet.final_hidden_t.device for packet in packets} |
| if ( |
| len(summary_positions) != 1 |
| or len(hidden_sizes) != 1 |
| or len(dtypes) != 1 |
| or len(devices) != 1 |
| ): |
| raise ValueError( |
| "frozen-backbone rows require one exact summary geometry" |
| ) |
| stacked = cls( |
| final_hidden_t=torch.cat( |
| tuple(packet.final_hidden_t for packet in packets), |
| dim=0, |
| ), |
| summary_hidden_t=torch.cat( |
| tuple(packet.summary_hidden_t for packet in packets), |
| dim=0, |
| ), |
| summary_mask_t=torch.cat( |
| tuple(packet.summary_mask_t for packet in packets), |
| dim=0, |
| ), |
| input_positions_t=torch.cat( |
| tuple(packet.input_positions_t for packet in packets), |
| dim=0, |
| ), |
| cache_key_sha256_t=torch.cat( |
| tuple(packet.cache_key_sha256_t for packet in packets), |
| dim=0, |
| ), |
| work_id_sha256_t=torch.cat( |
| tuple(packet.work_id_sha256_t for packet in packets), |
| dim=0, |
| ), |
| prompt_mask_sha256_t=torch.cat( |
| tuple(packet.prompt_mask_sha256_t for packet in packets), |
| dim=0, |
| ), |
| ) |
| stacked.validate_boundary() |
| return stacked |
|
|
| def to_device_boundary( |
| self, |
| device: torch.device, |
| ) -> FrozenBackbonePrefillPacket: |
| def move(value: torch.Tensor) -> torch.Tensor: |
| |
| |
| |
| return value.detach().to( |
| device=device, |
| non_blocking=device.type == "cuda", |
| copy=value.device == device, |
| ) |
|
|
| packet = FrozenBackbonePrefillPacket( |
| final_hidden_t=move(self.final_hidden_t), |
| summary_hidden_t=move(self.summary_hidden_t), |
| summary_mask_t=move(self.summary_mask_t), |
| input_positions_t=move(self.input_positions_t), |
| cache_key_sha256_t=move(self.cache_key_sha256_t), |
| work_id_sha256_t=move(self.work_id_sha256_t), |
| prompt_mask_sha256_t=move(self.prompt_mask_sha256_t), |
| ) |
| packet.validate_boundary() |
| return packet |
|
|
|
|
| @dataclass(frozen=True) |
| class _FrozenBackboneSidecarPaths: |
| root: Path |
| manifest: Path |
| progress: Path |
| lock: Path |
| final_hidden: Path |
| summary_hidden: Path |
| input_positions: Path |
|
|
| def artifact_paths_boundary(self) -> tuple[Path, Path, Path]: |
| return ( |
| self.final_hidden, |
| self.summary_hidden, |
| self.input_positions, |
| ) |
|
|
|
|
| def _frozen_backbone_file_stamp_boundary( |
| path: Path, |
| ) -> _FrozenBackboneFileStamp | None: |
| """Observe one regular file without reading or trusting its content.""" |
|
|
| try: |
| observed = path.stat() |
| except FileNotFoundError: |
| return None |
| if not stat.S_ISREG(observed.st_mode): |
| return None |
| return _FrozenBackboneFileStamp( |
| device=observed.st_dev, |
| inode=observed.st_ino, |
| mode=observed.st_mode, |
| size=observed.st_size, |
| mtime_ns=observed.st_mtime_ns, |
| ctime_ns=observed.st_ctime_ns, |
| ) |
|
|
|
|
| def _frozen_backbone_sidecar_file_stamps_boundary( |
| paths: _FrozenBackboneSidecarPaths, |
| ) -> ( |
| tuple[ |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| ] |
| | None |
| ): |
| """Return one complete immutable-object stamp, or no reusable object.""" |
|
|
| manifest_stamp = _frozen_backbone_file_stamp_boundary(paths.manifest) |
| artifact_stamps = tuple( |
| _frozen_backbone_file_stamp_boundary(path) |
| for path in paths.artifact_paths_boundary() |
| ) |
| if manifest_stamp is None or any( |
| artifact_stamp is None for artifact_stamp in artifact_stamps |
| ): |
| return None |
| final_stamp, summary_stamp, positions_stamp = artifact_stamps |
| assert final_stamp is not None |
| assert summary_stamp is not None |
| assert positions_stamp is not None |
| return ( |
| manifest_stamp, |
| final_stamp, |
| summary_stamp, |
| positions_stamp, |
| ) |
|
|
|
|
| def _build_legacy_frozen_backbone_manifest_index_boundary( |
| cache_root: Path, |
| ) -> dict[str, tuple[Path, ...]]: |
| """Build one read-only normalized index of legacy source-bound manifests. |
| |
| The index is read-only and contains manifest paths, not copied tensor |
| artifacts. Its key is the authority record after removing only code-source |
| diagnostics. Artifact geometry and hashes are deliberately validated later |
| by the selected sidecar reader, exactly as they are for current objects. |
| """ |
|
|
| objects_root = cache_root.expanduser().resolve() / "objects" |
| if not objects_root.is_dir(): |
| return {} |
| indexed: dict[str, list[Path]] = {} |
| for manifest_path in sorted(objects_root.glob("*/*/manifest.json")): |
| object_root = manifest_path.parent |
| object_sha256 = object_root.name |
| if ( |
| not _is_sha256_hex(object_sha256) |
| or object_root.parent.name != object_sha256[:2] |
| ): |
| continue |
| try: |
| raw = json.loads(manifest_path.read_text(encoding="utf-8")) |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError): |
| |
| |
| continue |
| manifest = _string_object_mapping_boundary(raw) |
| if ( |
| manifest is None |
| or manifest.get("schema") != FROZEN_BACKBONE_SIDECAR_SCHEMA |
| or manifest.get("authoritySha256") != object_sha256 |
| ): |
| continue |
| stored_authority = _string_object_mapping_boundary( |
| manifest.get("authority") |
| ) |
| if ( |
| stored_authority is None |
| or not any( |
| field in stored_authority |
| for field in _FROZEN_BACKBONE_SOURCE_DIAGNOSTIC_FIELDS |
| ) |
| or hashlib.sha256( |
| _canonical_json_bytes(stored_authority) |
| ).hexdigest() |
| != object_sha256 |
| ): |
| continue |
| normalized = _source_independent_authority_record_boundary( |
| stored_authority |
| ) |
| if normalized is None: |
| continue |
| normalized_sha256 = hashlib.sha256( |
| _canonical_json_bytes(normalized) |
| ).hexdigest() |
| indexed.setdefault(normalized_sha256, []).append(manifest_path) |
| return { |
| authority_sha256: tuple(paths) |
| for authority_sha256, paths in indexed.items() |
| } |
|
|
|
|
| def _legacy_frozen_backbone_manifest_index_boundary( |
| cache_root: Path, |
| ) -> dict[str, tuple[Path, ...]]: |
| """Return one single-flight index per resolved root and process.""" |
|
|
| resolved_root = cache_root.expanduser().resolve() |
| with _LEGACY_FROZEN_BACKBONE_INDEX_LOCK: |
| if resolved_root in _LEGACY_FROZEN_BACKBONE_INDEXES: |
| cached = _LEGACY_FROZEN_BACKBONE_INDEXES[resolved_root] |
| _LEGACY_FROZEN_BACKBONE_INDEXES.move_to_end(resolved_root) |
| return cached |
| built = _build_legacy_frozen_backbone_manifest_index_boundary( |
| resolved_root |
| ) |
| _LEGACY_FROZEN_BACKBONE_INDEXES[resolved_root] = built |
| while len(_LEGACY_FROZEN_BACKBONE_INDEXES) > 16: |
| _LEGACY_FROZEN_BACKBONE_INDEXES.popitem(last=False) |
| return built |
|
|
|
|
| class FrozenBackboneFeatureSidecar: |
| """One exact prompt feature object using packed-sidecar durability rules.""" |
|
|
| @staticmethod |
| def _paths_for_authority_boundary( |
| *, |
| cache_root: Path, |
| authority_sha256: str, |
| feature_dtype: str, |
| ) -> _FrozenBackboneSidecarPaths: |
| object_root = ( |
| cache_root.expanduser().resolve() |
| / "objects" |
| / authority_sha256[:2] |
| / authority_sha256 |
| ) |
| return _FrozenBackboneSidecarPaths( |
| root=object_root, |
| manifest=object_root / "manifest.json", |
| progress=object_root / "progress.json", |
| lock=object_root / ".writer.lock", |
| final_hidden=object_root |
| / f"final_hidden.{feature_dtype}.bin", |
| summary_hidden=object_root |
| / f"summary_hidden.{feature_dtype}.bin", |
| input_positions=object_root / "input_positions.i64le", |
| ) |
|
|
| def __init__( |
| self, |
| *, |
| cache_root: Path, |
| authority: FrozenBackboneCacheAuthority, |
| ) -> None: |
| self.authority = authority |
| self.authority_record = authority.record_boundary() |
| self.authority_sha256 = authority.sha256_boundary() |
| self.cache_root = cache_root.expanduser().resolve() |
| self.paths = self._paths_for_authority_boundary( |
| cache_root=self.cache_root, |
| authority_sha256=self.authority_sha256, |
| feature_dtype=authority.feature_dtype, |
| ) |
| self.legacy_authority_record = ( |
| authority.legacy_v1_record_boundary() |
| ) |
| self.legacy_authority_sha256 = ( |
| authority.legacy_v1_sha256_boundary() |
| ) |
| self.legacy_paths = ( |
| self._paths_for_authority_boundary( |
| cache_root=self.cache_root, |
| authority_sha256=self.legacy_authority_sha256, |
| feature_dtype=authority.feature_dtype, |
| ) |
| if self.legacy_authority_sha256 is not None |
| else None |
| ) |
|
|
| def _expected_elements_boundary(self) -> tuple[int, int, int]: |
| return ( |
| self.authority.hidden_size, |
| self.authority.summary_positions * self.authority.hidden_size, |
| 1, |
| ) |
|
|
| def _artifact_record_boundary( |
| self, |
| path: Path, |
| *, |
| elements: int, |
| dtype_name: str, |
| ) -> dict[str, object]: |
| stat = path.stat() |
| return { |
| "path": str(path), |
| "bytes": stat.st_size, |
| "elements": elements, |
| "dtype": dtype_name, |
| "sha256": _file_sha256_boundary(path), |
| "device": stat.st_dev, |
| "inode": stat.st_ino, |
| "mtimeNs": stat.st_mtime_ns, |
| } |
|
|
| def _validation_cache_key_boundary( |
| self, |
| *, |
| paths: _FrozenBackboneSidecarPaths, |
| ) -> tuple[int, Path, str]: |
| """Bind reusable validation to this process session and executable authority.""" |
|
|
| return ( |
| os.getpid(), |
| paths.manifest, |
| self.authority_sha256, |
| ) |
|
|
| def _cached_manifest_validation_boundary( |
| self, |
| *, |
| paths: _FrozenBackboneSidecarPaths, |
| ) -> dict[str, object] | None: |
| """Reuse exact digests only while all kernel-owned file stamps remain stable.""" |
|
|
| cache_key = self._validation_cache_key_boundary(paths=paths) |
| observed_stamps = _frozen_backbone_sidecar_file_stamps_boundary(paths) |
| with _FROZEN_BACKBONE_VALIDATION_CACHE_LOCK: |
| cached = _FROZEN_BACKBONE_VALIDATION_CACHE.get(cache_key) |
| if ( |
| cached is not None |
| and observed_stamps is not None |
| and cached.file_stamps == observed_stamps |
| ): |
| _FROZEN_BACKBONE_VALIDATION_CACHE.move_to_end(cache_key) |
| return cached.manifest |
| if cached is not None: |
| del _FROZEN_BACKBONE_VALIDATION_CACHE[cache_key] |
| return None |
|
|
| def _cache_manifest_validation_boundary( |
| self, |
| *, |
| paths: _FrozenBackboneSidecarPaths, |
| manifest: dict[str, object], |
| manifest_sha256: str, |
| artifact_sha256s: tuple[str, str, str], |
| file_stamps: tuple[ |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| _FrozenBackboneFileStamp, |
| ], |
| ) -> None: |
| """Retain a bounded process-session proof for later lookup boundaries.""" |
|
|
| cache_key = self._validation_cache_key_boundary(paths=paths) |
| validation = _FrozenBackboneManifestValidation( |
| manifest=manifest, |
| manifest_sha256=manifest_sha256, |
| artifact_sha256s=artifact_sha256s, |
| file_stamps=file_stamps, |
| ) |
| with _FROZEN_BACKBONE_VALIDATION_CACHE_LOCK: |
| _FROZEN_BACKBONE_VALIDATION_CACHE[cache_key] = validation |
| _FROZEN_BACKBONE_VALIDATION_CACHE.move_to_end(cache_key) |
| while ( |
| len(_FROZEN_BACKBONE_VALIDATION_CACHE) |
| > _FROZEN_BACKBONE_VALIDATION_CACHE_ENTRIES |
| ): |
| _FROZEN_BACKBONE_VALIDATION_CACHE.popitem(last=False) |
|
|
| def _validated_manifest_boundary( |
| self, |
| *, |
| paths: _FrozenBackboneSidecarPaths, |
| ) -> dict[str, object]: |
| cached = self._cached_manifest_validation_boundary(paths=paths) |
| if cached is not None: |
| return cached |
| file_stamps_before = _frozen_backbone_sidecar_file_stamps_boundary( |
| paths |
| ) |
| if file_stamps_before is None: |
| raise RuntimeError("frozen-backbone sidecar artifacts are absent") |
| manifest_bytes = paths.manifest.read_bytes() |
| raw = json.loads(manifest_bytes) |
| manifest = _string_object_mapping_boundary(raw) |
| if manifest is None: |
| raise RuntimeError("frozen-backbone sidecar manifest is not an object") |
| stored_authority = _string_object_mapping_boundary( |
| manifest.get("authority") |
| ) |
| stored_authority_sha256 = manifest.get("authoritySha256") |
| normalized_authority = ( |
| _source_independent_authority_record_boundary(stored_authority) |
| ) |
| if ( |
| manifest.get("schema") != FROZEN_BACKBONE_SIDECAR_SCHEMA |
| or not isinstance(stored_authority_sha256, str) |
| or not _is_sha256_hex(stored_authority_sha256) |
| or stored_authority is None |
| or hashlib.sha256( |
| _canonical_json_bytes(stored_authority) |
| ).hexdigest() |
| != stored_authority_sha256 |
| or paths.root.name != stored_authority_sha256 |
| or paths.root.parent.name != stored_authority_sha256[:2] |
| or normalized_authority != self.authority_record |
| or manifest.get("rows") != 1 |
| or manifest.get("backboneOnly") is not True |
| or manifest.get("targetEnteredForward") is not False |
| or manifest.get("additiveHeadCached") is not False |
| or manifest.get("historicalRboCached") is not False |
| or manifest.get("fabricCached") is not False |
| ): |
| raise RuntimeError("frozen-backbone sidecar authority differs") |
| artifacts = manifest.get("artifacts") |
| if not isinstance(artifacts, dict): |
| raise RuntimeError("frozen-backbone sidecar artifacts are absent") |
| expected_elements = self._expected_elements_boundary() |
| expected_paths = paths.artifact_paths_boundary() |
| expected_dtypes = ( |
| self.authority.feature_dtype, |
| self.authority.feature_dtype, |
| "int64", |
| ) |
| names = ("finalHidden", "summaryHidden", "inputPositions") |
| artifact_sha256s: list[str] = [] |
| for name, path, elements, dtype_name in zip( |
| names, |
| expected_paths, |
| expected_elements, |
| expected_dtypes, |
| strict=True, |
| ): |
| record = artifacts.get(name) |
| stat = path.stat() if path.is_file() else None |
| recorded_sha256 = ( |
| record.get("sha256") |
| if isinstance(record, dict) |
| else None |
| ) |
| observed_sha256 = ( |
| _file_sha256_boundary(path) |
| if stat is not None |
| and _is_sha256_hex(recorded_sha256) |
| else None |
| ) |
| if ( |
| not isinstance(record, dict) |
| or record.get("path") != str(path) |
| or record.get("elements") != elements |
| or record.get("dtype") != dtype_name |
| or stat is None |
| or record.get("bytes") != stat.st_size |
| or record.get("device") != stat.st_dev |
| or record.get("inode") != stat.st_ino |
| or record.get("mtimeNs") != stat.st_mtime_ns |
| or recorded_sha256 != observed_sha256 |
| ): |
| raise RuntimeError( |
| f"frozen-backbone {name} artifact authority differs" |
| ) |
| assert isinstance(recorded_sha256, str) |
| artifact_sha256s.append(recorded_sha256) |
| file_stamps_after = _frozen_backbone_sidecar_file_stamps_boundary(paths) |
| if file_stamps_after != file_stamps_before: |
| raise RuntimeError( |
| "frozen-backbone sidecar changed during content validation" |
| ) |
| final_sha256, summary_sha256, positions_sha256 = artifact_sha256s |
| self._cache_manifest_validation_boundary( |
| paths=paths, |
| manifest=manifest, |
| manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(), |
| artifact_sha256s=( |
| final_sha256, |
| summary_sha256, |
| positions_sha256, |
| ), |
| file_stamps=file_stamps_after, |
| ) |
| return manifest |
|
|
| def _lookup_paths_boundary( |
| self, |
| *, |
| paths: _FrozenBackboneSidecarPaths, |
| ) -> FrozenBackbonePrefillPacket | None: |
| if not paths.manifest.is_file(): |
| return None |
| self._validated_manifest_boundary(paths=paths) |
| dtype = _FROZEN_BACKBONE_DTYPES[self.authority.feature_dtype] |
| final_elements, summary_elements, input_elements = ( |
| self._expected_elements_boundary() |
| ) |
| final_hidden_t = torch.from_file( |
| str(paths.final_hidden), |
| shared=False, |
| size=final_elements, |
| dtype=dtype, |
| ).reshape(1, 1, self.authority.hidden_size) |
| summary_hidden_t = torch.from_file( |
| str(paths.summary_hidden), |
| shared=False, |
| size=summary_elements, |
| dtype=dtype, |
| ).reshape( |
| 1, |
| self.authority.summary_positions, |
| self.authority.hidden_size, |
| ) |
| input_positions_t = torch.from_file( |
| str(paths.input_positions), |
| shared=False, |
| size=input_elements, |
| dtype=torch.long, |
| ) |
| return FrozenBackbonePrefillPacket.from_features_boundary( |
| authority=self.authority, |
| final_hidden_t=final_hidden_t, |
| summary_hidden_t=summary_hidden_t, |
| input_positions_t=input_positions_t, |
| ) |
|
|
| def _lookup_current_boundary( |
| self, |
| ) -> FrozenBackbonePrefillPacket | None: |
| return self._lookup_paths_boundary(paths=self.paths) |
|
|
| def lookup_boundary(self) -> FrozenBackbonePrefillPacket | None: |
| """Return a current or legacy exact object without rewriting artifacts.""" |
|
|
| current = self._lookup_current_boundary() |
| if current is not None: |
| return current |
| if ( |
| self.legacy_paths is None |
| or self.legacy_authority_record is None |
| or self.legacy_authority_sha256 is None |
| ): |
| direct_legacy = None |
| else: |
| direct_legacy = self._lookup_paths_boundary( |
| paths=self.legacy_paths, |
| ) |
| if direct_legacy is not None: |
| return direct_legacy |
|
|
| |
| |
| |
| |
| |
| legacy_manifests = ( |
| _legacy_frozen_backbone_manifest_index_boundary( |
| self.cache_root |
| ).get(self.authority_sha256, ()) |
| ) |
| for manifest_path in legacy_manifests: |
| candidate = self._paths_for_authority_boundary( |
| cache_root=self.cache_root, |
| authority_sha256=manifest_path.parent.name, |
| feature_dtype=self.authority.feature_dtype, |
| ) |
| if candidate.root == self.paths.root: |
| continue |
| legacy = self._lookup_paths_boundary(paths=candidate) |
| if legacy is not None: |
| return legacy |
| return None |
|
|
| def _progress_payload_boundary( |
| self, |
| *, |
| completed_rows: int, |
| ) -> dict[str, object]: |
| final_elements, summary_elements, input_elements = ( |
| self._expected_elements_boundary() |
| ) |
| return { |
| "schema": FROZEN_BACKBONE_SIDECAR_PROGRESS_SCHEMA, |
| "authoritySha256": self.authority_sha256, |
| "authority": self.authority_record, |
| "completedRows": completed_rows, |
| "finalHiddenElements": final_elements if completed_rows else 0, |
| "summaryHiddenElements": summary_elements if completed_rows else 0, |
| "inputPositionElements": input_elements if completed_rows else 0, |
| "targetEnteredForward": False, |
| } |
|
|
| def _validated_progress_boundary(self) -> dict[str, object]: |
| if not self.paths.progress.is_file(): |
| return self._progress_payload_boundary(completed_rows=0) |
| raw = json.loads(self.paths.progress.read_text(encoding="utf-8")) |
| if not isinstance(raw, dict): |
| raise RuntimeError("frozen-backbone progress is not an object") |
| expected_empty = self._progress_payload_boundary(completed_rows=0) |
| expected_complete = self._progress_payload_boundary(completed_rows=1) |
| if raw != expected_empty and raw != expected_complete: |
| raise RuntimeError("frozen-backbone progress authority differs") |
| return { |
| str(key): value |
| for key, value in raw.items() |
| } |
|
|
| @staticmethod |
| def _write_tensor_boundary(path: Path, value: torch.Tensor) -> None: |
| with path.open("wb") as handle: |
| handle.write(_tensor_storage_bytes_boundary(value)) |
| handle.flush() |
| os.fsync(handle.fileno()) |
|
|
| def _validate_packet_boundary( |
| self, |
| packet: FrozenBackbonePrefillPacket, |
| ) -> None: |
| """Validate one row before any staged or durable object mutation.""" |
|
|
| packet.validate_boundary() |
| if ( |
| packet.final_hidden_t.shape |
| != (1, 1, self.authority.hidden_size) |
| or packet.summary_hidden_t.shape |
| != ( |
| 1, |
| self.authority.summary_positions, |
| self.authority.hidden_size, |
| ) |
| or packet.final_hidden_t.dtype |
| != _FROZEN_BACKBONE_DTYPES[self.authority.feature_dtype] |
| ): |
| raise ValueError("frozen-backbone packet does not match its authority") |
| expected_identity_t = self.authority.digest_tensors_boundary() |
| for actual_t, expected_t in zip( |
| ( |
| packet.cache_key_sha256_t, |
| packet.work_id_sha256_t, |
| packet.prompt_mask_sha256_t, |
| ), |
| expected_identity_t, |
| strict=True, |
| ): |
| if not torch.equal(actual_t.detach().to(device="cpu"), expected_t): |
| raise ValueError( |
| "frozen-backbone packet identity differs from its authority" |
| ) |
|
|
| def _manifest_payload_boundary(self) -> dict[str, object]: |
| """Build one manifest only after all three final artifacts exist.""" |
|
|
| final_elements, summary_elements, input_elements = ( |
| self._expected_elements_boundary() |
| ) |
| return { |
| "schema": FROZEN_BACKBONE_SIDECAR_SCHEMA, |
| "authoritySha256": self.authority_sha256, |
| "authority": self.authority_record, |
| "rows": 1, |
| "backboneOnly": True, |
| "additiveHeadCached": False, |
| "historicalRboCached": False, |
| "fabricCached": False, |
| "targetEnteredForward": False, |
| "artifacts": { |
| "finalHidden": self._artifact_record_boundary( |
| self.paths.final_hidden, |
| elements=final_elements, |
| dtype_name=self.authority.feature_dtype, |
| ), |
| "summaryHidden": self._artifact_record_boundary( |
| self.paths.summary_hidden, |
| elements=summary_elements, |
| dtype_name=self.authority.feature_dtype, |
| ), |
| "inputPositions": self._artifact_record_boundary( |
| self.paths.input_positions, |
| elements=input_elements, |
| dtype_name="int64", |
| ), |
| }, |
| } |
|
|
| @staticmethod |
| def store_batch_boundary( |
| captured: tuple[ |
| tuple[ |
| FrozenBackboneFeatureSidecar, |
| FrozenBackbonePrefillPacket, |
| ], |
| ..., |
| ], |
| ) -> int: |
| """Durably publish one producer outer wave with manifest-last phases. |
| |
| Every row keeps its independent content-addressed object and writer |
| lock. The optimization changes only when durability syscalls run: |
| artifact bytes for the complete outer wave are staged first, then |
| flushed as one filesystem transaction group; manifest bytes are staged |
| and flushed second; manifest renames are published last and their |
| object/prefix/cache directories are synchronized together. A failure |
| before the last phase leaves no manifest for that row, so lookup treats |
| it as an ordinary miss. The owning producer advances its durable cursor |
| only after this method returns. |
| """ |
|
|
| if not captured: |
| return 0 |
| ordered = tuple( |
| sorted( |
| captured, |
| key=lambda row: str(row[0].paths.root), |
| ) |
| ) |
| if len({sidecar.paths.root for sidecar, _packet in ordered}) != len( |
| ordered |
| ): |
| raise ValueError( |
| "frozen-backbone batch contains duplicate object authority" |
| ) |
| for sidecar, packet in ordered: |
| sidecar._validate_packet_boundary(packet) |
|
|
| lock_descriptors: list[int] = [] |
| artifact_descriptors: list[int] = [] |
| manifest_descriptors: list[int] = [] |
| staged_rows: list[ |
| tuple[ |
| FrozenBackboneFeatureSidecar, |
| tuple[Path, Path, Path], |
| ] |
| ] = [] |
| manifest_promotions: list[tuple[Path, Path]] = [] |
| try: |
| for sidecar, _packet in ordered: |
| sidecar.paths.root.mkdir(parents=True, exist_ok=True) |
| lock_descriptor = os.open( |
| sidecar.paths.lock, |
| os.O_RDWR | os.O_CREAT, |
| 0o600, |
| ) |
| fcntl.flock(lock_descriptor, fcntl.LOCK_EX) |
| lock_descriptors.append(lock_descriptor) |
|
|
| for sidecar, packet in ordered: |
| if sidecar._lookup_current_boundary() is not None: |
| continue |
| progress = sidecar._validated_progress_boundary() |
| completed_rows = progress["completedRows"] |
| if type(completed_rows) is not int: |
| raise RuntimeError( |
| "frozen-backbone progress row frontier is malformed" |
| ) |
| final_paths = sidecar.paths.artifact_paths_boundary() |
| partial_paths: tuple[Path, Path, Path] = ( |
| final_paths[0].with_name( |
| f"{final_paths[0].name}.partial" |
| ), |
| final_paths[1].with_name( |
| f"{final_paths[1].name}.partial" |
| ), |
| final_paths[2].with_name( |
| f"{final_paths[2].name}.partial" |
| ), |
| ) |
| for final_path, partial_path in zip( |
| final_paths, |
| partial_paths, |
| strict=True, |
| ): |
| if final_path.exists() and partial_path.exists(): |
| raise RuntimeError( |
| "frozen-backbone sidecar has duplicate " |
| "incomplete artifacts" |
| ) |
| if final_path.exists(): |
| os.replace(final_path, partial_path) |
| payloads = ( |
| _tensor_storage_bytes_boundary(packet.final_hidden_t), |
| _tensor_storage_bytes_boundary(packet.summary_hidden_t), |
| _tensor_storage_bytes_boundary(packet.input_positions_t), |
| ) |
| for partial_path, payload in zip( |
| partial_paths, |
| payloads, |
| strict=True, |
| ): |
| descriptor = os.open( |
| partial_path, |
| os.O_WRONLY | os.O_CREAT | os.O_TRUNC, |
| 0o600, |
| ) |
| artifact_descriptors.append(descriptor) |
| _write_all_fd_boundary(descriptor, payload) |
| staged_rows.append((sidecar, partial_paths)) |
|
|
| _synchronize_open_files_boundary(tuple(artifact_descriptors)) |
| for descriptor in artifact_descriptors: |
| os.close(descriptor) |
| artifact_descriptors.clear() |
|
|
| for sidecar, partial_paths in staged_rows: |
| for partial_path, final_path in zip( |
| partial_paths, |
| sidecar.paths.artifact_paths_boundary(), |
| strict=True, |
| ): |
| os.replace(partial_path, final_path) |
| temporary_manifest = sidecar.paths.manifest.with_name( |
| f".{sidecar.paths.manifest.name}.batch.partial" |
| ) |
| temporary_manifest.unlink(missing_ok=True) |
| descriptor = os.open( |
| temporary_manifest, |
| os.O_WRONLY | os.O_CREAT | os.O_TRUNC, |
| 0o600, |
| ) |
| manifest_descriptors.append(descriptor) |
| _write_all_fd_boundary( |
| descriptor, |
| _canonical_json_bytes( |
| sidecar._manifest_payload_boundary() |
| ) |
| + b"\n", |
| ) |
| manifest_promotions.append( |
| (temporary_manifest, sidecar.paths.manifest) |
| ) |
|
|
| _synchronize_open_files_boundary(tuple(manifest_descriptors)) |
| for descriptor in manifest_descriptors: |
| os.close(descriptor) |
| manifest_descriptors.clear() |
|
|
| for temporary_manifest, manifest_path in manifest_promotions: |
| os.replace(temporary_manifest, manifest_path) |
| for sidecar, _partial_paths in staged_rows: |
| sidecar.paths.progress.unlink(missing_ok=True) |
| _synchronize_directories_boundary( |
| tuple( |
| path |
| for sidecar, _partial_paths in staged_rows |
| for path in ( |
| sidecar.paths.root, |
| sidecar.paths.root.parent, |
| sidecar.cache_root / "objects", |
| sidecar.cache_root, |
| ) |
| ) |
| ) |
|
|
| for sidecar, _packet in ordered: |
| if sidecar._lookup_current_boundary() is None: |
| raise RuntimeError( |
| "frozen-backbone batch manifest published " |
| "without a readable packet" |
| ) |
| return len(captured) |
| finally: |
| for descriptor in artifact_descriptors: |
| os.close(descriptor) |
| for descriptor in manifest_descriptors: |
| os.close(descriptor) |
| for descriptor in reversed(lock_descriptors): |
| fcntl.flock(descriptor, fcntl.LOCK_UN) |
| os.close(descriptor) |
|
|
| def store_boundary( |
| self, |
| packet: FrozenBackbonePrefillPacket, |
| ) -> FrozenBackbonePrefillPacket: |
| """Durably publish one row and return its mmap-backed representation.""" |
|
|
| self._validate_packet_boundary(packet) |
|
|
| self.paths.root.mkdir(parents=True, exist_ok=True) |
| lock_fd = os.open(self.paths.lock, os.O_RDWR | os.O_CREAT, 0o600) |
| try: |
| fcntl.flock(lock_fd, fcntl.LOCK_EX) |
| |
| |
| |
| |
| existing = self._lookup_current_boundary() |
| if existing is not None: |
| return existing |
| progress = self._validated_progress_boundary() |
| completed_rows_value = progress["completedRows"] |
| if type(completed_rows_value) is not int: |
| raise RuntimeError( |
| "frozen-backbone progress row frontier is malformed" |
| ) |
| completed_rows = completed_rows_value |
| final_elements, summary_elements, input_elements = ( |
| self._expected_elements_boundary() |
| ) |
| dtype = _FROZEN_BACKBONE_DTYPES[self.authority.feature_dtype] |
| expected_bytes = ( |
| final_elements * torch.tensor([], dtype=dtype).element_size(), |
| summary_elements * torch.tensor([], dtype=dtype).element_size(), |
| input_elements |
| * torch.tensor([], dtype=torch.long).element_size(), |
| ) |
| final_paths = self.paths.artifact_paths_boundary() |
| partial_paths = tuple( |
| path.with_name(f"{path.name}.partial") |
| for path in final_paths |
| ) |
| if completed_rows == 0: |
| for final_path, partial_path in zip( |
| final_paths, |
| partial_paths, |
| strict=True, |
| ): |
| if final_path.exists() and partial_path.exists(): |
| raise RuntimeError( |
| "frozen-backbone sidecar has duplicate incomplete artifacts" |
| ) |
| if final_path.exists(): |
| os.replace(final_path, partial_path) |
| self._write_tensor_boundary( |
| partial_paths[0], |
| packet.final_hidden_t, |
| ) |
| self._write_tensor_boundary( |
| partial_paths[1], |
| packet.summary_hidden_t, |
| ) |
| self._write_tensor_boundary( |
| partial_paths[2], |
| packet.input_positions_t, |
| ) |
| _atomic_json_boundary( |
| self.paths.progress, |
| self._progress_payload_boundary(completed_rows=1), |
| ) |
| for final_path, partial_path, byte_count in zip( |
| final_paths, |
| partial_paths, |
| expected_bytes, |
| strict=True, |
| ): |
| selected = ( |
| final_path if final_path.is_file() else partial_path |
| ) |
| if not selected.is_file(): |
| raise RuntimeError( |
| "frozen-backbone durable progress lost an artifact" |
| ) |
| with selected.open("r+b") as handle: |
| handle.truncate(byte_count) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| if selected == partial_path: |
| os.replace(partial_path, final_path) |
| directory_fd = os.open( |
| self.paths.root, |
| os.O_RDONLY | os.O_DIRECTORY, |
| ) |
| try: |
| os.fsync(directory_fd) |
| finally: |
| os.close(directory_fd) |
| manifest = self._manifest_payload_boundary() |
| _atomic_json_boundary(self.paths.manifest, manifest) |
| self.paths.progress.unlink(missing_ok=True) |
| directory_fd = os.open( |
| self.paths.root, |
| os.O_RDONLY | os.O_DIRECTORY, |
| ) |
| try: |
| os.fsync(directory_fd) |
| finally: |
| os.close(directory_fd) |
| loaded = self._lookup_current_boundary() |
| if loaded is None: |
| raise RuntimeError( |
| "frozen-backbone manifest published without a readable packet" |
| ) |
| return loaded |
| finally: |
| fcntl.flock(lock_fd, fcntl.LOCK_UN) |
| os.close(lock_fd) |
|
|
|
|
| def base_forward_cache_enabled_boundary() -> bool: |
| """Return whether the frozen-parent prefill cache is active.""" |
|
|
| raw = os.environ.get("NNF_RESYNTHESIS_BASE_FORWARD_CACHE", "1") |
| return raw not in {"0", "false", "False", "no", "off"} |
|
|
|
|
| def base_forward_cache_capacity_boundary() -> int: |
| """Return the LRU capacity for cached prefill results.""" |
|
|
| raw = os.environ.get("NNF_RESYNTHESIS_BASE_FORWARD_CACHE_ENTRIES", "512") |
| try: |
| capacity = int(raw) |
| except ValueError as error: |
| raise RuntimeError( |
| "NNF_RESYNTHESIS_BASE_FORWARD_CACHE_ENTRIES must be an integer" |
| ) from error |
| if capacity < 1: |
| raise RuntimeError( |
| "NNF_RESYNTHESIS_BASE_FORWARD_CACHE_ENTRIES must be positive" |
| ) |
| return capacity |
|
|
|
|
| def base_forward_cache_allowed_boundary( |
| *, |
| input_mask: torch.Tensor | None, |
| use_past: bool, |
| parent_trainable: bool, |
| ) -> bool: |
| """Return whether a frozen-parent prefill may be served from cache. |
| |
| Cache is restricted to batched training prefills that supply an explicit |
| attention mask. Autoregressive decode arms (no mask, or KV continuation) |
| must always run the live parent forward so ``past_key_values`` state is |
| populated for the next token. |
| """ |
|
|
| return ( |
| base_forward_cache_enabled_boundary() |
| and input_mask is not None |
| and not use_past |
| and not parent_trainable |
| ) |
|
|
|
|
| def base_forward_cache_key_boundary( |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor | None, |
| ) -> tuple[str, str]: |
| """Return full SHA-256 hex and its 16-hex display prefix for one prompt.""" |
|
|
| if input_ids.ndim != 2 or input_ids.shape[1] < 1: |
| raise ValueError("base forward cache key requires [batch, sequence]") |
| digest = hashlib.sha256() |
| prompt = input_ids.detach().to(device="cpu", dtype=torch.long).contiguous() |
| digest.update(prompt.numpy().tobytes()) |
| if attention_mask is not None: |
| if attention_mask.shape != input_ids.shape: |
| raise ValueError("base forward cache attention mask geometry differs") |
| mask = attention_mask.detach().to(device="cpu", dtype=torch.long).contiguous() |
| digest.update(mask.numpy().tobytes()) |
| full_hex = digest.hexdigest() |
| return full_hex, full_hex[:16] |
|
|
|
|
| @dataclass(frozen=True) |
| class BaseForwardCacheEntry: |
| """CPU-resident frozen-parent prefill tensors keyed by prompt identity.""" |
|
|
| hidden: torch.Tensor |
| logits: torch.Tensor |
| parent_context_hidden: torch.Tensor |
| parent_expert_routes: torch.Tensor |
| parent_layer_routes: torch.Tensor |
| kv_prefix_positions: torch.Tensor |
| kv_new_positions: torch.Tensor |
| parent_prefill_hidden: torch.Tensor | None |
| parent_prefill_input_positions: torch.Tensor | None |
| sha16: str |
|
|
| @classmethod |
| def from_forward_boundary( |
| cls, |
| *, |
| hidden: torch.Tensor, |
| logits: torch.Tensor, |
| parent_context_hidden: torch.Tensor, |
| parent_expert_routes: torch.Tensor, |
| parent_layer_routes: torch.Tensor, |
| kv_prefix_positions: torch.Tensor, |
| kv_new_positions: torch.Tensor, |
| parent_prefill_hidden: torch.Tensor | None, |
| parent_prefill_input_positions: torch.Tensor | None, |
| sha16: str, |
| ) -> BaseForwardCacheEntry: |
| """Store detached CPU tensors, pinning when the host supports it.""" |
|
|
| def _pin(value: torch.Tensor) -> torch.Tensor: |
| detached = value.detach().to(device="cpu").contiguous() |
| if detached.is_floating_point() or detached.dtype in { |
| torch.int32, |
| torch.int64, |
| torch.bool, |
| }: |
| try: |
| return detached.pin_memory() |
| except RuntimeError: |
| return detached |
| return detached |
|
|
| prefill_hidden = ( |
| _pin(parent_prefill_hidden) |
| if isinstance(parent_prefill_hidden, torch.Tensor) |
| else None |
| ) |
| prefill_positions = ( |
| _pin(parent_prefill_input_positions) |
| if isinstance(parent_prefill_input_positions, torch.Tensor) |
| else None |
| ) |
| return cls( |
| hidden=_pin(hidden), |
| logits=_pin(logits), |
| parent_context_hidden=_pin(parent_context_hidden), |
| parent_expert_routes=_pin(parent_expert_routes), |
| parent_layer_routes=_pin(parent_layer_routes), |
| kv_prefix_positions=_pin(kv_prefix_positions), |
| kv_new_positions=_pin(kv_new_positions), |
| parent_prefill_hidden=prefill_hidden, |
| parent_prefill_input_positions=prefill_positions, |
| sha16=sha16, |
| ) |
|
|
| def to_device_boundary(self, device: torch.device) -> BaseForwardCacheEntry: |
| """Materialize one cache entry on the active compute device.""" |
|
|
| def _move(value: torch.Tensor) -> torch.Tensor: |
| return value.to(device=device, non_blocking=device.type == "cuda") |
|
|
| prefill_hidden = ( |
| _move(self.parent_prefill_hidden) |
| if isinstance(self.parent_prefill_hidden, torch.Tensor) |
| else None |
| ) |
| prefill_positions = ( |
| _move(self.parent_prefill_input_positions) |
| if isinstance(self.parent_prefill_input_positions, torch.Tensor) |
| else None |
| ) |
| return BaseForwardCacheEntry( |
| hidden=_move(self.hidden), |
| logits=_move(self.logits), |
| parent_context_hidden=_move(self.parent_context_hidden), |
| parent_expert_routes=_move(self.parent_expert_routes), |
| parent_layer_routes=_move(self.parent_layer_routes), |
| kv_prefix_positions=_move(self.kv_prefix_positions), |
| kv_new_positions=_move(self.kv_new_positions), |
| parent_prefill_hidden=prefill_hidden, |
| parent_prefill_input_positions=prefill_positions, |
| sha16=self.sha16, |
| ) |
|
|
|
|
| class BaseForwardCache: |
| """LRU cache for complete frozen-parent prefill outputs.""" |
|
|
| def __init__(self, *, capacity: int) -> None: |
| if capacity < 1: |
| raise ValueError("base forward cache capacity must be positive") |
| self._capacity = capacity |
| self._entries: OrderedDict[str, BaseForwardCacheEntry] = OrderedDict() |
| self._hits = 0 |
| self._misses = 0 |
|
|
| @property |
| def hits(self) -> int: |
| return self._hits |
|
|
| @property |
| def misses(self) -> int: |
| return self._misses |
|
|
| def clear(self) -> None: |
| self._entries.clear() |
|
|
| def lookup(self, cache_key: str) -> BaseForwardCacheEntry | None: |
| entry = self._entries.get(cache_key) |
| if entry is None: |
| self._misses += 1 |
| return None |
| self._hits += 1 |
| self._entries.move_to_end(cache_key) |
| return entry |
|
|
| def store(self, cache_key: str, entry: BaseForwardCacheEntry) -> None: |
| if cache_key in self._entries: |
| self._entries.move_to_end(cache_key) |
| self._entries[cache_key] = entry |
| while len(self._entries) > self._capacity: |
| self._entries.popitem(last=False) |
|
|
| def telemetry_boundary(self) -> dict[str, object]: |
| return { |
| "schema": "nnf.resynthesis.base_forward_cache.v1", |
| "enabled": base_forward_cache_enabled_boundary(), |
| "capacity": self._capacity, |
| "entries": len(self._entries), |
| "hits": self._hits, |
| "misses": self._misses, |
| } |
|
|
