OBLITERATUS/obliteratus/checkpoint_writer.py

"""Transactional canonical safetensors writer for validated neutral fragments."""

from __future__ import annotations

import json
import os
import shutil
import stat
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path, PurePosixPath
from typing import Any, Callable, Mapping, Sequence

import torch
from safetensors.torch import load_file, save_file

from obliteratus.checkpoint_errors import CheckpointContractError
from obliteratus.checkpoint_fragments import (
    TensorFragment,
    ValidationResult,
    reconstruct_logical_tensor,
    validate_fragments,
)
from obliteratus.checkpoint_provenance import (
    ProvenanceRecord,
    verify_provenance_record,
)
from obliteratus.persistence_contracts import atomic_checkpoint_directory


_DIGEST_PREFIX = "sha256:"
_INT64_MAX = (1 << 63) - 1
_WRITER_BUFFER_BYTES = 1 << 20
_COPY_KINDS = frozenset({"configuration", "tokenizer"})
_CANONICAL_DTYPES = frozenset(
    {
        "bool",
        "uint8",
        "uint16",
        "uint32",
        "uint64",
        "int8",
        "int16",
        "int32",
        "int64",
        "float8_e4m3fn",
        "float8_e5m2",
        "float16",
        "bfloat16",
        "float32",
        "float64",
        "complex64",
    }
)


@dataclass(frozen=True)
class WriterLimits:
    """Admission limits applied before a staging directory is created."""

    max_shard_bytes: int = 5 << 30
    max_output_bytes: int = 1 << 40
    max_temp_bytes: int = 1 << 40
    max_peak_ram_bytes: int = 16 << 30
    min_free_headroom_percent: int = 10

    def __post_init__(self) -> None:
        for name, value in vars(self).items():
            if type(value) is not int or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        if self.min_free_headroom_percent > 100:
            raise ValueError("min_free_headroom_percent cannot exceed 100")


@dataclass(frozen=True)
class VerifiedSourceFile:
    """One source artifact whose path and digest are independently rechecked."""

    path: Path
    relative_path: str
    expected_sha256: str

    def __post_init__(self) -> None:
        object.__setattr__(self, "path", Path(self.path))
        _safe_relative(self.relative_path)
        _require_digest(self.expected_sha256, "source digest")


@dataclass(frozen=True)
class ImmutableCopy:
    """A base configuration/tokenizer file copied only after exact hash match."""

    relative_path: str
    source_path: Path
    expected_sha256: str
    kind: str

    def __post_init__(self) -> None:
        object.__setattr__(self, "source_path", Path(self.source_path))
        _safe_basename(self.relative_path)
        _require_digest(self.expected_sha256, "copy digest")
        if self.kind not in _COPY_KINDS:
            raise ValueError("copy kind must be configuration or tokenizer")


@dataclass(frozen=True)
class CanonicalWriteResult:
    """Identity and immutable location of a successfully promoted checkpoint."""

    artifact_id: str
    output_path: Path
    weight_files: tuple[str, ...]
    manifest_digest: str


@dataclass(frozen=True)
class _Fingerprint:
    device: int
    inode: int
    mode: int
    size: int
    modified_ns: int

    @classmethod
    def from_stat(cls, value: os.stat_result) -> _Fingerprint:
        return cls(value.st_dev, value.st_ino, value.st_mode, value.st_size, value.st_mtime_ns)


@dataclass(frozen=True)
class _Snapshot:
    path: Path
    relative_path: str
    fingerprint: _Fingerprint
    sha256: str

    def file_record(self) -> dict[str, Any]:
        return {
            "relative_path": self.relative_path,
            "size_bytes": self.fingerprint.size,
            "sha256": self.sha256,
        }


@dataclass(frozen=True)
class _TensorSpec:
    name: str
    logical_tensor_id: str
    size_bytes: int


@dataclass(frozen=True)
class _TensorOracle:
    shape: tuple[int, ...]
    dtype: str
    sha256: str


@dataclass(frozen=True)
class _Admission:
    logical_bytes: int
    estimated_output_bytes: int
    estimated_temp_bytes: int
    estimated_peak_ram_bytes: int


def _safe_relative(value: str) -> None:
    pure = PurePosixPath(value)
    if (
        not isinstance(value, str)
        or not value
        or pure.is_absolute()
        or any(part in {"", ".", ".."} for part in pure.parts)
        or len(value) > 4096
    ):
        raise ValueError("relative path is unsafe")


def _safe_basename(value: str) -> None:
    _safe_relative(value)
    if PurePosixPath(value).name != value:
        raise ValueError("canonical copied artifacts must be root-level files")


def _require_digest(value: str, field: str) -> None:
    if (
        not isinstance(value, str)
        or not value.startswith(_DIGEST_PREFIX)
        or len(value) != len(_DIGEST_PREFIX) + 64
        or any(character not in "0123456789abcdef" for character in value[len(_DIGEST_PREFIX) :])
    ):
        raise ValueError(f"{field} must be a sha256 digest")


def _source_changed(reference: str) -> CheckpointContractError:
    return CheckpointContractError(
        "DCI_SOURCE_CHANGED",
        detail="source_digest_or_identity_changed",
        affected_refs=(reference,),
    )


def _open_flags() -> int:
    return os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0)


def _open_snapshot(path: Path, reference: str) -> int:
    try:
        return os.open(path, _open_flags())
    except OSError as error:
        raise _source_changed(reference) from error


def _hash_descriptor(descriptor: int, reference: str) -> str:
    digest = sha256()
    try:
        while True:
            chunk = os.read(descriptor, 1 << 20)
            if not chunk:
                return f"sha256:{digest.hexdigest()}"
            digest.update(chunk)
    except OSError as error:
        raise _source_changed(reference) from error


def _read_snapshot(descriptor: int, size: int, reference: str) -> bytes:
    try:
        return os.read(descriptor, size)
    except OSError as error:
        raise _source_changed(reference) from error


def _descriptor_fingerprint(descriptor: int, reference: str) -> _Fingerprint:
    try:
        return _Fingerprint.from_stat(os.fstat(descriptor))
    except OSError as error:
        raise _source_changed(reference) from error


def _snapshot(path: Path, relative_path: str, expected_digest: str) -> _Snapshot:
    try:
        observed = path.lstat()
    except OSError as error:
        raise _source_changed(relative_path) from error
    if not stat.S_ISREG(observed.st_mode):
        raise CheckpointContractError(
            "DCI_SOURCE_BOUNDARY_VIOLATION",
            detail="source_not_regular_file",
            affected_refs=(relative_path,),
        )
    try:
        if path.absolute() != path.resolve(strict=True):
            raise CheckpointContractError(
                "DCI_SOURCE_BOUNDARY_VIOLATION",
                detail="source_symlink",
                affected_refs=(relative_path,),
            )
    except OSError as error:
        raise _source_changed(relative_path) from error
    fingerprint = _Fingerprint.from_stat(observed)
    descriptor = _open_snapshot(path, relative_path)
    try:
        if _descriptor_fingerprint(descriptor, relative_path) != fingerprint:
            raise _source_changed(relative_path)
        digest = _hash_descriptor(descriptor, relative_path)
        if _descriptor_fingerprint(descriptor, relative_path) != fingerprint:
            raise _source_changed(relative_path)
    finally:
        os.close(descriptor)
    if digest != expected_digest:
        raise _source_changed(relative_path)
    try:
        if _Fingerprint.from_stat(path.lstat()) != fingerprint:
            raise _source_changed(relative_path)
    except OSError as error:
        raise _source_changed(relative_path) from error
    return _Snapshot(path, relative_path, fingerprint, digest)


def _revalidate_snapshot(snapshot: _Snapshot) -> None:
    try:
        observed = _Fingerprint.from_stat(snapshot.path.lstat())
    except OSError as error:
        raise _source_changed(snapshot.relative_path) from error
    if observed != snapshot.fingerprint:
        raise _source_changed(snapshot.relative_path)
    descriptor = _open_snapshot(snapshot.path, snapshot.relative_path)
    try:
        if _descriptor_fingerprint(descriptor, snapshot.relative_path) != snapshot.fingerprint:
            raise _source_changed(snapshot.relative_path)
        digest = _hash_descriptor(descriptor, snapshot.relative_path)
    finally:
        os.close(descriptor)
    if digest != snapshot.sha256:
        raise _source_changed(snapshot.relative_path)


def _tensor_bytes(shape: tuple[int, ...], element_size: int, reference: str) -> int:
    elements = 1
    for dimension in shape:
        if dimension and elements > _INT64_MAX // dimension:
            raise CheckpointContractError(
                "DCI_VALIDATION_FAILED",
                detail="integer_overflow",
                affected_refs=(reference,),
            )
        elements *= dimension
    if elements and elements > _INT64_MAX // element_size:
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="integer_overflow",
            affected_refs=(reference,),
        )
    return elements * element_size


def _tensor_specs(validation: ValidationResult) -> tuple[_TensorSpec, ...]:
    names = [logical.fqn for logical in validation.logical_tensors]
    if len(set(names)) != len(names):
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="canonical_tensor_name_duplicate",
            affected_refs=names,
        )
    specs = []
    for logical in validation.logical_tensors:
        if logical.dtype not in _CANONICAL_DTYPES:
            raise CheckpointContractError(
                "DCI_VALIDATION_FAILED",
                detail="canonical_dtype_unsupported",
                affected_refs=(logical.logical_tensor_id,),
            )
        payload = logical.fragments[0].payload
        if payload is None:
            raise CheckpointContractError(
                "DCI_VALIDATION_FAILED",
                detail="payload_unavailable",
                affected_refs=(logical.logical_tensor_id,),
            )
        specs.append(
            _TensorSpec(
                name=logical.fqn,
                logical_tensor_id=logical.logical_tensor_id,
                size_bytes=_tensor_bytes(
                    logical.global_shape,
                    payload.element_size(),
                    logical.logical_tensor_id,
                ),
            )
        )
    return tuple(sorted(specs, key=lambda item: item.name))


def _shard_plan(specs: tuple[_TensorSpec, ...], max_shard_bytes: int) -> tuple[tuple[_TensorSpec, ...], ...]:
    oversized = next((spec for spec in specs if spec.size_bytes > max_shard_bytes), None)
    if oversized is not None:
        raise CheckpointContractError(
            "DCI_ADMISSION_DENIED",
            detail="tensor_exceeds_max_shard_bytes",
            affected_refs=(oversized.logical_tensor_id,),
        )
    shards: list[list[_TensorSpec]] = []
    current: list[_TensorSpec] = []
    current_bytes = 0
    for spec in specs:
        if current and current_bytes + spec.size_bytes > max_shard_bytes:
            shards.append(current)
            current = []
            current_bytes = 0
        current.append(spec)
        current_bytes += spec.size_bytes
    if current:
        shards.append(current)
    return tuple(tuple(shard) for shard in shards)


def _admit(
    destination: Path,
    plan: tuple[tuple[_TensorSpec, ...], ...],
    copy_snapshots: tuple[_Snapshot, ...],
    limits: WriterLimits,
) -> _Admission:
    logical_bytes = sum(spec.size_bytes for shard in plan for spec in shard)
    copied_bytes = sum(snapshot.fingerprint.size for snapshot in copy_snapshots)
    header_allowance = (
        sum(len(spec.name.encode("utf-8")) + 256 for shard in plan for spec in shard)
        + 4096 * max(1, len(plan))
        + 8192
    )
    estimated_output = logical_bytes + copied_bytes + header_allowance
    estimated_temp = estimated_output
    largest_shard = max((sum(spec.size_bytes for spec in shard) for shard in plan), default=0)
    largest_tensor = max((spec.size_bytes for shard in plan for spec in shard), default=0)
    peak_ram = largest_shard + largest_tensor + _WRITER_BUFFER_BYTES
    checks = (
        (estimated_output, limits.max_output_bytes, "max_output_bytes"),
        (estimated_temp, limits.max_temp_bytes, "max_temp_bytes"),
        (peak_ram, limits.max_peak_ram_bytes, "max_peak_ram_bytes"),
    )
    for required, maximum, detail in checks:
        if required > maximum:
            raise CheckpointContractError(
                "DCI_ADMISSION_DENIED",
                detail=detail,
                affected_refs=(f"required:{required}", f"available:{maximum}"),
            )
    try:
        destination.parent.mkdir(parents=True, exist_ok=True)
    except OSError as error:
        raise CheckpointContractError(
            "DCI_ADMISSION_DENIED",
            detail="destination_parent_unavailable",
            affected_refs=("destination",),
        ) from error
    if destination.parent.is_symlink():
        raise CheckpointContractError(
            "DCI_SOURCE_BOUNDARY_VIOLATION",
            detail="destination_parent_symlink",
            affected_refs=("destination",),
        )
    try:
        available = shutil.disk_usage(destination.parent).free
    except OSError as error:
        raise CheckpointContractError(
            "DCI_ADMISSION_DENIED",
            detail="filesystem_capacity_unavailable",
            affected_refs=("destination",),
        ) from error
    required_with_headroom = estimated_temp * (100 + limits.min_free_headroom_percent)
    required_with_headroom = (required_with_headroom + 99) // 100
    if available < required_with_headroom:
        raise CheckpointContractError(
            "DCI_ADMISSION_DENIED",
            detail="filesystem_free_bytes",
            affected_refs=(f"required:{required_with_headroom}", f"available:{available}"),
        )
    return _Admission(logical_bytes, estimated_output, estimated_temp, peak_ram)


def _copy_snapshot(snapshot: _Snapshot, destination: Path) -> None:
    source = _open_snapshot(snapshot.path, snapshot.relative_path)
    try:
        if _descriptor_fingerprint(source, snapshot.relative_path) != snapshot.fingerprint:
            raise _source_changed(snapshot.relative_path)
        with destination.open("xb") as output:
            while True:
                chunk = _read_snapshot(source, 1 << 20, snapshot.relative_path)
                if not chunk:
                    break
                output.write(chunk)
            output.flush()
            os.fsync(output.fileno())
        if _descriptor_fingerprint(source, snapshot.relative_path) != snapshot.fingerprint:
            raise _source_changed(snapshot.relative_path)
    finally:
        os.close(source)
    if _file_digest(destination) != snapshot.sha256:
        raise _source_changed(snapshot.relative_path)


def _save_safetensors_file(tensors: Mapping[str, torch.Tensor], path: Path) -> None:
    save_file(dict(sorted(tensors.items())), path)


def _file_digest(path: Path) -> str:
    digest = sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1 << 20), b""):
            digest.update(chunk)
    return f"sha256:{digest.hexdigest()}"


def _tensor_digest(tensor: torch.Tensor) -> str:
    raw = tensor.detach().contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()
    return f"sha256:{sha256(raw).hexdigest()}"


def _write_json(path: Path, value: object) -> None:
    payload = json.dumps(
        value,
        indent=2,
        sort_keys=True,
        ensure_ascii=True,
        allow_nan=False,
    ) + "\n"
    with path.open("x", encoding="utf-8") as stream:
        stream.write(payload)
        stream.flush()
        os.fsync(stream.fileno())


def _output_record(path: Path) -> dict[str, Any]:
    observed = path.lstat()
    if not stat.S_ISREG(observed.st_mode):
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_not_regular_file",
            affected_refs=(path.name,),
        )
    return {
        "relative_path": path.name,
        "size_bytes": observed.st_size,
        "sha256": _file_digest(path),
    }


def _canonical_digest(value: object) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
    return f"sha256:{sha256(payload).hexdigest()}"


def _manifest(
    *,
    artifact_id: str,
    descriptor_digest: str,
    source_records: list[dict[str, Any]],
    source_topology: Mapping[str, Any],
    validation: ValidationResult,
    output_records: list[dict[str, Any]],
    index_name: str | None,
    admission: _Admission,
    provenance: dict[str, Any],
) -> dict[str, Any]:
    lost_state = provenance["state"]["lost_state"]
    return {
        "schema_id": "obliteratus.conversion-manifest",
        "schema_version": "1.0.0",
        "manifest_id": artifact_id,
        "descriptor": {"schema_version": "1.0.0", "digest": descriptor_digest},
        "source_inventory_digest": _canonical_digest(source_records),
        "source_files": source_records,
        "adapter": {
            "id": "producer-neutral-fragments",
            "version": "1.0.0",
            "capability_digest": validation.manifest_digest,
            "producer": "already-normalized-input",
            "producer_version": "1.0.0",
        },
        "source_topology": dict(source_topology),
        "state": {
            "source_classification": provenance["state"]["classification"],
            "output_classification": "weights_only",
            "observed_scopes": provenance["state"]["observed_scopes"],
        },
        "dropped_scopes": [
            {"scope": scope, "reason": "Canonical Wave 2 output contains weights only."}
            for scope in lost_state
        ],
        "canonical_output": {
            "format": "hf_safetensors",
            "dtype_policy": "preserve_exact",
            "files": output_records,
            "hf_index": index_name,
            "logical_tensor_count": len(validation.logical_tensors),
            "logical_bytes": admission.logical_bytes,
        },
        "resource_usage": {
            "estimated_peak_ram_bytes": admission.estimated_peak_ram_bytes,
            "actual_peak_ram_bytes": None,
            "estimated_temp_bytes": admission.estimated_temp_bytes,
            "actual_temp_bytes": None,
        },
        "validation": {
            "coverage": True,
            "replicas": True,
            "ties": True,
            "hashes": True,
            "index": True,
            "safe_reload": True,
            "source_unchanged": True,
            "result": "passed",
        },
        "provenance": {
            "obliteratus_commit": provenance["obliteratus_commit"],
            "configuration_digest": provenance["configuration_digest"],
            "tokenizer_digest": (
                provenance["tokenizer"]["digest"] if provenance["tokenizer"] else None
            ),
            "base_model": {
                "identity": provenance["base_model"]["identity"] if provenance["base_model"] else None,
                "revision": provenance["base_model"]["revision"] if provenance["base_model"] else None,
                "digest": provenance["base_model"]["digest"] if provenance["base_model"] else None,
            },
            "transformation_log": provenance["transformations"],
            "unknowns": provenance["unknowns"],
        },
        "publication": {
            "staging_validated": True,
            "promoted": True,
            "atomic_strategy": "sibling-stage-fsync-replace",
            "rollback_result": "not_required",
        },
    }


def _verify_provenance(
    record: dict[str, Any],
    *,
    output_digests: tuple[str, ...],
    source_digests: tuple[str, ...],
    configuration_digest: str,
    source_topology: Mapping[str, Any],
) -> dict[str, Any]:
    try:
        verified = verify_provenance_record(record)
    except (TypeError, ValueError) as error:
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="provenance_contract_invalid",
        ) from error
    if sorted(verified["output_digests"]) != sorted(set(output_digests)):
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="provenance_output_digest_mismatch",
        )
    if sorted(verified["input_digests"]) != sorted(set(source_digests)):
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="provenance_input_digest_mismatch",
        )
    if verified["configuration_digest"] != configuration_digest:
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="provenance_configuration_digest_mismatch",
        )
    try:
        topology_matches = dict(source_topology) == verified["source_topology"]
    except (TypeError, ValueError) as error:
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="provenance_source_topology_invalid",
        ) from error
    if not topology_matches:
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="provenance_source_topology_mismatch",
        )
    return verified


def _verify_json_object(path: Path) -> dict[str, Any]:
    if path.is_symlink() or not path.is_file() or path.stat().st_size > 8 << 20:
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_json_invalid",
            affected_refs=(path.name,),
        )
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (UnicodeError, json.JSONDecodeError) as error:
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_json_invalid",
            affected_refs=(path.name,),
        ) from error
    if not isinstance(value, dict):
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_json_invalid",
            affected_refs=(path.name,),
        )
    return value


def _verify_staging(
    staging: Path,
    *,
    weight_map: dict[str, str],
    oracles: dict[str, _TensorOracle],
    expected_files: set[str],
    snapshots: tuple[_Snapshot, ...],
    source_snapshots: tuple[_Snapshot, ...],
    configuration_digest: str,
    descriptor_digest: str,
    validation_digest: str,
    artifact_id: str,
    limits: WriterLimits,
) -> None:
    actual_files = {path.name for path in staging.iterdir() if path.is_file() and not path.is_symlink()}
    if actual_files != expected_files or any(path.is_symlink() or not path.is_file() for path in staging.iterdir()):
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_file_set_mismatch",
        )
    actual_bytes = sum((staging / name).lstat().st_size for name in expected_files)
    if actual_bytes > limits.max_output_bytes or actual_bytes > limits.max_temp_bytes:
        raise CheckpointContractError(
            "DCI_ADMISSION_DENIED",
            detail="post_write_size_limit",
            affected_refs=(f"required:{actual_bytes}",),
        )
    loaded_names: set[str] = set()
    loaded_bytes = 0
    for shard_name in sorted(set(weight_map.values())):
        shard = load_file(staging / shard_name, device="cpu")
        for name, tensor in shard.items():
            if weight_map.get(name) != shard_name or name not in oracles:
                raise CheckpointContractError(
                    "DCI_VALIDATION_FAILED",
                    detail="output_index_mismatch",
                    affected_refs=(name,),
                )
            oracle = oracles[name]
            if (
                tuple(tensor.shape) != oracle.shape
                or str(tensor.dtype).removeprefix("torch.") != oracle.dtype
                or _tensor_digest(tensor) != oracle.sha256
            ):
                raise CheckpointContractError(
                    "DCI_VALIDATION_FAILED",
                    detail="output_tensor_mismatch",
                    affected_refs=(name,),
                )
            loaded_names.add(name)
            loaded_bytes += tensor.numel() * tensor.element_size()
    if loaded_names != set(weight_map) or loaded_names != set(oracles):
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_tensor_set_mismatch",
        )
    index_name = (
        "model.safetensors.index.json"
        if "model.safetensors.index.json" in expected_files
        else None
    )
    if index_name is not None:
        index = _verify_json_object(staging / index_name)
        if index != {
            "metadata": {"total_size": loaded_bytes},
            "weight_map": dict(sorted(weight_map.items())),
        }:
            raise CheckpointContractError(
                "DCI_VALIDATION_FAILED",
                detail="output_index_mismatch",
            )
    for name in ("config.json", "tokenizer_config.json"):
        _verify_json_object(staging / name)
    provenance = _verify_json_object(staging / "checkpoint-provenance.json")
    manifest = _verify_json_object(staging / "conversion-manifest.json")
    metadata = _verify_json_object(staging / "abliteration_metadata.json")
    evidence_files = {
        "checkpoint-provenance.json",
        "conversion-manifest.json",
        "abliteration_metadata.json",
    }
    output_records = [
        _output_record(staging / name)
        for name in sorted(expected_files - evidence_files)
    ]
    output_digests = tuple(record["sha256"] for record in output_records)
    verified_provenance = _verify_provenance(
        provenance,
        output_digests=output_digests,
        source_digests=tuple(snapshot.sha256 for snapshot in source_snapshots),
        configuration_digest=configuration_digest,
        source_topology=provenance.get("source_topology", {}),
    )
    canonical_output = manifest.get("canonical_output")
    adapter = manifest.get("adapter")
    descriptor = manifest.get("descriptor")
    if (
        not isinstance(canonical_output, dict)
        or not isinstance(adapter, dict)
        or not isinstance(descriptor, dict)
        or canonical_output.get("files") != output_records
        or canonical_output.get("hf_index") != index_name
        or adapter.get("capability_digest") != validation_digest
        or descriptor.get("digest") != descriptor_digest
        or manifest.get("source_files")
        != [snapshot.file_record() for snapshot in source_snapshots]
        or manifest.get("source_topology") != verified_provenance["source_topology"]
    ):
        raise CheckpointContractError(
            "DCI_VALIDATION_FAILED",
            detail="output_manifest_mismatch",
        )
    if not (
        provenance.get("artifact_id")
        == manifest.get("manifest_id")
        == metadata.get("artifact_id")
        == artifact_id
    ):
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="artifact_identity_mismatch",
        )
    for snapshot in snapshots:
        _revalidate_snapshot(snapshot)


def write_canonical_checkpoint(
    destination: Path | str,
    fragments: Sequence[TensorFragment],
    *,
    source_files: Sequence[VerifiedSourceFile],
    copies: Sequence[ImmutableCopy],
    descriptor_digest: str,
    source_topology: Mapping[str, Any],
    provenance_factory: Callable[[tuple[str, ...]], ProvenanceRecord],
    tie_policy: str | None = None,
    limits: WriterLimits | None = None,
) -> CanonicalWriteResult:
    """Write validated CPU fragments without invoking a producer reader or adapter."""
    _require_digest(descriptor_digest, "descriptor digest")
    active_limits = limits or WriterLimits()
    validation = validate_fragments(tuple(fragments))
    if any(logical.tie_group_id is not None for logical in validation.logical_tensors):
        if tie_policy != "duplicate_validated":
            raise CheckpointContractError(
                "DCI_VALIDATION_FAILED",
                detail="tie_policy_required",
                affected_refs=(
                    logical.logical_tensor_id
                    for logical in validation.logical_tensors
                    if logical.tie_group_id is not None
                ),
            )
    if not source_files:
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="source_inventory_empty",
        )
    source_paths = [item.relative_path for item in source_files]
    if len(set(source_paths)) != len(source_paths):
        raise ValueError("source relative paths must be unique")
    copy_paths = [item.relative_path for item in copies]
    if len(set(copy_paths)) != len(copy_paths):
        raise ValueError("copy output paths must be unique")
    if not {"config.json", "tokenizer_config.json"} <= set(copy_paths):
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="canonical_config_or_tokenizer_missing",
        )
    copy_kinds = {item.relative_path: item.kind for item in copies}
    if (
        copy_kinds["config.json"] != "configuration"
        or copy_kinds["tokenizer_config.json"] != "tokenizer"
    ):
        raise CheckpointContractError(
            "DCI_EVIDENCE_UNAVAILABLE",
            detail="canonical_config_or_tokenizer_kind_mismatch",
        )
    source_snapshots = tuple(
        sorted(
            (
                _snapshot(item.path, item.relative_path, item.expected_sha256)
                for item in source_files
            ),
            key=lambda snapshot: snapshot.relative_path,
        )
    )
    copy_snapshots = tuple(
        _snapshot(item.source_path, item.relative_path, item.expected_sha256) for item in copies
    )
    specs = _tensor_specs(validation)
    plan = _shard_plan(specs, active_limits.max_shard_bytes)
    output_path = Path(destination)
    try:
        destination_unsafe = (
            output_path.is_symlink()
            or output_path.absolute() != output_path.resolve(strict=False)
        )
    except OSError as error:
        raise CheckpointContractError(
            "DCI_SOURCE_BOUNDARY_VIOLATION",
            detail="destination_unresolvable",
            affected_refs=("destination",),
        ) from error
    if destination_unsafe:
        raise CheckpointContractError(
            "DCI_SOURCE_BOUNDARY_VIOLATION",
            detail="destination_symlink",
            affected_refs=("destination",),
        )
    admission = _admit(output_path, plan, copy_snapshots, active_limits)
    shard_names = (
        ("model.safetensors",)
        if len(plan) == 1
        else tuple(
            f"model-{index:05d}-of-{len(plan):05d}.safetensors"
            for index in range(1, len(plan) + 1)
        )
    )
    weight_map = {
        spec.name: shard_name
        for shard_name, shard in zip(shard_names, plan, strict=True)
        for spec in shard
    }
    oracles: dict[str, _TensorOracle] = {}
    index_name = "model.safetensors.index.json" if len(plan) > 1 else None
    artifact_id = ""
    manifest_digest = ""
    all_snapshots = (*source_snapshots, *copy_snapshots)

    def verify(staging: Path) -> None:
        expected = {
            *copy_paths,
            *shard_names,
            "checkpoint-provenance.json",
            "conversion-manifest.json",
            "abliteration_metadata.json",
        }
        if index_name is not None:
            expected.add(index_name)
        _verify_staging(
            staging,
            weight_map=weight_map,
            oracles=oracles,
            expected_files=expected,
            snapshots=all_snapshots,
            source_snapshots=source_snapshots,
            configuration_digest=next(
                snapshot.sha256
                for snapshot in copy_snapshots
                if snapshot.relative_path == "config.json"
            ),
            descriptor_digest=descriptor_digest,
            validation_digest=validation.manifest_digest,
            artifact_id=artifact_id,
            limits=active_limits,
        )

    try:
        with atomic_checkpoint_directory(output_path, validate=verify) as staging:
            try:
                for copy, snapshot in zip(copies, copy_snapshots, strict=True):
                    _copy_snapshot(snapshot, staging / copy.relative_path)
                for shard_name, shard in zip(shard_names, plan, strict=True):
                    tensors: dict[str, torch.Tensor] = {}
                    for spec in shard:
                        tensor = reconstruct_logical_tensor(validation, spec.logical_tensor_id)
                        tensors[spec.name] = tensor
                        oracles[spec.name] = _TensorOracle(
                            shape=tuple(tensor.shape),
                            dtype=str(tensor.dtype).removeprefix("torch."),
                            sha256=_tensor_digest(tensor),
                        )
                    _save_safetensors_file(tensors, staging / shard_name)
                del tensors
                if index_name is not None:
                    _write_json(
                        staging / index_name,
                        {
                            "metadata": {"total_size": admission.logical_bytes},
                            "weight_map": dict(sorted(weight_map.items())),
                        },
                    )
                output_names = [*copy_paths, *shard_names]
                if index_name is not None:
                    output_names.append(index_name)
                output_records = [
                    _output_record(staging / name) for name in sorted(output_names)
                ]
                output_digests = tuple(sorted(item["sha256"] for item in output_records))
                try:
                    provenance_record = provenance_factory(output_digests)
                    provenance = _verify_provenance(
                        provenance_record.to_dict(),
                        output_digests=output_digests,
                        source_digests=tuple(
                            snapshot.sha256 for snapshot in source_snapshots
                        ),
                        configuration_digest=next(
                            snapshot.sha256
                            for snapshot in copy_snapshots
                            if snapshot.relative_path == "config.json"
                        ),
                        source_topology=source_topology,
                    )
                except CheckpointContractError:
                    raise
                except Exception as error:
                    raise CheckpointContractError(
                        "DCI_EVIDENCE_UNAVAILABLE",
                        detail="provenance_factory_failed",
                    ) from error
                artifact_id = provenance["artifact_id"]
                if getattr(provenance_record, "artifact_id", None) != artifact_id:
                    raise CheckpointContractError(
                        "DCI_EVIDENCE_UNAVAILABLE",
                        detail="provenance_artifact_identity_mismatch",
                    )
                _write_json(staging / "checkpoint-provenance.json", provenance)
                source_records = [snapshot.file_record() for snapshot in source_snapshots]
                manifest = _manifest(
                    artifact_id=artifact_id,
                    descriptor_digest=descriptor_digest,
                    source_records=source_records,
                    source_topology=provenance["source_topology"],
                    validation=validation,
                    output_records=output_records,
                    index_name=index_name,
                    admission=admission,
                    provenance=provenance,
                )
                _write_json(staging / "conversion-manifest.json", manifest)
                manifest_digest = _file_digest(staging / "conversion-manifest.json")
                _write_json(
                    staging / "abliteration_metadata.json",
                    {
                        "schema_version": 2,
                        "artifact_id": artifact_id,
                        "checkpoint_provenance": "checkpoint-provenance.json",
                        "conversion_manifest": "conversion-manifest.json",
                        "state_classification": "weights_only",
                        "lost_state": provenance["state"]["lost_state"],
                    },
                )
            except CheckpointContractError:
                raise
            except Exception as error:
                raise CheckpointContractError(
                    "DCI_MATERIALIZE_FAILED",
                    detail="canonical_write_failed",
                ) from error
    except CheckpointContractError:
        raise
    except Exception as error:
        raise CheckpointContractError(
            "DCI_PROMOTION_FAILED",
            detail="canonical_promotion_failed",
        ) from error
    return CanonicalWriteResult(
        artifact_id=artifact_id,
        output_path=output_path.resolve(),
        weight_files=shard_names,
        manifest_digest=manifest_digest,
    )