242 lines
9.7 KiB
Python
242 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import socket
|
|
import subprocess
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from ..config import ModelConfig, RuntimeConfig
|
|
from ..types import BackendError, ChatChunk, ChatParams
|
|
from .base import BaseBackend
|
|
|
|
|
|
def _free_port() -> int:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
|
sock.bind(("127.0.0.1", 0))
|
|
return int(sock.getsockname()[1])
|
|
|
|
|
|
class GGUFBackend(BaseBackend):
|
|
"""Run a GGUF model through an external CUDA-enabled llama-server."""
|
|
|
|
def __init__(self, model: ModelConfig, runtime: RuntimeConfig) -> None:
|
|
super().__init__(model, runtime)
|
|
self.port: int | None = None
|
|
self.process: subprocess.Popen[str] | None = None
|
|
self.log_path: Path | None = None
|
|
self._log_handle: Any = None
|
|
|
|
@property
|
|
def endpoint(self) -> str:
|
|
if self.port is None:
|
|
raise BackendError("GGUF 后端尚未启动")
|
|
return f"http://127.0.0.1:{self.port}"
|
|
|
|
async def load(self) -> None:
|
|
if self.loaded:
|
|
return
|
|
if not self.model.path:
|
|
raise BackendError(f"GGUF 模型缺少 path: {self.model.id}")
|
|
model_path = Path(self.model.path)
|
|
if not model_path.exists():
|
|
raise BackendError(f"GGUF 文件不存在: {model_path}")
|
|
server = Path(self.runtime.llama_server)
|
|
if not server.exists():
|
|
raise BackendError(
|
|
f"找不到 llama-server: {server}。请在 config.json 的 runtime.llama_server 中配置。"
|
|
)
|
|
|
|
self.port = _free_port()
|
|
runtime_dir = Path(self.runtime.runtime_dir)
|
|
runtime_dir.mkdir(parents=True, exist_ok=True)
|
|
self.log_path = runtime_dir / f"{self.model.id}.llama-server.log"
|
|
self._log_handle = self.log_path.open("a", encoding="utf-8", errors="replace")
|
|
|
|
options = self.model.options
|
|
args = [
|
|
str(server),
|
|
"--model",
|
|
str(model_path),
|
|
"--host",
|
|
"127.0.0.1",
|
|
"--port",
|
|
str(self.port),
|
|
"--alias",
|
|
self.model.id,
|
|
"--no-webui",
|
|
"--jinja",
|
|
"--cont-batching",
|
|
"--cache-prompt",
|
|
"--flash-attn",
|
|
str(options.get("flash_attn", self.runtime.flash_attn)),
|
|
"--gpu-layers",
|
|
str(options.get("n_gpu_layers", self.runtime.n_gpu_layers)),
|
|
"--ctx-size",
|
|
str(options.get("ctx_size", self.runtime.ctx_size)),
|
|
"--batch-size",
|
|
str(options.get("batch_size", self.runtime.batch_size)),
|
|
"--ubatch-size",
|
|
str(options.get("ubatch_size", self.runtime.ubatch_size)),
|
|
]
|
|
if self.runtime.fit_vram and options.get("fit_vram", True):
|
|
args.extend(["--fit", "on", "--fit-ctx", str(options.get("fit_ctx", 4096))])
|
|
if options.get("mlock", False):
|
|
args.append("--mlock")
|
|
if options.get("no_mmap", False):
|
|
args.append("--no-mmap")
|
|
if options.get("mmproj"):
|
|
args.extend(["--mmproj", str(options["mmproj"])])
|
|
if options.get("chat_template_file"):
|
|
args.extend(["--chat-template-file", str(options["chat_template_file"])])
|
|
elif self.model.template:
|
|
args.extend(["--chat-template", self.model.template])
|
|
|
|
creationflags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
|
|
self.process = subprocess.Popen(
|
|
args,
|
|
cwd=str(server.parent),
|
|
stdout=self._log_handle,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
creationflags=creationflags,
|
|
)
|
|
try:
|
|
await self._wait_ready()
|
|
except Exception:
|
|
await self.unload()
|
|
raise
|
|
self.loaded = True
|
|
|
|
async def _wait_ready(self) -> None:
|
|
deadline = asyncio.get_running_loop().time() + 90.0
|
|
last_error = ""
|
|
async with httpx.AsyncClient(timeout=2.0) as client:
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
if self.process and self.process.poll() is not None:
|
|
tail = ""
|
|
if self.log_path and self.log_path.exists():
|
|
tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-2000:]
|
|
if "check_tensor_dims" in tail or "wrong shape" in tail:
|
|
tail = (
|
|
"检测到模型张量与当前 llama.cpp runner 不兼容。若 path 指向 Ollama 的 "
|
|
"content-addressed blob，请改用 kind=ollama；否则换成与 runner 匹配的标准 GGUF 文件。\n"
|
|
+ tail
|
|
)
|
|
raise BackendError(f"llama-server 启动失败。日志尾部:\n{tail}")
|
|
try:
|
|
response = await client.get(f"{self.endpoint}/health")
|
|
if response.status_code == 200:
|
|
return
|
|
last_error = f"HTTP {response.status_code}"
|
|
except httpx.HTTPError as exc:
|
|
last_error = str(exc)
|
|
await asyncio.sleep(0.25)
|
|
raise BackendError(f"llama-server 在 90 秒内未就绪: {last_error}")
|
|
|
|
async def unload(self) -> None:
|
|
process, self.process = self.process, None
|
|
if process is not None and process.poll() is None:
|
|
process.terminate()
|
|
try:
|
|
await asyncio.to_thread(process.wait, 5)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
await asyncio.to_thread(process.wait)
|
|
if self._log_handle is not None:
|
|
self._log_handle.close()
|
|
self._log_handle = None
|
|
self.loaded = False
|
|
self.port = None
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
params: ChatParams,
|
|
stream: bool = False,
|
|
) -> AsyncIterator[ChatChunk]:
|
|
if not self.loaded:
|
|
await self.load()
|
|
payload: dict[str, Any] = {
|
|
"model": self.model.id,
|
|
"messages": messages,
|
|
"stream": stream,
|
|
"max_tokens": params.max_tokens,
|
|
"temperature": params.temperature,
|
|
"top_p": params.top_p,
|
|
"top_k": params.top_k,
|
|
"min_p": params.min_p,
|
|
"repeat_penalty": params.repeat_penalty,
|
|
}
|
|
if params.seed is not None:
|
|
payload["seed"] = params.seed
|
|
if params.stop:
|
|
payload["stop"] = params.stop
|
|
if params.enable_thinking is not None:
|
|
payload["chat_template_kwargs"] = {"enable_thinking": params.enable_thinking}
|
|
if params.response_format is not None:
|
|
payload["response_format"] = (
|
|
{"type": params.response_format} if isinstance(params.response_format, str) else params.response_format
|
|
)
|
|
if params.tools:
|
|
payload["tools"] = params.tools
|
|
if params.tool_choice is not None:
|
|
payload["tool_choice"] = params.tool_choice
|
|
payload.update(params.extra)
|
|
|
|
timeout = httpx.Timeout(600.0, connect=10.0)
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
if stream:
|
|
async with client.stream("POST", f"{self.endpoint}/v1/chat/completions", json=payload) as response:
|
|
if response.status_code >= 400:
|
|
detail = (await response.aread()).decode("utf-8", errors="replace")
|
|
raise BackendError(f"llama-server 请求失败 {response.status_code}: {detail}")
|
|
async for line in response.aiter_lines():
|
|
if not line.startswith("data:"):
|
|
continue
|
|
raw = line[5:].strip()
|
|
if raw == "[DONE]":
|
|
yield ChatChunk(done=True, finish_reason="stop")
|
|
return
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
choice = (data.get("choices") or [{}])[0]
|
|
delta = choice.get("delta") or {}
|
|
text = delta.get("content") or ""
|
|
finish = choice.get("finish_reason")
|
|
if text or finish or len(delta) > 1:
|
|
yield ChatChunk(text=text, finish_reason=finish, metadata=data, delta=delta)
|
|
yield ChatChunk(done=True, finish_reason="stop")
|
|
return
|
|
|
|
response = await client.post(f"{self.endpoint}/v1/chat/completions", json=payload)
|
|
if response.status_code >= 400:
|
|
raise BackendError(f"llama-server 请求失败 {response.status_code}: {response.text}")
|
|
data = response.json()
|
|
choice = (data.get("choices") or [{}])[0]
|
|
message = choice.get("message") or {}
|
|
yield ChatChunk(
|
|
text=message.get("content") or "",
|
|
done=True,
|
|
finish_reason=choice.get("finish_reason") or "stop",
|
|
metadata=data,
|
|
delta=message,
|
|
)
|
|
|
|
def info(self) -> dict[str, Any]:
|
|
return {
|
|
"id": self.model.id,
|
|
"kind": "gguf",
|
|
"path": self.model.path,
|
|
"loaded": self.loaded,
|
|
"endpoint": self.endpoint if self.port else None,
|
|
"pid": self.process.pid if self.process else None,
|
|
"log": str(self.log_path) if self.log_path else None,
|
|
}
|