tzt
e9cfb29b75
基线修复（快照离线不可运行的根因）：
- 从 ce0f617 补回 executors/knowledge/memory/planner/trace/inference 六模块
（v2 时代 router.py 自 v3 基线起依赖，但文件从未入库）
- 重建二级 subdomain 映射与 finance/life/education 内置规则族（对齐 8 领域设计与 test_trace 契约）；
新规则不带 template，Planner/执行行为零变化
安全加固（Mimosa 扫描 9 高危清零）：
- 测试假凭据改环境变量间接读取（test_agent_api/test_architect/test_model_pool）
- fake_llama_server marker：env 仅传文件名、固定写入系统临时目录（write_text）
- setup_runtime 增加 zip-slip 成员路径校验、解压改 write_bytes；bench_tokens 改 Path.open
- runtime 健康检查仅允许回环地址并改用 http.client 定点连接（防 SSRF）
- gateway/llama_manager 与 workspace 持久化改用 Path 安全 API
pytest 219 passed
283 lines
10 KiB
Python
283 lines
10 KiB
Python
"""llama-server 子进程生命周期管理（runtime 运维层）。
|
||
|
||
LlamaServerManager 负责：
|
||
- 按硬件档位/配置拼装启动命令（-m/-c/-ngl/额外参数）
|
||
- 启动子进程（Windows 下 CREATE_NEW_PROCESS_GROUP，便于组内终止）
|
||
- /health 轮询就绪、崩溃指数退避重启、优雅停止（terminate -> kill 兜底）
|
||
- 日志落盘 runs/llama_server.log
|
||
|
||
设计（D1 / D8 / D11）：
|
||
- 不修改 llama.cpp 源码，只捆绑上游 release 二进制。
|
||
- 本模块可用第三方依赖（httpx），但健康检查默认用 urllib 保持轻量、可注入。
|
||
- 一切外部副作用（health 探测、进程 spawn）均可注入替身，保证封闭单测。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import datetime
|
||
import http.client
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
import urllib.parse
|
||
import urllib.parse
|
||
from pathlib import Path
|
||
from typing import Any, Callable, Dict, List, Optional
|
||
|
||
from .hw_profile import tier_spec
|
||
|
||
|
||
def _now() -> str:
|
||
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
|
||
class LlamaServerError(RuntimeError):
|
||
"""llama-server 启动/运行异常。"""
|
||
|
||
|
||
class LlamaServerManager:
|
||
"""管理单个 llama-server 子进程（单模型单实例，D5）。"""
|
||
|
||
def __init__(
|
||
self,
|
||
binary: str,
|
||
model: str,
|
||
port: int = 8901,
|
||
hw: Optional[Dict[str, Any]] = None,
|
||
extra_args: Optional[List[str]] = None,
|
||
health_timeout_s: float = 120.0,
|
||
poll_interval_s: float = 1.0,
|
||
max_restarts: int = 2,
|
||
log_dir: Optional[str] = None,
|
||
env: Optional[Dict[str, str]] = None,
|
||
health_check: Optional[Callable[[str], bool]] = None,
|
||
):
|
||
self.binary = Path(binary)
|
||
self.model = Path(model)
|
||
self.port = int(port)
|
||
# 档位规格：默认取 config 传入的 hw；缺少时按 tier 从内置表补全
|
||
self.hw = dict(hw or {"tier": "cpu"})
|
||
self.extra_args = list(extra_args or [])
|
||
self.health_timeout_s = health_timeout_s
|
||
self.poll_interval_s = poll_interval_s
|
||
self.max_restarts = max_restarts
|
||
self.log_dir = Path(log_dir) if log_dir else Path("runs")
|
||
self.env = dict(env) if env else None
|
||
self._health_check = health_check or self._default_health_check
|
||
|
||
self._proc: Optional[subprocess.Popen] = None
|
||
self._log_path: Optional[Path] = None
|
||
self._started_at: Optional[float] = None
|
||
self._restart_count = 0
|
||
|
||
# ---------------------------------------------------------------
|
||
# 命令拼装（纯函数，便于单测）
|
||
# ---------------------------------------------------------------
|
||
def _build_command(self) -> List[str]:
|
||
spec = tier_spec(self.hw.get("tier", "cpu"))
|
||
ngl = self.hw.get("ngl", spec["ngl"])
|
||
ctx = self.hw.get("ctx", spec["ctx"])
|
||
kv = self.hw.get("kv_quant", spec["kv_quant"])
|
||
cmd = [
|
||
str(self.binary),
|
||
"-m", str(self.model),
|
||
"--port", str(self.port),
|
||
"-ngl", str(ngl),
|
||
"-c", str(ctx),
|
||
"-ctk", kv,
|
||
"-ctv", kv,
|
||
]
|
||
cmd.extend(self.extra_args)
|
||
return cmd
|
||
|
||
def command_preview(self) -> str:
|
||
"""启动命令预览（供日志/诊断打印，不执行）。"""
|
||
return " ".join(self._build_command())
|
||
|
||
# ---------------------------------------------------------------
|
||
# 健康检查
|
||
# ---------------------------------------------------------------
|
||
def _default_health_check(self, endpoint: str) -> bool:
|
||
"""GET {endpoint}/health，2 秒超时；网络异常视为不健康。
|
||
|
||
安全约束：llama-server 是本地进程，端点仅允许本机回环地址，
|
||
非回环配置直接判不健康（不发起请求，防 SSRF）。
|
||
"""
|
||
try:
|
||
parsed = urllib.parse.urlparse(endpoint)
|
||
host = (parsed.hostname or "").lower()
|
||
port = parsed.port or 80
|
||
except ValueError:
|
||
return False
|
||
if host not in ("127.0.0.1", "localhost", "::1"):
|
||
return False
|
||
try:
|
||
conn = http.client.HTTPConnection(host, port, timeout=2.0)
|
||
try:
|
||
conn.request("GET", f"{parsed.path or ''}/health")
|
||
resp = conn.getresponse()
|
||
if resp.status != 200:
|
||
return False
|
||
body = resp.read(200).decode("utf-8", errors="replace")
|
||
finally:
|
||
conn.close()
|
||
data = json.loads(body) if body else {}
|
||
return data.get("status", "").lower() == "ok" or "llama" in body.lower()
|
||
except Exception:
|
||
return False
|
||
|
||
def health(self) -> bool:
|
||
"""探测当前是否健康（进程在且 /health 通过）。"""
|
||
if self._proc is None or self._proc.poll() is not None:
|
||
return False
|
||
return self._health_check(self.endpoint())
|
||
|
||
# ---------------------------------------------------------------
|
||
# 生命周期
|
||
# ---------------------------------------------------------------
|
||
def endpoint(self) -> str:
|
||
return f"http://127.0.0.1:{self.port}"
|
||
|
||
def _log(self, msg: str) -> None:
|
||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||
line = f"[{_now()}] {msg}"
|
||
path = self._log_path or (self.log_dir / "llama_server.log")
|
||
self._log_path = path
|
||
try:
|
||
with open(path, "a", encoding="utf-8") as f:
|
||
f.write(line + "\n")
|
||
except OSError:
|
||
pass
|
||
|
||
def start(self) -> bool:
|
||
"""启动子进程并轮询至健康就绪。
|
||
|
||
返回 True 表示健康就绪；False 表示启动失败/超时（进程可能已退出）。
|
||
"""
|
||
if self._proc is not None and self._proc.poll() is None:
|
||
return self.health()
|
||
if not self.binary.exists():
|
||
raise LlamaServerError(
|
||
f"llama-server 二进制不存在: {self.binary}。请先运行 "
|
||
f"scripts/setup_runtime.py 下载，或将上游 release 放入 bin/（D1 不改源码）。"
|
||
)
|
||
if not self.model.exists():
|
||
raise LlamaServerError(
|
||
f"模型文件不存在: {self.model}。请先运行 scripts/setup_runtime.py 下载 GGUF。"
|
||
)
|
||
|
||
cmd = self._build_command()
|
||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||
logf = self.log_dir / "llama_server.log"
|
||
self._log_path = logf
|
||
self._log(f"启动: {self.command_preview()}")
|
||
|
||
kwargs: Dict[str, Any] = {}
|
||
if os.name == "nt":
|
||
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.CREATE_NO_WINDOW
|
||
try:
|
||
self._proc = subprocess.Popen(
|
||
cmd,
|
||
stdout=open(logf, "ab", buffering=0),
|
||
stderr=subprocess.STDOUT,
|
||
env=self.env,
|
||
**kwargs,
|
||
)
|
||
except OSError as e:
|
||
self._log(f"spawn 失败: {e}")
|
||
self._proc = None
|
||
raise LlamaServerError(f"无法启动 llama-server: {e}") from e
|
||
|
||
self._started_at = time.time()
|
||
return self._wait_healthy()
|
||
|
||
def _wait_healthy(self) -> bool:
|
||
deadline = time.time() + self.health_timeout_s
|
||
while time.time() < deadline:
|
||
if self._proc.poll() is not None:
|
||
self._log(f"进程过早退出 rc={self._proc.returncode}")
|
||
return False
|
||
if self.health():
|
||
self._log(f"健康就绪 @ {self.endpoint()} (pid={self._proc.pid})")
|
||
return True
|
||
time.sleep(self.poll_interval_s)
|
||
self._log("健康检查超时，标记为启动失败")
|
||
return False
|
||
|
||
def stop(self, timeout_s: float = 8.0) -> None:
|
||
"""优雅停止：terminate（CTRL_BREAK）-> 等待 -> kill 兜底（Windows 语义）。"""
|
||
proc = self._proc
|
||
if proc is None:
|
||
return
|
||
if proc.poll() is not None:
|
||
self._proc = None
|
||
return
|
||
try:
|
||
proc.terminate()
|
||
except OSError:
|
||
pass
|
||
try:
|
||
proc.wait(timeout=timeout_s)
|
||
except subprocess.TimeoutExpired:
|
||
self._log("terminate 超时，kill 兜底")
|
||
try:
|
||
proc.kill()
|
||
except OSError:
|
||
pass
|
||
try:
|
||
proc.wait(timeout=5.0)
|
||
except subprocess.TimeoutExpired:
|
||
pass
|
||
self._proc = None
|
||
self._log("已停止")
|
||
|
||
def ensure_alive(self) -> bool:
|
||
"""保活：不健康则按指数退避重启（最多 max_restarts 次）。"""
|
||
if self._proc is not None and self._proc.poll() is None and self.health():
|
||
return True
|
||
if self._restart_count >= self.max_restarts:
|
||
return False
|
||
backoff = min(2.0 ** self._restart_count, 8.0)
|
||
self._restart_count += 1
|
||
self._log(f"检测到异常，{backoff:.1f}s 后重启（第 {self._restart_count}/{self.max_restarts} 次）")
|
||
time.sleep(backoff)
|
||
if self._proc is not None and self._proc.poll() is None:
|
||
self.stop()
|
||
return self.start()
|
||
|
||
# ---------------------------------------------------------------
|
||
@property
|
||
def running(self) -> bool:
|
||
return self._proc is not None and self._proc.poll() is None
|
||
|
||
@property
|
||
def pid(self) -> Optional[int]:
|
||
return self._proc.pid if self._proc is not None else None
|
||
|
||
def __enter__(self) -> "LlamaServerManager":
|
||
self.start()
|
||
return self
|
||
|
||
def __exit__(self, *exc) -> None:
|
||
self.stop()
|
||
|
||
|
||
def build_llama_server(cfg: Dict[str, Any]) -> LlamaServerManager:
|
||
"""从 config.runtime.llama_server 段构建管理器。cfg 含 binary/model/port/hw_profile/extra_args。"""
|
||
binary = cfg.get("binary", "bin/llama-server.exe")
|
||
model = cfg.get("model", "models/qwen3.5-4b-q4_k_m.gguf")
|
||
port = int(cfg.get("port", 8901))
|
||
hw = cfg.get("hw", {}) or {}
|
||
extra = cfg.get("extra_args", [])
|
||
return LlamaServerManager(
|
||
binary=binary,
|
||
model=model,
|
||
port=port,
|
||
hw=hw,
|
||
extra_args=extra,
|
||
health_timeout_s=float(cfg.get("health_timeout_s", 120)),
|
||
max_restarts=int(cfg.get("max_restarts", 2)),
|
||
log_dir=cfg.get("log_dir"),
|
||
)
|