vdi-cloud/tools/gpu_dashboard.py

#!/usr/bin/env python3
"""服务器 GPU 状态看板（部署在 PVE 宿主机上，仅依赖 python3 标准库）。

数据源：
- lspci / /sys / lsmod：GPU 型号、当前驱动（nouveau/vfio-pci/nvidia）、IOMMU 状态
- qm config/status：VM 100 的直通配置（hostpci0）与运行状态
- 直通生效后：qm guest exec 在 VM 100 内执行 nvidia-smi 采集实时数据
  （利用率 / 显存 / 温度 / 驱动版本 / NVENC 编码会话）

用法：python3 gpu_dashboard.py [端口]   # 默认 8090，浏览器访问 http://<宿主机IP>:8090/
"""
import json
import subprocess
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

VMID = 100
CACHE_TTL = 2.0
GUEST_TIMEOUT = 25

_lock = threading.Lock()
_cache = {"t": 0.0, "data": None}


def run(cmd, timeout=15):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
        return (r.stdout or "").strip()
    except Exception as e:
        return "ERROR: %s" % e


def read(path):
    try:
        with open(path) as f:
            return f.read().strip()
    except Exception:
        return ""


def gpu_identity():
    out = run("lspci -nnk -s 01:00.0")
    name = ""
    driver = ""
    for line in out.splitlines():
        if "NVIDIA" in line and not name:
            idx = line.find("NVIDIA")
            name = line[idx:].strip()
        if "Kernel driver in use" in line:
            driver = line.split(":", 1)[-1].strip()
    audio = ""
    for line in run("lspci -nnk -s 01:00.1").splitlines():
        if "Kernel driver in use" in line:
            audio = line.split(":", 1)[-1].strip()
    return {
        "name": name or "NVIDIA GPU (0000:01:00.0)",
        "driver": driver or "（无，未绑定）",
        "audio_driver": audio or "（无）",
    }


def iommu_state():
    # 新内核（6.8+）默认启用 Intel DMAR，无需显式 intel_iommu=on；
    # 以 iommu_groups 是否存在且 GPU 已入组为准，内核参数仅作展示。
    cmdline = read("/proc/cmdline")
    flag = "intel_iommu=on" in cmdline or "iommu=on" in cmdline
    groups = run("find /sys/kernel/iommu_groups -maxdepth 1 -mindepth 1 -type d 2>/dev/null | wc -l")
    n = int(groups) if groups.isdigit() else 0
    gpu_group = run("for d in /sys/kernel/iommu_groups/*/devices/0000:01:00.0; do "
                    "[ -e \"$d\" ] && dirname \"$d\" | sed 's|.*iommu_groups/||;s|/devices||'; done")
    members = run("ls /sys/kernel/iommu_groups/%s/devices/ 2>/dev/null" % gpu_group) if gpu_group else ""
    if n == 0:
        state, ok = "未启用（iommu 组数为 0，需 BIOS 开 VT-d + 内核参数 intel_iommu=on）", False
    elif not gpu_group:
        state, ok = "已启用，但 GPU 不在任何组内（异常）", False
    else:
        state, ok = "已启用（DMAR 生效）", True
    return {"state": state, "ok": ok, "kernel_flag": flag, "groups": n,
            "gpu_group": gpu_group, "group_members": members}


def passthrough_state():
    cfg = run("qm config %d 2>/dev/null" % VMID)
    hostpci = ""
    for line in cfg.splitlines():
        if line.startswith("hostpci"):
            hostpci = line
            break
    status = run("qm status %d 2>/dev/null" % VMID).replace("status: ", "").replace("status:", "").strip()
    return {"hostpci": hostpci or "（未配置）", "vm_status": status or "unknown"}


def guest_gpu():
    """直通后在 VM 内经 guest agent 执行 nvidia-smi；未直通/无驱动时返回 None。"""
    q = ("nvidia-smi --query-gpu=name,utilization.gpu,memory.used,memory.total,"
         "temperature.gpu,driver_version,power.draw --format=csv,noheader")
    out = run("qm guest exec %d -- powershell -NoProfile -Command \"%s\"" % (VMID, q),
              timeout=GUEST_TIMEOUT)
    try:
        d = json.loads(out)
    except Exception:
        return None
    if d.get("returncode") != 0:
        return None
    text = (d.get("out-data") or "").strip()
    if not text or "nvidia-smi" in (d.get("err-data") or ""):
        return None
    parts = [p.strip() for p in text.split(",")]
    if len(parts) < 6:
        return None
    info = {
        "name": parts[0],
        "utilization": parts[1],
        "vram_used": parts[2],
        "vram_total": parts[3],
        "temperature": parts[4],
        "driver": parts[5],
        "power": parts[6] if len(parts) > 6 else "",
    }
    enc = run("qm guest exec %d -- powershell -NoProfile -Command \""
              "nvidia-smi --query-gpu=encoder.stats.sessionCount --format=csv,noheader\"" % VMID,
              timeout=GUEST_TIMEOUT)
    try:
        d2 = json.loads(enc)
        v = (d2.get("out-data") or "").strip()
        if d2.get("returncode") == 0 and v:
            info["enc_sessions"] = v
    except Exception:
        pass
    return info


def collect():
    ident = gpu_identity()
    driver = ident["driver"]
    if "vfio-pci" in driver:
        stage = "vfio 已接管（可直通/已解绑 nouveau）"
    elif "nouveau" in driver:
        stage = "未直通（GPU 由宿主 nouveau 占用，直通前需解绑）"
    elif "nvidia" in driver:
        stage = "未直通（GPU 由宿主 nvidia 驱动占用）"
    else:
        stage = "GPU 未绑定驱动"
    data = {
        "time": time.strftime("%Y-%m-%d %H:%M:%S"),
        "gpu": ident,
        "stage": stage,
        "iommu": iommu_state(),
        "pt": passthrough_state(),
        "guest": guest_gpu() if "hostpci" in passthrough_state()["hostpci"] else None,
    }
    return data


def get_data():
    with _lock:
        now = time.time()
        if _cache["data"] is not None and now - _cache["t"] < CACHE_TTL:
            return _cache["data"]
        data = collect()
        _cache["data"] = data
        _cache["t"] = now
        return data


HTML = """<!DOCTYPE html>
<html lang="zh-CN"><head><meta charset="utf-8">
<title>服务器 GPU 状态</title>
<style>
 body { background:#0f1420; color:#e6e9ef; font-family:"Microsoft YaHei",sans-serif;
        margin:0; padding:32px; }
 h1 { font-size:22px; margin:0 0 4px; }
 .sub { color:#8a94a6; font-size:12px; margin-bottom:24px; }
 .grid { display:grid; grid-template-columns:repeat(auto-fit,minmax(280px,1fr)); gap:14px; max-width:1100px; }
 .card { background:#1a2233; border:1px solid #2b3654; border-radius:10px; padding:16px 18px; }
 .card h3 { margin:0 0 10px; font-size:13px; color:#8a94a6; font-weight:normal; }
 .card .v { font-size:15px; line-height:1.7; word-break:break-all; }
 .ok { color:#3fb950; } .bad { color:#e05561; } .warn { color:#d29922; }
 .big { font-size:30px; font-weight:bold; }
</style></head><body>
<h1>服务器 GPU 状态</h1>
<div class="sub" id="time">加载中…（每 4 秒自动刷新）</div>
<div class="grid">
 <div class="card"><h3>显卡</h3><div class="v" id="gpu">-</div></div>
 <div class="card"><h3>直通阶段</h3><div class="v" id="stage">-</div></div>
 <div class="card"><h3>IOMMU / VT-d</h3><div class="v" id="iommu">-</div></div>
 <div class="card"><h3>VM %d 直通配置</h3><div class="v" id="pt">-</div></div>
 <div class="card"><h3>VM 内实时数据（nvidia-smi）</h3><div class="v" id="guest">直通生效后显示</div></div>
</div>
<script>
const $ = id => document.getElementById(id);
function esc(s){ return String(s).replace(/</g,"&lt;"); }
function render(d){
  $("time").textContent = "更新于 " + d.time + "（每 4 秒自动刷新）";
  $("gpu").innerHTML = "<b>" + esc(d.gpu.name) + "</b><br>当前驱动：" + esc(d.gpu.driver) +
                       "<br>HDMI 音频驱动：" + esc(d.gpu.audio_driver);
  const stColor = d.stage.includes("未直通") ? "warn" : "ok";
  $("stage").innerHTML = "<span class='" + stColor + "'>" + esc(d.stage) + "</span>";
  $("iommu").innerHTML = (d.iommu.ok ? "<span class='ok'>" : "<span class='bad'>") +
      esc(d.iommu.state) + "</span><br>内核参数 intel_iommu=on：" +
      (d.iommu.kernel_flag ? "已设置" : "未设置（新内核默认开启，非必须）") +
      (d.iommu.gpu_group ? "<br>GPU 所在组：" + esc(d.iommu.gpu_group) +
        (d.iommu.group_members ? "<br>组内设备：" + esc(d.iommu.group_members) : "") : "");
  $("pt").innerHTML = "运行状态：" + esc(d.pt.vm_status) + "<br>hostpci：" + esc(d.pt.hostpci);
  if (d.guest) {
    $("guest").innerHTML =
      "<span class='big'>" + esc(d.guest.utilization) + "</span> 利用率<br>" +
      "显存：" + esc(d.guest.vram_used) + " / " + esc(d.guest.vram_total) + "<br>" +
      "温度：" + esc(d.guest.temperature) +
      (d.guest.power ? "　功耗：" + esc(d.guest.power) : "") + "<br>" +
      "驱动：" + esc(d.guest.driver) +
      (d.guest.enc_sessions !== undefined ? "<br>NVENC 编码会话：" + esc(d.guest.enc_sessions) : "");
  } else {
    $("guest").innerHTML = "<span class='warn'>直通生效并在 VM 内装好驱动后，这里显示实时利用率/显存/温度/NVENC 会话</span>";
  }
}
function tick(){ fetch("/api").then(r=>r.json()).then(render).catch(()=>{}); }
tick(); setInterval(tick, 4000);
</script></body></html>""" % VMID


class Handler(BaseHTTPRequestHandler):
    def _send(self, code, body, ctype):
        raw = body.encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", ctype + "; charset=utf-8")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def do_GET(self):
        if self.path.startswith("/api"):
            self._send(200, json.dumps(get_data(), ensure_ascii=False), "application/json")
        elif self.path.startswith("/healthz"):
            self._send(200, "ok", "text/plain")
        else:
            self._send(200, HTML, "text/html")

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8090
    srv = ThreadingHTTPServer(("0.0.0.0", port), Handler)
    print("GPU dashboard on http://0.0.0.0:%d/ (VM %d)" % (port, VMID), flush=True)
    srv.serve_forever()