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