254 lines
10 KiB
Python
254 lines
10 KiB
Python
"""Hardware selection for PM-owned llama.cpp binaries; mutable runtime state stays separate."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import threading
|
|
from contextlib import contextmanager, suppress
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable, Iterator
|
|
|
|
import pm
|
|
from pm.downloader import ProgressFn
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BACKEND_PACKAGES = {
|
|
"cuda": "llamacpp-cuda",
|
|
"vulkan": "llamacpp-vulkan",
|
|
"metal": "llamacpp-metal",
|
|
"hip": "llamacpp-hip",
|
|
"cpu": "llamacpp-cpu",
|
|
}
|
|
|
|
|
|
class BinaryResolutionError(RuntimeError):
|
|
"""The requested backend has no pinned or installed engine."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Engine:
|
|
backend: str
|
|
tag: str
|
|
binary: Path
|
|
|
|
|
|
def runtimes_root() -> Path:
|
|
"""Machine-scoped presets and server state. Binaries belong to PM's store."""
|
|
from hermes_constants import get_default_hermes_root
|
|
|
|
return get_default_hermes_root() / "runtimes" / "llamacpp"
|
|
|
|
|
|
def pinned_tag(backend: str) -> str:
|
|
version = pm.Lockfile(pm.paths.lockfile_path()).version(BACKEND_PACKAGES[backend])
|
|
if version is None:
|
|
raise BinaryResolutionError(f"llama.cpp {backend} has no PM version pin")
|
|
return f"b{version}"
|
|
|
|
|
|
def select_backend(gpu_vendor: str | None, os_name: str | None = None) -> str:
|
|
if os_name is None:
|
|
os_name = "macos" if pm.current_target().startswith("darwin-") else "other"
|
|
if os_name == "macos":
|
|
return "metal"
|
|
vendor = (gpu_vendor or "").lower()
|
|
if "nvidia" in vendor:
|
|
return "cuda"
|
|
if any(name in vendor for name in ("amd", "intel", "radeon", "arc")):
|
|
return "vulkan"
|
|
return "cpu"
|
|
|
|
|
|
def _candidates(requested: str, gpu_vendor: str | None, target: str) -> tuple[str, ...]:
|
|
if requested != "auto":
|
|
if requested not in BACKEND_PACKAGES:
|
|
raise BinaryResolutionError(f"unknown backend {requested}")
|
|
return (requested,)
|
|
preferred = select_backend(gpu_vendor, "macos" if target.startswith("darwin-") else "other")
|
|
fallback = ("vulkan", "cpu") if gpu_vendor else ("cpu",)
|
|
return tuple(dict.fromkeys((preferred, *fallback)))
|
|
|
|
|
|
def unavailable_reason(backend: str, target: str | None = None) -> str | None:
|
|
name = BACKEND_PACKAGES.get(backend)
|
|
if name is None:
|
|
return f"unknown backend {backend}"
|
|
target = target or pm.current_target()
|
|
reason = pm.get_package(name).missing_reason(target)
|
|
if reason:
|
|
return reason
|
|
lock = pm.Lockfile(pm.paths.lockfile_path())
|
|
if not lock.version(name) or not lock.artifacts(name, target):
|
|
return f"{name} is not pinned for {target}"
|
|
return None
|
|
|
|
|
|
def resolve_backend(requested: str = "auto", *, gpu_vendor: str | None = None,
|
|
target: str | None = None) -> str:
|
|
"""Explicit choices are strict. Auto only falls back to compatible pinned builds."""
|
|
if requested == "auto" and target is None and gpu_vendor is None:
|
|
from hermes_cli.local_runtime.bootstrap import _detect_gpu_vendor
|
|
|
|
gpu_vendor = _detect_gpu_vendor()
|
|
target = target or pm.current_target()
|
|
reasons = []
|
|
for backend in _candidates(requested, gpu_vendor, target):
|
|
reason = unavailable_reason(backend, target)
|
|
if reason is None:
|
|
return backend
|
|
reasons.append(reason)
|
|
raise BinaryResolutionError("; ".join(reasons))
|
|
|
|
|
|
def _installed(candidates: tuple[str, ...], allow_outdated: bool) -> Engine | None:
|
|
for candidate in candidates:
|
|
found = pm.installed_package(BACKEND_PACKAGES[candidate], allow_outdated=allow_outdated)
|
|
if found is not None and found.binary is not None:
|
|
return Engine(candidate, f"b{found.version}", found.binary)
|
|
return None
|
|
|
|
|
|
def installed_engine(backend: str = "auto", *, allow_outdated: bool = True) -> Engine | None:
|
|
"""Boot may retain a prior PM pin and moves a pre-PM engine into PM's store once; it never downloads."""
|
|
vendor = None
|
|
if backend == "auto":
|
|
from hermes_cli.local_runtime.bootstrap import _detect_gpu_vendor
|
|
|
|
vendor = _detect_gpu_vendor()
|
|
candidates = _candidates(backend, vendor, pm.current_target())
|
|
engine = _installed(candidates, allow_outdated)
|
|
if engine is None and any(adopt_legacy_engine(candidate) for candidate in candidates):
|
|
engine = _installed(candidates, allow_outdated)
|
|
return engine
|
|
|
|
|
|
# Before PM owned binaries, Hermes installed each engine to runtimes/llamacpp/b<tag>/<backend>/
|
|
# and wrote a manifest.json with the archive digests and the llama-server --version it saw.
|
|
_SHA256 = re.compile(r"[0-9a-f]{64}")
|
|
_ADOPTION_LOCK = threading.Lock()
|
|
# Installs that failed to move stay put for this process; the pane polls status every few seconds.
|
|
_LEFT_IN_PLACE: set[Path] = set()
|
|
|
|
|
|
def _legacy_installs(backend: str) -> list[tuple[str, Path, dict]]:
|
|
"""Verified pre-PM installs of ``backend`` as (version, dir, manifest), newest first."""
|
|
found = []
|
|
for manifest_path in runtimes_root().glob(f"b*/{backend}/manifest.json"):
|
|
tag = manifest_path.parent.parent.name
|
|
try:
|
|
number = int(tag.removeprefix("b"))
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8-sig"))
|
|
except (OSError, ValueError):
|
|
continue
|
|
if (isinstance(manifest, dict) and manifest.get("tag") == tag and manifest.get("backend") == backend
|
|
and manifest.get("verified_version") and isinstance(manifest.get("assets"), dict)
|
|
and manifest_path.parent not in _LEFT_IN_PLACE):
|
|
found.append((number, str(number), manifest_path.parent, manifest))
|
|
return [(version, path, manifest) for _, version, path, manifest in sorted(found, reverse=True)]
|
|
|
|
|
|
def _legacy_artifacts(package, version: str, target: str, assets: dict) -> list[str] | None:
|
|
"""The manifest's archive digests in PM's archive order, or None when one is missing.
|
|
|
|
The pre-PM installer fetched only llama.cpp's own release assets. On Linux PM also pins the
|
|
libgomp .deb, so an adopted engine there records fewer archives than the lock and counts as
|
|
outdated: usable now (its manifest proves it ran against the host's libgomp), update offered."""
|
|
shas = [assets.get(name) for name in package._asset_names(version, target)]
|
|
if not shas or not all(isinstance(sha, str) and _SHA256.fullmatch(sha) for sha in shas):
|
|
return None
|
|
return shas
|
|
|
|
|
|
@contextmanager
|
|
def _store_lock(root: Path) -> Iterator[bool]:
|
|
"""PM's store lock, or False when another PM operation holds it (a download can take minutes)."""
|
|
from pm.filesystem import lock_fd
|
|
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
fd = os.open(root / ".install.lock", os.O_CREAT | os.O_RDWR, 0o600)
|
|
try:
|
|
yield lock_fd(fd, wait=True, timeout=2)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def adopt_legacy_engine(backend: str) -> bool:
|
|
"""Move the newest pre-PM install of ``backend`` into PM's store and record it as installed.
|
|
|
|
A machine that already ran a Hermes-installed engine keeps it across the update to PM. The
|
|
manifest's archive digests become the PM identity, so a tag that matches the pin counts as
|
|
current and an older tag counts as outdated, which offers the update. ``os.rename`` only:
|
|
instant on one volume, and a store on another volume leaves the engine where it is. Never
|
|
raises; returns whether an engine was moved.
|
|
"""
|
|
try:
|
|
candidates = _legacy_installs(backend)
|
|
if not candidates:
|
|
return False
|
|
from pm import paths as pm_paths
|
|
|
|
name = BACKEND_PACKAGES[backend]
|
|
package = pm.get_package(name)
|
|
target = pm.current_target()
|
|
root = pm_paths.writable_store_root()
|
|
facts_path = pm_paths.facts_path() if root == pm_paths.store_root() else root / "facts.json"
|
|
with _ADOPTION_LOCK, _store_lock(root) as held:
|
|
if not held or pm.installed_package(name, allow_outdated=True) is not None:
|
|
return False
|
|
for version, source, manifest in candidates:
|
|
if _adopt(package, version, target, source, manifest, root, facts_path):
|
|
return True
|
|
_LEFT_IN_PLACE.add(source)
|
|
except Exception as exc:
|
|
logger.warning("could not move the pre-PM llama.cpp %s engine into the PM store: %s", backend, exc)
|
|
return False
|
|
|
|
|
|
def _adopt(package, version: str, target: str, source: Path, manifest: dict, root: Path,
|
|
facts_path: Path) -> bool:
|
|
from pm.store import tree_digest
|
|
|
|
artifacts = _legacy_artifacts(package, version, target, manifest["assets"])
|
|
entry_name = package.store_entry(version, target)
|
|
entry = root / entry_name
|
|
if artifacts is None or entry.exists() or entry.is_symlink():
|
|
return False
|
|
reason = package.verify(source, target)
|
|
if reason:
|
|
logger.warning("pre-PM llama.cpp at %s left in place: %s", source, reason)
|
|
return False
|
|
try:
|
|
os.rename(source, entry)
|
|
except OSError as exc:
|
|
logger.warning("pre-PM llama.cpp at %s left in place: %s", source, exc)
|
|
return False
|
|
try:
|
|
pm.Facts(facts_path).record(package.name, version, entry_name, package.env(entry, target), root,
|
|
target=target, artifacts=artifacts, digest=tree_digest(entry))
|
|
except BaseException:
|
|
with suppress(OSError):
|
|
os.rename(entry, source)
|
|
raise
|
|
with suppress(OSError):
|
|
source.parent.rmdir()
|
|
logger.info("moved llama.cpp b%s (%s) from %s into the PM store", version, package.name, source)
|
|
return True
|
|
|
|
|
|
def ensure_engine(backend: str, *, progress: Callable[[str, int, int, str], None] | None = None,
|
|
pause_event: threading.Event | None = None,
|
|
download_progress: ProgressFn | None = None) -> Engine:
|
|
"""Only deliberate install/update jobs call this. PM verifies every archive and publishes."""
|
|
resolved = resolve_backend(backend)
|
|
pm.ensure(BACKEND_PACKAGES[resolved], explicit=True, progress=progress,
|
|
pause_event=pause_event, download_progress=download_progress)
|
|
engine = installed_engine(resolved, allow_outdated=False)
|
|
if engine is None:
|
|
raise BinaryResolutionError(f"llama.cpp {resolved} install has no usable pinned binary")
|
|
return engine
|