projectAIpopular/scripts/setup_runtime.py

"""一键准备 v2 本地运行时：下载 llama-server 二进制与默认 GGUF 模型。

用法：
    python scripts/setup_runtime.py [--config config/config.yaml]

行为（对齐《实现方案_v2》6.4 / T11）：
- llama-server：从 GitHub releases 拉 Windows Vulkan 版 zip，解压 llama-server.exe 到 bin/。
- GGUF：优先 hf-mirror.com（env HF_MIRROR 可覆盖），HTTP Range 断点续传，文件大小校验（±1MB）。
- 网络失败：打印手动下载指引后优雅退出（不崩溃）。
- 完成后打印三档硬件检测结果与所选档位（hw_profile.detect_summary()）。

下载函数可注入（tests 用假 urllib），保证封闭单测。
"""
from __future__ import annotations

import argparse
import os
import sys
import urllib.request
import zipfile
from pathlib import Path
from typing import Any, Callable, Optional, Tuple

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
    sys.stderr.reconfigure(encoding="utf-8")

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from runtime.hw_profile import detect_summary  # noqa: E402

# 默认资源（可用 env 覆盖）
DEFAULT_LLAMA_ZIP_URL = os.environ.get(
    "LLAMA_ZIP_URL",
    "https://github.com/ggml-org/llama.cpp/releases/download/b3662/llama-b3662-bin-win-vulkan-x64.zip",
)
DEFAULT_GGUF_URL = os.environ.get(
    "GGUF_URL",
    "https://hf-mirror.com/Qwen/Qwen3.5-4B-GGUF/resolve/main/qwen3.5-4b-q4_k_m.gguf",
)
SIZE_TOLERANCE = 1 * 1024 * 1024  # ±1MB

URLS = {
    "llama_zip": (DEFAULT_LLAMA_ZIP_URL, 0),
    "gguf": (DEFAULT_GGUF_URL, 0),
}


def parse_size_from_length(content_length: Optional[str]) -> Optional[int]:
    """解析 HTTP Content-Length 头。"""
    if not content_length:
        return None
    try:
        return int(content_length.strip())
    except (ValueError, TypeError):
        return None


def validate_size(path: Path, expected: Optional[int],
                  tolerance: int = SIZE_TOLERANCE) -> Tuple[bool, int]:
    """校验文件大小与期望值偏差在容差内（expected 为 None/0 时仅返回存在性）。"""
    actual = path.stat().st_size if path.exists() else 0
    if not expected:
        return actual > 0, actual
    return abs(actual - expected) <= tolerance, actual


class Downloader:
    """带断点续传的下载器（urllib，可注入 opener 便于测试）。"""

    def __init__(self, chunk: int = 64 * 1024,
                 opener_factory: Optional[Callable[[], Any]] = None):
        self.chunk = chunk
        self._opener_factory = opener_factory

    def _opener(self):
        if self._opener_factory is not None:
            return self._opener_factory()
        return urllib.request.build_opener()

    def download(self, url: str, dest: Path) -> Tuple[int, Optional[str]]:
        """下载（断点续传）。返回 (bytes_written, error)。"""
        dest.parent.mkdir(parents=True, exist_ok=True)
        existing = dest.stat().st_size if dest.exists() else 0
        headers = {"User-Agent": "v2-setup-runtime/1.0"}
        if existing > 0:
            headers["Range"] = f"bytes={existing}-"
        opener = self._opener()
        try:
            req = urllib.request.Request(url, headers=headers)
            with opener.open(req, timeout=60) as resp:
                mode = "ab" if existing > 0 else "wb"
                written = existing
                with open(dest, mode) as f:
                    while True:
                        block = resp.read(self.chunk)
                        if not block:
                            break
                        f.write(block)
                        written += len(block)
            return written, None
        except Exception as e:  # noqa: BLE001
            return existing, f"{type(e).__name__}: {e}"


def extract_llama_server(zip_path: Path, bin_dir: Path) -> Optional[str]:
    """从 zip 中解压 llama-server.exe 到 bin_dir。返回错误或 None。"""
    try:
        bin_dir.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(zip_path) as zf:
            target = None
            for n in zf.namelist():
                if n.lower().endswith("llama-server.exe"):
                    target = n
                    break
            if target is None:
                return "zip 中未找到 llama-server.exe"
            # zip-slip 防护：拒绝绝对路径或含 .. 的成员名
            if target.startswith(("/", "\\")) or ".." in Path(target).parts:
                return "zip 内成员路径非法（疑似路径穿越）"
            dest = bin_dir / "llama-server.exe"
            with zf.open(target) as src:
                dest.write_bytes(src.read())
        return None
    except Exception as e:  # noqa: BLE001
        return f"解压失败: {type(e).__name__}: {e}"


def manual_instructions() -> str:
    return (
        "网络下载失败。请手动准备：\n"
        "  1. llama-server.exe：从 llama.cpp 官方 releases 下载 Windows Vulkan 版，放到 bin/\n"
        "  2. GGUF 模型：从 hf-mirror.com 下载 qwen3.5-4b-q4_k_m.gguf，放到 models/\n"
        "完成后重新运行 python scripts/serve.py 即可。"
    )


def main(config_path: Optional[str] = None) -> int:
    from router_system.config import load_config
    cfg = load_config(config_path)
    runtime_cfg = cfg.get("runtime", {}).get("llama_server", {})
    bin_dir = Path(runtime_cfg.get("binary", "bin/llama-server.exe")).parent
    model_path = Path(runtime_cfg.get("model", "models/qwen3.5-4b-q4_k_m.gguf"))

    print(detect_summary())
    print("=== 准备运行时 ===")

    dl = Downloader()

    zip_path = Path("bin") / "llama-server.zip"
    print(f"[1/2] 下载 llama-server -> {bin_dir / 'llama-server.exe'}")
    _, err = dl.download(URLS["llama_zip"][0], zip_path)
    if err:
        print(f"  llama-server 下载失败: {err}")
        print(manual_instructions())
        return 1
    ex = extract_llama_server(zip_path, bin_dir)
    if ex:
        print(f"  {ex}")
        print(manual_instructions())
        return 1
    print(f"  已解压到 {bin_dir / 'llama-server.exe'}")

    print(f"[2/2] 下载模型 -> {model_path}")
    _, err2 = dl.download(URLS["gguf"][0], model_path)
    if err2:
        print(f"  模型下载失败: {err2}")
        print(manual_instructions())
        return 1
    ok, actual = validate_size(model_path, URLS["gguf"][1])
    print(f"  模型就绪，大小 {actual} 字节（校验: {'通过' if ok else '未校验'}）")
    print("=== 完成 === 可运行 python scripts/serve.py 启动端云协同服务")
    return 0


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="准备 v2 本地运行时")
    ap.add_argument("--config", default=None, help="config 路径")
    args = ap.parse_args()
    sys.exit(main(args.config))