"""llama.cpp backend: spawns the upstream `llama-server` binary natively. Install: - macOS: brew install llama.cpp (Metal build, MPS) - Linux: brew install llama.cpp OR github.com/ggml-org/llama.cpp/releases """ from __future__ import annotations import os import shutil import subprocess from pathlib import Path from typing import List, Optional from huggingface_hub import hf_hub_download from openai import OpenAI from surya.inference.backends.base import Backend, ServerHandle from surya.inference.backends.openai_client import chat_completions_batch from surya.inference.backends.spawn import ( SpawnHandle, SpawnError, attach_or_spawn, ) from surya.inference.schema import BatchInputItem, BatchOutputItem from surya.logging import get_logger from surya.settings import settings logger = get_logger() def _resolve_llama_server_binary() -> str: binary = settings.LLAMA_CPP_BINARY if binary and os.path.isfile(binary): return binary found = shutil.which(binary or "llama-server") if found: return found raise SpawnError( "llama-server binary not found. Install with:\n" " macOS: brew install llama.cpp\n" " Linux: brew install llama.cpp OR download from\n" " https://github.com/ggml-org/llama.cpp/releases\n" "Or set LLAMA_CPP_BINARY in your env to the binary path." ) def _download_gguf_files() -> tuple[str, str]: """Download model + mmproj GGUFs from HF Hub. Returns local paths.""" repo = settings.SURYA_GGUF_REPO model_file = settings.SURYA_GGUF_MODEL_FILE mmproj_file = settings.SURYA_GGUF_MMPROJ_FILE logger.info(f"Downloading {model_file} and {mmproj_file} from {repo}") model_path = hf_hub_download(repo_id=repo, filename=model_file) mmproj_path = hf_hub_download(repo_id=repo, filename=mmproj_file) return model_path, mmproj_path def _health_url(port: int) -> str: return f"http://{settings.SURYA_INFERENCE_HOST}:{port}" def _openai_url(port: int) -> str: return f"http://{settings.SURYA_INFERENCE_HOST}:{port}/v1" class LlamaCppBackend(Backend): name = "llamacpp" def __init__(self): self.handle: Optional[ServerHandle] = None self._client: Optional[OpenAI] = None def start(self) -> ServerHandle: if self.handle is not None: return self.handle # If user pinned an external server, attach without spawning. # No binary or GGUF download needed in that case. if settings.SURYA_INFERENCE_URL: spawned = attach_or_spawn( backend=self.name, expected_model_name=settings.SURYA_MODEL_CHECKPOINT, spawn_fn=lambda port: SpawnHandle( pid=None, cleanup_id="", cleanup_kind="process" ), # never called health_url_for=_health_url, openai_url_for=_openai_url, startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT, ) self.handle = ServerHandle( base_url=spawned.base_url, model_name=spawned.model_name, spawned_by_us=spawned.spawned_by_us, ) self._client = OpenAI(api_key="EMPTY", base_url=self.handle.base_url) return self.handle binary = _resolve_llama_server_binary() # Pre-download GGUFs so the spawn doesn't race the download if ( settings.SURYA_GGUF_LOCAL_MODEL_PATH and settings.SURYA_GGUF_LOCAL_MMPROJ_PATH ): model_path = settings.SURYA_GGUF_LOCAL_MODEL_PATH mmproj_path = settings.SURYA_GGUF_LOCAL_MMPROJ_PATH else: model_path, mmproj_path = _download_gguf_files() # Total KV-cache budget. llama-server divides --ctx-size across # --parallel slots, so a too-small total silently truncates outputs # once each slot's share fills. Scale with parallel by default; # SURYA_INFERENCE_CTX_SIZE overrides to a fixed value if set. parallel = settings.SURYA_INFERENCE_PARALLEL per_slot = settings.SURYA_INFERENCE_CTX_PER_SLOT ctx_size = settings.SURYA_INFERENCE_CTX_SIZE if ctx_size is None: ctx_size = max(16384, parallel * per_slot) effective_per_slot = ctx_size // max(parallel, 1) logger.info( f"llama-server ctx-size={ctx_size} " f"(~{effective_per_slot}/slot × {parallel} parallel slots)" ) if effective_per_slot < per_slot: logger.warning( f"per-slot ctx ({effective_per_slot}) is below recommended " f"{per_slot}; outputs may truncate. Raise " f"SURYA_INFERENCE_CTX_SIZE or SURYA_INFERENCE_CTX_PER_SLOT, " f"or lower SURYA_INFERENCE_PARALLEL." ) def spawn_fn(port: int) -> SpawnHandle: cmd = [ binary, "-m", model_path, "--mmproj", mmproj_path, "-ngl", str(settings.LLAMA_CPP_NGL), "--host", settings.SURYA_INFERENCE_HOST, "--port", str(port), "--parallel", str(parallel), "--ctx-size", str(ctx_size), "--no-mmproj-offload" if settings.LLAMA_CPP_NO_MMPROJ_OFFLOAD else "", "--alias", settings.SURYA_MODEL_CHECKPOINT, "--jinja", ] cmd = [c for c in cmd if c] for extra in (settings.LLAMA_CPP_EXTRA_ARGS or "").split(): cmd.append(extra) logger.info(f"Spawning: {' '.join(cmd)}") log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser() log_path.parent.mkdir(parents=True, exist_ok=True) log_fp = open(log_path, "ab") proc = subprocess.Popen( cmd, stdout=log_fp, stderr=subprocess.STDOUT, start_new_session=True, ) return SpawnHandle( pid=proc.pid, cleanup_id=str(proc.pid), cleanup_kind="process" ) spawned = attach_or_spawn( backend=self.name, expected_model_name=settings.SURYA_MODEL_CHECKPOINT, spawn_fn=spawn_fn, health_url_for=_health_url, openai_url_for=_openai_url, startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT, ) self.handle = ServerHandle( base_url=spawned.base_url, model_name=spawned.model_name, spawned_by_us=spawned.spawned_by_us, ) self._client = OpenAI( api_key="EMPTY", base_url=self.handle.base_url, ) return self.handle def stop(self) -> None: # atexit handler in spawn.py owns cleanup; nothing to do here. self.handle = None self._client = None def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]: if self.handle is None or self._client is None: self.start() return chat_completions_batch( batch, client=self._client, model_name=self.handle.model_name, timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS, max_workers=settings.SURYA_INFERENCE_PARALLEL, request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS, )