neuron-tai/packages/validator/meshnet_validator/audit.py

"""TOPLOC activation proof helpers for validator-side audits."""

from __future__ import annotations

from dataclasses import dataclass
from importlib import import_module
from typing import Any, Literal


ProofEncoding = Literal["base64", "bytes"]


@dataclass(frozen=True)
class ToplocAuditConfig:
    """Canonical audit parameters for one model preset."""

    dtype: str = "bfloat16"
    quantization: str = "bfloat16"
    decode_batching_size: int = 32
    topk: int = 8
    skip_prefill: bool = True
    encoding: ProofEncoding = "base64"
    # ADR-0018 §3: nodes retain boundary activations only briefly; a commitment
    # older than this can no longer be verified against a live node and must
    # fall back to the text-only audit path.
    commitment_ttl_seconds: float = 30.0


@dataclass(frozen=True)
class ToplocProofClaim:
    """Prover-provided TOPLOC proof and the parameters it was built with."""

    proofs: Any
    dtype: str
    quantization: str
    decode_batching_size: int
    topk: int
    skip_prefill: bool = True
    encoding: ProofEncoding = "base64"

    @classmethod
    def from_mapping(cls, value: dict[str, Any]) -> "ToplocProofClaim":
        return cls(
            proofs=value["proofs"],
            dtype=str(value.get("dtype", "bfloat16")),
            quantization=str(value.get("quantization", "bfloat16")),
            decode_batching_size=int(value.get("decode_batching_size", 32)),
            topk=int(value.get("topk", 8)),
            skip_prefill=bool(value.get("skip_prefill", True)),
            encoding=_proof_encoding(value.get("encoding", "base64")),
        )

    def as_mapping(self) -> dict[str, Any]:
        return {
            "proofs": self.proofs,
            "dtype": self.dtype,
            "quantization": self.quantization,
            "decode_batching_size": self.decode_batching_size,
            "topk": self.topk,
            "skip_prefill": self.skip_prefill,
            "encoding": self.encoding,
        }


def build_activation_proofs(
    activations: list[Any],
    *,
    config: ToplocAuditConfig | None = None,
    backend: Any | None = None,
) -> ToplocProofClaim:
    """Build a TOPLOC proof claim from captured activation tensors."""
    cfg = config or ToplocAuditConfig()
    module = backend or _load_toploc()
    function_name = f"build_proofs_{cfg.encoding}"
    build = getattr(module, function_name)
    proofs = _call_toploc(
        build,
        activations,
        decode_batching_size=cfg.decode_batching_size,
        topk=cfg.topk,
        skip_prefill=cfg.skip_prefill,
    )
    return ToplocProofClaim(
        proofs=proofs,
        dtype=cfg.dtype,
        quantization=cfg.quantization,
        decode_batching_size=cfg.decode_batching_size,
        topk=cfg.topk,
        skip_prefill=cfg.skip_prefill,
        encoding=cfg.encoding,
    )


def verify_activation_proofs(
    reference_activations: list[Any],
    claim: ToplocProofClaim,
    *,
    config: ToplocAuditConfig | None = None,
    backend: Any | None = None,
) -> bool:
    """Verify prover TOPLOC proofs against reference teacher-forced activations."""
    cfg = config or ToplocAuditConfig(
        dtype=claim.dtype,
        quantization=claim.quantization,
        decode_batching_size=claim.decode_batching_size,
        topk=claim.topk,
        skip_prefill=claim.skip_prefill,
        encoding=claim.encoding,
    )
    if claim.dtype != cfg.dtype or claim.quantization != cfg.quantization:
        return False
    if claim.decode_batching_size != cfg.decode_batching_size or claim.topk != cfg.topk:
        return False
    if claim.skip_prefill != cfg.skip_prefill or claim.encoding != cfg.encoding:
        return False

    module = backend or _load_toploc()
    function_name = f"verify_proofs_{claim.encoding}"
    verify = getattr(module, function_name)
    return bool(_call_toploc(
        verify,
        reference_activations,
        claim.proofs,
        decode_batching_size=claim.decode_batching_size,
        topk=claim.topk,
        skip_prefill=claim.skip_prefill,
    ))


def _load_toploc() -> Any:
    try:
        return import_module("toploc")
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "toploc is required for activation proof audits; install meshnet-validator with dependencies"
        ) from exc


def _call_toploc(function: Any, activations: list[Any], *args: Any, **kwargs: Any) -> Any:
    try:
        return function(activations, *args, **kwargs)
    except TypeError:
        if kwargs:
            ordered = [
                kwargs["decode_batching_size"],
                kwargs["topk"],
                kwargs["skip_prefill"],
            ]
            return function(activations, *args, *ordered)
        raise


def _proof_encoding(value: object) -> ProofEncoding:
    if value == "bytes":
        return "bytes"
    return "base64"


__all__ = [
    "ToplocAuditConfig",
    "ToplocProofClaim",
    "build_activation_proofs",
    "verify_activation_proofs",
]