#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业17 · GPU / 驱动 / Operator 盘点表（本环境 + 实验 K8s）

本环境现网为 Docker on 116（H20），非 K8s；脚本输出对照表 + 实验集群 kubectl 命令清单。

用法：
  python hw17_02_GPU与Operator盘点.py
  python hw17_02_GPU与Operator盘点.py --ssh-host 116.62.10.183

输出：data/homework17/gpu_inventory_<timestamp>.md
"""
from __future__ import annotations

import argparse
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path


KNOWN_ENV = """
## 本仓库已锚定环境（文档实录，非实时探测）

| 角色 | 主机 | GPU | 部署方式 | Operator |
|------|------|-----|----------|----------|
| 应用/网关 | 192.168.110.207 | 无 GPU | Docker Compose（Dify/LiteLLM/Xinference CPU） | 不适用 |
| 推理 | 116.62.10.183 (prod102) | **NVIDIA H20 ~96GB** | Docker `vllm-qwen3.6-27b` | **无 K8s Operator**（裸 Docker `--gpus`） |
| 实验 K8s | （自备 lab 集群） | 按节点 | Deployment + `nvidia.com/gpu` | **需装 GPU Operator** |

116 常用命令：
```bash
nvidia-smi
docker ps --filter name=vllm
curl -s http://116.62.10.183:8000/metrics | grep vllm_num_requests_waiting
```
"""


K8S_CHECKS = """
## 实验 K8s 集群 · 必跑命令

```bash
# 1) 节点 GPU 标签与可分配量
kubectl get nodes -o custom-columns=NAME:.metadata.name,GPU:.status.allocatable.nvidia\\.com/gpu

# 2) GPU Operator（命名空间以集群为准，常见 gpu-operator / nvidia-gpu-operator）
kubectl get pods -n gpu-operator
kubectl get clusterpolicy -n gpu-operator 2>/dev/null || true

# 3) 驱动 / CUDA 版本（节点上或 driver daemonset 日志）
kubectl get nodes -o json | jq '.items[] | {name:.metadata.name, labels:.metadata.labels}'

# 4) 目标 Pod NVIDIA 相关事件
kubectl describe pod -n ai-inference -l app=vllm-qwen
kubectl get events -n ai-inference --sort-by='.lastTimestamp' | grep -iE 'nvidia|gpu|cuda|device-plugin|Insufficient'
```

### 常见 Event 含义

| Event / Reason | 含义 | 动作 |
|----------------|------|------|
| `Insufficient nvidia.com/gpu` | 节点无可用 GPU | 缩其它 GPU 负载或换节点 |
| `FailedScheduling` + GPU | 调度器找不到 GPU | 查 allocatable、taint/toleration |
| `UnexpectedAdmissionError` | device plugin 异常 | 重启 plugin pod；查驱动 |
| `CUDA driver version is insufficient` | 驱动/CUDA 与镜像不匹配 | 升驱动或换镜像 tag |
| `nvidia-container-cli: initialization error` | runtime 未配置 nvidia | 查 containerd nvidia runtime |
"""


def ssh_run(host: str, cmd: str, timeout: int = 20) -> tuple[int, str]:
    try:
        r = subprocess.run(
            ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=8", host, cmd],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        return r.returncode, (r.stdout or "") + (r.stderr or "")
    except (FileNotFoundError, subprocess.TimeoutExpired) as e:
        return 127, str(e)


def probe_nvidia_smi(host: str | None) -> str:
    if not host:
        return "（未指定 --ssh-host，跳过实时 nvidia-smi）"
    code, out = ssh_run(host, "nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader")
    if code != 0:
        return f"SSH/nvidia-smi 失败 (code={code}): {out.strip()[:500]}"
    return out.strip()


def main() -> None:
    parser = argparse.ArgumentParser(description="GPU/Operator 盘点")
    parser.add_argument("--ssh-host", default=os.environ.get("GPU_SSH_HOST"))
    args = parser.parse_args()

    od = Path(os.environ.get("HW17_OUT", Path(__file__).resolve().parents[2] / "data" / "homework17"))
    od.mkdir(parents=True, exist_ok=True)

    smi = probe_nvidia_smi(args.ssh_host)
    lines = [
        "# GPU / 驱动 / Operator 盘点（作业17）",
        "",
        f"- 生成时间：`{datetime.now().isoformat(timespec='seconds')}`",
        f"- SSH 探测主机：`{args.ssh_host or '未指定'}`",
        "",
        KNOWN_ENV.strip(),
        "",
        "## 实时 nvidia-smi（若 SSH 成功）",
        "",
        "```",
        smi,
        "```",
        "",
        K8S_CHECKS.strip(),
        "",
        "## Operator 是否安装 · 判定标准",
        "",
        "| 检查项 | 已安装 | 未安装 |",
        "|--------|--------|--------|",
        "| `kubectl get pods -n gpu-operator` | 有 Running 的 operator/driver/plugin | Namespace 不存在或 CrashLoop |",
        "| Node allocatable | 显示 `nvidia.com/gpu: N` | 无该字段 |",
        "| Pod 请求 GPU | 可调度到 GPU 节点 | FailedScheduling Insufficient gpu |",
        "",
        "**本环境 116**：Docker 直跑 vLLM，**不需要** Operator；上 K8s 才需要。",
        "",
    ]
    out = od / f"gpu_inventory_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
    out.write_text("\n".join(lines), encoding="utf-8")
    print(f"Wrote {out}")
    if args.ssh_host and "失败" in smi:
        sys.exit(1)


if __name__ == "__main__":
    main()