File size: 3,081 Bytes
35514b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """Managed, loopback-only llama.cpp server used by the validation pipeline."""
import contextlib
import socket
import subprocess
import time
import unicodedata
import requests
from test_equivalence import ROOT
class Server:
def __init__(self, model, gpu_layers=99, context=4096, mmproj=None):
self.model, self.gpu_layers, self.context, self.mmproj = model, gpu_layers, context, mmproj
def __enter__(self):
(ROOT / "logs").mkdir(exist_ok=True)
with socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
self.port = sock.getsockname()[1]
self.url = f"http://127.0.0.1:{self.port}"
self.log = (ROOT / "logs" / f"server-{self.model.stem}.log").open("a")
args = [str(ROOT / "llama.cpp/build/bin/llama-server"), "-m", str(self.model),
"--alias", "agnes", "--host", "127.0.0.1", "--port", str(self.port), "-c", str(self.context),
"-ngl", str(self.gpu_layers), "--parallel", "1", "--no-warmup", "--jinja",
"--flash-attn", "on", "--split-mode", "layer", "--tensor-split", "1,1",
"--threads", "6", "--batch-size", "256", "--ubatch-size", "128"]
if self.mmproj:
args += ["--mmproj", str(self.mmproj)]
self.process = subprocess.Popen(args, stdout=self.log, stderr=subprocess.STDOUT)
try:
deadline = time.monotonic() + 900
while time.monotonic() < deadline:
if self.process.poll() is not None:
raise RuntimeError(f"llama-server exited; see {self.log.name}")
try:
if requests.get(self.url + "/health", timeout=3).status_code == 200:
return self
except requests.RequestException:
pass
time.sleep(1)
raise TimeoutError("llama-server startup timed out")
except BaseException:
self.__exit__(None, None, None)
raise
def post(self, endpoint, payload):
response = requests.post(self.url + endpoint, json=payload, timeout=1800)
response.raise_for_status()
return response.json()
def complete(self, prompt, tokens=96, **extra):
if isinstance(prompt, str):
prompt = unicodedata.normalize("NFC", prompt)
return self.post("/completion", {"prompt": prompt, "n_predict": tokens,
"temperature": 0, "seed": 20260912, "cache_prompt": False, **extra})
def __exit__(self, *args):
if hasattr(self, "process") and self.process.poll() is None:
self.process.terminate()
try:
self.process.wait(timeout=30)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait(timeout=15)
if hasattr(self, "log"):
self.log.close()
def chat(tokenizer, text):
return unicodedata.normalize("NFC", tokenizer.apply_chat_template([{"role": "user", "content": text}],
tokenize=False, add_generation_prompt=True, enable_thinking=False))
|