132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
"""Core profiling handler using ComfyUI's ProgressHandler API"""
|
|
import time
|
|
import logging
|
|
from comfy_execution.progress import ProgressHandler
|
|
from .metrics import snapshot_memory, reset_peak, get_peak
|
|
from .storage import StorageManager
|
|
|
|
logger = logging.getLogger('ComfyUI-ProfilerX')
|
|
|
|
|
|
class ProfilerXProgressHandler(ProgressHandler):
|
|
def __init__(self, storage: StorageManager):
|
|
super().__init__("profilerx")
|
|
self.storage = storage
|
|
self.registry = None
|
|
# Current run state
|
|
self._current_run = None
|
|
self._started_nodes = set() # Track which nodes got start_handler
|
|
|
|
def set_registry(self, registry):
|
|
self.registry = registry
|
|
|
|
def reset(self):
|
|
"""Called when progress registry resets (new execution starting).
|
|
Finalize previous run if any, then clear state."""
|
|
if self._current_run is not None:
|
|
self._finalize_run()
|
|
self._current_run = None
|
|
self._started_nodes = set()
|
|
|
|
def _init_run(self, prompt_id: str):
|
|
"""Lazily initialize a new run on first node event."""
|
|
if self._current_run is not None and 'endTime' not in self._current_run:
|
|
return
|
|
# Either no run, or previous run was already flushed — start fresh
|
|
self._current_run = None
|
|
self._started_nodes = set()
|
|
self._current_run = {
|
|
'promptId': prompt_id,
|
|
'startTime': time.time() * 1000,
|
|
'nodes': {},
|
|
'executionOrder': [],
|
|
'totalVramPeak': 0,
|
|
'totalRamPeak': 0,
|
|
'cacheHits': 0,
|
|
'cacheMisses': 0,
|
|
}
|
|
|
|
def _finalize_run(self):
|
|
"""Finalize and save the current run."""
|
|
run = self._current_run
|
|
if run is None:
|
|
return
|
|
if 'endTime' in run:
|
|
return # Already finalized
|
|
run['endTime'] = time.time() * 1000
|
|
vram_peak, ram_peak = get_peak()
|
|
run['totalVramPeak'] = vram_peak
|
|
run['totalRamPeak'] = ram_peak
|
|
logger.info(f"Finalizing run {run.get('promptId', '?')}: {len(run.get('nodes', {}))} nodes, "
|
|
f"{run.get('cacheMisses', 0)} executed, {run.get('cacheHits', 0)} cached")
|
|
self.storage.save_run(run)
|
|
|
|
def flush(self):
|
|
"""Flush current run to storage if it has data. Called by stats endpoint."""
|
|
if self._current_run is not None and self._current_run.get('nodes'):
|
|
self._finalize_run()
|
|
|
|
def _get_class_type(self, node_id: str) -> str:
|
|
"""Get class_type for a node from the dynamic prompt."""
|
|
try:
|
|
if self.registry and self.registry.dynprompt:
|
|
node_info = self.registry.dynprompt.get_node(node_id)
|
|
if node_info:
|
|
return node_info.get('class_type', 'unknown')
|
|
except Exception:
|
|
pass
|
|
return 'unknown'
|
|
|
|
def start_handler(self, node_id: str, state, prompt_id: str):
|
|
"""Called when a node starts executing (not called for cached nodes)."""
|
|
self._init_run(prompt_id)
|
|
self._started_nodes.add(node_id)
|
|
|
|
class_type = self._get_class_type(node_id)
|
|
vram_before, ram_before = snapshot_memory()
|
|
reset_peak()
|
|
|
|
self._current_run['nodes'][node_id] = {
|
|
'nodeId': node_id,
|
|
'nodeType': class_type,
|
|
'startTime': time.time() * 1000,
|
|
'vramBefore': vram_before,
|
|
'ramBefore': ram_before,
|
|
'cacheHit': False,
|
|
}
|
|
self._current_run['executionOrder'].append(node_id)
|
|
|
|
def finish_handler(self, node_id: str, state, prompt_id: str):
|
|
"""Called when a node finishes. If start_handler was never called, it's a cache hit."""
|
|
self._init_run(prompt_id)
|
|
|
|
if node_id not in self._started_nodes:
|
|
# Cache hit — finish called without start
|
|
class_type = self._get_class_type(node_id)
|
|
self._current_run['nodes'][node_id] = {
|
|
'nodeId': node_id,
|
|
'nodeType': class_type,
|
|
'startTime': 0,
|
|
'endTime': 0,
|
|
'vramBefore': 0,
|
|
'vramAfter': 0,
|
|
'vramPeak': 0,
|
|
'ramBefore': 0,
|
|
'ramAfter': 0,
|
|
'cacheHit': True,
|
|
}
|
|
self._current_run['executionOrder'].append(node_id)
|
|
self._current_run['cacheHits'] += 1
|
|
else:
|
|
# Normal execution finish
|
|
node = self._current_run['nodes'].get(node_id)
|
|
if node is None:
|
|
return
|
|
node['endTime'] = time.time() * 1000
|
|
vram_after, ram_after = snapshot_memory()
|
|
vram_peak, _ = get_peak()
|
|
node['vramAfter'] = vram_after
|
|
node['vramPeak'] = vram_peak - node['vramBefore']
|
|
node['ramAfter'] = ram_after
|
|
self._current_run['cacheMisses'] += 1
|