Spaces:

User1342
/

distinct

Sleeping

Drawn mockups instead of recorded demos; delete conversations; drop Llama; fix the weights cache re-downloading
1d4a48e about 2 months ago
27.5 kB
"""Fetch model weights on demand from Hugging Face, into a capped cache.
Why this module replaced ``--model ID=PATH``
--------------------------------------------
Agent setup used to hand the worker a filesystem path to a GGUF file the
operator had already downloaded by some route the worker never saw. That is
backwards on two counts.
**It makes the operator the integration.** The server publishes a catalogue of
models it may ask an agent to run. The operator's decision is which of those
they are willing to host on their hardware. Turning that decision into "find
the right file, put it somewhere, and type its path correctly" is asking a
volunteer to do a build step, and every step of it can go wrong silently: a
requantisation from a different publisher has the same filename and different
weights.
**It makes identity unverifiable in practice.** A digest is only meaningful
against a repository and a revision. A path says nothing about where the file
came from, so a path-configured worker could only ever report "no digest
recorded", which is the state this project treats as an unanswered question
rather than a passing check.
So the flow is now: the operator names *models*, never paths. The worker
resolves each to the pinned repository, revision and filename, downloads it
when a run first needs it, verifies it against the published digest before it
is ever loaded, and keeps it in a cache whose size the operator sets at run.
The rule that follows, and it is deliberate
-------------------------------------------
**A model that cannot be fetched this way is not offered.** Fetchable means all
three of repository, revision and digest are recorded. A manifest missing any
of them stays in the catalogue as a documented absence and is never advertised,
because the alternatives are worse: downloading from a repository with no
pinned revision fetches whatever is at the head of a branch today, and
verifying nothing means the worker cannot say what it ran.
What this module never does
---------------------------
* It never substitutes. If the pinned file cannot be obtained or does not match
  its digest, the run fails naming the model. That is the no-fallback rule.
* It never writes a file into place before verifying it. Downloads land in a
  temporary part-file, are hashed as they are written, and are renamed only
  after the digest matches.
* It never evicts to make room for something that would not fit anyway.
"""

from __future__ import annotations

import os
import shutil
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

from distinct_protocol import ModelManifest

from .models import (
    KNOWN_MODEL_MANIFESTS,
    DiscoveredModel,
    ModelVerificationError,
    sha256_file,
)

#: The only host this module will fetch from. Not configurable: a worker that
#: could be pointed at an arbitrary origin by configuration is a worker whose
#: weights provenance is decided by whoever wrote the configuration.
WEIGHTS_HOST = "huggingface.co"

#: Where a weights download may be redirected to, and nowhere else.
#:
#: Hugging Face answers a ``/resolve/`` URL with a redirect to its CDN, so
#: refusing redirects outright would refuse every download. Following them
#: without checking was the other extreme, and it was a blind request
#: primitive: whoever could answer that first response could point this at any
#: host and any scheme, including a cloud metadata service on the volunteer's
#: own network. The pinned digest stops a substituted file; it does not stop
#: the request being made, and the request is the finding.
#:
#: So redirects are followed, to https only, and only to these hosts or a
#: subdomain of one.
WEIGHTS_REDIRECT_HOSTS = frozenset({"huggingface.co", "hf.co", "cdn-lfs.huggingface.co", "cdn-lfs.hf.co", "cdn-lfs-us-1.hf.co", "cdn-lfs-eu-1.hf.co", "xet.huggingface.co", "transfer.xethub.hf.co", "cas-bridge.xethub.hf.co", "cas-server.xethub.hf.co"})


class WeightsRedirects(urllib.request.HTTPRedirectHandler):
    """Follow a redirect only to Hugging Face, only over TLS.
    A redirect anywhere else is refused with a sentence rather than followed,
    because a download that has been bounced to another host is no longer the
    download the operator asked for even when the digest would have caught a
    substitution.
    """

    #: A legitimate chain is one or two hops. Ten is generous and finite.
    max_redirections = 10

    @staticmethod
    def _permitted(url: str) -> bool:
        parsed = urllib.parse.urlsplit(url)
        if parsed.scheme != "https":
            return False
        host = (parsed.hostname or "").casefold().rstrip(".")
        if not host:
            return False
        return any(
            host == allowed or host.endswith("." + allowed)
            for allowed in WEIGHTS_REDIRECT_HOSTS
        )

    def redirect_request(self, req, fp, code, msg, headers, newurl):  # noqa: D102
        if not self._permitted(newurl):
            raise urllib.error.HTTPError(
                newurl,
                code,
                "the weights download was redirected off Hugging Face and was refused",
                headers,
                fp,
            )
        return super().redirect_request(req, fp, code, msg, headers, newurl)


def _weights_opener() -> Callable[..., Any]:
    """The default fetcher: redirects allowed, but only to the expected place."""

    return urllib.request.build_opener(WeightsRedirects).open

#: Default cache ceiling. Small enough to be a polite default on a volunteer's
#: machine, large enough to hold the models in the shipped catalogue.
#: The cache ceiling, in gigabytes.
#:
#: A DEFAULT SMALLER THAN THE CATALOGUE IS A DEFAULT THAT THRASHES.
#:
#: This was 8, and the two models this project ships come to about nine
#: together. So a worker offering both downloaded the first, evicted it to fit
#: the second, and re-downloaded it on the next start — five gigabytes over
#: the wire, every run, for ever, on a project whose entire subject is the
#: energy cost of running models. Nothing reported it, because eviction was
#: silent and the usage line only counted the models that run offered.
#:
#: Sized to hold the shipped catalogue with room for one more, so the common
#: case never evicts anything. An operator with a small disk lowers it and is
#: told what that costs.
DEFAULT_CACHE_GB = 16.0

#: Anything above this is almost certainly a typo rather than an intention.
MAX_CACHE_GB = 512.0

_CHUNK = 1024 * 1024


def shared_cache_directory(environ: Mapping[str, str] | None = None) -> Path:
    """Where weights live so that everything on the machine can share them.
    A model file is hundreds of megabytes and is not specific to this project.
    Putting the cache inside the checkout meant a second agent, a second clone,
    or any other tool on the same machine downloaded its own copy of the same
    bytes: wasteful on disk, wasteful on bandwidth, and a strange thing for a
    project about energy to make people do.
    Resolution order, most specific first:
    1. ``DISTINCT_MODELS_DIR``, for an operator who wants to say exactly where.
    2. ``HF_HOME``/``HUGGINGFACE_HUB_CACHE``, because a machine that already
       has a Hugging Face cache has already answered this question, and a
       model fetched by anything else is a model this worker should reuse.
    3. The platform's own user cache: ``%LOCALAPPDATA%`` on Windows,
       ``XDG_CACHE_HOME`` or ``~/.cache`` elsewhere.
    ``--cache-dir`` still overrides all of it. This only decides the default.
    """

    values = os.environ if environ is None else environ
    explicit = values.get("DISTINCT_MODELS_DIR", "").strip()
    if explicit:
        return Path(explicit).expanduser()
    for name in ("HUGGINGFACE_HUB_CACHE", "HF_HOME"):
        value = values.get(name, "").strip()
        if value:
            return Path(value).expanduser() / "distinct-models"
    if os.name == "nt":
        base = values.get("LOCALAPPDATA", "").strip()
        root = Path(base).expanduser() if base else Path.home() / "AppData" / "Local"
        return root / "distinct" / "models"
    base = values.get("XDG_CACHE_HOME", "").strip()
    root = Path(base).expanduser() if base else Path.home() / ".cache"
    return root / "distinct" / "models"


class WeightsUnavailable(RuntimeError):
    """The pinned weights could not be obtained, and nothing was substituted."""


class CacheTooSmall(WeightsUnavailable):
    """The model does not fit in the cache the operator configured.
    Kept separate because the remedy is different: this is not a network
    problem and retrying will not fix it. The operator raises the cap or
    stops offering the model.
    """


def is_fetchable(manifest: ModelManifest) -> bool:
    """Can this manifest be obtained and verified without a human in the loop?"""

    return bool(manifest.source_repo and manifest.source_revision and manifest.digest_recorded)


def unfetchable_reason(manifest: ModelManifest) -> str:
    """Why a manifest is not offered. Named parts, so it can be fixed."""

    missing = []
    if not manifest.source_repo:
        missing.append("upstream repository")
    if not manifest.source_revision:
        missing.append("pinned revision")
    if not manifest.digest_recorded:
        missing.append("SHA-256 digest")
    if not missing:
        return ""
    if len(missing) == 1:
        parts = missing[0]
    else:
        parts = f"{', '.join(missing[:-1])} or {missing[-1]}"
    return (
        f"{manifest.id}: no {parts} recorded, so this worker cannot obtain the "
        "exact published file or confirm what it ran"
    )


def fetchable_manifests(
    manifests: Sequence[ModelManifest] = KNOWN_MODEL_MANIFESTS,
) -> tuple[ModelManifest, ...]:
    """The models a worker may offer. Everything else is a documented absence."""

    return tuple(manifest for manifest in manifests if is_fetchable(manifest))


def weights_url(manifest: ModelManifest) -> str:
    """The exact revision-pinned download URL for one manifest.
    ``/resolve/<revision>/<filename>`` rather than ``/resolve/main/...``: main
    moves, and a digest pinned against a moving reference is a digest pinned
    against nothing.
    """

    if not is_fetchable(manifest):
        raise WeightsUnavailable(unfetchable_reason(manifest))
    quoted = urllib.parse.quote(manifest.filename)
    return (
        f"https://{WEIGHTS_HOST}/{manifest.source_repo}"
        f"/resolve/{manifest.source_revision}/{quoted}?download=true"
    )


@dataclass(frozen=True)
class CacheUsage:
    used_bytes: int
    limit_bytes: int
    model_ids: tuple[str, ...]

    @property
    def free_bytes(self) -> int:
        return max(0, self.limit_bytes - self.used_bytes)


class WeightsCache:
    """A size-capped directory of verified weights, filled on demand.
    ``limit_bytes`` is a ceiling on the *cache*, not on one model. A model
    larger than the ceiling is refused outright rather than being allowed to
    evict every other model and then still not fit.
    """

    def __init__(
        self,
        directory: Path | str,
        *,
        limit_bytes: int,
        manifests: Sequence[ModelManifest] = KNOWN_MODEL_MANIFESTS,
        opener: Optional[Callable[..., Any]] = None,
        clock: Callable[[], float] = time.time,
    ) -> None:
        if isinstance(limit_bytes, bool) or not isinstance(limit_bytes, int):
            raise ValueError("limit_bytes must be an integer")
        if limit_bytes < 1:
            raise ValueError("limit_bytes must be positive")
        self.directory = Path(directory).expanduser()
        self.limit_bytes = limit_bytes
        self.manifests = {manifest.id: manifest for manifest in manifests}
        self._opener = opener if opener is not None else _weights_opener()
        self._clock = clock
        self._lock = threading.RLock()

    # Inspection ---------------------------------------------------------

    def path_for(self, manifest: ModelManifest) -> Path:
        """Where this model lives when cached.
        Keyed by model id, not by filename: two manifests can legitimately
        share a filename across publishers, and the id is what the run asked
        for.
        """

        return self.directory / manifest.id / manifest.filename

    def is_cached(self, manifest: ModelManifest) -> bool:
        path = self.path_for(manifest)
        try:
            return path.is_file() and path.stat().st_size > 0
        except OSError:
            return False

    def adopt_loose_file(self, manifest: ModelManifest) -> bool:
        """Take over a correctly named file sitting loose in the cache root.
        Weights used to be placed by hand, typically as
        ``models/Qwen3-0.6B-Q4_K_M.gguf``; the cache keys them by model id, so
        the same file now belongs at ``models/qwen3-0.6b/...``. Without this,
        every existing installation would re-download several hundred megabytes
        of a file it already has, which is a poor way to repay someone for
        lending their machine, and a strange thing for a project about energy
        to do.
        The file is **verified before it is adopted**, not after and not
        instead. A loose file that does not match the pinned digest is left
        exactly where it is and the normal download proceeds: this is a
        convenience, and a convenience is never allowed to become a way for an
        unverified file to enter the cache.
        """

        if not manifest.digest_recorded or self.is_cached(manifest):
            return False
        for candidate in self._loose_candidates(manifest):
            try:
                if not candidate.is_file() or candidate.stat().st_size <= 0:
                    continue
                if sha256_file(candidate) != manifest.sha256:
                    continue
                target = self.path_for(manifest)
                target.parent.mkdir(parents=True, exist_ok=True)
                candidate.replace(target)
            except OSError:
                continue
            return True
        return False

    def _loose_candidates(self, manifest: ModelManifest) -> list[Path]:
        """Places a copy of this file may already be sitting.
        The cache root first, then any directory an operator listed in
        ``DISTINCT_EXTRA_MODEL_DIRS``. Each is still verified against the
        pinned digest before it is adopted, so a wrong file in a shared
        location is skipped rather than trusted.
        """

        names = {manifest.filename, manifest.filename.lower()}
        roots = [self.directory]
        extra = os.environ.get("DISTINCT_EXTRA_MODEL_DIRS", "")
        roots += [
            Path(part.strip()).expanduser()
            for part in extra.split(os.pathsep)
            if part.strip()
        ]
        return [root / name for root in roots for name in names]

    def _cached_entries(self) -> list[tuple[str, Path, int]]:
        """Everything actually in the cache, as ``(model_id, path, bytes)``.
        THE CACHE IS ON DISK, NOT IN THIS RUN'S ARGUMENTS.
        This used to be derived from ``self.manifests``, which is only the
        models the current invocation offers. A worker started with
        ``--models olmoe`` therefore reported an empty cache while several
        gigabytes of another model sat right there in it — the operator was
        told "0.00 GB used, ready: none" about a directory that was nearly
        full — and the eviction arithmetic below inherited the same blind
        spot, so it believed there was room it did not have.
        """

        entries: list[tuple[str, Path, int]] = []
        try:
            folders = sorted(self.directory.iterdir())
        except OSError:
            return entries
        for folder in folders:
            if not folder.is_dir():
                continue
            try:
                for path in folder.iterdir():
                    # `.part` files are downloads in flight; they occupy space
                    # but they are not a cached model and must never be counted
                    # as one, or an interrupted fetch looks like a hit.
                    if not path.is_file() or path.suffix == ".part":
                        continue
                    entries.append((folder.name, path, path.stat().st_size))
            except OSError:
                continue
        return entries

    def usage(self) -> CacheUsage:
        entries = self._cached_entries()
        used = sum(size for _, _, size in entries)
        present = sorted({model_id for model_id, _, size in entries if size > 0})
        return CacheUsage(used, self.limit_bytes, tuple(present))

    def installed(self) -> tuple[DiscoveredModel, ...]:
        """Cached models, each verified against its digest.
        Verification happens here rather than at load time so that a corrupted
        or replaced cache entry is found when the worker starts, not when a
        user is waiting for an answer.
        """

        found: list[DiscoveredModel] = []
        for manifest in self.manifests.values():
            self.adopt_loose_file(manifest)
            path = self.path_for(manifest)
            try:
                if not path.is_file():
                    continue
                size = path.stat().st_size
            except OSError:
                continue
            if size <= 0:
                continue
            if manifest.digest_recorded and sha256_file(path) != manifest.sha256:
                raise ModelVerificationError(
                    f"cached weights for {manifest.id!r} do not match the pinned "
                    f"SHA-256; delete {path} and let the worker fetch it again"
                )
            found.append(DiscoveredModel(manifest, path, size, manifest.digest_recorded))
        return tuple(found)

    # Acquisition --------------------------------------------------------

    def ensure(
        self,
        manifest: ModelManifest,
        *,
        progress: Callable[[float, str], None] | None = None,
        cancel_event: threading.Event | None = None,
    ) -> DiscoveredModel:
        """Return the verified local file for ``manifest``, fetching if needed."""

        if not is_fetchable(manifest):
            raise WeightsUnavailable(unfetchable_reason(manifest))
        with self._lock:
            self.adopt_loose_file(manifest)
            path = self.path_for(manifest)
            if self.is_cached(manifest):
                actual = sha256_file(path)
                if actual == manifest.sha256:
                    self._touch(path)
                    return DiscoveredModel(manifest, path, path.stat().st_size, True)
                # A cached file that no longer matches is not repaired quietly.
                # It is removed and re-fetched, and the operator is told.
                path.unlink(missing_ok=True)
            return self._download(manifest, progress=progress, cancel_event=cancel_event)

    def _download(
        self,
        manifest: ModelManifest,
        *,
        progress: Callable[[float, str], None] | None,
        cancel_event: threading.Event | None,
    ) -> DiscoveredModel:
        url = weights_url(manifest)
        target = self.path_for(manifest)
        target.parent.mkdir(parents=True, exist_ok=True)
        part = target.with_suffix(target.suffix + ".part")
        part.unlink(missing_ok=True)

        import hashlib

        digest = hashlib.sha256()
        written = 0
        try:
            request = urllib.request.Request(
                url,
                headers={"User-Agent": "distinct-agent", "Accept": "application/octet-stream"},
            )
            with self._opener(request, timeout=60) as response:
                declared = _content_length(response)
                if declared:
                    evicted = self._make_room(declared, keep=manifest.id)
                    if evicted and progress is not None:
                        # Deleting gigabytes somebody already waited for is
                        # not a detail. Silent, it looks like the download
                        # simply never persisted, which is exactly how it was
                        # reported to us.
                        progress(
                            0.0,
                            "Evicted from the weights cache to make room: "
                            + ", ".join(evicted)
                            + f" — raise --model-cache-gb above {_gb(self.limit_bytes):.0f}"
                            " to keep them",
                        )
                with part.open("wb") as handle:
                    while True:
                        if cancel_event is not None and cancel_event.is_set():
                            raise WeightsUnavailable(
                                f"weights download for {manifest.id!r} was cancelled"
                            )
                        chunk = response.read(_CHUNK)
                        if not chunk:
                            break
                        handle.write(chunk)
                        digest.update(chunk)
                        written += len(chunk)
                        if written > self.limit_bytes:
                            raise CacheTooSmall(
                                f"{manifest.id!r} exceeds the {_gb(self.limit_bytes):.1f} GB "
                                "weights cache; raise --model-cache-gb or stop offering "
                                "this model"
                            )
                        if progress is not None and declared:
                            progress(
                                min(0.99, written / declared),
                                f"Fetching {manifest.id} weights",
                            )
        except (OSError, urllib.error.URLError, ValueError) as exc:
            part.unlink(missing_ok=True)
            raise WeightsUnavailable(
                f"could not fetch {manifest.id!r} from {WEIGHTS_HOST}: "
                f"{type(exc).__name__}: {exc}"
            ) from exc
        except BaseException:
            part.unlink(missing_ok=True)
            raise

        actual = digest.hexdigest()
        if actual != manifest.sha256:
            part.unlink(missing_ok=True)
            raise ModelVerificationError(
                f"downloaded weights for {manifest.id!r} do not match the pinned "
                f"SHA-256: expected {manifest.sha256}, computed {actual}. Nothing "
                "was installed and no other model was substituted."
            )
        # Only now does the file take its real name, so a partial or wrong
        # download can never be found by discovery.
        part.replace(target)
        self._touch(target)
        if progress is not None:
            progress(1.0, f"{manifest.id} weights verified")
        return DiscoveredModel(manifest, target, target.stat().st_size, True)

    # Eviction -----------------------------------------------------------

    def _make_room(self, needed_bytes: int, *, keep: str) -> list[str]:
        """Evict least-recently-used models until ``needed_bytes`` fits.
        Returns the ids evicted, so the caller can say what happened. Eviction
        is a real cost: the next run that wants an evicted model pays the
        download again, and a worker that silently churned its cache would
        look slow for no visible reason.
        """

        if needed_bytes > self.limit_bytes:
            raise CacheTooSmall(
                f"{_gb(needed_bytes):.1f} GB of weights cannot fit in a "
                f"{_gb(self.limit_bytes):.1f} GB cache"
            )
        evicted: list[str] = []
        while True:
            usage = self.usage()
            if usage.used_bytes + needed_bytes <= self.limit_bytes:
                return evicted
            # Any model in the cache is a candidate, not only the ones this
            # run happens to offer: they take up the same disk either way, and
            # a run that could not see them could not free them.
            candidates = [
                (self._last_used(path), model_id)
                for model_id, path, _ in self._cached_entries()
                if model_id != keep
            ]
            if not candidates:
                raise CacheTooSmall(
                    f"the weights cache is full at {_gb(self.limit_bytes):.1f} GB and "
                    "nothing else can be evicted to make room"
                )
            candidates.sort()
            _, victim = candidates[0]
            self.evict(victim)
            evicted.append(victim)

    def evict(self, model_id: str) -> bool:
        """Remove one cached model. Returns whether anything was removed."""

        manifest = self.manifests.get(model_id)
        if manifest is None:
            return False
        path = self.path_for(manifest)
        directory = path.parent
        removed = False
        try:
            if path.is_file():
                path.unlink()
                removed = True
            if directory.is_dir() and not any(directory.iterdir()):
                shutil.rmtree(directory, ignore_errors=True)
        except OSError:
            return removed
        return removed

    def _touch(self, path: Path) -> None:
        """Record use, so eviction can be least-recently-used rather than random."""

        try:
            now = self._clock()
            os.utime(path, (now, now))
        except OSError:
            pass

    def _last_used(self, path: Path) -> float:
        try:
            return path.stat().st_mtime
        except OSError:
            return 0.0


def _content_length(response: Any) -> int:
    for getter in ("headers", "info"):
        source = getattr(response, getter, None)
        if source is None:
            continue
        headers = source() if callable(source) else source
        try:
            value = headers.get("Content-Length")
        except AttributeError:
            continue
        try:
            return max(0, int(value))
        except (TypeError, ValueError):
            continue
    return 0


def _gb(value: int) -> float:
    return value / (1024**3)


def cache_bytes(gigabytes: float) -> int:
    if not isinstance(gigabytes, (int, float)) or isinstance(gigabytes, bool):
        raise ValueError("cache size must be a number of gigabytes")
    if not 0.1 <= float(gigabytes) <= MAX_CACHE_GB:
        raise ValueError(f"cache size must be between 0.1 and {MAX_CACHE_GB:g} GB")
    return int(float(gigabytes) * (1024**3))


def describe_catalogue(
    manifests: Iterable[ModelManifest] = KNOWN_MODEL_MANIFESTS,
) -> tuple[list[str], list[str]]:
    """``(offerable, withheld)`` lines for the setup console.
    The withheld list is printed rather than filtered away in silence. An
    operator who wonders why a model they expected is missing should be able to
    read the reason on the same screen.
    """

    offerable: list[str] = []
    withheld: list[str] = []
    for manifest in manifests:
        if is_fetchable(manifest):
            offerable.append(f"{manifest.id}  {manifest.label}")
        else:
            withheld.append(unfetchable_reason(manifest))
    return offerable, withheld