kk8bit
31f20a5041
psutil's virtual_memory().used counts only active+wired on macOS, while .percent is defined as (total - available) / total. Showing one next to the other produced readings like 10.1/24GB (68%). Derive the used figure from total - available so both describe the same thing; it also lines up with Activity Monitor's Memory Used.
127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
import logging
|
||
import platform
|
||
import time
|
||
import psutil
|
||
from server import PromptServer
|
||
from aiohttp import web
|
||
|
||
PLATFORM = platform.system()
|
||
IS_MACOS = PLATFORM == "Darwin"
|
||
IS_LINUX = PLATFORM == "Linux"
|
||
IS_WINDOWS = PLATFORM == "Windows"
|
||
|
||
pynvml_instance = None
|
||
pynvml_available = False
|
||
CPU_NAME = None
|
||
|
||
class KayResourceCollector:
|
||
def __init__(self):
|
||
self.cpu_count = psutil.cpu_count()
|
||
self._initialize_hardware()
|
||
|
||
def _initialize_hardware(self):
|
||
global pynvml_instance, pynvml_available, CPU_NAME
|
||
if CPU_NAME is None:
|
||
try:
|
||
import cpuinfo
|
||
CPU_NAME = cpuinfo.get_cpu_info().get('brand_raw', "Unknown")
|
||
except Exception:
|
||
CPU_NAME = "Unknown"
|
||
|
||
if pynvml_instance is None and (IS_LINUX or IS_WINDOWS):
|
||
try:
|
||
import pynvml
|
||
pynvml.nvmlInit()
|
||
pynvml_instance = pynvml
|
||
pynvml_available = True
|
||
except Exception:
|
||
# 不能写成 except (ImportError, pynvml.NVMLError)：import 失败时 pynvml
|
||
# 这个名字根本没绑定，求值 except 子句本身就会抛 NameError，而它不在
|
||
# 捕获范围内 —— 于是「没装 pynvml」这个本该被兜住的场景反而把监控打死。
|
||
pynvml_instance = None
|
||
pynvml_available = False
|
||
|
||
def get_status(self):
|
||
try:
|
||
cpu_percent = psutil.cpu_percent(interval=None)
|
||
ram = psutil.virtual_memory()
|
||
ram_total = round(ram.total / (1024 ** 3), 1)
|
||
# 不能用 ram.used：在 macOS 上它只算 active+wired，而 ram.percent 的定义是
|
||
# (total - available) / total，两者口径不同，显示出来就是「10.1/24GB (68%)」这种自相矛盾。
|
||
# 用 total - available 才和百分比自洽，也接近活动监视器的「已使用内存」。
|
||
ram_used = round((ram.total - ram.available) / (1024 ** 3), 1)
|
||
ram_percent = ram.percent
|
||
gpu_info = self.get_gpu_info() if not IS_MACOS and pynvml_available else []
|
||
return {
|
||
"cpu_percent": cpu_percent,
|
||
"cpu_count": self.cpu_count,
|
||
"cpu_name": CPU_NAME or "Unknown",
|
||
"ram_total": ram_total,
|
||
"ram_used": ram_used,
|
||
"ram_percent": ram_percent,
|
||
"gpu": gpu_info
|
||
}
|
||
except Exception:
|
||
return {
|
||
"cpu_percent": 0,
|
||
"cpu_count": self.cpu_count,
|
||
"cpu_name": CPU_NAME or "Unknown",
|
||
"ram_total": 0,
|
||
"ram_used": 0,
|
||
"ram_percent": 0,
|
||
"gpu": "Error"
|
||
}
|
||
|
||
def get_gpu_info(self):
|
||
try:
|
||
device_count = pynvml_instance.nvmlDeviceGetCount()
|
||
gpu_info = []
|
||
for i in range(device_count):
|
||
handle = pynvml_instance.nvmlDeviceGetHandleByIndex(i)
|
||
name = pynvml_instance.nvmlDeviceGetName(handle)
|
||
if isinstance(name, bytes):
|
||
name = name.decode('utf-8')
|
||
util = pynvml_instance.nvmlDeviceGetUtilizationRates(handle)
|
||
mem_info = pynvml_instance.nvmlDeviceGetMemoryInfo(handle)
|
||
gpu_info.append({
|
||
"index": i,
|
||
"name": name,
|
||
"load": float(util.gpu),
|
||
"memory_used": float(mem_info.used) / (1024 ** 3),
|
||
"memory_total": float(mem_info.total) / (1024 ** 3),
|
||
"memory_percent": (float(mem_info.used) / float(mem_info.total)) * 100,
|
||
"temperature": float(pynvml_instance.nvmlDeviceGetTemperature(handle, pynvml_instance.NVML_TEMPERATURE_GPU))
|
||
})
|
||
return gpu_info
|
||
except Exception:
|
||
return "GPU information unavailable"
|
||
|
||
# 前端按需轮询，这里只负责按需采集。多个标签页同时拉时靠一个很短的缓存去重，
|
||
# 也避免 psutil.cpu_percent(interval=None) 的「距上次调用」基准被多个客户端互相打乱。
|
||
_collector = None
|
||
_cache = {"at": 0.0, "data": None}
|
||
CACHE_TTL_SECONDS = 0.2
|
||
|
||
|
||
def _get_status_cached():
|
||
global _collector
|
||
now = time.monotonic()
|
||
if _cache["data"] is None or now - _cache["at"] > CACHE_TTL_SECONDS:
|
||
if _collector is None:
|
||
_collector = KayResourceCollector()
|
||
_cache["data"] = _collector.get_status()
|
||
_cache["at"] = now
|
||
return _cache["data"]
|
||
|
||
|
||
routes = PromptServer.instance.routes
|
||
|
||
|
||
@routes.get("/kaytool/resources")
|
||
async def resources_endpoint(request):
|
||
try:
|
||
return web.json_response(_get_status_cached())
|
||
except Exception:
|
||
logging.exception("[KayTool] Failed to collect resource status")
|
||
return web.json_response({"error": "unavailable"}, status=500)
|