scripts/record_llama_pin.py

"""Record a llama.cpp build for `llama_runtime.json`, by downloading and hashing it.

The digest in that file has to be produced by someone who fetched the bytes and
hashed them. Copying a hash out of a release page or an API response would pin
the same origin that served the download, which proves nothing TLS had not
already said -- and this project withholds a *model* on exactly that reasoning.

Usage::

    python scripts/record_llama_pin.py windows-x64            # runs anywhere
    python scripts/record_llama_pin.py windows-x64-cuda       # CUDA, if present
    python scripts/record_llama_pin.py linux-x64 --tag b10742

An accelerated key records more than one archive, because an accelerated build
is not self-contained: llama.cpp ships the CUDA binaries in one zip and the CUDA
runtime libraries they link against in another, and a worker that fetched only
the first gets ``cudart64_12.dll was not found`` instead of a model. Both are
downloaded, both are hashed, and the fragment names both.

Prints a JSON fragment to paste into `llama_runtime.json`, and exits non-zero if
no matching asset was found. It never writes the file itself: adding a pin
should be a diff somebody reads.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
import urllib.request

RELEASES = "https://api.github.com/repos/ggml-org/llama.cpp/releases?per_page=%d"

#: What a platform's asset name looks like, and what must not be in it.
#:
#: The plain keys still exclude every accelerator: they are the build that runs
#: anywhere, and they remain the fallback every machine can reach. The
#: accelerated keys are additions beside them, never replacements, and are only
#: ever installed on a machine that has been checked for the hardware.
PATTERNS = {
    "windows-x64": (("win", "x64"), ("cuda", "hip", "sycl", "arm64", "vulkan")),
    "windows-arm64": (("win", "arm64"), ("cuda", "hip", "sycl")),
    "linux-x64": (("ubuntu", "x64"), ("cuda", "hip", "sycl", "arm64", "vulkan")),
    "macos-arm64": (("macos", "arm64"), ()),
    "macos-x64": (("macos", "x64"), ()),
    "windows-x64-cuda": (("win", "cuda", "x64"), ("arm64", "hip", "sycl")),
    "linux-x64-cuda": (("ubuntu", "cuda", "x64"), ("arm64", "hip", "sycl")),
}

#: The companion archive an accelerated build cannot start without.
COMPANIONS = {
    "windows-x64-cuda": ("cudart", "win", "x64"),
    "linux-x64-cuda": ("cudart", "ubuntu", "x64"),
}

#: ``llama-b10742-bin-win-cuda-12.4-x64.zip`` -> the CUDA version it needs.
_CUDA_IN_NAME = re.compile(r"cuda[-_]?(\d+)\.(\d+)")


def _get(url: str) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": "distinct-pin"})
    with urllib.request.urlopen(request, timeout=300) as response:  # noqa: S310
        return response.read()


def _hash(url: str) -> tuple[str, int]:
    """Stream the asset through a hash rather than holding it in memory.

    A CUDA build and its runtime libraries come to several hundred megabytes.
    Reading that into a string to hash it once works until the machine doing
    the recording is the volunteer's laptop, which is the only machine that
    ever does it.
    """

    digest = hashlib.sha256()
    size = 0
    request = urllib.request.Request(url, headers={"User-Agent": "distinct-pin"})
    with urllib.request.urlopen(request, timeout=300) as response:  # noqa: S310
        while True:
            chunk = response.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
            size += len(chunk)
            if size % (32 * 1024 * 1024) < 1024 * 1024:
                print(f"  ... {size // (1024 * 1024)} MB", file=sys.stderr, flush=True)
    return digest.hexdigest(), size


def _matches(name: str, wanted: tuple[str, ...], refused: tuple[str, ...]) -> bool:
    return (
        name.endswith(".zip")
        and all(part in name for part in wanted)
        and not any(part in name for part in refused)
    )


def _record(asset: dict) -> dict:
    url = asset["browser_download_url"]
    print(f"downloading {asset['name']} ({asset.get('size', 0)} bytes)", file=sys.stderr)
    digest, size = _hash(url)
    return {"asset": asset["name"], "url": url, "sha256": digest, "bytes": size}


def _cuda_version(*names: str) -> int:
    """The CUDA version an asset name declares, in NVML's integer form."""

    for name in names:
        found = _CUDA_IN_NAME.search(name.casefold())
        if found:
            return int(found.group(1)) * 1000 + int(found.group(2)) * 10
    return 0


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("platform", choices=sorted(PATTERNS))
    parser.add_argument("--tag", default="", help="a specific build tag, e.g. b10742")
    parser.add_argument("--scan", type=int, default=15, help="how many releases to look through")
    args = parser.parse_args(argv)

    wanted, refused = PATTERNS[args.platform]
    companion = COMPANIONS.get(args.platform)
    releases = json.loads(_get(RELEASES % args.scan))

    for release in releases:
        if args.tag and release.get("tag_name") != args.tag:
            continue
        assets = release.get("assets", [])
        main_asset = next(
            (item for item in assets if _matches(str(item.get("name", "")).casefold(), wanted, refused)),
            None,
        )
        if main_asset is None:
            continue

        archives = [_record(main_asset)]
        entry: dict = {"tag": release["tag_name"], "archives": archives}

        if companion is not None:
            # A CUDA build without its runtime libraries does not run slowly,
            # it does not start. Refusing to record half a pin is the whole
            # reason this script knows about companions at all.
            wanted_cuda = _cuda_version(main_asset["name"])
            partner = next(
                (
                    item
                    for item in assets
                    if _matches(str(item.get("name", "")).casefold(), companion, ())
                    and (not wanted_cuda or _cuda_version(str(item.get("name", ""))) == wanted_cuda)
                ),
                None,
            )
            if partner is None:
                print(
                    f"{main_asset['name']} was found but its {companion[0]} companion was not, "
                    "so nothing was recorded: half an accelerated pin is a worker that will "
                    "not start.",
                    file=sys.stderr,
                )
                return 1
            archives.append(_record(partner))
            entry["accelerator"] = "cuda"
            entry["requires_cuda"] = wanted_cuda

        print(json.dumps({args.platform: entry}, indent=2))
        return 0

    print(
        f"no asset matching {wanted} found in the last {args.scan} releases",
        file=sys.stderr,
    )
    return 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))