| """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)) |
|
|
