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

"""Resumable, hash-verifying provisioning of exact split-GGUF artifacts (DGR-026).

Model artifacts must use configured mounted-drive storage and never `/home`
(RALPH-CONTEXT). This module is the enforcement point: every entry point here
resolves and rejects a destination under `/home` before touching disk, mirroring
the existing `artifact_storage_root` check in
`meshnet_node.recipe_drivers._validate_config`.

Provisioning never trusts a partially-downloaded file. Each split is staged as
`<name>.partial` so an interrupted run resumes from the exact byte offset
already on disk — a `SplitFetcher` is handed that offset and is responsible for
continuing from it — and a partial is promoted to its final name only after its
SHA-256 matches the manifest exactly. A short, truncated, or hash-mismatched
split is deleted and raises rather than being silently accepted or left on disk
to be mistaken for complete later.
"""

from __future__ import annotations

import hashlib
import shutil
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from .manifest import SplitArtifactManifest, SplitFile

_CHUNK_SIZE = 4 * 1024 * 1024
_HOME_ROOT = Path("/home")


class SplitProvisionError(ValueError):
    """Raised when provisioning cannot produce a manifest-conformant local artifact."""


def reject_home_path(root: Path | str) -> Path:
    """Resolve *root* and fail closed if it is (or is under) `/home`.

    Does not require *root* to exist yet — provisioning creates it — so this
    performs the same structural check as
    `meshnet_node.recipe_drivers._validate_config` without `strict=True`.
    """
    resolved = Path(root).expanduser().resolve()
    if not resolved.is_absolute() or resolved == _HOME_ROOT or _HOME_ROOT in resolved.parents:
        raise SplitProvisionError(
            f"refusing to provision split-GGUF artifacts under {resolved}: model artifacts "
            "must use configured mounted-drive storage, never /home"
        )
    return resolved


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(_CHUNK_SIZE), b""):
            digest.update(chunk)
    return digest.hexdigest()


# fetch(split, partial_dest, resume_from_bytes) must, on success, leave
# partial_dest containing exactly the bytes of `split` starting from byte 0,
# with total length equal to split.size_bytes; resume_from_bytes bytes are
# already present at the start of partial_dest and must not be re-fetched.
SplitFetcher = Callable[[SplitFile, Path, int], None]


def local_directory_fetcher(source_dir: Path) -> SplitFetcher:
    """A fetcher that copies split bytes from files already present in *source_dir*.

    No network access. Used by deterministic tests against tiny local
    fixtures, and for provisioning from splits already staged on another local
    or mounted path (e.g. a pre-synced mirror).
    """
    source_dir = Path(source_dir)

    def fetch(split: SplitFile, dest: Path, resume_from_bytes: int) -> None:
        source_path = source_dir / split.name
        if not source_path.is_file():
            raise SplitProvisionError(f"split source is missing: {source_path}")
        mode = "r+b" if resume_from_bytes else "wb"
        dest.parent.mkdir(parents=True, exist_ok=True)
        if not dest.exists():
            dest.touch()
        with source_path.open("rb") as src, dest.open(mode) as out:
            src.seek(resume_from_bytes)
            out.seek(resume_from_bytes)
            out.truncate(resume_from_bytes)
            shutil.copyfileobj(src, out, length=_CHUNK_SIZE)

    return fetch


def http_split_fetcher(url_for: Callable[[SplitFile], str], timeout: float = 30.0) -> SplitFetcher:
    """A fetcher that downloads each split over HTTP(S) with Range-header resume.

    Falls back to a full restart if the server ignores the `Range` request
    (some static hosts return `200` with the whole body instead of `206`).
    """

    def fetch(split: SplitFile, dest: Path, resume_from_bytes: int) -> None:
        request = urllib.request.Request(url_for(split))
        if resume_from_bytes:
            request.add_header("Range", f"bytes={resume_from_bytes}-")
        dest.parent.mkdir(parents=True, exist_ok=True)
        with urllib.request.urlopen(request, timeout=timeout) as resp:
            resumed = bool(resume_from_bytes) and getattr(resp, "status", 200) == 206
            with dest.open("ab" if resumed else "wb") as out:
                shutil.copyfileobj(resp, out, length=_CHUNK_SIZE)

    return fetch


@dataclass(frozen=True)
class ProvisionResult:
    dest_dir: Path
    verified_splits: tuple[str, ...]

    def to_dict(self) -> dict:
        return {"dest_dir": str(self.dest_dir), "verified_splits": list(self.verified_splits)}


def provision_split_artifact(
    manifest: SplitArtifactManifest,
    dest_dir: Path,
    fetch: SplitFetcher,
) -> ProvisionResult:
    """Provision every split in *manifest* under *dest_dir*: resumable, hash-verified.

    Refuses any destination under `/home`. A split already present at the
    correct size and hash is left untouched (a re-run is a no-op); a file
    present with the wrong size or hash is deleted and re-fetched rather than
    trusted. On success every split is byte- and hash-verified against the
    manifest before this function returns.
    """
    dest_dir = reject_home_path(dest_dir)
    dest_dir.mkdir(parents=True, exist_ok=True)

    verified: list[str] = []
    for split in manifest.splits:
        final_path = dest_dir / split.name
        if (
            final_path.is_file()
            and final_path.stat().st_size == split.size_bytes
            and _sha256_file(final_path) == split.sha256
        ):
            verified.append(split.name)
            continue
        if final_path.is_file():
            final_path.unlink()

        partial_path = dest_dir / f"{split.name}.partial"
        resume_from = partial_path.stat().st_size if partial_path.is_file() else 0
        if resume_from > split.size_bytes:
            partial_path.unlink()
            resume_from = 0

        if resume_from < split.size_bytes:
            fetch(split, partial_path, resume_from)

        actual_size = partial_path.stat().st_size if partial_path.is_file() else 0
        if actual_size != split.size_bytes:
            raise SplitProvisionError(
                f"split {split.name!r} is incomplete after provisioning: "
                f"got {actual_size} of {split.size_bytes} bytes"
            )

        actual_sha256 = _sha256_file(partial_path)
        if actual_sha256 != split.sha256:
            partial_path.unlink()
            raise SplitProvisionError(
                f"split {split.name!r} hash mismatch: expected {split.sha256}, got {actual_sha256}"
            )

        partial_path.replace(final_path)
        verified.append(split.name)

    verify_provisioned_split_artifact(manifest, dest_dir)
    return ProvisionResult(dest_dir=dest_dir, verified_splits=tuple(verified))


def verify_provisioned_split_artifact(manifest: SplitArtifactManifest, dest_dir: Path) -> None:
    """Fail closed unless every manifest split is present, complete, and hash-exact.

    This is the check a downstream loader — or a resumed provisioning run —
    should call before trusting *dest_dir*, so a partially-provisioned
    directory is never mistaken for a ready artifact.
    """
    dest_dir = reject_home_path(dest_dir)
    missing: list[str] = []
    mismatched: list[str] = []
    for split in manifest.splits:
        path = dest_dir / split.name
        if not path.is_file():
            missing.append(split.name)
            continue
        if path.stat().st_size != split.size_bytes:
            mismatched.append(split.name)
            continue
        if _sha256_file(path) != split.sha256:
            mismatched.append(split.name)

    if missing:
        raise SplitProvisionError(f"missing split(s) in {dest_dir}: {sorted(missing)}")
    if mismatched:
        raise SplitProvisionError(f"hash/size mismatch for split(s) in {dest_dir}: {sorted(mismatched)}")