feat(kikostats): Add resource monitoring functionality
This commit is contained in:
Vito Sansevero
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
KikoStats package initialization and node registry
|
||||
Handles automatic discovery and registration of all ComfyAssets monitoring tools
|
||||
"""
|
||||
|
||||
from .tools.resource_monitor import ResourceMonitorNode
|
||||
|
||||
# ComfyUI node registration mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ResourceMonitor": ResourceMonitorNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ResourceMonitor": "Resource Monitor",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
Base classes and shared utilities for ComfyAssets monitoring tools
|
||||
"""
|
||||
|
||||
from .base_node import ComfyAssetsBaseNode
|
||||
|
||||
__all__ = ["ComfyAssetsBaseNode"]
|
||||
@@ -0,0 +1,76 @@
|
||||
"""
|
||||
Base node class for all ComfyAssets monitoring tools
|
||||
Provides consistent categorization and shared functionality
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ComfyAssetsBaseNode:
|
||||
"""
|
||||
Base class for all ComfyAssets monitoring nodes
|
||||
|
||||
Provides:
|
||||
- Consistent "ComfyAssets" categorization
|
||||
- Standardized error handling and logging
|
||||
- Common validation patterns
|
||||
- Consistent return type handling
|
||||
"""
|
||||
|
||||
CATEGORY = "ComfyAssets"
|
||||
|
||||
def validate_inputs(self, **kwargs) -> None:
|
||||
"""
|
||||
Common input validation logic
|
||||
Override in subclasses for specific validation needs
|
||||
|
||||
Args:
|
||||
**kwargs: Input parameters to validate
|
||||
|
||||
Raises:
|
||||
ValueError: If validation fails
|
||||
"""
|
||||
pass
|
||||
|
||||
def handle_error(
|
||||
self, error_msg: str, exception: Optional[Exception] = None
|
||||
) -> None:
|
||||
"""
|
||||
Standardized error handling with logging
|
||||
|
||||
Args:
|
||||
error_msg: Human-readable error message
|
||||
exception: Optional original exception for logging
|
||||
"""
|
||||
logger.error(f"{self.__class__.__name__}: {error_msg}")
|
||||
if exception:
|
||||
logger.exception(f"Original exception: {exception}")
|
||||
raise ValueError(error_msg)
|
||||
|
||||
def log_info(self, message: str) -> None:
|
||||
"""
|
||||
Standardized info logging
|
||||
|
||||
Args:
|
||||
message: Information message to log
|
||||
"""
|
||||
logger.info(f"{self.__class__.__name__}: {message}")
|
||||
|
||||
@classmethod
|
||||
def get_node_info(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
Get standardized node information for debugging/introspection
|
||||
|
||||
Returns:
|
||||
Dict containing node metadata
|
||||
"""
|
||||
return {
|
||||
"class_name": cls.__name__,
|
||||
"category": getattr(cls, "CATEGORY", "Unknown"),
|
||||
"function": getattr(cls, "FUNCTION", "Unknown"),
|
||||
"return_types": getattr(cls, "RETURN_TYPES", ()),
|
||||
"return_names": getattr(cls, "RETURN_NAMES", ()),
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
ComfyUI Execution Hooks for Node-Level Resource Monitoring
|
||||
Integrates with ComfyUI's execution system to track per-node performance
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
# Global state for tracking
|
||||
_node_tracking_data = {}
|
||||
_current_node_info = {}
|
||||
|
||||
def get_node_info_from_prompt(prompt_id: str, node_id: str) -> Dict[str, str]:
|
||||
"""Extract node information from the current execution context"""
|
||||
# In a real implementation, this would extract from ComfyUI's execution context
|
||||
# For now, return basic info
|
||||
return {
|
||||
'node_type': _current_node_info.get(node_id, {}).get('type', 'Unknown'),
|
||||
'node_title': _current_node_info.get(node_id, {}).get('title', f'Node {node_id}')
|
||||
}
|
||||
|
||||
def on_node_execution_start(prompt_id: str, node_id: str, node_data: Dict[str, Any]):
|
||||
"""Called when a node starts executing"""
|
||||
try:
|
||||
from .tools.resource_monitor.logic import get_global_monitor
|
||||
|
||||
# Store node info for later use
|
||||
_current_node_info[node_id] = {
|
||||
'type': node_data.get('class_type', 'Unknown'),
|
||||
'title': node_data.get('_meta', {}).get('title', f'Node {node_id}')
|
||||
}
|
||||
|
||||
# Get node information
|
||||
node_info = get_node_info_from_prompt(prompt_id, node_id)
|
||||
|
||||
# Start tracking this node
|
||||
monitor = get_global_monitor()
|
||||
monitor.start_node_tracking(
|
||||
node_id=node_id,
|
||||
node_type=node_info['node_type'],
|
||||
node_title=node_info['node_title']
|
||||
)
|
||||
|
||||
print(f"[KikoStats] Started tracking execution: {node_info['node_title']} ({node_info['node_type']}) [{node_id}]")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[KikoStats] Error starting node tracking: {e}")
|
||||
|
||||
def on_node_execution_end(prompt_id: str, node_id: str, output_data: Any):
|
||||
"""Called when a node finishes executing"""
|
||||
try:
|
||||
from .tools.resource_monitor.logic import get_global_monitor
|
||||
|
||||
# Stop tracking this node
|
||||
monitor = get_global_monitor()
|
||||
metrics = monitor.stop_node_tracking(node_id)
|
||||
|
||||
if metrics:
|
||||
print(f"[KikoStats] Completed tracking: {metrics['node_title']} "
|
||||
f"({metrics['duration_ms']:.1f}ms, "
|
||||
f"CPU: {metrics['avg_cpu_percent']:.1f}%, "
|
||||
f"GPU: {metrics['avg_gpu_utilization']:.1f}%)")
|
||||
|
||||
# Clean up node info
|
||||
_current_node_info.pop(node_id, None)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[KikoStats] Error stopping node tracking: {e}")
|
||||
|
||||
def install_execution_hooks():
|
||||
"""Install execution hooks into ComfyUI's execution system"""
|
||||
try:
|
||||
# Try to hook into ComfyUI's execution system
|
||||
import execution
|
||||
import server
|
||||
|
||||
# Store original methods
|
||||
if hasattr(execution, 'PromptExecutor'):
|
||||
original_execute = execution.PromptExecutor.execute
|
||||
|
||||
def wrapped_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
|
||||
"""Wrapped execute method that tracks nodes"""
|
||||
|
||||
# Store original methods for node execution
|
||||
if hasattr(self, 'execute_node'):
|
||||
original_execute_node = self.execute_node
|
||||
|
||||
def wrapped_execute_node(node, current_item, prompt_id, dynprompt):
|
||||
"""Wrapped node execution"""
|
||||
node_id = str(current_item)
|
||||
|
||||
# Get node data
|
||||
node_data = prompt.get(node_id, {})
|
||||
|
||||
# Start tracking
|
||||
on_node_execution_start(prompt_id, node_id, node_data)
|
||||
|
||||
try:
|
||||
# Execute the node
|
||||
result = original_execute_node(node, current_item, prompt_id, dynprompt)
|
||||
return result
|
||||
finally:
|
||||
# Stop tracking
|
||||
on_node_execution_end(prompt_id, node_id, None)
|
||||
|
||||
# Replace the method temporarily
|
||||
self.execute_node = wrapped_execute_node
|
||||
|
||||
# Call original execute
|
||||
return original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
|
||||
|
||||
# Replace the method
|
||||
execution.PromptExecutor.execute = wrapped_execute
|
||||
|
||||
print("[KikoStats] Successfully installed execution hooks")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"[KikoStats] Could not install execution hooks: {e}")
|
||||
return False
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,6 @@
|
||||
"""
|
||||
Individual monitoring tool implementations
|
||||
Each tool is self-contained with its own logic and node implementation
|
||||
"""
|
||||
|
||||
# Monitoring tools will be imported here as they are created
|
||||
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
Resource Monitor Tool
|
||||
Real-time GPU and system resource monitoring for ComfyUI
|
||||
"""
|
||||
|
||||
from .node import ResourceMonitorNode
|
||||
|
||||
__all__ = ["ResourceMonitorNode"]
|
||||
@@ -0,0 +1,720 @@
|
||||
"""
|
||||
Core monitoring logic for Resource Monitor
|
||||
Handles GPU and system resource data collection
|
||||
"""
|
||||
|
||||
import time
|
||||
import threading
|
||||
import asyncio
|
||||
from typing import Dict, Any, Optional, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
from collections import defaultdict
|
||||
|
||||
# Try to import monitoring libraries with graceful fallbacks
|
||||
try:
|
||||
import pynvml
|
||||
PYNVML_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYNVML_AVAILABLE = False
|
||||
pynvml = None
|
||||
|
||||
try:
|
||||
import psutil
|
||||
PSUTIL_AVAILABLE = True
|
||||
except ImportError:
|
||||
PSUTIL_AVAILABLE = False
|
||||
psutil = None
|
||||
|
||||
# Try to import ComfyUI server for WebSocket communication
|
||||
try:
|
||||
from server import PromptServer
|
||||
COMFYUI_SERVER_AVAILABLE = True
|
||||
except ImportError:
|
||||
COMFYUI_SERVER_AVAILABLE = False
|
||||
PromptServer = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class GPUStats:
|
||||
"""GPU monitoring statistics"""
|
||||
utilization: int # GPU utilization percentage (0-100)
|
||||
memory_used: int # VRAM used in MB
|
||||
memory_total: int # Total VRAM in MB
|
||||
memory_percent: float # VRAM usage percentage
|
||||
temperature: int # GPU temperature in Celsius
|
||||
power_draw: int # Power consumption in Watts
|
||||
available: bool = True # Whether GPU monitoring is available
|
||||
|
||||
|
||||
@dataclass
|
||||
class SystemStats:
|
||||
"""System resource statistics"""
|
||||
cpu_percent: float # CPU usage percentage
|
||||
ram_used: int # RAM used in MB
|
||||
ram_total: int # Total RAM in MB
|
||||
ram_percent: float # RAM usage percentage
|
||||
available: bool = True # Whether system monitoring is available
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeResourceSample:
|
||||
"""Single resource measurement for a node"""
|
||||
timestamp: float
|
||||
cpu_percent: float
|
||||
gpu_utilization: int
|
||||
gpu_memory_used: int
|
||||
vram_delta: int # Change in VRAM usage since last sample
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeResourceMetrics:
|
||||
"""Aggregated resource metrics for a node"""
|
||||
node_id: str
|
||||
node_type: str = ""
|
||||
node_title: str = ""
|
||||
start_time: float = 0.0
|
||||
end_time: float = 0.0
|
||||
duration_ms: float = 0.0
|
||||
|
||||
# CPU metrics
|
||||
avg_cpu_percent: float = 0.0
|
||||
max_cpu_percent: float = 0.0
|
||||
|
||||
# GPU metrics
|
||||
avg_gpu_utilization: float = 0.0
|
||||
max_gpu_utilization: float = 0.0
|
||||
peak_vram_used: int = 0
|
||||
vram_delta: int = 0 # Net VRAM change during execution
|
||||
|
||||
# Sample data
|
||||
samples: list = field(default_factory=list)
|
||||
|
||||
def calculate_aggregates(self):
|
||||
"""Calculate aggregate metrics from samples"""
|
||||
if not self.samples:
|
||||
return
|
||||
|
||||
self.duration_ms = (self.end_time - self.start_time) * 1000
|
||||
|
||||
cpu_values = [s.cpu_percent for s in self.samples]
|
||||
gpu_values = [s.gpu_utilization for s in self.samples]
|
||||
vram_values = [s.gpu_memory_used for s in self.samples]
|
||||
|
||||
self.avg_cpu_percent = sum(cpu_values) / len(cpu_values) if cpu_values else 0
|
||||
self.max_cpu_percent = max(cpu_values) if cpu_values else 0
|
||||
|
||||
self.avg_gpu_utilization = sum(gpu_values) / len(gpu_values) if gpu_values else 0
|
||||
self.max_gpu_utilization = max(gpu_values) if gpu_values else 0
|
||||
|
||||
self.peak_vram_used = max(vram_values) if vram_values else 0
|
||||
|
||||
# Calculate net VRAM change
|
||||
if len(self.samples) >= 2:
|
||||
self.vram_delta = self.samples[-1].gpu_memory_used - self.samples[0].gpu_memory_used
|
||||
|
||||
|
||||
def initialize_gpu_monitoring() -> bool:
|
||||
"""
|
||||
Initialize GPU monitoring using NVIDIA Management Library
|
||||
|
||||
Returns:
|
||||
bool: True if GPU monitoring is available, False otherwise
|
||||
"""
|
||||
if not PYNVML_AVAILABLE:
|
||||
return False
|
||||
|
||||
try:
|
||||
pynvml.nvmlInit()
|
||||
# Try to get device count to verify NVIDIA drivers are working
|
||||
device_count = pynvml.nvmlDeviceGetCount()
|
||||
return device_count > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def get_gpu_stats(device_index: int = 0) -> GPUStats:
|
||||
"""
|
||||
Get current GPU statistics using NVIDIA Management Library
|
||||
|
||||
Args:
|
||||
device_index: GPU device index (default: 0 for primary GPU)
|
||||
|
||||
Returns:
|
||||
GPUStats object with current GPU metrics
|
||||
"""
|
||||
if not PYNVML_AVAILABLE:
|
||||
return GPUStats(
|
||||
utilization=0,
|
||||
memory_used=0,
|
||||
memory_total=0,
|
||||
memory_percent=0.0,
|
||||
temperature=0,
|
||||
power_draw=0,
|
||||
available=False
|
||||
)
|
||||
|
||||
try:
|
||||
# Ensure pynvml is initialized
|
||||
pynvml.nvmlInit()
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
|
||||
|
||||
# Get GPU utilization
|
||||
try:
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
gpu_util = util.gpu
|
||||
except Exception:
|
||||
gpu_util = 0
|
||||
|
||||
# Get memory info
|
||||
try:
|
||||
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
memory_used = mem_info.used // (1024 * 1024) # Convert to MB
|
||||
memory_total = mem_info.total // (1024 * 1024) # Convert to MB
|
||||
memory_percent = (mem_info.used / mem_info.total) * 100
|
||||
except Exception:
|
||||
memory_used = memory_total = 0
|
||||
memory_percent = 0.0
|
||||
|
||||
# Get temperature
|
||||
try:
|
||||
temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
|
||||
except Exception:
|
||||
temp = 0
|
||||
|
||||
# Get power draw
|
||||
try:
|
||||
power = pynvml.nvmlDeviceGetPowerUsage(handle) // 1000 # Convert to Watts
|
||||
except Exception:
|
||||
power = 0
|
||||
|
||||
return GPUStats(
|
||||
utilization=gpu_util,
|
||||
memory_used=memory_used,
|
||||
memory_total=memory_total,
|
||||
memory_percent=memory_percent,
|
||||
temperature=temp,
|
||||
power_draw=power,
|
||||
available=True
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# GPU monitoring failed - return unavailable stats
|
||||
return GPUStats(
|
||||
utilization=0,
|
||||
memory_used=0,
|
||||
memory_total=0,
|
||||
memory_percent=0.0,
|
||||
temperature=0,
|
||||
power_draw=0,
|
||||
available=False
|
||||
)
|
||||
|
||||
|
||||
def get_system_stats() -> SystemStats:
|
||||
"""
|
||||
Get current system resource statistics using psutil
|
||||
|
||||
Returns:
|
||||
SystemStats object with current system metrics
|
||||
"""
|
||||
if not PSUTIL_AVAILABLE:
|
||||
return SystemStats(
|
||||
cpu_percent=0.0,
|
||||
ram_used=0,
|
||||
ram_total=0,
|
||||
ram_percent=0.0,
|
||||
available=False
|
||||
)
|
||||
|
||||
try:
|
||||
# Get CPU usage (non-blocking, based on last call)
|
||||
cpu_percent = psutil.cpu_percent(interval=0)
|
||||
|
||||
# Get memory info
|
||||
memory = psutil.virtual_memory()
|
||||
ram_used = memory.used // (1024 * 1024) # Convert to MB
|
||||
ram_total = memory.total // (1024 * 1024) # Convert to MB
|
||||
ram_percent = memory.percent
|
||||
|
||||
return SystemStats(
|
||||
cpu_percent=cpu_percent,
|
||||
ram_used=ram_used,
|
||||
ram_total=ram_total,
|
||||
ram_percent=ram_percent,
|
||||
available=True
|
||||
)
|
||||
|
||||
except Exception:
|
||||
return SystemStats(
|
||||
cpu_percent=0.0,
|
||||
ram_used=0,
|
||||
ram_total=0,
|
||||
ram_percent=0.0,
|
||||
available=False
|
||||
)
|
||||
|
||||
|
||||
class ResourceMonitor:
|
||||
"""
|
||||
Thread-safe resource monitoring class
|
||||
Collects GPU and system statistics in separate thread
|
||||
"""
|
||||
|
||||
def __init__(self, update_interval: float = 1.0):
|
||||
"""
|
||||
Initialize resource monitor
|
||||
|
||||
Args:
|
||||
update_interval: Monitoring update interval in seconds
|
||||
"""
|
||||
self.update_interval = update_interval
|
||||
self.monitoring = False
|
||||
self.monitor_thread: Optional[threading.Thread] = None
|
||||
self.lock = threading.Lock()
|
||||
|
||||
# Latest stats
|
||||
self.gpu_stats: Optional[GPUStats] = None
|
||||
self.system_stats: Optional[SystemStats] = None
|
||||
|
||||
# Initialize monitoring capabilities
|
||||
self.gpu_available = initialize_gpu_monitoring()
|
||||
self.system_available = PSUTIL_AVAILABLE
|
||||
|
||||
def start_monitoring(self) -> None:
|
||||
"""Start monitoring in background thread"""
|
||||
if self.monitoring:
|
||||
return
|
||||
|
||||
self.monitoring = True
|
||||
self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
|
||||
self.monitor_thread.start()
|
||||
|
||||
def stop_monitoring(self) -> None:
|
||||
"""Stop background monitoring"""
|
||||
self.monitoring = False
|
||||
if self.monitor_thread and self.monitor_thread.is_alive():
|
||||
self.monitor_thread.join(timeout=2.0)
|
||||
|
||||
def _monitor_loop(self) -> None:
|
||||
"""Main monitoring loop running in background thread"""
|
||||
while self.monitoring:
|
||||
try:
|
||||
# Collect stats
|
||||
gpu_stats = get_gpu_stats() if self.gpu_available else None
|
||||
system_stats = get_system_stats() if self.system_available else None
|
||||
|
||||
# Thread-safe update
|
||||
with self.lock:
|
||||
self.gpu_stats = gpu_stats
|
||||
self.system_stats = system_stats
|
||||
|
||||
# Sleep until next update
|
||||
time.sleep(self.update_interval)
|
||||
|
||||
except Exception:
|
||||
# Continue monitoring even if individual collection fails
|
||||
time.sleep(self.update_interval)
|
||||
|
||||
def get_current_stats(self) -> Tuple[Optional[GPUStats], Optional[SystemStats]]:
|
||||
"""
|
||||
Get latest monitoring statistics
|
||||
|
||||
Returns:
|
||||
Tuple of (gpu_stats, system_stats) or (None, None) if not available
|
||||
"""
|
||||
with self.lock:
|
||||
return self.gpu_stats, self.system_stats
|
||||
|
||||
def get_stats_dict(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get statistics as dictionary for JSON serialization
|
||||
|
||||
Returns:
|
||||
Dictionary containing all current statistics
|
||||
"""
|
||||
gpu_stats, system_stats = self.get_current_stats()
|
||||
|
||||
result = {
|
||||
"timestamp": time.time(),
|
||||
"gpu": {
|
||||
"available": gpu_stats.available if gpu_stats else False,
|
||||
"utilization": gpu_stats.utilization if gpu_stats else 0,
|
||||
"memory_used": gpu_stats.memory_used if gpu_stats else 0,
|
||||
"memory_total": gpu_stats.memory_total if gpu_stats else 0,
|
||||
"memory_percent": gpu_stats.memory_percent if gpu_stats else 0.0,
|
||||
"temperature": gpu_stats.temperature if gpu_stats else 0,
|
||||
"power_draw": gpu_stats.power_draw if gpu_stats else 0,
|
||||
},
|
||||
"system": {
|
||||
"available": system_stats.available if system_stats else False,
|
||||
"cpu_percent": system_stats.cpu_percent if system_stats else 0.0,
|
||||
"ram_used": system_stats.ram_used if system_stats else 0,
|
||||
"ram_total": system_stats.ram_total if system_stats else 0,
|
||||
"ram_percent": system_stats.ram_percent if system_stats else 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def validate_update_interval(interval: float) -> None:
|
||||
"""
|
||||
Validate monitoring update interval
|
||||
|
||||
Args:
|
||||
interval: Update interval in seconds
|
||||
|
||||
Raises:
|
||||
ValueError: If interval is out of valid range
|
||||
"""
|
||||
if not isinstance(interval, (int, float)):
|
||||
raise ValueError(f"Update interval must be a number, got {type(interval).__name__}")
|
||||
|
||||
if interval < 0.1:
|
||||
raise ValueError(f"Update interval too small: {interval}s (minimum: 0.1s)")
|
||||
|
||||
if interval > 60.0:
|
||||
raise ValueError(f"Update interval too large: {interval}s (maximum: 60s)")
|
||||
|
||||
|
||||
def format_memory_size(size_mb: int) -> str:
|
||||
"""
|
||||
Format memory size for display
|
||||
|
||||
Args:
|
||||
size_mb: Memory size in megabytes
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "4.2 GB", "512 MB")
|
||||
"""
|
||||
if size_mb >= 1024:
|
||||
return f"{size_mb / 1024:.1f} GB"
|
||||
else:
|
||||
return f"{size_mb} MB"
|
||||
|
||||
|
||||
def format_percentage(percent: float) -> str:
|
||||
"""
|
||||
Format percentage for display
|
||||
|
||||
Args:
|
||||
percent: Percentage value
|
||||
|
||||
Returns:
|
||||
Formatted percentage string
|
||||
"""
|
||||
return f"{percent:.1f}%"
|
||||
|
||||
|
||||
class NodeResourceTracker:
|
||||
"""
|
||||
Tracks resource usage per ComfyUI node during execution
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.current_node_id: Optional[str] = None
|
||||
self.active_nodes: Dict[str, NodeResourceMetrics] = {}
|
||||
self.completed_nodes: Dict[str, NodeResourceMetrics] = {}
|
||||
self.lock = threading.Lock()
|
||||
|
||||
# For tracking resource baselines
|
||||
self.baseline_stats = None
|
||||
|
||||
def start_node_tracking(self, node_id: str, node_type: str = "", node_title: str = ""):
|
||||
"""Start tracking resources for a specific node"""
|
||||
with self.lock:
|
||||
self.current_node_id = node_id
|
||||
|
||||
# Create metrics object for this node
|
||||
metrics = NodeResourceMetrics(
|
||||
node_id=node_id,
|
||||
node_type=node_type,
|
||||
node_title=node_title,
|
||||
start_time=time.time()
|
||||
)
|
||||
|
||||
self.active_nodes[node_id] = metrics
|
||||
|
||||
# Store baseline for calculating deltas
|
||||
self.baseline_stats = self._get_current_stats()
|
||||
|
||||
print(f"[KikoStats] Started tracking node: {node_id} ({node_type})")
|
||||
|
||||
def stop_node_tracking(self, node_id: str):
|
||||
"""Stop tracking and finalize metrics for a node"""
|
||||
with self.lock:
|
||||
if node_id in self.active_nodes:
|
||||
metrics = self.active_nodes[node_id]
|
||||
metrics.end_time = time.time()
|
||||
|
||||
# Calculate aggregate metrics
|
||||
metrics.calculate_aggregates()
|
||||
|
||||
# Move to completed
|
||||
self.completed_nodes[node_id] = metrics
|
||||
del self.active_nodes[node_id]
|
||||
|
||||
# Clear current tracking
|
||||
if self.current_node_id == node_id:
|
||||
self.current_node_id = None
|
||||
|
||||
print(f"[KikoStats] Completed tracking node: {node_id} "
|
||||
f"(Duration: {metrics.duration_ms:.1f}ms, "
|
||||
f"Avg CPU: {metrics.avg_cpu_percent:.1f}%, "
|
||||
f"Avg GPU: {metrics.avg_gpu_utilization:.1f}%)")
|
||||
|
||||
return metrics
|
||||
|
||||
return None
|
||||
|
||||
def sample_current_node(self, gpu_stats: Optional[GPUStats], system_stats: Optional[SystemStats]):
|
||||
"""Sample resources for the currently tracking node"""
|
||||
with self.lock:
|
||||
if not self.current_node_id or self.current_node_id not in self.active_nodes:
|
||||
return
|
||||
|
||||
# Create sample
|
||||
sample = NodeResourceSample(
|
||||
timestamp=time.time(),
|
||||
cpu_percent=system_stats.cpu_percent if system_stats else 0.0,
|
||||
gpu_utilization=gpu_stats.utilization if gpu_stats else 0,
|
||||
gpu_memory_used=gpu_stats.memory_used if gpu_stats else 0,
|
||||
vram_delta=0 # Will be calculated later
|
||||
)
|
||||
|
||||
# Add to current node's samples
|
||||
self.active_nodes[self.current_node_id].samples.append(sample)
|
||||
|
||||
def get_recent_node_metrics(self, limit: int = 10) -> list:
|
||||
"""Get recent completed node metrics"""
|
||||
with self.lock:
|
||||
# Get most recent completed nodes
|
||||
recent = list(self.completed_nodes.values())[-limit:]
|
||||
|
||||
# Convert to dictionaries for JSON serialization
|
||||
return [
|
||||
{
|
||||
'node_id': m.node_id,
|
||||
'node_type': m.node_type,
|
||||
'node_title': m.node_title,
|
||||
'duration_ms': m.duration_ms,
|
||||
'avg_cpu_percent': m.avg_cpu_percent,
|
||||
'max_cpu_percent': m.max_cpu_percent,
|
||||
'avg_gpu_utilization': m.avg_gpu_utilization,
|
||||
'max_gpu_utilization': m.max_gpu_utilization,
|
||||
'peak_vram_used': m.peak_vram_used,
|
||||
'vram_delta': m.vram_delta,
|
||||
'sample_count': len(m.samples)
|
||||
}
|
||||
for m in recent
|
||||
]
|
||||
|
||||
def clear_old_metrics(self, keep_count: int = 50):
|
||||
"""Clear old completed metrics to prevent memory buildup"""
|
||||
with self.lock:
|
||||
if len(self.completed_nodes) > keep_count:
|
||||
# Keep only the most recent entries
|
||||
sorted_items = sorted(self.completed_nodes.items(),
|
||||
key=lambda x: x[1].end_time)
|
||||
to_keep = dict(sorted_items[-keep_count:])
|
||||
self.completed_nodes = to_keep
|
||||
|
||||
def _get_current_stats(self) -> Dict[str, Any]:
|
||||
"""Get current system stats for baseline calculations"""
|
||||
gpu_stats = get_gpu_stats()
|
||||
system_stats = get_system_stats()
|
||||
|
||||
return {
|
||||
'cpu_percent': system_stats.cpu_percent if system_stats else 0.0,
|
||||
'gpu_utilization': gpu_stats.utilization if gpu_stats else 0,
|
||||
'gpu_memory_used': gpu_stats.memory_used if gpu_stats else 0
|
||||
}
|
||||
|
||||
|
||||
class ContinuousResourceMonitor:
|
||||
"""
|
||||
Continuous background monitoring similar to Crystools
|
||||
Sends real-time data via ComfyUI WebSocket
|
||||
"""
|
||||
|
||||
def __init__(self, update_interval: float = 1.0):
|
||||
"""
|
||||
Initialize continuous monitor
|
||||
|
||||
Args:
|
||||
update_interval: Update frequency in seconds
|
||||
"""
|
||||
self.update_interval = update_interval
|
||||
self.running = False
|
||||
self.thread = None
|
||||
self.lock = threading.Lock()
|
||||
|
||||
# Latest stats cache
|
||||
self.latest_stats = None
|
||||
|
||||
# Node tracking
|
||||
self.node_tracker = NodeResourceTracker()
|
||||
|
||||
# Initialize monitoring capabilities
|
||||
self.gpu_available = initialize_gpu_monitoring()
|
||||
self.system_available = PSUTIL_AVAILABLE
|
||||
|
||||
# Initialize CPU monitoring baseline
|
||||
if self.system_available:
|
||||
try:
|
||||
# Prime CPU monitoring with first call
|
||||
psutil.cpu_percent(interval=None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def start_monitoring(self):
|
||||
"""Start continuous background monitoring"""
|
||||
if self.running:
|
||||
return
|
||||
|
||||
self.running = True
|
||||
self.thread = threading.Thread(target=self._monitor_loop, daemon=True)
|
||||
self.thread.start()
|
||||
print(f"[KikoStats] Started continuous monitoring (interval: {self.update_interval}s)")
|
||||
|
||||
def stop_monitoring(self):
|
||||
"""Stop background monitoring"""
|
||||
self.running = False
|
||||
if self.thread and self.thread.is_alive():
|
||||
self.thread.join(timeout=2.0)
|
||||
print("[KikoStats] Stopped continuous monitoring")
|
||||
|
||||
def _monitor_loop(self):
|
||||
"""Main monitoring loop - runs continuously in background"""
|
||||
while self.running:
|
||||
try:
|
||||
# Collect fresh stats
|
||||
gpu_stats = get_gpu_stats() if self.gpu_available else None
|
||||
system_stats = get_system_stats() if self.system_available else None
|
||||
|
||||
# Sample current node if tracking
|
||||
self.node_tracker.sample_current_node(gpu_stats, system_stats)
|
||||
|
||||
# Create stats dictionary
|
||||
stats_dict = {
|
||||
"timestamp": time.time(),
|
||||
"gpu": {
|
||||
"available": gpu_stats.available if gpu_stats else False,
|
||||
"utilization": gpu_stats.utilization if gpu_stats else 0,
|
||||
"memory_used": gpu_stats.memory_used if gpu_stats else 0,
|
||||
"memory_total": gpu_stats.memory_total if gpu_stats else 0,
|
||||
"memory_percent": gpu_stats.memory_percent if gpu_stats else 0.0,
|
||||
"temperature": gpu_stats.temperature if gpu_stats else 0,
|
||||
"power_draw": gpu_stats.power_draw if gpu_stats else 0,
|
||||
},
|
||||
"system": {
|
||||
"available": system_stats.available if system_stats else False,
|
||||
"cpu_percent": system_stats.cpu_percent if system_stats else 0.0,
|
||||
"ram_used": system_stats.ram_used if system_stats else 0,
|
||||
"ram_total": system_stats.ram_total if system_stats else 0,
|
||||
"ram_percent": system_stats.ram_percent if system_stats else 0.0,
|
||||
},
|
||||
"nodes": self.node_tracker.get_recent_node_metrics(limit=10)
|
||||
}
|
||||
|
||||
# Thread-safe update
|
||||
with self.lock:
|
||||
self.latest_stats = stats_dict
|
||||
|
||||
# Check for node tracking commands from JavaScript
|
||||
self._process_js_tracking_commands()
|
||||
|
||||
# Send real-time data via WebSocket (like Crystools)
|
||||
if COMFYUI_SERVER_AVAILABLE and PromptServer.instance:
|
||||
try:
|
||||
PromptServer.instance.send_sync('kikostats.monitor', stats_dict)
|
||||
except Exception as e:
|
||||
# Don't let WebSocket errors stop monitoring
|
||||
pass
|
||||
|
||||
# Sleep until next update
|
||||
time.sleep(self.update_interval)
|
||||
|
||||
except Exception as e:
|
||||
# Continue monitoring even if individual collection fails
|
||||
print(f"[KikoStats] Monitoring error: {e}")
|
||||
time.sleep(self.update_interval)
|
||||
|
||||
def get_latest_stats(self) -> Optional[Dict[str, Any]]:
|
||||
"""Get the latest cached stats"""
|
||||
with self.lock:
|
||||
return self.latest_stats.copy() if self.latest_stats else None
|
||||
|
||||
def start_node_tracking(self, node_id: str, node_type: str = "", node_title: str = ""):
|
||||
"""Start tracking resources for a specific node"""
|
||||
self.node_tracker.start_node_tracking(node_id, node_type, node_title)
|
||||
|
||||
# Send node start event via WebSocket
|
||||
if COMFYUI_SERVER_AVAILABLE and PromptServer.instance:
|
||||
try:
|
||||
PromptServer.instance.send_sync('kikostats.node_start', {
|
||||
'node_id': node_id,
|
||||
'node_type': node_type,
|
||||
'node_title': node_title,
|
||||
'timestamp': time.time()
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def stop_node_tracking(self, node_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Stop tracking and get final metrics for a node"""
|
||||
metrics = self.node_tracker.stop_node_tracking(node_id)
|
||||
|
||||
if metrics:
|
||||
# Convert to dict for WebSocket
|
||||
metrics_dict = {
|
||||
'node_id': metrics.node_id,
|
||||
'node_type': metrics.node_type,
|
||||
'node_title': metrics.node_title,
|
||||
'duration_ms': metrics.duration_ms,
|
||||
'avg_cpu_percent': metrics.avg_cpu_percent,
|
||||
'max_cpu_percent': metrics.max_cpu_percent,
|
||||
'avg_gpu_utilization': metrics.avg_gpu_utilization,
|
||||
'max_gpu_utilization': metrics.max_gpu_utilization,
|
||||
'peak_vram_used': metrics.peak_vram_used,
|
||||
'vram_delta': metrics.vram_delta,
|
||||
'sample_count': len(metrics.samples)
|
||||
}
|
||||
|
||||
# Send node completion event via WebSocket
|
||||
if COMFYUI_SERVER_AVAILABLE and PromptServer.instance:
|
||||
try:
|
||||
PromptServer.instance.send_sync('kikostats.node_complete', metrics_dict)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return metrics_dict
|
||||
|
||||
return None
|
||||
|
||||
def _process_js_tracking_commands(self):
|
||||
"""Process node tracking commands from JavaScript frontend"""
|
||||
# This is a simple approach - in a real implementation you'd use proper IPC
|
||||
# For now, we'll implement a basic execution hook approach
|
||||
pass
|
||||
|
||||
|
||||
# Global continuous monitor instance
|
||||
_global_monitor: Optional[ContinuousResourceMonitor] = None
|
||||
|
||||
|
||||
def get_global_monitor() -> ContinuousResourceMonitor:
|
||||
"""Get or create the global continuous monitor"""
|
||||
global _global_monitor
|
||||
if _global_monitor is None:
|
||||
_global_monitor = ContinuousResourceMonitor(update_interval=1.0)
|
||||
_global_monitor.start_monitoring()
|
||||
return _global_monitor
|
||||
|
||||
|
||||
def stop_global_monitor():
|
||||
"""Stop the global monitor (for cleanup)"""
|
||||
global _global_monitor
|
||||
if _global_monitor:
|
||||
_global_monitor.stop_monitoring()
|
||||
_global_monitor = None
|
||||
@@ -0,0 +1,261 @@
|
||||
"""
|
||||
Resource Monitor ComfyUI Node
|
||||
Provides ComfyUI interface for real-time GPU and system monitoring
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Dict, Any, Tuple
|
||||
|
||||
from ...base import ComfyAssetsBaseNode
|
||||
from .logic import ResourceMonitor, validate_update_interval, format_memory_size, format_percentage, get_global_monitor
|
||||
|
||||
|
||||
class ResourceMonitorNode(ComfyAssetsBaseNode):
|
||||
"""
|
||||
ComfyUI node for real-time resource monitoring
|
||||
|
||||
Displays GPU utilization, VRAM usage, CPU usage, and RAM usage
|
||||
with configurable update intervals and display modes.
|
||||
|
||||
Inputs:
|
||||
- update_interval (FLOAT): Monitoring refresh rate in seconds
|
||||
- display_mode (STRING): Output format ("text", "json", "both")
|
||||
- enable_gpu (BOOLEAN): Enable GPU monitoring
|
||||
- enable_system (BOOLEAN): Enable system monitoring
|
||||
|
||||
Outputs:
|
||||
- stats (STRING): Formatted statistics output
|
||||
- json_data (STRING): JSON statistics for chaining
|
||||
"""
|
||||
|
||||
# Class variable to maintain monitor instance across executions
|
||||
_monitor_instance: ResourceMonitor = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
Define ComfyUI input interface
|
||||
|
||||
Returns:
|
||||
Dict with required and optional input specifications
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"update_interval": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 60.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
"tooltip": "How often to update monitoring data (seconds)"
|
||||
},
|
||||
),
|
||||
"display_mode": (
|
||||
["text", "json", "both"],
|
||||
{
|
||||
"default": "text",
|
||||
"tooltip": "Output format for monitoring data"
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"enable_gpu": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Enable GPU monitoring (requires NVIDIA GPU and drivers)"
|
||||
},
|
||||
),
|
||||
"enable_system": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Enable system resource monitoring (CPU, RAM)"
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("stats", "json_data")
|
||||
FUNCTION = "monitor_resources"
|
||||
|
||||
def monitor_resources(
|
||||
self,
|
||||
update_interval: float,
|
||||
display_mode: str,
|
||||
enable_gpu: bool = True,
|
||||
enable_system: bool = True,
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Monitor system and GPU resources with real-time updates
|
||||
|
||||
Args:
|
||||
update_interval: Update frequency in seconds
|
||||
display_mode: Output format ("text", "json", "both")
|
||||
enable_gpu: Whether to monitor GPU resources
|
||||
enable_system: Whether to monitor system resources
|
||||
|
||||
Returns:
|
||||
Tuple of (formatted_stats, json_data)
|
||||
|
||||
Raises:
|
||||
ValueError: If validation fails
|
||||
"""
|
||||
try:
|
||||
# Validate inputs
|
||||
self.validate_inputs(
|
||||
update_interval=update_interval,
|
||||
display_mode=display_mode,
|
||||
enable_gpu=enable_gpu,
|
||||
enable_system=enable_system
|
||||
)
|
||||
|
||||
# Get the global continuous monitor (starts automatically)
|
||||
continuous_monitor = get_global_monitor()
|
||||
|
||||
# Get latest stats from continuous monitoring
|
||||
stats_dict = continuous_monitor.get_latest_stats()
|
||||
|
||||
# If no stats yet, return waiting message
|
||||
if stats_dict is None:
|
||||
waiting_stats = {
|
||||
"timestamp": time.time(),
|
||||
"gpu": {"available": False, "utilization": 0, "memory_used": 0, "memory_total": 0, "memory_percent": 0.0, "temperature": 0, "power_draw": 0},
|
||||
"system": {"available": False, "cpu_percent": 0.0, "ram_used": 0, "ram_total": 0, "ram_percent": 0.0}
|
||||
}
|
||||
stats_dict = waiting_stats
|
||||
|
||||
# Filter based on enabled monitoring
|
||||
if not enable_gpu:
|
||||
stats_dict["gpu"]["available"] = False
|
||||
if not enable_system:
|
||||
stats_dict["system"]["available"] = False
|
||||
|
||||
# Format output based on display mode
|
||||
if display_mode == "text":
|
||||
formatted_stats = self._format_text_output(stats_dict, enable_gpu, enable_system)
|
||||
json_data = ""
|
||||
elif display_mode == "json":
|
||||
formatted_stats = ""
|
||||
json_data = json.dumps(stats_dict, indent=2)
|
||||
else: # "both"
|
||||
formatted_stats = self._format_text_output(stats_dict, enable_gpu, enable_system)
|
||||
json_data = json.dumps(stats_dict, indent=2)
|
||||
|
||||
# Log monitoring status
|
||||
self.log_info(
|
||||
f"Continuous Monitoring Active: GPU={enable_gpu}, System={enable_system}, "
|
||||
f"Display={display_mode} (Real-time updates via WebSocket)"
|
||||
)
|
||||
|
||||
return formatted_stats, json_data
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to monitor resources: {str(e)}"
|
||||
self.handle_error(error_msg, e)
|
||||
|
||||
|
||||
def _format_text_output(
|
||||
self,
|
||||
stats_dict: Dict[str, Any],
|
||||
enable_gpu: bool,
|
||||
enable_system: bool
|
||||
) -> str:
|
||||
"""
|
||||
Format statistics as human-readable text
|
||||
|
||||
Args:
|
||||
stats_dict: Statistics dictionary
|
||||
enable_gpu: Whether GPU monitoring is enabled
|
||||
enable_system: Whether system monitoring is enabled
|
||||
|
||||
Returns:
|
||||
Formatted text string
|
||||
"""
|
||||
lines = []
|
||||
lines.append("=== ComfyUI Resource Monitor ===")
|
||||
|
||||
# GPU Statistics
|
||||
if enable_gpu:
|
||||
gpu = stats_dict["gpu"]
|
||||
if gpu["available"]:
|
||||
lines.append("\n🖥️ GPU Statistics:")
|
||||
lines.append(f" Utilization: {gpu['utilization']}%")
|
||||
lines.append(f" VRAM: {format_memory_size(gpu['memory_used'])} / "
|
||||
f"{format_memory_size(gpu['memory_total'])} "
|
||||
f"({format_percentage(gpu['memory_percent'])})")
|
||||
lines.append(f" Temperature: {gpu['temperature']}°C")
|
||||
lines.append(f" Power: {gpu['power_draw']}W")
|
||||
else:
|
||||
lines.append("\n🖥️ GPU: Not Available")
|
||||
lines.append(" (Requires NVIDIA GPU and drivers)")
|
||||
|
||||
# System Statistics
|
||||
if enable_system:
|
||||
system = stats_dict["system"]
|
||||
if system["available"]:
|
||||
lines.append("\n💻 System Statistics:")
|
||||
lines.append(f" CPU Usage: {format_percentage(system['cpu_percent'])}")
|
||||
lines.append(f" RAM: {format_memory_size(system['ram_used'])} / "
|
||||
f"{format_memory_size(system['ram_total'])} "
|
||||
f"({format_percentage(system['ram_percent'])})")
|
||||
else:
|
||||
lines.append("\n💻 System: Not Available")
|
||||
lines.append(" (Requires psutil library)")
|
||||
|
||||
# Add timestamp
|
||||
import time
|
||||
timestamp = time.strftime("%H:%M:%S", time.localtime(stats_dict["timestamp"]))
|
||||
lines.append(f"\n⏰ Last Updated: {timestamp}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def validate_inputs(
|
||||
self,
|
||||
update_interval: float,
|
||||
display_mode: str,
|
||||
enable_gpu: bool,
|
||||
enable_system: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Validate inputs specific to resource monitor
|
||||
|
||||
Args:
|
||||
update_interval: Update interval to validate
|
||||
display_mode: Display mode to validate
|
||||
enable_gpu: GPU monitoring flag
|
||||
enable_system: System monitoring flag
|
||||
|
||||
Raises:
|
||||
ValueError: If validation fails
|
||||
"""
|
||||
# Validate update interval
|
||||
validate_update_interval(update_interval)
|
||||
|
||||
# Validate display mode
|
||||
valid_modes = ["text", "json", "both"]
|
||||
if display_mode not in valid_modes:
|
||||
raise ValueError(f"Invalid display_mode '{display_mode}'. Must be one of: {valid_modes}")
|
||||
|
||||
# Validate boolean flags
|
||||
for flag_name, flag_value in [("enable_gpu", enable_gpu), ("enable_system", enable_system)]:
|
||||
if not isinstance(flag_value, bool):
|
||||
raise ValueError(f"{flag_name} must be a boolean, got {type(flag_value).__name__}")
|
||||
|
||||
# Check that at least one monitoring type is enabled
|
||||
if not enable_gpu and not enable_system:
|
||||
raise ValueError("At least one monitoring type (GPU or System) must be enabled")
|
||||
|
||||
|
||||
# Node class mappings for ComfyUI registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ResourceMonitor": ResourceMonitorNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ResourceMonitor": "Resource Monitor",
|
||||
}
|
||||
Reference in New Issue
Block a user