neuron-tai/packages/node/meshnet_node/split_gguf/manifest.py

"""Exact split-GGUF artifact manifest (DGR-026).

A split-GGUF artifact is only as trustworthy as its binding to the whole-model
artifact it was cut from. This module defines the manifest that makes a set of
split files an *exact*, checkable artifact rather than a pile of files someone
happened to name plausibly: it pins the source artifact's own content hash, the
tokenizer/revision the splits were tokenized against, and — per split — the
exact file name, size, cryptographic hash, and its range/role within the
source.

Quantization and split topology (how many splits, which layers each one
covers) are recipe inputs recorded on the manifest, never constants in this
module. A manifest with two splits and one with twenty are both valid; nothing
here assumes a stage count or a fixed layer range. Provisioning
(:mod:`meshnet_node.split_gguf.provision`) consumes whatever this manifest
declares.

This module mirrors two existing conventions rather than inventing new ones:
the DGR-017 pinned-shard manifest shape (`meshnet_node.glm_alpha.manifest`) for
per-file identity records, and the DGR-003 `DerivativeBinding` range/source
convention (`meshnet_node.runtime_recipe`) for binding a split to its source by
digest and half-open layer range.
"""

from __future__ import annotations

import hashlib
import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping

SPLIT_ARTIFACT_MANIFEST_SCHEMA_VERSION = 1

_SHA256_RE = re.compile(r"\A[0-9a-f]{64}\Z")
_REVISION_RE = re.compile(r"\A[0-9a-f]{40}\Z")


class SplitArtifactManifestError(ValueError):
    """Raised when a split-GGUF manifest is missing, malformed, or self-inconsistent."""


def canonical_sha256(value: Any) -> str:
    """SHA-256 over canonical JSON — the repository's digest convention."""
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def _require_mapping(value: Any, what: str, origin: str) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise SplitArtifactManifestError(
            f"{what} in {origin} must be a JSON object, got {type(value).__name__}"
        )
    return value


def _require_text(value: Any, what: str, origin: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise SplitArtifactManifestError(f"{what} in {origin} must be a non-empty string")
    return value


def _require_int(value: Any, what: str, origin: str, minimum: int = 0) -> int:
    if not isinstance(value, int) or isinstance(value, bool):
        raise SplitArtifactManifestError(
            f"{what} in {origin} must be an integer, got {type(value).__name__}"
        )
    if value < minimum:
        raise SplitArtifactManifestError(f"{what} in {origin} must be >= {minimum}, got {value}")
    return value


def _require_sha256(value: Any, what: str, origin: str) -> str:
    text = _require_text(value, what, origin)
    if not _SHA256_RE.match(text):
        raise SplitArtifactManifestError(
            f"{what} in {origin} must be a lowercase 64-character hex SHA-256, got {text!r}"
        )
    return text


def _require_revision(value: Any, what: str, origin: str) -> str:
    text = _require_text(value, what, origin)
    if not _REVISION_RE.match(text):
        raise SplitArtifactManifestError(
            f"{what} in {origin} must be a full 40-character commit revision, got {text!r}; "
            "a branch name, tag, or short SHA is not an immutable pin"
        )
    return text


@dataclass(frozen=True)
class SourceArtifact:
    """The whole-model artifact every split in this manifest was cut from."""

    artifact_id: str
    repo_id: str
    revision: str
    sha256: str
    size_bytes: int

    def to_dict(self) -> dict:
        return {
            "artifact_id": self.artifact_id,
            "repo_id": self.repo_id,
            "revision": self.revision,
            "sha256": self.sha256,
            "size_bytes": self.size_bytes,
        }


@dataclass(frozen=True)
class TokenizerRef:
    """The exact tokenizer/revision the split artifact's routing assumes."""

    repo_id: str
    revision: str
    sha256: str

    def to_dict(self) -> dict:
        return {"repo_id": self.repo_id, "revision": self.revision, "sha256": self.sha256}


@dataclass(frozen=True)
class SplitFile:
    """One split-GGUF file: name, size, hash, and its role/range in the source.

    `shard_start`/`shard_end` are half-open (end-exclusive), matching the
    `DerivativeBinding` protocol convention in `meshnet_node.runtime_recipe`.
    They are optional because not every split is a layer range — a shared
    embedding or tokenizer-adjacent split may carry only a `role` label — but
    when present they must describe a real, non-empty range.
    """

    name: str
    size_bytes: int
    sha256: str
    role: str
    url: str = ""
    shard_start: int | None = None
    shard_end: int | None = None

    def __post_init__(self) -> None:
        if (self.shard_start is None) != (self.shard_end is None):
            raise SplitArtifactManifestError(
                f"split {self.name!r} must declare both shard_start and shard_end, or neither"
            )
        if self.shard_start is not None and self.shard_end is not None:
            if self.shard_start < 0:
                raise SplitArtifactManifestError(f"split {self.name!r} shard_start must be >= 0")
            if self.shard_end <= self.shard_start:
                raise SplitArtifactManifestError(
                    f"split {self.name!r} shard_end ({self.shard_end}) must be greater than "
                    f"shard_start ({self.shard_start}); an empty range covers nothing"
                )

    @property
    def has_range(self) -> bool:
        return self.shard_start is not None

    def to_dict(self) -> dict:
        doc: dict[str, Any] = {
            "name": self.name,
            "size_bytes": self.size_bytes,
            "sha256": self.sha256,
            "role": self.role,
            "url": self.url,
        }
        if self.has_range:
            doc["shard_start"] = self.shard_start
            doc["shard_end"] = self.shard_end
        return doc


@dataclass(frozen=True)
class SplitArtifactManifest:
    """A parsed, self-consistent exact split-GGUF artifact manifest."""

    schema_version: int
    manifest_id: str
    manifest_version: str
    quantization: str
    source: SourceArtifact
    tokenizer: TokenizerRef
    total_bytes: int
    splits: tuple[SplitFile, ...]
    raw: Mapping[str, Any]
    origin: str = "<memory>"

    @property
    def digest(self) -> str:
        """Stable identity of this manifest, for binding into the DGR-003 recipe identity."""
        return canonical_sha256(self.raw)

    def split(self, name: str) -> SplitFile:
        for split in self.splits:
            if split.name == name:
                return split
        raise SplitArtifactManifestError(f"split {name!r} is not in {self.origin}")

    def to_dict(self) -> dict:
        return dict(self.raw)


def _parse_splits(raw: Any, expected_total: int, origin: str) -> tuple[SplitFile, ...]:
    if not isinstance(raw, list) or not raw:
        raise SplitArtifactManifestError(f"'splits' in {origin} must be a non-empty JSON array")

    splits: list[SplitFile] = []
    seen_names: set[str] = set()
    seen_sha: set[str] = set()
    for position, entry in enumerate(raw):
        item = _require_mapping(entry, f"splits[{position}]", origin)
        name = _require_text(item.get("name"), f"splits[{position}].name", origin)
        if name in seen_names:
            raise SplitArtifactManifestError(f"duplicate split name {name!r} in {origin}")
        seen_names.add(name)

        size_bytes = _require_int(item.get("size_bytes"), f"splits[{name}].size_bytes", origin, minimum=1)
        sha256 = _require_sha256(item.get("sha256"), f"splits[{name}].sha256", origin)
        if sha256 in seen_sha:
            raise SplitArtifactManifestError(
                f"split {name!r} repeats SHA-256 {sha256}; two distinct splits cannot "
                "have the same content digest"
            )
        seen_sha.add(sha256)

        role = _require_text(item.get("role"), f"splits[{name}].role", origin)
        url = item.get("url", "")
        if not isinstance(url, str):
            raise SplitArtifactManifestError(f"splits[{name}].url in {origin} must be a string")

        shard_start = item.get("shard_start")
        shard_end = item.get("shard_end")
        if shard_start is not None:
            shard_start = _require_int(shard_start, f"splits[{name}].shard_start", origin, minimum=0)
        if shard_end is not None:
            shard_end = _require_int(shard_end, f"splits[{name}].shard_end", origin, minimum=1)

        splits.append(
            SplitFile(
                name=name,
                size_bytes=size_bytes,
                sha256=sha256,
                role=role,
                url=url,
                shard_start=shard_start,
                shard_end=shard_end,
            )
        )

    summed = sum(split.size_bytes for split in splits)
    if summed != expected_total:
        raise SplitArtifactManifestError(
            f"declared total_bytes {expected_total} in {origin} does not equal the sum of "
            f"the split sizes {summed}; the manifest is not self-consistent"
        )

    return tuple(splits)


def parse_split_artifact_manifest(data: Any, origin: str = "<memory>") -> SplitArtifactManifest:
    """Validate an already-decoded split-artifact manifest document, failing closed."""
    doc = _require_mapping(data, "manifest root", origin)

    schema_version = _require_int(doc.get("schema_version"), "'schema_version'", origin, minimum=1)
    if schema_version != SPLIT_ARTIFACT_MANIFEST_SCHEMA_VERSION:
        raise SplitArtifactManifestError(
            f"{origin} declares split-artifact manifest schema version {schema_version}, "
            f"but this reader understands version {SPLIT_ARTIFACT_MANIFEST_SCHEMA_VERSION}"
        )

    manifest_id = _require_text(doc.get("manifest_id"), "'manifest_id'", origin)
    manifest_version = _require_text(doc.get("manifest_version"), "'manifest_version'", origin)
    quantization = _require_text(doc.get("quantization"), "'quantization'", origin)

    source_doc = _require_mapping(doc.get("source"), "'source'", origin)
    source = SourceArtifact(
        artifact_id=_require_text(source_doc.get("artifact_id"), "source.artifact_id", origin),
        repo_id=_require_text(source_doc.get("repo_id"), "source.repo_id", origin),
        revision=_require_revision(source_doc.get("revision"), "source.revision", origin),
        sha256=_require_sha256(source_doc.get("sha256"), "source.sha256", origin),
        size_bytes=_require_int(source_doc.get("size_bytes"), "source.size_bytes", origin, minimum=1),
    )

    tokenizer_doc = _require_mapping(doc.get("tokenizer"), "'tokenizer'", origin)
    tokenizer = TokenizerRef(
        repo_id=_require_text(tokenizer_doc.get("repo_id"), "tokenizer.repo_id", origin),
        revision=_require_revision(tokenizer_doc.get("revision"), "tokenizer.revision", origin),
        sha256=_require_sha256(tokenizer_doc.get("sha256"), "tokenizer.sha256", origin),
    )

    total_bytes = _require_int(doc.get("total_bytes"), "'total_bytes'", origin, minimum=1)
    splits = _parse_splits(doc.get("splits"), total_bytes, origin)

    return SplitArtifactManifest(
        schema_version=schema_version,
        manifest_id=manifest_id,
        manifest_version=manifest_version,
        quantization=quantization,
        source=source,
        tokenizer=tokenizer,
        total_bytes=total_bytes,
        splits=splits,
        raw=doc,
        origin=origin,
    )


def load_split_artifact_manifest(path: Path) -> SplitArtifactManifest:
    """Load and validate a split-artifact manifest from *path*."""
    try:
        raw = path.read_text(encoding="utf-8")
    except OSError as exc:
        raise SplitArtifactManifestError(f"cannot read split-artifact manifest {path}: {exc.strerror or exc}") from exc
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise SplitArtifactManifestError(
            f"{path} is not valid JSON: {exc.msg} at line {exc.lineno} column {exc.colno}"
        ) from exc
    return parse_split_artifact_manifest(data, origin=str(path))