226 lines
8.2 KiB
Python
226 lines
8.2 KiB
Python
"""
|
|
ComfyUI Process Monitor
|
|
=======================
|
|
Monitors ComfyUI process during image generation.
|
|
Auto-detects generation start/end and records peak memory/CPU.
|
|
|
|
Usage:
|
|
from monitor import get_process_monitor
|
|
monitor = get_process_monitor()
|
|
monitor.start_monitoring()
|
|
# ... generate images ...
|
|
logger.info("%s", monitor.get_summary())
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import Optional, List
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class GenerationRecord:
|
|
"""Stats for a single generation run"""
|
|
timestamp: float = 0.0
|
|
duration: float = 0.0
|
|
vram_peak_mb: float = 0.0
|
|
vram_start_mb: float = 0.0
|
|
vram_end_mb: float = 0.0
|
|
vram_delta_mb: float = 0.0
|
|
ram_peak_mb: float = 0.0
|
|
ram_start_mb: float = 0.0
|
|
ram_end_mb: float = 0.0
|
|
cpu_peak: float = 0.0
|
|
error: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class ProcessSnapshot:
|
|
"""Point-in-time process metrics"""
|
|
timestamp: float
|
|
memory_rss_mb: float
|
|
memory_vms_mb: float
|
|
cpu_percent: float
|
|
num_threads: int
|
|
|
|
|
|
class ComfyProcessMonitor:
|
|
"""Monitors ComfyUI process for generation metrics"""
|
|
|
|
def __init__(self, poll_interval: float = 0.5):
|
|
self.poll_interval = poll_interval
|
|
self._process = None
|
|
self._process_pid = None
|
|
self._running = False
|
|
self._thread = None
|
|
self._is_generating = False
|
|
self._gen_start_time = 0.0
|
|
self._current_record: Optional[GenerationRecord] = None
|
|
self._prev_memory_mb = 0.0
|
|
self.history: List[GenerationRecord] = []
|
|
self.current: Optional[GenerationRecord] = None
|
|
self._snapshots: List[ProcessSnapshot] = []
|
|
self._lock = threading.Lock()
|
|
|
|
# Detection thresholds (tuned)
|
|
self._MEMORY_START_THRESHOLD_MB = 50
|
|
self._CPU_START_THRESHOLD = 15
|
|
self._MEMORY_END_THRESHOLD_MB = -30
|
|
self._CPU_END_THRESHOLD = 8
|
|
|
|
def find_comfy_process(self) -> bool:
|
|
"""Auto-detect ComfyUI Python process"""
|
|
try:
|
|
import psutil
|
|
pid = None
|
|
current = psutil.Process()
|
|
cmdline = ' '.join(current.cmdline()).lower()
|
|
|
|
# Check if WE are the ComfyUI process
|
|
if 'comfy' in cmdline or 'main.py' in cmdline:
|
|
pid = current.pid
|
|
|
|
# Search by name
|
|
if not pid:
|
|
for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
|
|
try:
|
|
name = (proc.info['name'] or '').lower()
|
|
cl = ' '.join(proc.info['cmdline'] or []).lower()
|
|
if ('python' in name or 'python' in cl) and ('comfy' in cl or 'main.py' in cl):
|
|
pid = proc.info['pid']
|
|
break
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
|
continue
|
|
|
|
if pid:
|
|
self._process_pid = pid
|
|
self._process = psutil.Process(pid)
|
|
logger.info(f"Process Monitor: found PID {pid}")
|
|
return True
|
|
|
|
logger.warning("Process Monitor: ComfyUI not found")
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"Process Monitor: find error: {e}")
|
|
return False
|
|
|
|
def start_monitoring(self):
|
|
"""Begin background monitoring thread"""
|
|
if self._running:
|
|
return
|
|
if not self.find_comfy_process():
|
|
return
|
|
self._running = True
|
|
self._thread = threading.Thread(target=self._poll_loop, daemon=True)
|
|
self._thread.start()
|
|
logger.info("Process Monitor: started")
|
|
|
|
def stop_monitoring(self):
|
|
"""Stop background monitoring"""
|
|
self._running = False
|
|
if self._thread:
|
|
self._thread.join(timeout=2)
|
|
|
|
def _poll_loop(self):
|
|
try:
|
|
import psutil
|
|
except ImportError:
|
|
self._running = False
|
|
return
|
|
|
|
while self._running:
|
|
try:
|
|
if not self._process or not self._process.is_running():
|
|
if not self.find_comfy_process():
|
|
time.sleep(self.poll_interval * 10)
|
|
continue
|
|
|
|
mem = self._process.memory_info()
|
|
cpu = self._process.cpu_percent(interval=0)
|
|
rss_mb = mem.rss / (1024 * 1024)
|
|
delta = rss_mb - self._prev_memory_mb
|
|
snap = ProcessSnapshot(time.time(), rss_mb, mem.vms / (1024*1024), cpu, self._process.num_threads())
|
|
|
|
with self._lock:
|
|
if not self._is_generating:
|
|
if delta > self._MEMORY_START_THRESHOLD_MB or cpu > self._CPU_START_THRESHOLD:
|
|
self._is_generating = True
|
|
self._gen_start_time = time.time()
|
|
self._current_record = GenerationRecord(
|
|
timestamp=self._gen_start_time, duration=0,
|
|
vram_start_mb=rss_mb, vram_peak_mb=rss_mb,
|
|
ram_start_mb=rss_mb, ram_peak_mb=rss_mb,
|
|
)
|
|
self._snapshots = [snap]
|
|
logger.info(f"Process Monitor: START (RAM: {rss_mb:.0f}MB)")
|
|
else:
|
|
if self._current_record:
|
|
self._current_record.vram_peak_mb = max(self._current_record.vram_peak_mb, rss_mb)
|
|
self._current_record.ram_peak_mb = max(self._current_record.ram_peak_mb, rss_mb)
|
|
self._current_record.cpu_peak = max(self._current_record.cpu_peak, cpu)
|
|
self._snapshots.append(snap)
|
|
|
|
if delta < self._MEMORY_END_THRESHOLD_MB and cpu < self._CPU_END_THRESHOLD:
|
|
self._finalize(rss_mb)
|
|
logger.info(f"Process Monitor: END ({self.current.duration:.1f}s)")
|
|
|
|
self._prev_memory_mb = rss_mb
|
|
except Exception as e:
|
|
logger.debug(f"Process Monitor: {e}")
|
|
time.sleep(self.poll_interval)
|
|
|
|
def _finalize(self, end_mb: float):
|
|
if not self._current_record:
|
|
return
|
|
self._current_record.duration = time.time() - self._gen_start_time
|
|
self._current_record.vram_end_mb = end_mb
|
|
self._current_record.ram_end_mb = end_mb
|
|
self._current_record.vram_delta_mb = self._current_record.vram_peak_mb - self._current_record.vram_start_mb
|
|
self.current = self._current_record
|
|
self.history.append(self._current_record)
|
|
if len(self.history) > 50:
|
|
self.history = self.history[-50:]
|
|
self._is_generating = False
|
|
self._current_record = None
|
|
self._snapshots = []
|
|
|
|
@property
|
|
def is_generating(self) -> bool:
|
|
return self._is_generating
|
|
|
|
def get_current_generation(self) -> GenerationRecord:
|
|
with self._lock:
|
|
if self._is_generating and self._current_record:
|
|
return self._current_record
|
|
return GenerationRecord()
|
|
|
|
def get_last_generation(self) -> Optional[GenerationRecord]:
|
|
if self.history:
|
|
return self.history[-1]
|
|
return None
|
|
|
|
def get_summary(self) -> str:
|
|
if not self.history:
|
|
return "No generations recorded."
|
|
lines = ["=== ComfyUI Process Monitor ==="]
|
|
total_time = peak_ram = 0
|
|
for gen in self.history[-5:]:
|
|
total_time += gen.duration
|
|
peak_ram = max(peak_ram, gen.ram_peak_mb)
|
|
lines.append(f" #{len(lines)}: {gen.duration:.1f}s | RAM: {gen.ram_start_mb:.0f}→{gen.ram_peak_mb:.0f}MB")
|
|
lines.append(f" --- Total: {total_time:.1f}s | Peak RAM: {peak_ram:.0f}MB")
|
|
return "\n".join(lines)
|
|
|
|
|
|
# Singleton
|
|
_process_monitor = None
|
|
|
|
|
|
def get_process_monitor() -> ComfyProcessMonitor:
|
|
global _process_monitor
|
|
if _process_monitor is None:
|
|
_process_monitor = ComfyProcessMonitor()
|
|
return _process_monitor |