Add GGUF model inference via llama-server backend
This commit is contained in:
Roland Tannous
parent
de1303f10c
commit
2f985ccbb5
13 changed files with 790 additions and 39 deletions
414
studio/backend/core/inference/llama_cpp.py
Normal file
414
studio/backend/core/inference/llama_cpp.py
Normal file
|
|
@ -0,0 +1,414 @@
|
|||
"""
|
||||
llama-server inference backend for GGUF models.
|
||||
|
||||
Manages a llama-server subprocess and proxies chat completions
|
||||
through its /v1/completions endpoint.
|
||||
"""
|
||||
import atexit
|
||||
import json
|
||||
import logging
|
||||
import shutil
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Generator, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LlamaCppBackend:
|
||||
"""
|
||||
Manages a llama-server subprocess for GGUF model inference.
|
||||
|
||||
Lifecycle:
|
||||
1. load_model() — starts llama-server with the GGUF file
|
||||
2. generate_chat_completion() — formats prompt, proxies to /v1/completions, streams back
|
||||
3. unload_model() — terminates llama-server subprocess
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._process: Optional[subprocess.Popen] = None
|
||||
self._port: Optional[int] = None
|
||||
self._model_identifier: Optional[str] = None
|
||||
self._gguf_path: Optional[str] = None
|
||||
self._healthy = False
|
||||
self._lock = threading.Lock()
|
||||
self._chat_template: Optional[str] = None
|
||||
|
||||
atexit.register(self._cleanup)
|
||||
|
||||
# ── Properties ────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def is_loaded(self) -> bool:
|
||||
return self._process is not None and self._healthy
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return f"http://127.0.0.1:{self._port}"
|
||||
|
||||
@property
|
||||
def model_identifier(self) -> Optional[str]:
|
||||
return self._model_identifier
|
||||
|
||||
# ── Binary discovery ──────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _find_llama_server_binary() -> Optional[str]:
|
||||
"""
|
||||
Locate the llama-server binary.
|
||||
|
||||
Search order:
|
||||
1. LLAMA_SERVER_PATH environment variable
|
||||
2. ./bin/llama-server (built by setup.sh)
|
||||
3. llama-server on PATH (system install)
|
||||
4. ./llama.cpp/llama-server (unsloth-zoo build output)
|
||||
"""
|
||||
import os
|
||||
|
||||
# 1. Env var
|
||||
env_path = os.environ.get("LLAMA_SERVER_PATH")
|
||||
if env_path and Path(env_path).is_file():
|
||||
return env_path
|
||||
|
||||
# 2. Project bin/ directory (setup.sh output)
|
||||
project_root = Path(__file__).resolve().parents[3] # core/inference/ → backend/ → studio/ → root
|
||||
bin_path = project_root / "bin" / "llama-server"
|
||||
if bin_path.is_file():
|
||||
return str(bin_path)
|
||||
|
||||
# 3. System PATH
|
||||
system_path = shutil.which("llama-server")
|
||||
if system_path:
|
||||
return system_path
|
||||
|
||||
# 4. unsloth-zoo build output (from GGUF export)
|
||||
llama_cpp_path = project_root / "llama.cpp" / "llama-server"
|
||||
if llama_cpp_path.is_file():
|
||||
return str(llama_cpp_path)
|
||||
|
||||
return None
|
||||
|
||||
# ── Port allocation ───────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _find_free_port() -> int:
|
||||
"""Find an available TCP port."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
# ── Lifecycle ─────────────────────────────────────────────────
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
gguf_path: str,
|
||||
model_identifier: str,
|
||||
n_ctx: int = 4096,
|
||||
n_gpu_layers: int = -1,
|
||||
n_threads: Optional[int] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Start llama-server with the given GGUF file.
|
||||
|
||||
Args:
|
||||
gguf_path: Path to the .gguf file
|
||||
model_identifier: Display identifier for the model
|
||||
n_ctx: Context window size
|
||||
n_gpu_layers: Number of layers to offload to GPU (-1 = all)
|
||||
n_threads: Number of CPU threads (None = auto)
|
||||
|
||||
Returns:
|
||||
True if server started and health check passed.
|
||||
"""
|
||||
with self._lock:
|
||||
# Kill existing process if any
|
||||
self._kill_process()
|
||||
|
||||
binary = self._find_llama_server_binary()
|
||||
if not binary:
|
||||
raise RuntimeError(
|
||||
"llama-server binary not found. "
|
||||
"Run setup.sh to build it, install llama.cpp, "
|
||||
"or set LLAMA_SERVER_PATH environment variable."
|
||||
)
|
||||
|
||||
if not Path(gguf_path).is_file():
|
||||
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
||||
|
||||
self._port = self._find_free_port()
|
||||
cmd = [
|
||||
binary,
|
||||
"-m", gguf_path,
|
||||
"--port", str(self._port),
|
||||
"-c", str(n_ctx),
|
||||
"-ngl", str(n_gpu_layers),
|
||||
]
|
||||
if n_threads is not None:
|
||||
cmd.extend(["--threads", str(n_threads)])
|
||||
|
||||
logger.info(f"Starting llama-server: {' '.join(cmd)}")
|
||||
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
|
||||
self._gguf_path = gguf_path
|
||||
self._model_identifier = model_identifier
|
||||
|
||||
# Wait for health
|
||||
if not self._wait_for_health(timeout=120.0):
|
||||
self._kill_process()
|
||||
raise RuntimeError(
|
||||
"llama-server failed to start. "
|
||||
"Check that the GGUF file is valid and you have enough memory."
|
||||
)
|
||||
|
||||
self._healthy = True
|
||||
|
||||
# Try to read chat template from GGUF metadata
|
||||
self._chat_template = self._read_gguf_chat_template(gguf_path)
|
||||
|
||||
logger.info(
|
||||
f"llama-server ready on port {self._port} "
|
||||
f"for model '{model_identifier}'"
|
||||
)
|
||||
return True
|
||||
|
||||
def unload_model(self) -> bool:
|
||||
"""Terminate the llama-server subprocess and clean up state."""
|
||||
with self._lock:
|
||||
self._kill_process()
|
||||
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
|
||||
self._model_identifier = None
|
||||
self._gguf_path = None
|
||||
self._port = None
|
||||
self._healthy = False
|
||||
self._chat_template = None
|
||||
return True
|
||||
|
||||
def _kill_process(self):
|
||||
"""Terminate the subprocess if running."""
|
||||
if self._process is None:
|
||||
return
|
||||
try:
|
||||
self._process.terminate()
|
||||
self._process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
|
||||
self._process.kill()
|
||||
self._process.wait(timeout=5)
|
||||
except Exception as e:
|
||||
logger.warning(f"Error killing llama-server process: {e}")
|
||||
finally:
|
||||
self._process = None
|
||||
|
||||
def _cleanup(self):
|
||||
"""atexit handler to ensure llama-server is terminated."""
|
||||
self._kill_process()
|
||||
|
||||
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
|
||||
"""
|
||||
Poll llama-server's /health endpoint until it responds 200.
|
||||
|
||||
Also monitors subprocess for early exit/crash.
|
||||
"""
|
||||
deadline = time.monotonic() + timeout
|
||||
url = f"http://127.0.0.1:{self._port}/health"
|
||||
|
||||
while time.monotonic() < deadline:
|
||||
# Check if process crashed
|
||||
if self._process.poll() is not None:
|
||||
# Read remaining output for error info
|
||||
output = self._process.stdout.read() if self._process.stdout else ""
|
||||
logger.error(
|
||||
f"llama-server exited with code {self._process.returncode}. "
|
||||
f"Output: {output[:2000]}"
|
||||
)
|
||||
return False
|
||||
|
||||
try:
|
||||
resp = httpx.get(url, timeout=2.0)
|
||||
if resp.status_code == 200:
|
||||
return True
|
||||
except (httpx.ConnectError, httpx.TimeoutException):
|
||||
pass
|
||||
|
||||
time.sleep(interval)
|
||||
|
||||
logger.error(f"llama-server health check timed out after {timeout}s")
|
||||
return False
|
||||
|
||||
# ── Chat template ─────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _read_gguf_chat_template(gguf_path: str) -> Optional[str]:
|
||||
"""
|
||||
Try to read the chat_template from GGUF file metadata.
|
||||
|
||||
Uses the gguf Python library if available.
|
||||
Returns the Jinja2 template string, or None.
|
||||
"""
|
||||
try:
|
||||
from gguf import GGUFReader
|
||||
|
||||
reader = GGUFReader(gguf_path)
|
||||
for field_name in reader.fields:
|
||||
if field_name == "tokenizer.chat_template":
|
||||
field = reader.fields[field_name]
|
||||
# Field data is an array of bytes
|
||||
template_bytes = bytes(field.parts[field.data[0]])
|
||||
template = template_bytes.decode("utf-8")
|
||||
logger.info(f"Read chat template from GGUF metadata ({len(template)} chars)")
|
||||
return template
|
||||
except ImportError:
|
||||
logger.debug("gguf library not available, cannot read chat template from GGUF metadata")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not read chat template from GGUF: {e}")
|
||||
|
||||
return None
|
||||
|
||||
def format_prompt(self, messages: list[dict], system_prompt: str = "") -> str:
|
||||
"""
|
||||
Format chat messages into a raw prompt string for /v1/completions.
|
||||
|
||||
Attempts to:
|
||||
1. Render the GGUF's embedded chat_template with Jinja2
|
||||
2. Fallback to ChatML format
|
||||
"""
|
||||
# Build full message list with system prompt
|
||||
full_messages = []
|
||||
if system_prompt:
|
||||
full_messages.append({"role": "system", "content": system_prompt})
|
||||
full_messages.extend(messages)
|
||||
|
||||
# Try Jinja2 rendering if we have a template
|
||||
if self._chat_template:
|
||||
try:
|
||||
return self._render_jinja_template(full_messages)
|
||||
except Exception as e:
|
||||
logger.warning(f"Jinja2 template rendering failed, falling back to ChatML: {e}")
|
||||
|
||||
# Fallback: ChatML format
|
||||
return self._format_chatml(full_messages)
|
||||
|
||||
def _render_jinja_template(self, messages: list[dict]) -> str:
|
||||
"""Render messages using the GGUF's Jinja2 chat template."""
|
||||
from jinja2 import BaseLoader, Environment
|
||||
|
||||
env = Environment(loader=BaseLoader(), keep_trailing_newline=True)
|
||||
# Add common template globals
|
||||
env.globals["raise_exception"] = lambda msg: (_ for _ in ()).throw(ValueError(msg))
|
||||
|
||||
template = env.from_string(self._chat_template)
|
||||
rendered = template.render(
|
||||
messages=messages,
|
||||
add_generation_prompt=True,
|
||||
bos_token="<s>",
|
||||
eos_token="</s>",
|
||||
)
|
||||
return rendered
|
||||
|
||||
@staticmethod
|
||||
def _format_chatml(messages: list[dict]) -> str:
|
||||
"""Format messages using ChatML template (universal fallback)."""
|
||||
parts = []
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
|
||||
parts.append("<|im_start|>assistant")
|
||||
return "\n".join(parts) + "\n"
|
||||
|
||||
# ── Generation (proxy to llama-server) ────────────────────────
|
||||
|
||||
def generate_chat_completion(
|
||||
self,
|
||||
prompt: str,
|
||||
temperature: float = 0.7,
|
||||
top_p: float = 0.9,
|
||||
top_k: int = 40,
|
||||
min_p: float = 0.0,
|
||||
max_tokens: int = 512,
|
||||
repetition_penalty: float = 1.1,
|
||||
stop: Optional[list[str]] = None,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
) -> Generator[str, None, None]:
|
||||
"""
|
||||
Send a completion request to llama-server and stream tokens back.
|
||||
|
||||
Uses /v1/completions (NOT /v1/chat/completions) so we control
|
||||
the prompt format entirely.
|
||||
|
||||
Yields cumulative text (matching InferenceBackend's convention).
|
||||
"""
|
||||
if not self.is_loaded:
|
||||
raise RuntimeError("llama-server is not loaded")
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"stream": True,
|
||||
"temperature": temperature,
|
||||
"top_p": top_p,
|
||||
"top_k": top_k if top_k >= 0 else 0,
|
||||
"min_p": min_p,
|
||||
"n_predict": max_tokens,
|
||||
"repeat_penalty": repetition_penalty,
|
||||
}
|
||||
if stop:
|
||||
payload["stop"] = stop
|
||||
|
||||
url = f"{self.base_url}/v1/completions"
|
||||
cumulative = ""
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=None) as client:
|
||||
with client.stream("POST", url, json=payload) as response:
|
||||
if response.status_code != 200:
|
||||
error_body = response.read().decode()
|
||||
raise RuntimeError(
|
||||
f"llama-server returned {response.status_code}: {error_body}"
|
||||
)
|
||||
|
||||
buffer = ""
|
||||
for raw_chunk in response.iter_text():
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
break
|
||||
|
||||
buffer += raw_chunk
|
||||
while "\n" in buffer:
|
||||
line, buffer = buffer.split("\n", 1)
|
||||
line = line.strip()
|
||||
|
||||
if not line:
|
||||
continue
|
||||
if line == "data: [DONE]":
|
||||
return
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
|
||||
try:
|
||||
data = json.loads(line[6:])
|
||||
choices = data.get("choices", [])
|
||||
if choices:
|
||||
token = choices[0].get("text", "")
|
||||
if token:
|
||||
cumulative += token
|
||||
yield cumulative
|
||||
except json.JSONDecodeError:
|
||||
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
|
||||
|
||||
except httpx.ConnectError:
|
||||
raise RuntimeError("Lost connection to llama-server")
|
||||
except Exception as e:
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return
|
||||
raise
|
||||
Loading…
Add table
Add a link
Reference in a new issue