100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
|
# GPU监控模块 · 采集nvidia-smi数据
|
||
|
|
# HLDP://zhuyuan-agent/gpu-monitor
|
||
|
|
|
||
|
|
import subprocess
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
from typing import List, Dict, Optional
|
||
|
|
|
||
|
|
|
||
|
|
def collect_gpu_metrics() -> Dict:
|
||
|
|
"""采集所有GPU的实时指标
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
{
|
||
|
|
"gpus": [
|
||
|
|
{
|
||
|
|
"index": 0,
|
||
|
|
"name": "NVIDIA GeForce RTX 3090",
|
||
|
|
"uuid": "GPU-xxxx",
|
||
|
|
"utilization_gpu": 85, # %
|
||
|
|
"memory_used_mb": 18432, # MB
|
||
|
|
"memory_total_mb": 24576, # MB
|
||
|
|
"temperature_gpu": 72, # °C
|
||
|
|
"power_draw_w": 285.5, # W
|
||
|
|
"fan_speed": 65 # %
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"error": null
|
||
|
|
}
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
# 查询GPU关键指标
|
||
|
|
result = subprocess.run(
|
||
|
|
[
|
||
|
|
"nvidia-smi",
|
||
|
|
"--query-gpu=index,name,uuid,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw,fan.speed",
|
||
|
|
"--format=csv,noheader,nounits"
|
||
|
|
],
|
||
|
|
capture_output=True, text=True, timeout=10
|
||
|
|
)
|
||
|
|
|
||
|
|
if result.returncode != 0:
|
||
|
|
return {"gpus": [], "error": f"nvidia-smi failed: {result.stderr.strip()}"}
|
||
|
|
|
||
|
|
gpus = []
|
||
|
|
for line in result.stdout.strip().split("\n"):
|
||
|
|
if not line.strip():
|
||
|
|
continue
|
||
|
|
parts = [p.strip() for p in line.split(",")]
|
||
|
|
if len(parts) < 9:
|
||
|
|
continue
|
||
|
|
|
||
|
|
try:
|
||
|
|
gpu = {
|
||
|
|
"index": int(parts[0]),
|
||
|
|
"name": parts[1],
|
||
|
|
"uuid": parts[2],
|
||
|
|
"utilization_gpu": int(parts[3]) if parts[3] != "[Not Supported]" else 0,
|
||
|
|
"memory_used_mb": int(parts[4]) if parts[4] != "[Not Supported]" else 0,
|
||
|
|
"memory_total_mb": int(parts[5]) if parts[5] != "[Not Supported]" else 0,
|
||
|
|
"temperature_gpu": int(parts[6]) if parts[6] != "[Not Supported]" else 0,
|
||
|
|
"power_draw_w": float(parts[7]) if parts[7] not in ("[Not Supported]", "") else 0.0,
|
||
|
|
"fan_speed": int(parts[8]) if parts[8] not in ("[Not Supported]", "") else 0,
|
||
|
|
}
|
||
|
|
gpus.append(gpu)
|
||
|
|
except (ValueError, IndexError):
|
||
|
|
continue
|
||
|
|
|
||
|
|
return {"gpus": gpus, "error": None}
|
||
|
|
|
||
|
|
except FileNotFoundError:
|
||
|
|
return {"gpus": [], "error": "nvidia-smi not found - not a GPU machine?"}
|
||
|
|
except subprocess.TimeoutExpired:
|
||
|
|
return {"gpus": [], "error": "nvidia-smi timed out"}
|
||
|
|
except Exception as e:
|
||
|
|
return {"gpus": [], "error": str(e)}
|
||
|
|
|
||
|
|
|
||
|
|
def gpu_summary(gpus: List[Dict]) -> str:
|
||
|
|
"""生成GPU状态的一行摘要"""
|
||
|
|
if not gpus:
|
||
|
|
return "无GPU"
|
||
|
|
parts = []
|
||
|
|
for g in gpus:
|
||
|
|
util = g.get("utilization_gpu", 0)
|
||
|
|
temp = g.get("temperature_gpu", 0)
|
||
|
|
mem = g.get("memory_used_mb", 0)
|
||
|
|
mem_total = g.get("memory_total_mb", 0)
|
||
|
|
mem_pct = int(mem / mem_total * 100) if mem_total > 0 else 0
|
||
|
|
parts.append(f"GPU{g['index']}: {util}%/{temp}°C/{mem_pct}%VRAM")
|
||
|
|
return " | ".join(parts)
|
||
|
|
|
||
|
|
|
||
|
|
# 快速测试
|
||
|
|
if __name__ == "__main__":
|
||
|
|
data = collect_gpu_metrics()
|
||
|
|
print(json.dumps(data, indent=2, ensure_ascii=False))
|
||
|
|
if data["gpus"]:
|
||
|
|
print(f"\n摘要: {gpu_summary(data['gpus'])}")
|