"""Local inference backends behind the OpenAI-compatible surface. ``LlamaCppBackend`` is the real one: an in-process ``llama-cpp-python`` model loaded from a GGUF pulled off the Hub. ``EchoBackend`` is a deterministic stand-in that lets the server, registration, heartbeat and routing paths be tested in CI without downloading weights or needing a GPU — the mesh contract is what those tests are about, and a real model would only make them slow and flaky. """ from __future__ import annotations import logging import os import threading from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Any, Iterable, Mapping, Sequence from .errors import BackendError from .gpu import resolve_gpu_layers logger = logging.getLogger("thox.backend") @dataclass class Completion: """One generated completion plus its token accounting.""" text: str prompt_tokens: int completion_tokens: int @property def total_tokens(self) -> int: return self.prompt_tokens + self.completion_tokens class Backend(ABC): """A loaded model that can answer chat completions.""" model_id: str = "unknown" @abstractmethod def generate( self, messages: Sequence[Mapping[str, Any]], *, max_tokens: int, temperature: float, stop: Sequence[str] | None = None, ) -> Completion: """Produce a single completion. Must be safe to call concurrently.""" def close(self) -> None: """Release model resources. Idempotent.""" class EchoBackend(Backend): """Deterministic backend used by tests and by ``--dry-run``. It is intentionally not a language model: it returns a stable, inspectable string so a test can assert that a prompt travelled router -> node -> back without asserting anything about model quality. """ def __init__(self, model_id: str = "thox-echo") -> None: self.model_id = model_id def generate( self, messages: Sequence[Mapping[str, Any]], *, max_tokens: int, temperature: float, stop: Sequence[str] | None = None, ) -> Completion: if not messages: raise BackendError("at least one message is required") last_user = next( (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), messages[-1].get("content", ""), ) text = f"[{self.model_id}] {last_user}" words = text.split() if max_tokens > 0: words = words[:max_tokens] rendered = " ".join(words) return Completion( text=rendered, prompt_tokens=sum(len(str(m.get("content", "")).split()) for m in messages), completion_tokens=len(words), ) class LlamaCppBackend(Backend): """GGUF model served in-process via ``llama-cpp-python``. A single process-wide lock serialises generation. ``llama.cpp``'s context is not safe to use from several threads at once, and on the free CPU tier there is no headroom for parallel decoding anyway — one queue with honest latency telemetry ranks better in the mesh than several that all thrash. """ def __init__( self, model_path: str, *, model_id: str, n_ctx: int = 2048, n_threads: int = 0, n_gpu_layers: str | int | None = "auto", chat_format: str | None = None, ) -> None: try: from llama_cpp import Llama except ImportError as exc: # pragma: no cover - env dependent raise BackendError( "llama-cpp-python is not installed. Install with:\n" " pip install llama-cpp-python\n" "There is no manylinux wheel, so it compiles from source and needs a " "C/C++ toolchain (build-essential + cmake)." ) from exc if not os.path.exists(model_path): raise BackendError(f"GGUF not found at {model_path}") try: size = os.path.getsize(model_path) except OSError as exc: raise BackendError(f"cannot stat GGUF at {model_path}: {exc}") from exc if size < 1_000_000: # A truncated download or a saved error page presents as a tiny # "GGUF" and otherwise fails much later inside the loader with an # opaque message. raise BackendError( f"GGUF at {model_path} is only {size} bytes - the download was " "truncated or returned an error page. Delete it and retry." ) self.model_id = model_id self._lock = threading.Lock() threads = n_threads or _default_threads() # Resolved here rather than by callers so every entry point (agent, # Space, notebook, CLI) gets the same CPU-safe default. gpu_layers = resolve_gpu_layers(n_gpu_layers) self.n_gpu_layers = gpu_layers logger.info( "loading %s (n_ctx=%d, threads=%d, n_gpu_layers=%d)", model_path, n_ctx, threads, gpu_layers, ) try: self._llama = self._load(Llama, model_path, n_ctx, threads, gpu_layers, chat_format) except Exception as exc: # noqa: BLE001 message = str(exc) if gpu_layers != 0 and _looks_like_gpu_failure(message): # A CPU-only llama.cpp build raises when asked to offload. Falling # back beats refusing to serve: the node is still a useful CPU # member of the mesh, just slower. logger.warning( "GPU offload failed (%s); retrying CPU-only. Rebuild " "llama-cpp-python with CMAKE_ARGS='-DGGML_CUDA=ON' for GPU support.", message.splitlines()[0] if message else exc, ) self.n_gpu_layers = 0 try: self._llama = self._load(Llama, model_path, n_ctx, threads, 0, chat_format) except Exception as retry_exc: # noqa: BLE001 raise BackendError( f"failed to load GGUF {model_path} even CPU-only: {retry_exc}" ) from retry_exc else: raise BackendError(f"failed to load GGUF {model_path}: {exc}") from exc @staticmethod def _load(llama_cls, model_path, n_ctx, threads, gpu_layers, chat_format): """Instantiate llama.cpp. Split out so the CPU retry cannot drift.""" return llama_cls( model_path=model_path, n_ctx=n_ctx, n_threads=threads, n_gpu_layers=gpu_layers, verbose=False, **({"chat_format": chat_format} if chat_format else {}), ) def generate( self, messages: Sequence[Mapping[str, Any]], *, max_tokens: int, temperature: float, stop: Sequence[str] | None = None, ) -> Completion: if not messages: raise BackendError("at least one message is required") payload = [ {"role": str(m.get("role", "user")), "content": str(m.get("content", ""))} for m in messages ] try: with self._lock: result = self._llama.create_chat_completion( messages=payload, max_tokens=max_tokens, temperature=temperature, stop=list(stop) if stop else None, ) except Exception as exc: # noqa: BLE001 raise BackendError(f"generation failed: {exc}") from exc choices = result.get("choices") or [] if not choices: raise BackendError("model returned no choices") text = (choices[0].get("message") or {}).get("content") or "" usage = result.get("usage") or {} return Completion( text=text, prompt_tokens=int(usage.get("prompt_tokens", 0)), completion_tokens=int(usage.get("completion_tokens", 0)), ) def close(self) -> None: llama = getattr(self, "_llama", None) if llama is not None and hasattr(llama, "close"): try: llama.close() except Exception: # noqa: BLE001 - best effort logger.debug("llama close failed", exc_info=True) def _default_threads() -> int: """Thread count honouring the container's CPU quota, not the host's cores. ``os.cpu_count()`` reports the *host* CPUs inside a cgroup-limited container, which on a 2-vCPU free Space means spawning many more threads than the quota allows and losing throughput to contention. The cgroup v2 quota is read first and the host count is only a fallback. """ try: with open("/sys/fs/cgroup/cpu.max", "r", encoding="utf-8") as handle: quota_text, period_text = handle.read().split() if quota_text != "max": quota = int(quota_text) / int(period_text) if quota >= 1: return max(1, int(quota)) except (OSError, ValueError): pass return max(1, (os.cpu_count() or 2)) def download_gguf(repo_id: str, filename: str, *, token: str | None = None) -> str: """Fetch a GGUF from the Hub and return its local path. The token is read from the environment by the caller and passed through; it is never written to disk or logged. """ try: from huggingface_hub import hf_hub_download except ImportError as exc: # pragma: no cover - env dependent raise BackendError("huggingface_hub is required to download a GGUF") from exc if not repo_id or not filename: raise BackendError("both THOX_MODEL_REPO and THOX_MODEL_FILE must be set") logger.info("downloading %s/%s", repo_id, filename) try: return hf_hub_download(repo_id=repo_id, filename=filename, token=token or None) except Exception as exc: # noqa: BLE001 # The raw hub error is accurate but rarely actionable. The two failures # that actually happen here are a missing token on a private repo and a # mistyped GGUF filename, so name both explicitly. detail = str(exc) lowered = detail.lower() hint = "" if any(k in lowered for k in ("401", "403", "gated", "unauthorized", "authentication")): hint = ( "\nTHOX model repos are private/gated. Set HF_TOKEN to a token with read " "access (Colab: Secrets panel; Space: Settings -> Secrets)." ) elif any(k in lowered for k in ("404", "not found", "entrynotfound")): hint = ( f"\nCheck that '{filename}' exists in '{repo_id}'. GGUF filenames are " "case-sensitive and differ between repos " "(e.g. thoxmini-3b-Q4_K_M.gguf vs thoxmythos-9b-q4_k_m.gguf)." ) raise BackendError(f"could not download {repo_id}/{filename}: {detail}{hint}") from exc def build_backend( *, model_id: str, repo_id: str, filename: str, n_ctx: int, n_threads: int, n_gpu_layers: str | int | None = "auto", token: str | None = None, echo: bool = False, ) -> Backend: """Construct the configured backend, falling back to echo when asked.""" if echo or not repo_id: if not echo: logger.warning("no THOX_MODEL_REPO set; serving the echo backend") return EchoBackend(model_id=model_id) path = download_gguf(repo_id, filename, token=token) return LlamaCppBackend( path, model_id=model_id, n_ctx=n_ctx, n_threads=n_threads, n_gpu_layers=n_gpu_layers, ) def _looks_like_gpu_failure(message: str) -> bool: """Whether a loader error is plausibly about GPU offload. Matched on text because llama.cpp raises a generic exception for these. Kept narrow on purpose: misclassifying a genuine model problem as a GPU problem would hide it behind a confusing CPU retry. """ lowered = message.lower() return any( token in lowered for token in ( "cuda", "cublas", "gpu", "vram", "out of memory", "no kernel image", "device", ) )