Edouard Durand
4fdbc035f4
Floating cards share one base (float.js) with magnet docking; docked cards stack in the corner. Durations come from ComfyUI's history timestamps, the start of the current run from the execution_start event.
130 lines
4.3 KiB
Python
130 lines
4.3 KiB
Python
"""Server resource usage for the panel: CPU, RAM, GPUs and the prompt queue.
|
|
|
|
psutil and pynvml ship with most ComfyUI installs (ComfyUI-Crystools and many
|
|
others depend on them); every metric is optional and skipped if unavailable.
|
|
"""
|
|
|
|
import logging
|
|
import warnings
|
|
|
|
log = logging.getLogger("ComfyUI-RigShare")
|
|
|
|
try:
|
|
import psutil
|
|
except ImportError:
|
|
psutil = None
|
|
|
|
try:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
import pynvml
|
|
pynvml.nvmlInit()
|
|
except Exception:
|
|
pynvml = None
|
|
|
|
GB = 1024 ** 3
|
|
|
|
|
|
def _gpus():
|
|
if not pynvml:
|
|
return []
|
|
out = []
|
|
for i in range(pynvml.nvmlDeviceGetCount()):
|
|
h = pynvml.nvmlDeviceGetHandleByIndex(i)
|
|
name = pynvml.nvmlDeviceGetName(h)
|
|
if isinstance(name, bytes):
|
|
name = name.decode()
|
|
mem = pynvml.nvmlDeviceGetMemoryInfo(h)
|
|
gpu = {"index": i, "name": name.replace("NVIDIA ", "").replace("GeForce ", ""),
|
|
"util": pynvml.nvmlDeviceGetUtilizationRates(h).gpu,
|
|
"vram_used": round(mem.used / GB, 2), "vram_total": round(mem.total / GB, 2)}
|
|
try:
|
|
gpu["temp"] = pynvml.nvmlDeviceGetTemperature(h, pynvml.NVML_TEMPERATURE_GPU)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
gpu["power"] = round(pynvml.nvmlDeviceGetPowerUsage(h) / 1000)
|
|
gpu["power_limit"] = round(pynvml.nvmlDeviceGetEnforcedPowerLimit(h) / 1000)
|
|
except Exception:
|
|
pass
|
|
out.append(gpu)
|
|
return out
|
|
|
|
|
|
def queue_items(server):
|
|
"""Running and pending prompts: [{id, number, status, nodes}], running first."""
|
|
try:
|
|
running, pending = server.prompt_queue.get_current_queue_volatile()
|
|
except Exception:
|
|
return None
|
|
items = []
|
|
for status, entries in (("running", running), ("pending", sorted(pending, key=lambda e: e[0]))):
|
|
for entry in entries:
|
|
try:
|
|
prompt = entry[2] if isinstance(entry[2], dict) else {}
|
|
items.append({"id": entry[1], "number": entry[0], "status": status, "nodes": len(prompt)})
|
|
except (IndexError, TypeError):
|
|
continue
|
|
return items
|
|
|
|
|
|
RESULTS = {"execution_success": "done", "execution_error": "error", "execution_interrupted": "stopped"}
|
|
|
|
|
|
def run_history(server, n=40):
|
|
"""Finished runs from ComfyUI's history, newest first: [{id, ms, end, result}].
|
|
|
|
ComfyUI stamps execution_start and the final event with millisecond times, so
|
|
durations are exact. Only each entry's status is read (no outputs copied).
|
|
"""
|
|
try:
|
|
history = server.prompt_queue.get_history(max_items=n, map_function=lambda p: p.get("status"))
|
|
except Exception:
|
|
return None
|
|
runs = []
|
|
for prompt_id, status in history.items():
|
|
start = end = result = None
|
|
for message in (status or {}).get("messages") or []:
|
|
try:
|
|
event, data = message[0], message[1] or {}
|
|
except (IndexError, TypeError):
|
|
continue
|
|
if event == "execution_start":
|
|
start = data.get("timestamp")
|
|
elif event in RESULTS:
|
|
end, result = data.get("timestamp"), RESULTS[event]
|
|
if start and end and end >= start:
|
|
runs.append({"id": prompt_id, "ms": end - start, "end": end, "result": result})
|
|
runs.sort(key=lambda r: r["end"], reverse=True)
|
|
return runs
|
|
|
|
|
|
def _queue(server):
|
|
try:
|
|
running, pending = server.prompt_queue.get_current_queue_volatile()
|
|
return {"running": len(running), "pending": len(pending)}
|
|
except Exception:
|
|
try:
|
|
return {"running": 0, "pending": server.prompt_queue.get_tasks_remaining()}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def collect_stats(server=None):
|
|
stats = {}
|
|
if psutil:
|
|
stats["cpu"] = psutil.cpu_percent(interval=None)
|
|
vm = psutil.virtual_memory()
|
|
stats["ram_used"] = round((vm.total - vm.available) / GB, 1)
|
|
stats["ram_total"] = round(vm.total / GB, 1)
|
|
try:
|
|
stats["gpus"] = _gpus()
|
|
except Exception as e:
|
|
log.debug(f"[RigShare] GPU stats failed: {e}")
|
|
stats["gpus"] = []
|
|
if server is not None:
|
|
stats["queue"] = _queue(server)
|
|
stats["queue_items"] = queue_items(server)
|
|
stats["run_history"] = run_history(server)
|
|
return stats
|