distinct_agent/offload.py

"""Remember how many layers this machine should actually put on its GPU.

WHY THIS IS MEASURED AND NOT CALCULATED.

Two estimates were tried and both were wrong on the same laptop, in opposite
directions. Offloading everything was wrong because Windows does not refuse an
oversubscribed allocation -- it backs the excess with system memory and every
token then crawls across the PCIe bus, which took a benchmark that had been
completing to twenty timeouts. Sizing the offload to free VRAM was wrong too,
more subtly: nine of thirty-two layers fitted, started cleanly, reported itself
healthy, and was still slower than the CPU, because a split model pays a
round-trip per token and a quarter of the layers does not earn it back.

The variables are the card, how much of it the desktop is already using, the
width of the link, the model, the context length and the CPU on the other side
of the split. Nobody can hold that in a formula. But it takes about a minute to
*try* three settings and see which is fastest, and the answer is stable for as
long as that machine and that model stay the same. So it is measured once,
written down, and read thereafter.

Zero is always one of the settings tried, so the worst outcome of measuring is
the CPU-only behaviour it replaced -- and unlike an estimate, it is impossible
for this to leave a machine running slower than that.
"""

from __future__ import annotations

import json
import os
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping, Optional

#: Where the answers live. Beside the installed runtime, which is already
#: per-machine and already ignored by version control.
CACHE_NAME = "offload-cache.json"

#: An answer older than this is re-measured. Drivers change, cards get
#: replaced, and a figure nobody can explain is worse than a minute of
#: measuring.
MAX_AGE_SECONDS = 90 * 24 * 3600

#: How much faster than the CPU an offload has to be before it is worth the
#: split. Below this the two are the same speed within noise, and the CPU has
#: the advantage of leaving the card free for whatever else the volunteer is
#: doing with their computer.
WORTH_IT_MARGIN = 1.08


@dataclass(frozen=True)
class Measurement:
    """What one setting achieved, in tokens per second."""

    layers: int
    tokens_per_second: float
    note: str = ""


def cache_path(root: Optional[Path] = None) -> Path:
    from .runtime import runtime_directory  # noqa: PLC0415 - avoids a cycle

    return runtime_directory(root) / CACHE_NAME


def machine_key(
    *, model_id: str, gpu: str, vram_bytes: int, build: str, context: int
) -> str:
    """Everything that would change the answer, in one string.

    Deliberately includes the context length: the key/value cache is per
    token, so the same model at 4096 and at 32768 are different problems.
    """

    return f"{model_id}|{gpu}|{vram_bytes}|{build}|{context}"


def _read(path: Path) -> dict:
    try:
        loaded = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}
    return loaded if isinstance(loaded, dict) else {}


def remembered(key: str, root: Optional[Path] = None) -> Optional[int]:
    """The layer count measured for this machine and model, if it is still fresh."""

    entry = _read(cache_path(root)).get(key)
    if not isinstance(entry, Mapping):
        return None
    try:
        layers = int(entry["layers"])
        measured_at = float(entry.get("measured_at", 0))
    except (KeyError, TypeError, ValueError):
        return None
    if layers < 0 or time.time() - measured_at > MAX_AGE_SECONDS:
        return None
    return layers


def remember(
    key: str,
    layers: int,
    measurements: tuple[Measurement, ...] = (),
    root: Optional[Path] = None,
) -> None:
    """Write the answer down, with the measurements that produced it.

    The measurements are kept because the number on its own is unfalsifiable.
    A volunteer who wonders why their card is idle can read the row that says
    the GPU was tried and was slower.
    """

    path = cache_path(root)
    store = _read(path)
    store[key] = {
        "layers": int(layers),
        "measured_at": time.time(),
        "measurements": [
            {
                "layers": item.layers,
                "tokens_per_second": round(item.tokens_per_second, 3),
                "note": item.note,
            }
            for item in measurements
        ],
    }
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        # Written whole and moved into place: a half-written cache read by the
        # next start-up would be a permanent, silent fallback to the CPU.
        handle, temporary = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
        with os.fdopen(handle, "w", encoding="utf-8") as out:
            json.dump(store, out, indent=2)
        os.replace(temporary, path)
    except OSError:
        # A machine that cannot write the cache measures again next time. That
        # is a cost, not a failure, and not worth ending a run over.
        return


def choose(measurements: tuple[Measurement, ...]) -> int:
    """The layer count worth using, given what each one achieved.

    The CPU wins ties and near-ties. An offload has to be meaningfully faster
    to justify taking a volunteer's graphics card for the duration of a run,
    and a figure inside the noise is not a reason to.
    """

    usable = [item for item in measurements if item.tokens_per_second > 0]
    if not usable:
        return 0
    baseline = next((item.tokens_per_second for item in usable if item.layers == 0), 0.0)
    best = max(usable, key=lambda item: item.tokens_per_second)
    if best.layers == 0:
        return 0
    if baseline > 0 and best.tokens_per_second < baseline * WORTH_IT_MARGIN:
        return 0
    return best.layers


def candidates(plan: tuple[int, ...]) -> tuple[int, ...]:
    """Which settings to actually time, from the ladder that would fit.

    Three, not the whole ladder: each one costs a model load, and the shape of
    the curve -- flat, rising, or falling -- is legible from the top of what
    fits, half of it, and none of it.
    """

    top = plan[0] if plan else 0
    return tuple(sorted({top, top // 2, 0}, reverse=True)) if top > 0 else (0,)