Download distinct_agent/server_runner.py from User1342/distinct: direct link, hf CLI and curl.
- Browser
- Download file 46.3 kB
-
https://huggingface.co/spaces/User1342/distinct/resolve/main/distinct_agent/server_runner.py
- Command line
-
hf download hf://spaces/User1342/distinct/distinct_agent/server_runner.py
-
curl -L -o server_runner.py https://huggingface.co/spaces/User1342/distinct/resolve/main/distinct_agent/server_runner.py
46.3 kB
| """Run inference against a supervised ``llama-server`` over loopback HTTP. | |
| This replaces spawning ``llama-cli`` per request, and it is a better design | |
| rather than a way around a bug. | |
| **Why the server, not the CLI.** From llama.cpp b10451 the CLI is a chat front | |
| end: it prints a logo, build metadata and a slash-command list to stdout, then | |
| blocks on stdin. ``-st`` stops the hang but the furniture still lands in the | |
| response, and parsing it out means depending on a format that changes between | |
| builds. The server's banner goes to its own log and never enters the response | |
| body, so there is nothing to parse. | |
| Three things follow that the CLI could not give us at all: | |
| * **token counts**, which per-turn usage reporting needs; | |
| * **streaming**, which live agent turns and traces need; | |
| * **an OpenAI-compatible endpoint**, which is what DSPy speaks natively, so | |
| the RLM harness points at a URL instead of wrapping a subprocess. | |
| And the model stays resident across queued requests instead of reloading per | |
| request, which on a project measuring energy per run is not a footnote. | |
| **The costs, handled rather than waved past.** It is a long-lived child, so it | |
| is supervised and killed with the worker. It binds a port, so it binds | |
| ``127.0.0.1`` on an ephemeral port, consistent with | |
| :mod:`distinct_protocol.netpolicy`. It is spawned through | |
| :func:`distinct_agent.isolation.spawn_restricted`, so it sits inside the Job | |
| Object rather than beside it. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import re | |
| import socket | |
| import subprocess | |
| import threading | |
| import time | |
| import urllib.error | |
| import urllib.request | |
| from collections import deque | |
| from collections.abc import Callable, Mapping | |
| from pathlib import Path | |
| from typing import Any | |
| from distinct_protocol import PHASE_GENERATING, PHASE_LOADING_MODEL, STEP_PHASE | |
| from .energy import EnergyMeter, UnavailableEnergyMeter | |
| from .isolation import ProcessRestriction, spawn_restricted | |
| from .models import DiscoveredModel | |
| from .runners import ( | |
| InferenceResult, | |
| RunnerCancelled, | |
| RunnerError, | |
| RunnerTimedOut, | |
| RunnerUnavailable, | |
| build_llama_environment, | |
| ) | |
| LOOPBACK = "127.0.0.1" | |
| #: The share of *free* VRAM an offload is allowed to plan for. The rest is | |
| #: left to the display, to whatever else is on the card, and to llama.cpp's own | |
| #: scratch buffers, which are not in the arithmetic below. | |
| VRAM_BUDGET_SHARE = 0.72 | |
| #: Fixed cost of having a CUDA context and llama.cpp's compute buffers on the | |
| #: card at all, before a single layer is placed there. | |
| VRAM_FIXED_OVERHEAD = 320 * 1024 * 1024 | |
| def plan_offload( | |
| *, | |
| model_bytes: int, | |
| shape: Any, | |
| free_vram_bytes: int, | |
| context_tokens: int, | |
| ) -> tuple[int, ...]: | |
| """How many layers to put on the GPU, most first, never more than fits. | |
| ASKING FOR MORE THAN FITS IS NOT A FAILURE ON WINDOWS, AND THAT IS THE | |
| PROBLEM. | |
| The obvious design is to ask for everything and step down when the | |
| allocation fails. It is what this did first, and on the benchmark machine | |
| -- a 4 GB laptop card with 2.3 GB free and a 4.5 GB model -- it made every | |
| single workload time out where they had previously completed. Nothing | |
| failed. The WDDM driver backed the excess with system memory, the | |
| allocation succeeded, the server reported itself healthy, and then every | |
| token dragged weights back across the PCIe bus. The result was several | |
| times slower than the CPU-only build it replaced, and it looked like a | |
| hang rather than a misconfiguration. | |
| So the ceiling is worked out before anything is asked for. Two costs per | |
| offloaded layer: its weights, and its slice of the key/value cache, which | |
| at a 4096-token context can exceed the weights and is the term that a | |
| weights-only estimate misses. Both come out of the model file's own header. | |
| A model whose header cannot be read plans nothing and runs on the CPU, | |
| which is where it ran before any of this existed. | |
| """ | |
| if not shape.known or model_bytes <= 0 or free_vram_bytes <= 0: | |
| return (0,) | |
| usable = int(free_vram_bytes * VRAM_BUDGET_SHARE) - VRAM_FIXED_OVERHEAD | |
| if usable <= 0: | |
| return (0,) | |
| # The +1 stands in for the embedding and output tensors, which are not a | |
| # layer but do take room when they are offloaded with one. | |
| weights_per_layer = model_bytes / (shape.layers + 1) | |
| per_layer = weights_per_layer + shape.kv_bytes_per_layer(context_tokens) | |
| if per_layer <= 0: | |
| return (0,) | |
| top = max(0, min(shape.layers, int(usable // per_layer))) | |
| if top <= 0: | |
| return (0,) | |
| rungs: list[int] = [] | |
| for value in (top, top * 3 // 4, top // 2, top // 4, 0): | |
| if value not in rungs: | |
| rungs.append(value) | |
| return tuple(rungs) | |
| #: What llama.cpp says when the accelerator has not got the room. Matched on | |
| #: the stderr tail, and only to decide whether offloading *less* is worth | |
| #: trying: a failure that is not about memory will fail identically on fewer | |
| #: layers, so the ladder stops rather than walking every rung to the same end. | |
| _SHORT_OF_MEMORY = re.compile( | |
| r"out of memory|failed to allocate|cudaMalloc|unable to allocate|" | |
| r"insufficient memory|ggml_backend_.*alloc.*failed|not enough (?:VRAM|memory)", | |
| re.IGNORECASE, | |
| ) | |
| def _short_of_memory(said: str) -> bool: | |
| return bool(said) and bool(_SHORT_OF_MEMORY.search(said)) | |
| #: What a per-run energy figure covers. | |
| #: | |
| #: Keeping the model resident across a queue makes this a real question rather | |
| #: than a pedantic one. Loading the weights costs energy once; the naive | |
| #: measurement charges the whole of it to whichever run happened to arrive | |
| #: first and nothing to the rest. Two runs of the identical prompt then differ | |
| #: by the load cost purely because of queue position, and the figure misleads | |
| #: in opposite directions depending on which one you look at. | |
| #: | |
| #: So load is never folded into a run. Every run reports | |
| #: ``inference-only``, and the load cost is measured once, attributed to the | |
| #: model's residency, and reported beside the run that triggered it. | |
| ENERGY_BOUNDARY_INFERENCE_ONLY = "inference-only" | |
| class LlamaServerUnavailable(RunnerUnavailable): | |
| """The server binary is missing, or it never became healthy.""" | |
| def _free_port() -> int: | |
| with socket.socket() as probe: | |
| probe.bind((LOOPBACK, 0)) | |
| return int(probe.getsockname()[1]) | |
| class _ProcessLog: | |
| """The last few kilobytes llama-server said, kept for its epitaph. | |
| A thread rather than a poll, because the pipe has to be drained: a child | |
| whose stderr buffer fills blocks on write, and a blocked inference server | |
| looks exactly like a hung one. | |
| """ | |
| #: Enough to hold a start-up banner or a stack of allocator complaints, | |
| #: little enough that a week-long residency cannot grow into it. | |
| LIMIT = 8192 | |
| def __init__(self) -> None: | |
| self._chunks: deque[str] = deque() | |
| self._size = 0 | |
| self._lock = threading.Lock() | |
| self._thread: threading.Thread | None = None | |
| def follow(self, stream: Any) -> None: | |
| if stream is None: | |
| return | |
| def pump() -> None: | |
| try: | |
| for raw in iter(stream.readline, b""): | |
| text = raw.decode("utf-8", "replace") | |
| with self._lock: | |
| self._chunks.append(text) | |
| self._size += len(text) | |
| while self._size > self.LIMIT and len(self._chunks) > 1: | |
| self._size -= len(self._chunks.popleft()) | |
| except (ValueError, OSError): | |
| # The pipe closed under us, which is what happens when the | |
| # process exits. Not an error: the exit is the news. | |
| pass | |
| self._thread = threading.Thread(target=pump, name="llama-stderr", daemon=True) | |
| self._thread.start() | |
| def tail(self, *, limit: int = 600) -> str: | |
| with self._lock: | |
| text = "".join(self._chunks) | |
| lines = [line.strip() for line in text.splitlines() if line.strip()] | |
| joined = " | ".join(lines[-6:]) | |
| return joined[-limit:] if joined else "" | |
| class LlamaServerRunner: | |
| """Supervise one ``llama-server`` and answer requests against it. | |
| One server per model. ``base_url`` is exposed so the DSPy RLM harness can | |
| point at the same endpoint rather than starting a second copy. | |
| """ | |
| name = "llama.cpp-server" | |
| #: The calibration generation. Short, deterministic and long enough for | |
| #: llama.cpp's own tokens-per-second figure to settle. | |
| PROBE_PROMPT = "Count from one to twenty, separated by commas." | |
| PROBE_TOKENS = 96 | |
| PROBE_TIMEOUT = 180.0 | |
| def __init__( | |
| self, | |
| executable: Path | str, | |
| *, | |
| context_length: int = 4096, | |
| startup_timeout_seconds: float = 240.0, | |
| recovery_seconds: float = 20.0, | |
| gpu_layers: tuple[int, ...] | None = None, | |
| restriction: ProcessRestriction | None = None, | |
| clock: Callable[[], float] = time.monotonic, | |
| energy_meter: EnergyMeter | None = None, | |
| ) -> None: | |
| self.executable = str(executable) | |
| self.context_length = int(context_length) | |
| self.startup_timeout_seconds = float(startup_timeout_seconds) | |
| #: How long a server that has stopped answering is given to come back | |
| #: before it is assumed dead and restarted. Long enough to cover a | |
| #: cancelled task being unwound, short enough that a genuinely dead | |
| #: server is not waited on. | |
| self.recovery_seconds = float(recovery_seconds) | |
| #: An explicit ladder, or ``None`` to work one out per model from the | |
| #: card's free memory and the model's own header. | |
| self.gpu_layers = tuple(gpu_layers) if gpu_layers is not None else None | |
| self.restriction = restriction or ProcessRestriction(mode="auto") | |
| self._clock = clock | |
| # Used only to price model loading, which is a property of residency | |
| # rather than of any one run. Runs are metered by the caller. | |
| self._energy_meter = energy_meter or UnavailableEnergyMeter("no meter supplied") | |
| self._lock = threading.Lock() | |
| self._process: subprocess.Popen | None = None | |
| self._guard: Any = None | |
| self._port: int | None = None | |
| self._model_path: Path | None = None | |
| self._model_id: str | None = None | |
| self.isolation: Any = None | |
| #: Energy spent loading the resident model, or ``None`` if unmeasured. | |
| #: Never zero for "unknown": absence and a measured zero are different. | |
| self.load_energy_joules: float | None = None | |
| self.load_seconds: float | None = None | |
| #: How many layers the resident server actually got onto the GPU. | |
| #: ``None`` until something is loaded; ``0`` is a real answer meaning | |
| #: everything is on the CPU, and is not the same as "not measured". | |
| self.gpu_layers_used: int | None = None | |
| #: Whether this build honours a JSON-schema constraint. ``None`` until | |
| #: one has been tried; ``False`` after a build refuses one, which stops | |
| #: every later request paying for the same refusal. | |
| self.supports_schema: bool | None = None | |
| #: Why it refused, when it did. Empty otherwise. | |
| self.schema_refusal: str = "" | |
| #: What calibration found, when it ran in this process. Empty when the | |
| #: answer came from the cache, which is the usual case. | |
| self.offload_measurements: tuple[Any, ...] = () | |
| #: Called with a sentence while calibration runs, so a worker that is | |
| #: busy for a minute says what it is doing. | |
| self._notify: Callable[[str], None] | None = None | |
| self._load_charged = False | |
| # -- lifecycle ---------------------------------------------------- | |
| def available(self) -> bool: | |
| return bool(self.executable and Path(self.executable).is_file()) | |
| def base_url(self) -> str: | |
| if self._port is None: | |
| raise LlamaServerUnavailable("server is not running") | |
| return f"http://{LOOPBACK}:{self._port}" | |
| def ensure_started(self, model: DiscoveredModel) -> None: | |
| """Start the server for ``model``, restarting if the model changed. | |
| Serialised: two jobs arriving together must not race two servers onto | |
| two ports and leave one orphaned. | |
| """ | |
| with self._lock: | |
| if self._process is not None and self._process.poll() is None: | |
| if self._model_path == model.path and self._answering(): | |
| return | |
| # ALIVE IS NOT THE SAME AS ANSWERING, AND THE DIFFERENCE COST A | |
| # RUN EVERY TIME. | |
| # | |
| # This checked only that the process existed. When llama-server | |
| # lost its listener but stayed up -- which is what it does on | |
| # some failures rather than exiting -- the check passed, the | |
| # request went to a closed port, and the run failed with | |
| # "connection actively refused". The run AFTER a crash was | |
| # therefore lost too: the first one to see the dead socket paid | |
| # for discovering it. Asking the health endpoint costs a | |
| # millisecond on loopback and turns that into a restart. | |
| self._stop_locked() | |
| self._start_locked(model) | |
| def _await_health(self, seconds: float) -> bool: | |
| """Poll ``/health`` until the server answers, or the time runs out. | |
| Distinct from :meth:`_wait_for_health`, which is start-up: this is for | |
| a server that was healthy a moment ago and may only be busy. | |
| """ | |
| deadline = self._clock() + max(0.0, seconds) | |
| while True: | |
| if self._process is not None and self._process.poll() is not None: | |
| return False | |
| if self._answering(): | |
| return True | |
| if self._clock() >= deadline: | |
| return False | |
| time.sleep(0.25) | |
| def _retry_after_transport_loss( | |
| self, | |
| model: DiscoveredModel, | |
| prompt: str, | |
| lost: _TransportLost, | |
| *, | |
| max_tokens: int, | |
| temperature: float, | |
| timeout: float, | |
| schema: Mapping[str, Any] | None = None, | |
| ) -> tuple[Mapping[str, Any], str]: | |
| """One request lost its connection. Get the server back and ask again. | |
| A DROPPED CONNECTION IS NOT A WRONG ANSWER, AND WAS BEING TREATED AS | |
| ONE. | |
| ``llama-server became unreachable: [WinError 10054]`` ended a run and | |
| threw away every tool result it had already earned. But the reset does | |
| not mean the model was wrong or the request was bad; it means the | |
| socket died, usually because the previous request was cancelled and | |
| the server was still releasing the slot, occasionally because the | |
| server crashed and has to be brought back. Neither is a reason to lose | |
| the work. | |
| So: wait for the server to answer ``/health`` again, and if it will | |
| not, restart it -- and either way ask once more. Once, not in a loop. | |
| A second loss with a healthy server in between is a real fault and | |
| deserves to be reported rather than papered over, and a retry loop | |
| around an expensive generation is how a worker burns an evening of | |
| somebody's electricity on a request that was never going to succeed. | |
| """ | |
| if not self._await_health(self.recovery_seconds): | |
| with self._lock: | |
| self._stop_locked() | |
| self._start_locked(model) | |
| try: | |
| return self._generate( | |
| prompt, | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| timeout=timeout, | |
| schema=schema, | |
| ) | |
| except _TransportLost as again: | |
| raise RunnerError(f"{again} (it had already been recovered once)") from again | |
| def _answering(self) -> bool: | |
| """Whether the server on our port is actually serving.""" | |
| if self._port is None: | |
| return False | |
| try: | |
| with urllib.request.urlopen( | |
| f"http://{LOOPBACK}:{self._port}/health", timeout=2 | |
| ) as response: | |
| return response.status == 200 | |
| except (urllib.error.URLError, OSError): | |
| return False | |
| def evict(self, model_id: str | None = None) -> bool: | |
| """Shut the server down when its model is evicted from the cache. | |
| Without this, a cache that evicts a model leaves its server resident, | |
| the worker accumulates loaded models, and the energy figures the | |
| product exists to report become nonsense. Returns whether anything was | |
| stopped, so an eviction policy can tell a hit from a miss. | |
| """ | |
| with self._lock: | |
| if self._process is None: | |
| return False | |
| if model_id is not None and self._model_id != model_id: | |
| return False | |
| self._stop_locked() | |
| return True | |
| def resident_model_id(self) -> str | None: | |
| """Which model is currently loaded, if any.""" | |
| return self._model_id | |
| def _start_locked(self, model: DiscoveredModel) -> None: | |
| """Bring the server up, on the GPU if the GPU will have it. | |
| WHY THIS TRIES SEVERAL TIMES RATHER THAN CALCULATING THE ANSWER. | |
| The benchmark that prompted this ran at 3.1 tokens per second with the | |
| GPU sitting at its 7 W idle draw: a CUDA-capable laptop doing every | |
| matrix multiply on its CPU. Two separate causes, and this handles the | |
| second. | |
| The first is which build was installed, and belongs to | |
| :mod:`distinct_agent.runtime`. The second is that a build with GPU | |
| support still offloads nothing unless it is told how many layers to | |
| offload, and the right number is a property of *this model on this | |
| card*: 4 GB of laptop VRAM takes every layer of a 0.6 B model and about | |
| half of a 7 B one. Guessing it needs the layer count and per-layer size | |
| from the GGUF, the free VRAM at this instant, and however much the | |
| display is about to take -- three numbers, two of which move. | |
| :func:`plan_offload` works out what fits from the card's free memory | |
| and the model's own header, and this starts there. The ladder below it | |
| is insurance, not the plan: an estimate can still be a little | |
| optimistic, so an allocation failure steps down rather than giving up, | |
| and a failure that is *not* about memory stops the ladder at once | |
| because a missing model file will not load on four layers either. The | |
| last rung is zero, which is the CPU-only behaviour this replaces, so | |
| the worst case is what used to be the only case. | |
| """ | |
| if not self.available: | |
| raise LlamaServerUnavailable(f"llama-server not found at {self.executable!r}") | |
| if not model.path.is_file(): | |
| raise RunnerUnavailable(f"model file is missing: {model.path}") | |
| # Price the load separately. It is paid once per residency and must | |
| # never be charged to a run, or queue position changes the answer. | |
| load_token = self._energy_meter.start() | |
| load_started = time.monotonic() | |
| rungs = self.gpu_layers or self._offload_ladder(model) | |
| attempts: list[str] = [] | |
| healthy = False | |
| for index, layers in enumerate(rungs): | |
| healthy = self._spawn_locked(model, layers) | |
| if healthy: | |
| self.gpu_layers_used = layers | |
| break | |
| said = self._log.tail() | |
| code = self._process.returncode if self._process is not None and self._process.poll() is not None else None | |
| attempts.append( | |
| f"-ngl {layers}: " + ("exited with %s" % code if code is not None else "timed out") | |
| ) | |
| last = index == len(rungs) - 1 | |
| self._stop_locked() | |
| if last or not _short_of_memory(said): | |
| if said: | |
| attempts[-1] += f". It said: {said}" | |
| break | |
| load_usage = self._energy_meter.stop(load_token) | |
| self.load_seconds = round(time.monotonic() - load_started, 3) | |
| self.load_energy_joules = ( | |
| float(load_usage.joules) | |
| if getattr(load_usage, "available", False) and load_usage.joules is not None | |
| else None | |
| ) | |
| self._load_charged = False | |
| if not healthy: | |
| raise LlamaServerUnavailable( | |
| "llama-server did not become healthy on " | |
| f"{LOOPBACK} ({'; '.join(attempts) or 'no attempt was made'})" | |
| ) | |
| def _offload_ladder(self, model: DiscoveredModel) -> tuple[int, ...]: | |
| """What to give this machine's card, measured once and remembered. | |
| :func:`plan_offload` says what would *fit*. That is not the same | |
| question as what is *faster*, and the difference is the whole reason | |
| this method exists: nine of thirty-two layers fitted on the benchmark | |
| laptop, started cleanly, and was slower than not using the card at | |
| all. So the settings that fit are timed, once, and the winner is | |
| written down. See :mod:`distinct_agent.offload`. | |
| """ | |
| from . import offload as _offload # noqa: PLC0415 | |
| from . import runtime as _runtime # noqa: PLC0415 - optional at import time | |
| from .gguf import read_shape # noqa: PLC0415 | |
| free, total = _runtime.accelerator_memory() | |
| if not free: | |
| return (0,) | |
| try: | |
| model_bytes = model.path.stat().st_size | |
| except OSError: | |
| return (0,) | |
| context = min(self.context_length, model.manifest.context_length) | |
| plan = plan_offload( | |
| model_bytes=model_bytes, | |
| shape=read_shape(model.path), | |
| free_vram_bytes=free, | |
| context_tokens=context, | |
| ) | |
| if plan == (0,): | |
| return plan | |
| key = _offload.machine_key( | |
| model_id=model.manifest.id, | |
| gpu=_runtime.accelerator_name(), | |
| vram_bytes=total, | |
| build=Path(self.executable).parent.name, | |
| context=context, | |
| ) | |
| known = _offload.remembered(key) | |
| if known is None: | |
| measured = self._time_offloads(model, _offload.candidates(plan)) | |
| known = _offload.choose(measured) | |
| _offload.remember(key, known, measured) | |
| self.offload_measurements = measured | |
| return (known, 0) if known else (0,) | |
| def _time_offloads( | |
| self, model: DiscoveredModel, settings: tuple[int, ...] | |
| ) -> tuple[Any, ...]: | |
| """Start the server at each setting and time one short generation. | |
| Runs before the model is resident for real, so nothing it does is | |
| charged to a job. It costs one model load per setting, once per | |
| machine and model, and it is the only thing here that can tell the | |
| difference between an offload that helps and one that does not. | |
| """ | |
| from .offload import Measurement # noqa: PLC0415 | |
| results: list[Any] = [] | |
| for layers in settings: | |
| if self._notify is not None: | |
| self._notify(f"Timing {layers} layer(s) on the GPU") | |
| healthy = self._spawn_locked(model, layers) | |
| if not healthy: | |
| results.append(Measurement(layers, 0.0, "did not start")) | |
| self._stop_locked() | |
| continue | |
| try: | |
| body, mode = self._generate( | |
| self.PROBE_PROMPT, | |
| max_tokens=self.PROBE_TOKENS, | |
| temperature=0.0, | |
| timeout=self.PROBE_TIMEOUT, | |
| ) | |
| speed = float((body.get("timings") or {}).get("predicted_per_second") or 0.0) | |
| results.append(Measurement(layers, speed)) | |
| except Exception as exc: # noqa: BLE001 - a setting that fails is a result | |
| results.append(Measurement(layers, 0.0, f"{type(exc).__name__}")) | |
| finally: | |
| self._stop_locked() | |
| return tuple(results) | |
| def _spawn_locked(self, model: DiscoveredModel, gpu_layers: int) -> bool: | |
| """Start one attempt and say whether it came up. Never raises.""" | |
| port = _free_port() | |
| command = [ | |
| self.executable, | |
| "-m", str(model.path), | |
| "--host", LOOPBACK, | |
| "--port", str(port), | |
| "-c", str(min(self.context_length, model.manifest.context_length)), | |
| "--no-webui", | |
| # A build without GPU support prints one warning and carries on, | |
| # so this is safe to pass unconditionally and saves the worker | |
| # having to know which build it installed. | |
| "-ngl", str(gpu_layers), | |
| ] | |
| # A CRASHED RUNNER USED TO TAKE ITS REASON WITH IT. | |
| # | |
| # Both streams went to DEVNULL, so when llama-server died mid-request | |
| # the operator got "llama-server became unreachable: [WinError 10054]" | |
| # and nothing else -- not the out-of-memory line, not the context | |
| # overflow, not the unsupported-flag complaint that had made it exit | |
| # during start-up. The one component whose failures a volunteer cannot | |
| # debug from the outside was the one component whose output was | |
| # discarded. | |
| # | |
| # Its last output is now kept in a bounded ring and attached to the | |
| # error. Bounded because llama-server is chatty per token and an | |
| # unbounded pipe on a long residency is a slow leak; the tail is the | |
| # part that says why it stopped. | |
| self._log = _ProcessLog() | |
| process, isolation, guard = spawn_restricted( | |
| command, | |
| self.restriction, | |
| stdin=subprocess.DEVNULL, | |
| stdout=subprocess.DEVNULL, | |
| stderr=subprocess.PIPE, | |
| shell=False, | |
| creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), | |
| env=build_llama_environment(), | |
| ) | |
| self._log.follow(process.stderr) | |
| self._process, self._guard, self._port = process, guard, port | |
| self._model_path, self.isolation = model.path, isolation | |
| self._model_id = model.manifest.id | |
| return self._wait_for_health() | |
| def _wait_for_health(self) -> bool: | |
| assert self._process is not None | |
| deadline = self._clock() + self.startup_timeout_seconds | |
| url = f"http://{LOOPBACK}:{self._port}/health" | |
| while self._clock() < deadline: | |
| if self._process.poll() is not None: | |
| return False | |
| try: | |
| with urllib.request.urlopen(url, timeout=3) as response: | |
| if response.status == 200: | |
| return True | |
| except (urllib.error.URLError, OSError): | |
| time.sleep(0.5) | |
| return False | |
| # Generation ----------------------------------------------------- | |
| #: Chat first, raw completion only if the server has no chat endpoint. | |
| CHAT_PATH = "/v1/chat/completions" | |
| COMPLETION_PATH = "/completion" | |
| PROMPT_MODE_CHAT = "chat-template" | |
| PROMPT_MODE_RAW = "raw-completion" | |
| def _generate( | |
| self, | |
| prompt: str, | |
| *, | |
| max_tokens: int, | |
| temperature: float, | |
| timeout: float, | |
| schema: Mapping[str, Any] | None = None, | |
| ) -> tuple[Mapping[str, Any], str]: | |
| """Ask the model, through its own chat template wherever possible. | |
| THIS IS THE DIFFERENCE BETWEEN AN ANSWER AND A CONTINUATION. | |
| llama.cpp's ``/completion`` endpoint is raw text completion: whatever | |
| you send is continued, verbatim, with no template applied. Every model | |
| in this catalogue is an *instruct* model, which means it was trained to | |
| respond to a specific framing (for Qwen3, ``<|im_start|>user`` ... | |
| ``<|im_end|><|im_start|>assistant``) that lives in the GGUF's own chat | |
| template. Send it bare prose and it does exactly what a base language | |
| model does: it continues the document. | |
| That is not a hypothetical. Asked "In one sentence, why can a smaller | |
| language model use less energy per answer?" through ``/completion``, | |
| Qwen3-0.6B replied: | |
| A) Because it has more neurons. B) Because it has fewer neurons. | |
| C) Because it has more connections. D) Because it has fewer | |
| connections. Answer: B **Process:** 1. Identify the key elements... | |
| It had not misunderstood the question. It was never asked one. It saw | |
| the opening of a quiz and wrote the rest of the quiz. | |
| ``/v1/chat/completions`` hands the message to llama-server as a *turn*, | |
| and llama-server applies the template the model was published with. The | |
| same model then answers. | |
| The raw endpoint remains as a fallback for a server too old to have the | |
| chat route, and when it is used the run says so in ``prompt_mode``, | |
| because an answer produced without a template is a different artefact | |
| and a run log that could not distinguish them would be lying by | |
| omission. | |
| """ | |
| # ASKING NICELY FOR JSON DOES NOT WORK ON A SEVEN BILLION PARAMETER | |
| # MODEL, AND IT DOES NOT HAVE TO. | |
| # | |
| # The harness needs one JSON object per turn. The prompt said so in | |
| # capitals, gave a worked example, and re-asked once when prose came | |
| # back; the model still wrote a perfectly good paragraph *about* the | |
| # document it had been asked to create, called nothing, and produced | |
| # no file. Eleven of fourteen benchmark failures were that. | |
| # | |
| # llama.cpp can constrain the sampler to a JSON schema, which turns | |
| # "please reply in this shape" from an instruction the model may | |
| # ignore into a shape it cannot leave. A build that does not support | |
| # it says so with a 400, and the flag below stops it being asked | |
| # again, so this degrades to the behaviour it replaces rather than | |
| # failing. | |
| constrain = schema is not None and self.supports_schema is not False | |
| chat_payload = { | |
| "messages": [{"role": "user", "content": prompt}], | |
| "max_tokens": max_tokens, | |
| "temperature": temperature, | |
| "stream": False, | |
| # Qwen3 and the other hybrid reasoning models interleave a | |
| # <think> block by default. It is not the answer, the user did not | |
| # ask for it, and it consumes the output budget. Asking the | |
| # template to omit it is the clean route; _answer_text strips any | |
| # block that arrives anyway, because not every build honours this. | |
| "chat_template_kwargs": {"enable_thinking": False}, | |
| } | |
| if constrain: | |
| chat_payload["response_format"] = { | |
| "type": "json_schema", | |
| "json_schema": {"name": "action", "strict": True, "schema": dict(schema or {})}, | |
| } | |
| try: | |
| body = self._post(self.CHAT_PATH, chat_payload, timeout) | |
| if constrain: | |
| self.supports_schema = True | |
| return body, self.PROMPT_MODE_CHAT | |
| except RunnerTimedOut: | |
| raise | |
| except _NoSuchEndpoint: | |
| pass | |
| except RunnerError as refused: | |
| if not constrain: | |
| raise | |
| # Refused with a schema attached. Assume the schema is why, and | |
| # carry on without it rather than losing the run -- but SAY SO. | |
| # | |
| # This was silent, and silence cost an evening. A build that | |
| # refuses the grammar falls back to asking the model nicely, which | |
| # is the behaviour every one of these fixes exists to replace, and | |
| # from the outside the run looks identical to one where the | |
| # constraint worked and the model ignored it. Those need different | |
| # fixes, so they have to be distinguishable. | |
| self.supports_schema = False | |
| self.schema_refusal = str(refused)[:400] | |
| if self._notify is not None: | |
| self._notify( | |
| "llama-server refused the action schema, so the model is " | |
| f"only being asked to follow it: {self.schema_refusal}" | |
| ) | |
| return self._generate( | |
| prompt, | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| timeout=timeout, | |
| schema=None, | |
| ) | |
| raw_payload = { | |
| "prompt": prompt, | |
| "n_predict": max_tokens, | |
| "temperature": temperature, | |
| "stream": False, | |
| "cache_prompt": True, | |
| } | |
| if constrain: | |
| raw_payload["json_schema"] = dict(schema or {}) | |
| return self._post(self.COMPLETION_PATH, raw_payload, timeout), self.PROMPT_MODE_RAW | |
| def _post( | |
| self, path: str, payload: Mapping[str, Any], timeout: float | |
| ) -> Mapping[str, Any]: | |
| request = urllib.request.Request( | |
| f"{self.base_url}{path}", | |
| data=json.dumps(payload).encode("utf-8"), | |
| headers={"Content-Type": "application/json"}, | |
| ) | |
| try: | |
| with urllib.request.urlopen(request, timeout=timeout) as response: | |
| body = json.loads(response.read().decode("utf-8")) | |
| except urllib.error.HTTPError as exc: | |
| if exc.code in {404, 501}: | |
| raise _NoSuchEndpoint(path) from exc | |
| # THE BODY IS THE DIAGNOSIS. llama-server refuses a request that | |
| # does not fit the context window with a 400 and a sentence saying | |
| # so by name. Reporting only "HTTP 400" turned the one failure a | |
| # volunteer can actually fix -- shorten the transcript -- into an | |
| # anonymous number. | |
| raise RunnerError( | |
| f"llama-server rejected the request: HTTP {exc.code}{_http_detail(exc)}" | |
| ) from exc | |
| except TimeoutError as exc: | |
| raise RunnerTimedOut(f"inference exceeded {timeout:g} seconds") from exc | |
| except (urllib.error.URLError, OSError) as exc: | |
| # Attach what it said on the way out. "became unreachable" names | |
| # the symptom; the tail names the cause, and an operator can only | |
| # act on the second. | |
| said = self._log.tail() if getattr(self, "_log", None) else "" | |
| gone = self._process is not None and self._process.poll() is not None | |
| detail = f"llama-server became unreachable: {exc}" | |
| if gone: | |
| detail += f" (it exited with {self._process.returncode})" | |
| if said: | |
| detail += f". It said: {said}" | |
| raise _TransportLost(detail) from exc | |
| if not isinstance(body, Mapping): | |
| raise RunnerError("llama-server returned a malformed response") | |
| return body | |
| def stop(self) -> None: | |
| with self._lock: | |
| self._stop_locked() | |
| def _stop_locked(self) -> None: | |
| process, guard = self._process, self._guard | |
| self._process = self._guard = self._port = self._model_path = None | |
| self._model_id = None | |
| self.load_energy_joules = self.load_seconds = None | |
| self.gpu_layers_used = None | |
| self._load_charged = False | |
| if process is not None and process.poll() is None: | |
| process.terminate() | |
| try: | |
| process.communicate(timeout=10) | |
| except subprocess.TimeoutExpired: | |
| process.kill() | |
| if guard is not None: | |
| guard.close() | |
| def __enter__(self) -> LlamaServerRunner: | |
| return self | |
| def __exit__(self, *exc: object) -> None: | |
| self.stop() | |
| # -- inference ---------------------------------------------------- | |
| def run( | |
| self, | |
| model: DiscoveredModel, | |
| prompt: str, | |
| *, | |
| cancel_event: threading.Event | None = None, | |
| progress: Callable[[float, str], None] | None = None, | |
| limits: Mapping[str, Any] | None = None, | |
| ) -> InferenceResult: | |
| if not isinstance(prompt, str) or not prompt: | |
| raise RunnerError("prompt must be non-empty text") | |
| limits = limits or {} | |
| max_tokens = int(limits.get("max_output_tokens", 512)) | |
| timeout = float(limits.get("timeout_seconds", 600.0)) | |
| was_resident = self.resident_model_id == model.manifest.id | |
| if progress is not None: | |
| from .harness import emit | |
| emit( | |
| progress, | |
| 0.05, | |
| "Model already resident" if was_resident else "Loading the model", | |
| phase=PHASE_LOADING_MODEL, | |
| step={ | |
| "kind": STEP_PHASE, | |
| "text": ( | |
| f"{model.manifest.id} is already loaded" | |
| if was_resident | |
| else f"Loading {model.manifest.id} into memory" | |
| ), | |
| }, | |
| ) | |
| self.ensure_started(model) | |
| # Only the run that actually caused the load hears about it, and it | |
| # hears about it as a separate figure rather than inside its own. | |
| triggered_load = not was_resident and not self._load_charged | |
| if triggered_load: | |
| self._load_charged = True | |
| if cancel_event is not None and cancel_event.is_set(): | |
| raise RunnerCancelled("cancelled before inference started") | |
| if progress is not None: | |
| from .harness import emit | |
| emit(progress, 0.2, "Generating", phase=PHASE_GENERATING) | |
| started = time.monotonic() | |
| temperature = float(limits.get("temperature", 0.2)) | |
| # Carried in limits rather than in the signature, so a runner that | |
| # cannot constrain its output ignores it instead of refusing the call. | |
| schema = limits.get("response_schema") | |
| schema = schema if isinstance(schema, Mapping) else None | |
| try: | |
| body, prompt_mode = self._generate( | |
| prompt, | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| timeout=timeout, | |
| schema=schema, | |
| ) | |
| except _TransportLost as lost: | |
| body, prompt_mode = self._retry_after_transport_loss( | |
| model, | |
| prompt, | |
| lost, | |
| max_tokens=max_tokens, | |
| temperature=temperature, | |
| timeout=timeout, | |
| schema=schema, | |
| ) | |
| except RunnerTimedOut: | |
| # LET IT FINISH TIDYING UP BEFORE THE NEXT RUN ARRIVES. | |
| # | |
| # A client-side timeout closes the socket mid-generation. | |
| # llama-server notices, logs `srv stop: cancel task`, and unwinds | |
| # the slot -- which takes a moment, during which it will reset a | |
| # new connection. That reset landed on whichever run came next, so | |
| # a single slow workload failed two: itself, and the innocent one | |
| # behind it. Waiting here means the run that caused the timeout is | |
| # the run that pays for it. | |
| self._await_health(self.recovery_seconds) | |
| raise | |
| text = _answer_text(body, prompt_mode) | |
| if not text: | |
| raise RunnerError("llama-server returned an empty completion") | |
| if progress is not None: | |
| progress(1.0, "Local inference complete") | |
| timings = body.get("timings") or {} | |
| return InferenceResult( | |
| text=text, | |
| usage={ | |
| "runner": self.name, | |
| "input_characters": len(prompt), | |
| "output_characters": len(text), | |
| # Real token counts, which the CLI could not report. | |
| "input_tokens": _token_count(body, prompt_mode, "prompt"), | |
| "output_tokens": _token_count(body, prompt_mode, "completion"), | |
| "tokens_per_second": timings.get("predicted_per_second"), | |
| "stop_reason": _stop_reason(body, prompt_mode), | |
| # Which endpoint answered. A raw completion is a materially | |
| # different thing from a templated chat turn, and a run log has | |
| # to be able to say which one it got. | |
| "prompt_mode": prompt_mode, | |
| # Whether the sampler was held to the action schema. An answer | |
| # produced under a constraint is a different artefact from one | |
| # produced freely, in the same way a templated turn is. | |
| "schema_constrained": bool(schema) and self.supports_schema is not False, | |
| "max_output_tokens": max_tokens, | |
| "elapsed_seconds": round(time.monotonic() - started, 3), | |
| "isolation": self.isolation.to_dict() if self.isolation else None, | |
| # State the boundary, exactly as scope='cpu-package' is stated. | |
| # The caller's meter brackets this call only, so the figure is | |
| # inference alone whether the model was already resident or | |
| # loaded a moment ago. Cold and warm runs are comparable. | |
| "energy_boundary": ENERGY_BOUNDARY_INFERENCE_ONLY, | |
| "model_was_resident": was_resident, | |
| "triggered_model_load": triggered_load, | |
| # Reported beside the run, never inside it. None means the load | |
| # was not measured, which is not the same as a load costing | |
| # nothing. | |
| "model_load_joules": self.load_energy_joules if triggered_load else None, | |
| "model_load_seconds": self.load_seconds if triggered_load else None, | |
| # Whether this answer came off the GPU, and how much of it did. | |
| # On a project whose figures are joules per run, "which device | |
| # did the arithmetic" is not a diagnostic detail, it is the | |
| # single largest term in the number being reported. | |
| "gpu_layers": self.gpu_layers_used, | |
| }, | |
| ) | |
| class _NoSuchEndpoint(RuntimeError): | |
| """This llama-server build does not implement the route that was tried.""" | |
| class _TransportLost(RuntimeError): | |
| """The connection to llama-server broke mid-request. | |
| Internal, and deliberately not a :class:`RunnerError`: a broken connection | |
| is the one inference failure that is often not the run's fault and is | |
| frequently survivable, so it has to be distinguishable from a refusal | |
| before anything can decide whether to try again. | |
| """ | |
| def _http_detail(exc: urllib.error.HTTPError) -> str: | |
| """The server's own words about why it said no, if it offered any.""" | |
| try: | |
| raw = exc.read().decode("utf-8", "replace") | |
| except Exception: # noqa: BLE001 - a body we cannot read is simply absent | |
| return "" | |
| if not raw: | |
| return "" | |
| try: | |
| parsed = json.loads(raw) | |
| except ValueError: | |
| parsed = None | |
| message = "" | |
| if isinstance(parsed, Mapping): | |
| error = parsed.get("error") | |
| if isinstance(error, Mapping): | |
| message = str(error.get("message") or "") | |
| elif isinstance(error, str): | |
| message = error | |
| if not message: | |
| message = str(parsed.get("message") or "") | |
| text = (message or raw).strip().replace("\n", " ") | |
| return f". It said: {text[:400]}" if text else "" | |
| #: A reasoning block, which is working-out rather than an answer. | |
| _THINK_BLOCK = re.compile(r"<think\b[^>]*>.*?</think\s*>", re.DOTALL | re.IGNORECASE) | |
| #: An unterminated one, which happens when the budget ran out mid-thought. | |
| _OPEN_THINK = re.compile(r"<think\b[^>]*>.*\Z", re.DOTALL | re.IGNORECASE) | |
| def strip_reasoning(text: str) -> str: | |
| """Remove a hybrid reasoning model's ``<think>`` working-out. | |
| The block is not the answer. Leaving it in puts several hundred tokens of | |
| the model talking to itself at the top of a reply that was asked for in one | |
| sentence, which is the single loudest way an interface can look broken. | |
| An *unterminated* block is removed too, and that case is why this is more | |
| than a tidy-up: when the output budget runs out inside the thought, the | |
| entire reply is working-out with no answer after it. Stripping it leaves an | |
| empty string, and an empty completion is refused upstream with a reason, | |
| which is a far better outcome than presenting a monologue as an answer. | |
| """ | |
| cleaned = _THINK_BLOCK.sub("", text) | |
| cleaned = _OPEN_THINK.sub("", cleaned) | |
| return cleaned.strip() | |
| def _answer_text(body: Mapping[str, Any], prompt_mode: str) -> str: | |
| """The assistant's reply, from whichever response shape arrived.""" | |
| if prompt_mode == LlamaServerRunner.PROMPT_MODE_CHAT: | |
| choices = body.get("choices") | |
| if not isinstance(choices, (list, tuple)) or not choices: | |
| return "" | |
| first = choices[0] if isinstance(choices[0], Mapping) else {} | |
| message = first.get("message") if isinstance(first.get("message"), Mapping) else {} | |
| content = message.get("content") | |
| return strip_reasoning(content) if isinstance(content, str) else "" | |
| content = body.get("content") | |
| return strip_reasoning(content) if isinstance(content, str) else "" | |
| def _token_count(body: Mapping[str, Any], prompt_mode: str, which: str) -> Any: | |
| if prompt_mode == LlamaServerRunner.PROMPT_MODE_CHAT: | |
| usage = body.get("usage") | |
| if isinstance(usage, Mapping): | |
| return usage.get(f"{which}_tokens") | |
| return None | |
| return body.get("tokens_evaluated" if which == "prompt" else "tokens_predicted") | |
| def _stop_reason(body: Mapping[str, Any], prompt_mode: str) -> Any: | |
| if prompt_mode == LlamaServerRunner.PROMPT_MODE_CHAT: | |
| choices = body.get("choices") | |
| if isinstance(choices, (list, tuple)) and choices and isinstance(choices[0], Mapping): | |
| return choices[0].get("finish_reason") | |
| return None | |
| return body.get("stop_type") | |