Teknium
326 lines
13 KiB
Python
326 lines
13 KiB
Python
"""Binary acquisition for the managed llama.cpp runtime."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import urllib.request
|
|
import zipfile
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
RELEASE_URL = "https://github.com/ggml-org/llama.cpp/releases/download/{tag}/{asset}"
|
|
|
|
# Windows CUDA zips ship per CUDA major; the runtime zip must be paired with
|
|
# its cudart zip so end users need no toolkit. 13.3 verified on 13.1 and
|
|
# 13.2 drivers.
|
|
_WIN_CUDA_VERSION = "13.3"
|
|
# arm64 Windows CUDA prebuilts landed upstream (~b1036x) on CUDA 13.4 —
|
|
# verified against live asset lists (b10362, b10630, b10679). Tags at or before
|
|
# b10290 don't have them; resolution succeeds and the download 404s
|
|
# honestly on such tags, which only arises if a user pins backward.
|
|
_WIN_CUDA_VERSION_ARM64 = "13.4"
|
|
|
|
|
|
# Fallback when the config section is missing entirely (deep-merge normally
|
|
# guarantees the key). Single source: DEFAULT_CONFIG owns the shipped tag.
|
|
def default_tag() -> str:
|
|
from hermes_cli.config_defaults import DEFAULT_CONFIG
|
|
|
|
return DEFAULT_CONFIG["local_runtime"]["tag"]
|
|
|
|
|
|
class BinaryResolutionError(RuntimeError):
|
|
"""No usable asset combination for this platform/backend."""
|
|
|
|
|
|
@dataclass
|
|
class AssetPlan:
|
|
"""The exact zips one runtime install needs, in extraction order."""
|
|
|
|
tag: str
|
|
backend: str # cuda | metal | vulkan | hip | cpu
|
|
assets: list[str] = field(default_factory=list)
|
|
|
|
@property
|
|
def install_dir(self) -> Path:
|
|
return runtimes_root() / self.tag / self.backend
|
|
|
|
|
|
def runtimes_root() -> Path:
|
|
"""Machine-scoped, deliberately NOT profile-scoped. Engine binaries, presets, and server state
|
|
describe this machine's hardware and its one managed server (stable port) — a second profile
|
|
re-downloading the engine or fighting over the port would be the bug. Profile-scoped things
|
|
(which model is the default, enabled) live in each profile's config.yaml as ever.
|
|
"""
|
|
from hermes_constants import get_default_hermes_root
|
|
|
|
return get_default_hermes_root() / "runtimes" / "llamacpp"
|
|
|
|
|
|
def installed_tags() -> list[str]:
|
|
"""Tags with a verified install (manifest carries verified_version),
|
|
newest first by release number. The boot ladder and the update check
|
|
both read installed-ness from here — one resolver, every caller."""
|
|
root = runtimes_root()
|
|
if not root.exists():
|
|
return []
|
|
found: list[str] = []
|
|
for entry in root.iterdir():
|
|
if not entry.is_dir() or entry.name == "downloads":
|
|
continue
|
|
for manifest in entry.glob("*/manifest.json"):
|
|
try:
|
|
if json.loads(manifest.read_text(encoding="utf-8")).get("verified_version"):
|
|
found.append(entry.name)
|
|
break
|
|
except (json.JSONDecodeError, OSError):
|
|
continue
|
|
|
|
def _release_number(tag: str) -> int:
|
|
digits = "".join(ch for ch in tag if ch.isdigit())
|
|
return int(digits) if digits else 0
|
|
|
|
return sorted(set(found), key=_release_number, reverse=True)
|
|
|
|
|
|
def _host_os_arch() -> tuple[str, str]:
|
|
"""(os, arch) normalized to release-asset vocabulary.
|
|
|
|
PITFALL: PROCESSOR_ARCHITECTURE lies under x64 emulation on ARM64 Windows. platform.machine()
|
|
reads the same env on some Pythons, so on Windows prefer PROCESSOR_IDENTIFIER's text when
|
|
present.
|
|
"""
|
|
system = platform.system().lower()
|
|
os_name = {"windows": "win", "darwin": "macos", "linux": "ubuntu"}.get(system, system)
|
|
machine = platform.machine().lower()
|
|
arch = "arm64" if machine in ("arm64", "aarch64") else "x64"
|
|
if os_name == "win":
|
|
import os as _os
|
|
ident = _os.environ.get("PROCESSOR_IDENTIFIER", "").lower()
|
|
if "armv8" in ident or "arm " in ident:
|
|
arch = "arm64"
|
|
return os_name, arch
|
|
|
|
|
|
def select_backend(gpu_vendor: str | None, os_name: str | None = None) -> str:
|
|
"""Backend choice per design: CUDA if NVIDIA, Metal on macOS, Vulkan if
|
|
a non-NVIDIA GPU is present, else CPU. ``--list-devices`` validates the
|
|
choice post-install; the supervisor's touch generation is ground truth."""
|
|
if os_name is None:
|
|
os_name, _ = _host_os_arch()
|
|
if os_name == "macos":
|
|
return "metal"
|
|
vendor = (gpu_vendor or "").lower()
|
|
if "nvidia" in vendor:
|
|
return "cuda"
|
|
if vendor in ("amd", "intel") or "radeon" in vendor or "arc" in vendor:
|
|
return "vulkan"
|
|
return "cpu"
|
|
|
|
|
|
# Per-OS (human label, {backend: asset-name templates}). Windows CUDA pairs the
|
|
# runtime zip with its cudart zip; ubuntu ships tarballs, win ships zips.
|
|
_ASSET_TEMPLATES = {
|
|
"ubuntu": ("linux", {
|
|
"vulkan": ["llama-{tag}-bin-ubuntu-vulkan-{arch}.tar.gz"],
|
|
"hip": ["llama-{tag}-bin-ubuntu-rocm-7.2-{arch}.tar.gz"],
|
|
"cpu": ["llama-{tag}-bin-ubuntu-{arch}.tar.gz"],
|
|
}),
|
|
"win": ("windows", {
|
|
"cuda": ["llama-{tag}-bin-win-cuda-{cuda_ver}-{arch}.zip",
|
|
"cudart-llama-bin-win-cuda-{cuda_ver}-{arch}.zip"],
|
|
"vulkan": ["llama-{tag}-bin-win-vulkan-x64.zip"],
|
|
"hip": ["llama-{tag}-bin-win-hip-radeon-x64.zip"],
|
|
"cpu": ["llama-{tag}-bin-win-cpu-{arch}.zip"],
|
|
}),
|
|
}
|
|
|
|
|
|
def resolve_assets(tag: str, backend: str, os_name: str | None = None,
|
|
arch: str | None = None) -> AssetPlan:
|
|
"""Compose the asset list for (tag, backend, platform).
|
|
|
|
Raises BinaryResolutionError for platform/backend pairs the release ships no artifact for;
|
|
callers fall back down the backend ladder cuda -> vulkan -> cpu.
|
|
"""
|
|
host_os, host_arch = _host_os_arch()
|
|
os_name = os_name or host_os
|
|
arch = arch or host_arch
|
|
if os_name == "macos":
|
|
# macOS tarballs are unified (Metal built in).
|
|
return AssetPlan(tag, backend, [f"llama-{tag}-bin-macos-{arch}.tar.gz"])
|
|
if os_name not in _ASSET_TEMPLATES:
|
|
raise BinaryResolutionError(f"unsupported platform {os_name}-{arch}")
|
|
if os_name == "ubuntu" and backend == "cuda":
|
|
# No prebuilt Linux CUDA zips at current tags — Linux CUDA users
|
|
# build from source or use vulkan; resolver is honest about it.
|
|
raise BinaryResolutionError(
|
|
f"no prebuilt linux CUDA asset at {tag}; use vulkan/cpu or a source build")
|
|
if os_name == "win" and backend == "vulkan" and arch == "arm64":
|
|
raise BinaryResolutionError(f"no win-vulkan-arm64 asset at {tag}")
|
|
label, templates = _ASSET_TEMPLATES[os_name]
|
|
if backend not in templates:
|
|
raise BinaryResolutionError(f"unsupported {label} backend {backend}")
|
|
cuda_ver = _WIN_CUDA_VERSION_ARM64 if arch == "arm64" else _WIN_CUDA_VERSION
|
|
return AssetPlan(tag, backend, [t.format(tag=tag, arch=arch, cuda_ver=cuda_ver)
|
|
for t in templates[backend]])
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 22), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def _download(url: str, dest: Path,
|
|
progress: "Callable[[int, int], None] | None" = None) -> None:
|
|
"""Stream url -> dest. ``progress(done_bytes, total_bytes)`` ticks per
|
|
chunk (total 0 when the server sends no Content-Length) — a several-
|
|
hundred-MB archive on a slow line must never look hung."""
|
|
logger.info("downloading %s", url)
|
|
tmp = dest.with_suffix(dest.suffix + ".part")
|
|
with urllib.request.urlopen(url, timeout=120) as r, open(tmp, "wb") as f:
|
|
total = int(r.headers.get("Content-Length") or 0)
|
|
done = 0
|
|
while True:
|
|
chunk = r.read(1 << 20)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
done += len(chunk)
|
|
if progress is not None:
|
|
progress(done, total)
|
|
tmp.replace(dest)
|
|
|
|
|
|
def _extract(archive: Path, dest: Path,
|
|
progress: "Callable[[int, int], None] | None" = None) -> None:
|
|
"""Extract member by member so ``progress(done, total)`` can tick in
|
|
uncompressed bytes — big archives take real time on laptop disks."""
|
|
if archive.name.endswith(".zip"):
|
|
opener, list_members, size = zipfile.ZipFile, "infolist", "file_size"
|
|
kwargs = {}
|
|
else:
|
|
import tarfile
|
|
opener, list_members, size = tarfile.open, "getmembers", "size"
|
|
kwargs = {"filter": "data"}
|
|
with opener(archive) as ar:
|
|
members = getattr(ar, list_members)()
|
|
total = sum(getattr(m, size) for m in members)
|
|
done = 0
|
|
for m in members:
|
|
ar.extract(m, dest, **kwargs)
|
|
done += getattr(m, size)
|
|
if progress is not None:
|
|
progress(done, total)
|
|
|
|
|
|
def server_binary(install_dir: Path) -> Path:
|
|
"""Locate llama-server within an extracted runtime (zips differ in
|
|
whether they nest a build/bin directory)."""
|
|
names = ("llama-server.exe", "llama-server")
|
|
for name in names:
|
|
direct = install_dir / name
|
|
if direct.exists():
|
|
return direct
|
|
for name in names:
|
|
hits = sorted(install_dir.rglob(name))
|
|
if hits:
|
|
return hits[0]
|
|
raise BinaryResolutionError(f"llama-server not found under {install_dir}")
|
|
|
|
|
|
def verify_install(install_dir: Path, tag: str) -> str:
|
|
"""Run --version; require the tag's build number in the output.
|
|
(The binary prints the tag WITHOUT the 'b' prefix.)"""
|
|
exe = server_binary(install_dir)
|
|
out = subprocess.run([str(exe), "--version"], capture_output=True,
|
|
text=True, encoding="utf-8", errors="replace",
|
|
timeout=60, cwd=str(exe.parent))
|
|
text = (out.stdout + out.stderr).strip()
|
|
if tag.lstrip("b") not in text:
|
|
raise BinaryResolutionError(
|
|
f"version check failed for {exe}: expected {tag}, got: {text[:120]}")
|
|
return text.splitlines()[0] if text else ""
|
|
|
|
|
|
def prune_old_tags(keep: list[str]) -> None:
|
|
"""Retain only the tags in ``keep`` (current + previous — N-1 rollback).
|
|
The shared ``downloads/`` archive cache is not a tag and always survives."""
|
|
root = runtimes_root()
|
|
if not root.exists():
|
|
return
|
|
for entry in root.iterdir():
|
|
if entry.is_dir() and entry.name != "downloads" and entry.name not in keep:
|
|
shutil.rmtree(entry, ignore_errors=True)
|
|
logger.info("pruned old runtime %s", entry.name)
|
|
|
|
|
|
def ensure_runtime_installed(tag: str, backend: str,
|
|
expected_sha256: dict[str, str] | None = None,
|
|
progress: "Callable[[str, int, int, str], None] | None" = None) -> Path:
|
|
"""Idempotent: resolve, download, verify, extract, version-check.
|
|
|
|
``expected_sha256`` pins hashes per asset when the catalog has them; without pins the computed
|
|
hash is recorded in the manifest (trust on first download, verified on every reinstall).
|
|
``progress(stage, done, total, label)`` ticks through download/extract/verify. Returns the
|
|
install dir containing llama-server.
|
|
"""
|
|
plan = resolve_assets(tag, backend)
|
|
install_dir = plan.install_dir
|
|
manifest_path = install_dir / "manifest.json"
|
|
if manifest_path.exists():
|
|
try:
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
if manifest.get("verified_version"):
|
|
return install_dir
|
|
except (json.JSONDecodeError, OSError):
|
|
pass # damaged manifest -> reinstall
|
|
|
|
install_dir.mkdir(parents=True, exist_ok=True)
|
|
downloads = runtimes_root() / "downloads"
|
|
downloads.mkdir(parents=True, exist_ok=True)
|
|
|
|
recorded: dict[str, str] = {}
|
|
n_assets = len(plan.assets)
|
|
for i, asset in enumerate(plan.assets, 1):
|
|
label = f"{i}/{n_assets}" if n_assets > 1 else ""
|
|
archive = downloads / asset
|
|
if not archive.exists():
|
|
_download(RELEASE_URL.format(tag=tag, asset=asset), archive,
|
|
progress=(lambda d, t, _l=label: progress("download", d, t, _l))
|
|
if progress is not None else None)
|
|
if progress is not None:
|
|
progress("verify", 0, 0, label)
|
|
digest = _sha256(archive)
|
|
expected = (expected_sha256 or {}).get(asset)
|
|
if expected and digest != expected:
|
|
archive.unlink(missing_ok=True)
|
|
raise BinaryResolutionError(
|
|
f"sha256 mismatch for {asset}: expected {expected}, got {digest}")
|
|
recorded[asset] = digest
|
|
_extract(archive, install_dir,
|
|
progress=(lambda d, t, _l=label: progress("extract", d, t, _l))
|
|
if progress is not None else None)
|
|
|
|
if progress is not None:
|
|
progress("verify", 0, 0, "")
|
|
version = verify_install(install_dir, tag)
|
|
manifest_path.write_text(json.dumps({
|
|
"tag": tag, "backend": plan.backend, "assets": recorded,
|
|
"verified_version": version,
|
|
}, indent=2), encoding="utf-8")
|
|
logger.info("installed llama.cpp %s (%s): %s", tag, backend, version)
|
|
return install_dir
|