454 lines
17 KiB
Python
454 lines
17 KiB
Python
"""llama-server 进程管理 & 模型下载。
|
||
|
||
职责：
|
||
- 启动/停止本地 llama-server 子进程（Windows 兼容）
|
||
- 探测已有 .gguf 模型文件
|
||
- 从 HuggingFace URL 下载模型（支持 huggingface.co 路径别名）
|
||
- 下载进度可通过 SSE /llama/download/stream 订阅
|
||
|
||
用法：
|
||
from gateway.llama_manager import get_llama_manager
|
||
lm = get_llama_manager()
|
||
await lm.start(model="models/qwen3.5-4b-q4_k_m.gguf")
|
||
await lm.stop()
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any, AsyncGenerator, Optional
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 路径配置（与 config.yaml runtime.llama_server 段保持一致）
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
_ROOT = Path(__file__).resolve().parent.parent # E:\projectAIpopular
|
||
BIN_DIR = _ROOT / "bin"
|
||
MODELS_DIR = _ROOT / "models"
|
||
PID_FILE = _ROOT / "data" / "llama-server.pid"
|
||
LOG_FILE = _ROOT / "data" / "llama-server.log"
|
||
|
||
# 确保目录存在
|
||
BIN_DIR.mkdir(parents=True, exist_ok=True)
|
||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||
PID_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 数据模型
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
@dataclass
|
||
class LlamaServerStatus:
|
||
running: bool
|
||
pid: Optional[int] = None
|
||
model: Optional[str] = None
|
||
port: Optional[int] = None
|
||
base_url: Optional[str] = None
|
||
started_at: Optional[float] = None
|
||
error: Optional[str] = None
|
||
|
||
|
||
@dataclass
|
||
class DownloadProgress:
|
||
url: str
|
||
dest: str
|
||
total_bytes: Optional[int] = None
|
||
downloaded_bytes: int = 0
|
||
progress_pct: float = 0.0
|
||
speed: str = ""
|
||
eta: str = ""
|
||
done: bool = False
|
||
error: Optional[str] = None
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# llama_manager 单例
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
class LlamaManager:
|
||
_instance: Optional["LlamaManager"] = None
|
||
|
||
def __init__(self):
|
||
self._proc: Optional[subprocess.Popen] = None
|
||
self._pid: Optional[int] = None
|
||
self._model: Optional[str] = None
|
||
self._port: Optional[int] = None
|
||
self._started_at: Optional[float] = None
|
||
self._downloading: dict[str, DownloadProgress] = {} # url -> progress
|
||
self._dl_lock = threading.Lock()
|
||
# 加载已有进程
|
||
self._load_pid()
|
||
|
||
# ── 进程持久化 ─────────────────────────────────────────────────────────
|
||
|
||
def _load_pid(self) -> None:
|
||
"""从 pid 文件恢复进程引用（进程仍在运行时）。"""
|
||
if not PID_FILE.exists():
|
||
return
|
||
try:
|
||
pid = int(PID_FILE.read_text().strip())
|
||
os.kill(pid, 0) # 检查进程是否存活
|
||
# 进程还在，尝试接管（通过 cmdline 判断是否是 llama-server）
|
||
self._pid = pid
|
||
self._proc = self._attach_to_process(pid)
|
||
except (ValueError, FileNotFoundError, OSError):
|
||
PID_FILE.unlink(missing_ok=True)
|
||
|
||
def _attach_to_process(self, pid: int) -> Optional[subprocess.Popen]:
|
||
"""通过 pid 重新关联到 Popen（仅作状态恢复，不拥有 stdout）。"""
|
||
try:
|
||
return subprocess.Popen(
|
||
[sys.executable, "-c",
|
||
f"import os; os.kill({pid}, 0)"], # 存活检查
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.DEVNULL,
|
||
)
|
||
except Exception:
|
||
return None
|
||
|
||
def _write_pid(self, pid: int) -> None:
|
||
PID_FILE.write_text(str(pid), encoding="utf-8")
|
||
|
||
def _clear_pid(self) -> None:
|
||
PID_FILE.unlink(missing_ok=True)
|
||
|
||
# ── 进程管理 ─────────────────────────────────────────────────────────
|
||
|
||
def find_binary(self) -> Optional[Path]:
|
||
"""查找 llama-server 可执行文件。"""
|
||
candidates = [
|
||
BIN_DIR / "llama-server.exe",
|
||
BIN_DIR / "llama-server",
|
||
_ROOT / "llama-server.exe",
|
||
_ROOT / "llama-server",
|
||
]
|
||
for p in candidates:
|
||
if p.exists():
|
||
return p
|
||
# PATH 中查找
|
||
import shutil
|
||
found = shutil.which("llama-server") or shutil.which("llama-server.exe")
|
||
if found:
|
||
return Path(found)
|
||
return None
|
||
|
||
def status(self) -> LlamaServerStatus:
|
||
"""返回当前服务状态。"""
|
||
if self._proc is None or self._pid is None:
|
||
return LlamaServerStatus(running=False)
|
||
try:
|
||
# 检查进程是否存活
|
||
os.kill(self._pid, 0)
|
||
except OSError:
|
||
# 进程已死
|
||
self._proc = None
|
||
self._pid = None
|
||
self._model = None
|
||
self._port = None
|
||
self._started_at = None
|
||
self._clear_pid()
|
||
return LlamaServerStatus(running=False)
|
||
return LlamaServerStatus(
|
||
running=True,
|
||
pid=self._pid,
|
||
model=self._model,
|
||
port=self._port,
|
||
base_url=f"http://127.0.0.1:{self._port}/v1",
|
||
started_at=self._started_at,
|
||
)
|
||
|
||
async def start(
|
||
self,
|
||
model: str,
|
||
port: int = 8901,
|
||
ngl: int = 99,
|
||
ctx: int = 4096,
|
||
extra_args: Optional[list] = None,
|
||
) -> LlamaServerStatus:
|
||
"""启动 llama-server，阻塞直到监听就绪或超时。"""
|
||
if self.status().running:
|
||
s = self.status()
|
||
if s.model == model and s.port == port:
|
||
return s # 已是同一模型，无需重启
|
||
await self.stop()
|
||
|
||
binary = self.find_binary()
|
||
if binary is None:
|
||
return LlamaServerStatus(
|
||
running=False,
|
||
error="未找到 llama-server 可执行文件。"
|
||
"请将 llama-server.exe 放入 bin/ 目录，"
|
||
"或从 https://github.com/ggerganov/llama.cpp/releases 下载。",
|
||
)
|
||
|
||
model_path = Path(model)
|
||
if not model_path.is_absolute():
|
||
model_path = MODELS_DIR / model
|
||
if not model_path.exists():
|
||
return LlamaServerStatus(
|
||
running=False,
|
||
error=f"模型文件不存在：{model_path}。"
|
||
"请先下载模型，或在设置页填写 HuggingFace URL 下载。",
|
||
)
|
||
|
||
args = [
|
||
str(binary),
|
||
"-m", str(model_path),
|
||
"-c", str(ctx),
|
||
"-ngl", str(ngl),
|
||
"--port", str(port),
|
||
"--host", "127.0.0.1",
|
||
]
|
||
if extra_args:
|
||
args.extend(extra_args)
|
||
|
||
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||
log_f = open(LOG_FILE, "w", encoding="utf-8", buffering=1)
|
||
|
||
try:
|
||
self._proc = subprocess.Popen(
|
||
args,
|
||
stdout=log_f,
|
||
stderr=subprocess.STDOUT,
|
||
cwd=str(_ROOT),
|
||
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0,
|
||
)
|
||
except OSError as e:
|
||
log_f.close()
|
||
return LlamaServerStatus(running=False, error=f"启动失败：{e}")
|
||
|
||
self._pid = self._proc.pid
|
||
self._model = str(model_path)
|
||
self._port = port
|
||
self._started_at = time.time()
|
||
self._write_pid(self._pid)
|
||
|
||
# 等待服务就绪
|
||
ok = await self._wait_until_ready(port, timeout=30)
|
||
if not ok:
|
||
await self.stop()
|
||
return LlamaServerStatus(
|
||
running=False,
|
||
error=f"llama-server 启动后 {port} 端口在 30 秒内未响应",
|
||
)
|
||
|
||
return self.status()
|
||
|
||
async def stop(self) -> None:
|
||
"""优雅停止 llama-server。"""
|
||
if self._pid is None:
|
||
self._proc = None
|
||
return
|
||
|
||
try:
|
||
if sys.platform == "win32":
|
||
# Windows: CTRL_BREAK_EVENT 或 taskkill
|
||
subprocess.run(
|
||
["taskkill", "/PID", str(self._pid), "/T", "/F"],
|
||
capture_output=True,
|
||
timeout=10,
|
||
)
|
||
else:
|
||
os.kill(self._pid, 15) # SIGTERM
|
||
time.sleep(1)
|
||
try:
|
||
os.kill(self._pid, 0)
|
||
os.kill(self._pid, 9)
|
||
except OSError:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
self._proc = None
|
||
self._pid = None
|
||
self._model = None
|
||
self._port = None
|
||
self._started_at = None
|
||
self._clear_pid()
|
||
|
||
async def _wait_until_ready(self, port: int, timeout: float = 30) -> bool:
|
||
"""轮询检查端口是否开始监听。"""
|
||
import httpx
|
||
url = f"http://127.0.0.1:{port}/v1/models"
|
||
deadline = time.time() + timeout
|
||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||
while time.time() < deadline:
|
||
try:
|
||
r = await client.get(url)
|
||
if r.status_code < 500:
|
||
return True
|
||
except (httpx.ConnectError, httpx.ReadTimeout, OSError):
|
||
pass
|
||
await asyncio.sleep(0.5)
|
||
return False
|
||
|
||
# ── 模型列表 ─────────────────────────────────────────────────────────
|
||
|
||
def list_local_models(self) -> list[dict[str, str]]:
|
||
"""列出 models/ 目录下所有 .gguf 文件。"""
|
||
models = []
|
||
for p in MODELS_DIR.glob("*.gguf"):
|
||
size_mb = p.stat().st_size // (1024 * 1024)
|
||
models.append({
|
||
"id": p.name,
|
||
"name": p.name,
|
||
"size_mb": size_mb,
|
||
"path": str(p),
|
||
})
|
||
return sorted(models, key=lambda m: m["name"])
|
||
|
||
# ── 模型下载 ─────────────────────────────────────────────────────────
|
||
|
||
async def download_model(
|
||
self,
|
||
url: str,
|
||
dest: Optional[str] = None,
|
||
) -> DownloadProgress:
|
||
"""从 HuggingFace 或直链下载 .gguf 模型文件。
|
||
|
||
HuggingFace 路径别名：用户输入 "Qwen/Qwen3-4B-GGUF/Qwen3-4B-Q4_K_M.gguf"
|
||
自动转换为 "https://huggingface.co/<repo>/resolve/main/<file>"
|
||
|
||
支持断点续传（Content-Range）。
|
||
|
||
返回 DownloadProgress 对象（含当前进度），进度通过 get_download_progress() 查询。
|
||
"""
|
||
import httpx
|
||
|
||
# 路径别名转换
|
||
if not url.startswith("http"):
|
||
url = f"https://huggingface.co/{url}/resolve/main"
|
||
|
||
# 解析文件名
|
||
filename = url.rstrip("/").split("/")[-1]
|
||
if not filename.endswith(".gguf"):
|
||
filename += ".gguf"
|
||
|
||
if dest:
|
||
dest_path = Path(dest)
|
||
else:
|
||
dest_path = MODELS_DIR / filename
|
||
|
||
# 构造 HTTP 头
|
||
headers = {}
|
||
resume_bytes = 0
|
||
if dest_path.exists():
|
||
resume_bytes = dest_path.stat().st_size
|
||
headers["Range"] = f"bytes={resume_bytes}-"
|
||
|
||
# 获取文件大小
|
||
total_bytes: Optional[int] = None
|
||
try:
|
||
async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, connect=10.0, read=60.0, write=30.0, pool=10.0), follow_redirects=True) as client:
|
||
head = await client.head(url, headers={"Range": "bytes=0-0"})
|
||
total_raw = head.headers.get("Content-Length")
|
||
if total_raw:
|
||
total_bytes = int(total_raw)
|
||
# Content-Range 响应时 total_bytes 在 Content-Range 头里
|
||
cr = head.headers.get("Content-Range", "")
|
||
m = re.search(r"/(\d+)", cr)
|
||
if m:
|
||
total_bytes = int(m.group(1))
|
||
except Exception:
|
||
pass
|
||
|
||
prog = DownloadProgress(
|
||
url=url,
|
||
dest=str(dest_path),
|
||
total_bytes=total_bytes,
|
||
downloaded_bytes=resume_bytes,
|
||
)
|
||
with self._dl_lock:
|
||
self._downloading[url] = prog
|
||
|
||
try:
|
||
mode = "ab" if resume_bytes > 0 else "wb"
|
||
async with httpx.AsyncClient(
|
||
timeout=httpx.Timeout(300.0, connect=10.0, read=300.0, write=30.0, pool=10.0),
|
||
follow_redirects=True,
|
||
) as client:
|
||
t0 = time.time()
|
||
last_bytes = resume_bytes
|
||
async with client.stream("GET", url, headers=headers) as resp:
|
||
if resp.status_code not in (200, 206):
|
||
raise RuntimeError(f"HTTP {resp.status_code}")
|
||
with open(dest_path, mode) as f:
|
||
async for chunk in resp.aiter_bytes(chunk_size=8192):
|
||
f.write(chunk)
|
||
prog.downloaded_bytes += len(chunk)
|
||
|
||
# 速度 & ETA
|
||
elapsed = time.time() - t0
|
||
if elapsed > 0.5:
|
||
speed_bps = (prog.downloaded_bytes - last_bytes) / elapsed
|
||
speed_str = _format_speed(speed_bps)
|
||
if prog.total_bytes and speed_bps > 0:
|
||
remain = prog.total_bytes - prog.downloaded_bytes
|
||
eta_s = remain / speed_bps
|
||
prog.eta = _format_eta(eta_s)
|
||
else:
|
||
prog.eta = ""
|
||
prog.speed = speed_str
|
||
last_bytes = prog.downloaded_bytes
|
||
t0 = time.time()
|
||
|
||
if prog.total_bytes:
|
||
prog.progress_pct = min(prog.downloaded_bytes / prog.total_bytes * 100, 100)
|
||
except Exception as e:
|
||
prog.error = str(e)
|
||
finally:
|
||
prog.done = True
|
||
with self._dl_lock:
|
||
self._downloading[url] = prog
|
||
|
||
return prog
|
||
|
||
def get_download_progress(self, url: str) -> Optional[DownloadProgress]:
|
||
"""查询下载进度。"""
|
||
with self._dl_lock:
|
||
return self._downloading.get(url)
|
||
|
||
def list_downloads(self) -> list[DownloadProgress]:
|
||
"""列出所有活跃下载。"""
|
||
with self._dl_lock:
|
||
return list(self._downloading.values())
|
||
|
||
|
||
def _format_speed(bps: float) -> str:
|
||
if bps >= 1e9:
|
||
return f"{bps/1e9:.1f} GB/s"
|
||
if bps >= 1e6:
|
||
return f"{bps/1e6:.1f} MB/s"
|
||
if bps >= 1e3:
|
||
return f"{bps/1e3:.1f} KB/s"
|
||
return f"{bps:.0f} B/s"
|
||
|
||
|
||
def _format_eta(seconds: float) -> str:
|
||
if seconds < 60:
|
||
return f"{seconds:.0f}s"
|
||
if seconds < 3600:
|
||
return f"{seconds/60:.0f}m"
|
||
return f"{seconds/3600:.1f}h"
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 全局单例
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
_lm: Optional[LlamaManager] = None
|
||
|
||
|
||
def get_llama_manager() -> LlamaManager:
|
||
global _lm
|
||
if _lm is None:
|
||
_lm = LlamaManager()
|
||
return _lm
|