From e7947dbfd104e4cb4b9535de01ea1d006325ad3b Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Fri, 20 Jun 2025 05:20:51 -0700 Subject: [PATCH] feat(kikostats): Add resource monitoring functionality --- kikostats/__init__.py | 17 + kikostats/base/__init__.py | 7 + kikostats/base/base_node.py | 76 ++ kikostats/execution_hooks.py | 122 ++++ kikostats/tools/__init__.py | 6 + kikostats/tools/resource_monitor/__init__.py | 8 + kikostats/tools/resource_monitor/logic.py | 720 +++++++++++++++++++ kikostats/tools/resource_monitor/node.py | 261 +++++++ 8 files changed, 1217 insertions(+) create mode 100644 kikostats/__init__.py create mode 100644 kikostats/base/__init__.py create mode 100644 kikostats/base/base_node.py create mode 100644 kikostats/execution_hooks.py create mode 100644 kikostats/tools/__init__.py create mode 100644 kikostats/tools/resource_monitor/__init__.py create mode 100644 kikostats/tools/resource_monitor/logic.py create mode 100644 kikostats/tools/resource_monitor/node.py diff --git a/kikostats/__init__.py b/kikostats/__init__.py new file mode 100644 index 0000000..fa0cbb4 --- /dev/null +++ b/kikostats/__init__.py @@ -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"] \ No newline at end of file diff --git a/kikostats/base/__init__.py b/kikostats/base/__init__.py new file mode 100644 index 0000000..ae8264b --- /dev/null +++ b/kikostats/base/__init__.py @@ -0,0 +1,7 @@ +""" +Base classes and shared utilities for ComfyAssets monitoring tools +""" + +from .base_node import ComfyAssetsBaseNode + +__all__ = ["ComfyAssetsBaseNode"] \ No newline at end of file diff --git a/kikostats/base/base_node.py b/kikostats/base/base_node.py new file mode 100644 index 0000000..c33d91a --- /dev/null +++ b/kikostats/base/base_node.py @@ -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", ()), + } \ No newline at end of file diff --git a/kikostats/execution_hooks.py b/kikostats/execution_hooks.py new file mode 100644 index 0000000..46f83be --- /dev/null +++ b/kikostats/execution_hooks.py @@ -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 \ No newline at end of file diff --git a/kikostats/tools/__init__.py b/kikostats/tools/__init__.py new file mode 100644 index 0000000..50a6d99 --- /dev/null +++ b/kikostats/tools/__init__.py @@ -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 \ No newline at end of file diff --git a/kikostats/tools/resource_monitor/__init__.py b/kikostats/tools/resource_monitor/__init__.py new file mode 100644 index 0000000..9f94cda --- /dev/null +++ b/kikostats/tools/resource_monitor/__init__.py @@ -0,0 +1,8 @@ +""" +Resource Monitor Tool +Real-time GPU and system resource monitoring for ComfyUI +""" + +from .node import ResourceMonitorNode + +__all__ = ["ResourceMonitorNode"] \ No newline at end of file diff --git a/kikostats/tools/resource_monitor/logic.py b/kikostats/tools/resource_monitor/logic.py new file mode 100644 index 0000000..ffa5136 --- /dev/null +++ b/kikostats/tools/resource_monitor/logic.py @@ -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 \ No newline at end of file diff --git a/kikostats/tools/resource_monitor/node.py b/kikostats/tools/resource_monitor/node.py new file mode 100644 index 0000000..e8701d6 --- /dev/null +++ b/kikostats/tools/resource_monitor/node.py @@ -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", +} \ No newline at end of file