141 lines
5.5 KiB
Python
141 lines
5.5 KiB
Python
"""硬件档位检测（runtime 运维层，纯标准库）。
|
||
|
||
把真实机器映射到三档保守模板之一：
|
||
|
||
- gpu12 : 约 ≥12GB 显存（NVIDIA / Vulkan 可探测） -> ngl 99, ctx 32768
|
||
- gpu8 : 约 ≥8GB 显存 -> ngl 14, ctx 16384
|
||
- cpu : 无独显或探测失败（保守兜底） -> ngl 0, ctx 8192
|
||
|
||
探测来源：nvidia-smi（NVIDIA 显存）优先；其次 vulkaninfo（AMD/Intel/通用，
|
||
只能判断是否存在 Vulkan 设备，无法可靠拿到显存 -> 保守回退 cpu，并在结果标注
|
||
probe:"conservative"）。总系统内存仅作为 cpu 档提示参考，不作为分档依据。
|
||
|
||
任何探测失败都回退到 cpu 保守档，保证不崩、可离线运行（D5 / D8）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import shutil
|
||
import subprocess
|
||
from typing import Any, Callable, Dict, List, Optional
|
||
|
||
# 三档硬件模板（保守默认，可被 config.tiers 手动覆盖）
|
||
TIER_SPECS: Dict[str, Dict[str, Any]] = {
|
||
"gpu12": {"tier": "gpu12", "ngl": 99, "ctx": 32768, "kv_quant": "q8_0"},
|
||
"gpu8": {"tier": "gpu8", "ngl": 14, "ctx": 16384, "kv_quant": "q8_0"},
|
||
"cpu": {"tier": "cpu", "ngl": 0, "ctx": 8192, "kv_quant": "q8_0"},
|
||
}
|
||
|
||
_GPU12_THRESHOLD_GB = 12.0
|
||
_GPU8_THRESHOLD_GB = 8.0
|
||
|
||
|
||
def _run(cmd: List[str], timeout: float = 10.0,
|
||
runner: Optional[Callable[[List[str], float], subprocess.CompletedProcess]] = None
|
||
) -> Optional[subprocess.CompletedProcess]:
|
||
"""执行命令并捕获输出；失败/超时返回 None（不抛异常）。"""
|
||
if runner is not None:
|
||
try:
|
||
return runner(cmd, timeout)
|
||
except Exception:
|
||
return None
|
||
try:
|
||
return subprocess.run(
|
||
cmd, capture_output=True, text=True, timeout=timeout,
|
||
creationflags=subprocess.CREATE_NO_WINDOW,
|
||
)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def nvidia_vram_gb(runner: Optional[Callable[..., subprocess.CompletedProcess]] = None) -> Optional[float]:
|
||
"""通过 nvidia-smi 读取显存总量（GB）；无 NVIDIA 返回 None。"""
|
||
exe = shutil.which("nvidia-smi")
|
||
if not exe:
|
||
return None
|
||
out = _run([exe, "--query-gpu=memory.total", "--format=csv,noheader,nounits"], runner=runner)
|
||
if out is None or out.returncode != 0 or not out.stdout.strip():
|
||
return None
|
||
try:
|
||
# 多卡取最大值（第一行也接受，但保守起见取最大以保证模板内存够用）
|
||
vals = [float(v.strip()) for v in out.stdout.strip().splitlines() if v.strip().isdigit()]
|
||
if not vals:
|
||
return None
|
||
return max(vals) / 1024.0
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def vulkan_present(runner: Optional[Callable[..., subprocess.CompletedProcess]] = None) -> bool:
|
||
"""检测是否存在 Vulkan 设备（无法可靠拿显存 -> 只用于判定非 cpu 的候选）。"""
|
||
exe = shutil.which("vulkaninfo")
|
||
if not exe:
|
||
return False
|
||
out = _run([exe, "--summary"], timeout=15.0, runner=runner)
|
||
if out is None or out.returncode != 0:
|
||
return False
|
||
low = out.stdout.lower()
|
||
# 出现 deviceName 且非 "llvmpipe"/"software" 视为有真实设备
|
||
return ("devicename" in low or "gpu" in low) and "llvmpipe" not in low and "lavapipe" not in low
|
||
|
||
|
||
def pick_tier(vram_gb: Optional[float]) -> str:
|
||
"""按显存选择档位；None/未知 -> cpu 保守档。"""
|
||
if vram_gb is None:
|
||
return "cpu"
|
||
if vram_gb >= _GPU12_THRESHOLD_GB:
|
||
return "gpu12"
|
||
if vram_gb >= _GPU8_THRESHOLD_GB:
|
||
return "gpu8"
|
||
return "cpu"
|
||
|
||
|
||
def detect(override: Optional[Dict[str, Any]] = None,
|
||
runner: Optional[Callable[..., subprocess.CompletedProcess]] = None) -> Dict[str, Any]:
|
||
"""检测并返回当前档位规格。
|
||
|
||
override（可选）：{"tier": "gpu12"} 强制指定档位；或覆盖单个字段如 {"ctx": 16384}。
|
||
|
||
返回形如 {"tier": "cpu", "ngl": 0, "ctx": 8192, "kv_quant": "q8_0",
|
||
"probe": "nvidia|vulkan|cpu|override", "note": str}
|
||
"""
|
||
if override and override.get("tier") in TIER_SPECS:
|
||
spec = dict(TIER_SPECS[override["tier"]])
|
||
spec.update({k: v for k, v in override.items() if k in spec})
|
||
spec["probe"] = "override"
|
||
spec["note"] = f"手动指定档位 {override['tier']}"
|
||
return spec
|
||
|
||
vram = nvidia_vram_gb(runner=runner)
|
||
probe = "nvidia"
|
||
if vram is None:
|
||
if vulkan_present(runner=runner):
|
||
probe = "vulkan"
|
||
note = "检测到 Vulkan 设备但无法读取显存，按保守档 cpu 运行（可在 config 手动覆盖 tier）"
|
||
else:
|
||
probe = "cpu"
|
||
note = "未检测到 GPU，按 cpu 档运行（-ngl 0，速度受限）"
|
||
else:
|
||
note = f"nvidia-smi 探测显存 {vram:.1f}GB"
|
||
|
||
tier = pick_tier(vram)
|
||
spec = dict(TIER_SPECS[tier])
|
||
spec["probe"] = probe
|
||
spec["note"] = note if probe != "nvidia" else f"{note} -> 档位 {tier}"
|
||
return spec
|
||
|
||
|
||
def tier_spec(tier: str) -> Dict[str, Any]:
|
||
"""返回指定档位的规格副本（供 config.tiers 兜底）。"""
|
||
if tier not in TIER_SPECS:
|
||
raise ValueError(f"未知硬件档位: {tier}（支持: {sorted(TIER_SPECS)}）")
|
||
return dict(TIER_SPECS[tier])
|
||
|
||
|
||
def detect_summary() -> str:
|
||
"""人类可读的检测摘要（setup_runtime / serve 启动时打印）。"""
|
||
spec = detect()
|
||
return (
|
||
f"硬件档位: {spec['tier']} (ngl={spec['ngl']}, ctx={spec['ctx']}, "
|
||
f"kv_quant={spec['kv_quant']}) [{spec.get('note', '')}]"
|
||
)
|