406 lines
16 KiB
Python
406 lines
16 KiB
Python
import time
|
|
import psutil
|
|
import torch
|
|
import threading
|
|
import json
|
|
import os
|
|
import logging
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Optional
|
|
|
|
logger = logging.getLogger('ComfyUI-ProfilerX')
|
|
logger.setLevel(logging.ERROR)
|
|
|
|
# Add console handler if not already added
|
|
if not logger.handlers:
|
|
ch = logging.StreamHandler()
|
|
ch.setLevel(logging.ERROR)
|
|
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
ch.setFormatter(formatter)
|
|
logger.addHandler(ch)
|
|
|
|
class ProfilerManager:
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __init__(self):
|
|
self.active_profiles: Dict[str, Dict] = {}
|
|
self.history: List[Dict] = []
|
|
self.max_history = 10000
|
|
self.process = psutil.Process()
|
|
|
|
# Run benchmark
|
|
logger.setLevel(logging.DEBUG) # Temporarily set to debug to see benchmark
|
|
self._benchmark_reset_stats()
|
|
logger.setLevel(logging.ERROR) # Reset to normal level
|
|
|
|
# Rolling averages for nodes and workflow
|
|
self.node_averages: Dict[str, Dict] = defaultdict(lambda: {
|
|
'total_time': 0.0,
|
|
'count': 0,
|
|
'vram_usage': 0.0,
|
|
'ram_usage': 0.0
|
|
})
|
|
self.workflow_averages = {
|
|
'total_time': 0.0,
|
|
'count': 0,
|
|
'vram_peak': 0.0,
|
|
'ram_peak': 0.0
|
|
}
|
|
|
|
# Create data directory if it doesn't exist
|
|
self.data_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data")
|
|
os.makedirs(self.data_dir, exist_ok=True)
|
|
logger.debug(f"Data directory created/verified at: {self.data_dir}")
|
|
|
|
# Load existing history
|
|
self.history_file = os.path.join(self.data_dir, "profiling_history.json")
|
|
if os.path.exists(self.history_file):
|
|
try:
|
|
with open(self.history_file, 'r') as f:
|
|
self.history = json.load(f)
|
|
logger.debug(f"Loaded {len(self.history)} profiles from history file")
|
|
except Exception as e:
|
|
logger.error(f"Failed to load history file: {e}")
|
|
self.history = []
|
|
else:
|
|
logger.debug("No existing history file found")
|
|
self.history = []
|
|
|
|
# Batching and async I/O for performance
|
|
self._save_counter = 0
|
|
self._save_batch_size = 5 # Only save every 5 workflows
|
|
self._save_executor = None
|
|
import concurrent.futures
|
|
self._save_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="ProfilerSave")
|
|
|
|
def _update_node_average(self, node_type: str, execution_time: float, vram_used: float, ram_used: float) -> Dict:
|
|
"""Update rolling average for a node type"""
|
|
avg = self.node_averages[node_type]
|
|
avg['count'] += 1
|
|
avg['total_time'] = (avg['total_time'] * (avg['count'] - 1) + execution_time) / avg['count']
|
|
avg['vram_usage'] = (avg['vram_usage'] * (avg['count'] - 1) + vram_used) / avg['count']
|
|
avg['ram_usage'] = (avg['ram_usage'] * (avg['count'] - 1) + ram_used) / avg['count']
|
|
return avg
|
|
|
|
def _update_workflow_average(self, execution_time: float, vram_peak: float, ram_peak: float) -> Dict:
|
|
"""Update rolling average for workflow execution"""
|
|
self.workflow_averages['count'] += 1
|
|
self.workflow_averages['total_time'] = (
|
|
self.workflow_averages['total_time'] * (self.workflow_averages['count'] - 1) + execution_time
|
|
) / self.workflow_averages['count']
|
|
self.workflow_averages['vram_peak'] = (
|
|
self.workflow_averages['vram_peak'] * (self.workflow_averages['count'] - 1) + vram_peak
|
|
) / self.workflow_averages['count']
|
|
self.workflow_averages['ram_peak'] = (
|
|
self.workflow_averages['ram_peak'] * (self.workflow_averages['count'] - 1) + ram_peak
|
|
) / self.workflow_averages['count']
|
|
return self.workflow_averages
|
|
|
|
def _save_history(self, force=False):
|
|
"""Save profiling history to disk (batched and async for performance)"""
|
|
if not force:
|
|
self._save_counter += 1
|
|
if self._save_counter < self._save_batch_size:
|
|
return # Skip saving until batch size reached
|
|
self._save_counter = 0 # Reset counter
|
|
|
|
# Submit async save operation
|
|
if self._save_executor:
|
|
self._save_executor.submit(self._async_save_history)
|
|
else:
|
|
self._sync_save_history()
|
|
|
|
def _async_save_history(self):
|
|
"""Async save operation"""
|
|
try:
|
|
self._sync_save_history()
|
|
except Exception as e:
|
|
logger.error(f"Async save failed: {e}")
|
|
|
|
def _sync_save_history(self):
|
|
"""Synchronous save operation"""
|
|
try:
|
|
with open(self.history_file, 'w') as f:
|
|
json.dump(self.history[-self.max_history:], f, indent=2)
|
|
logger.debug(f"Saved {len(self.history)} profiles to history file")
|
|
except Exception as e:
|
|
logger.error(f"Failed to save history file: {e}")
|
|
|
|
@classmethod
|
|
def get_instance(cls) -> 'ProfilerManager':
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
logger.debug("Created new ProfilerManager instance")
|
|
return cls._instance
|
|
|
|
def start_workflow(self, prompt_id: str) -> None:
|
|
"""Start profiling a workflow execution"""
|
|
logger.debug(f"Starting workflow profiling for prompt_id: {prompt_id}")
|
|
self.active_profiles[prompt_id] = {
|
|
'promptId': prompt_id,
|
|
'startTime': time.time() * 1000, # Convert to ms
|
|
'nodes': {},
|
|
'executionOrder': [],
|
|
'totalVramPeak': 0,
|
|
'totalRamPeak': 0,
|
|
'cacheHits': 0,
|
|
'cacheMisses': 0
|
|
}
|
|
|
|
def end_workflow(self, prompt_id: str) -> Optional[Dict]:
|
|
"""End profiling a workflow execution"""
|
|
if prompt_id not in self.active_profiles:
|
|
logger.warning(f"Attempted to end non-existent workflow profile: {prompt_id}")
|
|
return None
|
|
|
|
logger.debug(f"Ending workflow profiling for prompt_id: {prompt_id}")
|
|
profile = self.active_profiles[prompt_id]
|
|
profile['endTime'] = time.time() * 1000
|
|
|
|
# Update peak memory usage
|
|
profile['totalVramPeak'] = torch.cuda.max_memory_allocated()
|
|
profile['totalRamPeak'] = self.process.memory_info().rss
|
|
|
|
# Calculate and update workflow averages
|
|
execution_time = profile['endTime'] - profile['startTime']
|
|
avg = self._update_workflow_average(execution_time, profile['totalVramPeak'], profile['totalRamPeak'])
|
|
profile['averages'] = {
|
|
'execution_time': avg['total_time'],
|
|
'vram_peak': avg['vram_peak'],
|
|
'ram_peak': avg['ram_peak'],
|
|
'count': avg['count']
|
|
}
|
|
|
|
# Store in history
|
|
self.history.append(profile)
|
|
if len(self.history) >= self.max_history:
|
|
# Auto-archive when limit is reached
|
|
logger.info(f"History limit of {self.max_history} reached. Auto-archiving...")
|
|
self.archive_history()
|
|
# History is now empty after archiving
|
|
else:
|
|
# Only save to current history file if we haven't archived
|
|
self._save_history()
|
|
|
|
# Cleanup
|
|
del self.active_profiles[prompt_id]
|
|
return profile
|
|
|
|
def start_node(self, prompt_id: str, node_id: str, node_type: str, inputs: Dict) -> None:
|
|
"""Start profiling a node execution"""
|
|
if prompt_id not in self.active_profiles:
|
|
logger.warning(f"Attempted to start node profiling for non-existent workflow: {prompt_id}")
|
|
return
|
|
|
|
logger.debug(f"Starting node profiling - prompt: {prompt_id}, node: {node_id}, type: {node_type}")
|
|
profile = self.active_profiles[prompt_id]
|
|
|
|
# Reset peak stats to track this node's peak specifically
|
|
torch.cuda.reset_peak_memory_stats()
|
|
base_vram = torch.cuda.memory_allocated() # Store base VRAM to calculate true peak increase
|
|
|
|
profile['nodes'][node_id] = {
|
|
'nodeId': node_id,
|
|
'nodeType': node_type,
|
|
'startTime': time.time() * 1000,
|
|
'vramBefore': base_vram, # This is also our base VRAM
|
|
'ramBefore': self.process.memory_info().rss,
|
|
'inputSizes': self._get_tensor_sizes(inputs),
|
|
'outputSizes': {},
|
|
'cacheHit': False
|
|
}
|
|
profile['executionOrder'].append(node_id)
|
|
|
|
def end_node(self, prompt_id: str, node_id: str, outputs: Dict, cache_hit: bool = False) -> None:
|
|
"""End profiling a node execution"""
|
|
if prompt_id not in self.active_profiles:
|
|
logger.warning(f"Attempted to end node profiling for non-existent workflow: {prompt_id}")
|
|
return
|
|
if node_id not in self.active_profiles[prompt_id]['nodes']:
|
|
logger.warning(f"Attempted to end non-existent node profile: {node_id}")
|
|
return
|
|
|
|
logger.debug(f"Ending node profiling - prompt: {prompt_id}, node: {node_id}, cache_hit: {cache_hit}")
|
|
profile = self.active_profiles[prompt_id]
|
|
node = profile['nodes'][node_id]
|
|
node['endTime'] = time.time() * 1000
|
|
node['vramAfter'] = torch.cuda.memory_allocated()
|
|
total_peak = torch.cuda.max_memory_allocated()
|
|
node['vramPeak'] = total_peak - node['vramBefore'] # Calculate the actual peak increase from base
|
|
node['ramAfter'] = self.process.memory_info().rss
|
|
node['outputSizes'] = self._get_tensor_sizes(outputs)
|
|
node['cacheHit'] = cache_hit
|
|
|
|
# Calculate and update averages
|
|
execution_time = node['endTime'] - node['startTime']
|
|
vram_used = node['vramAfter'] - node['vramBefore']
|
|
ram_used = node['ramAfter'] - node['ramBefore']
|
|
|
|
# Update rolling averages with the true peak
|
|
avg = self._update_node_average(node['nodeType'], execution_time, vram_used, ram_used)
|
|
node['averages'] = {
|
|
'execution_time': avg['total_time'],
|
|
'vram_usage': avg['vram_usage'],
|
|
'ram_usage': avg['ram_usage'],
|
|
'count': avg['count']
|
|
}
|
|
|
|
if cache_hit:
|
|
profile['cacheHits'] += 1
|
|
else:
|
|
profile['cacheMisses'] += 1
|
|
|
|
def record_error(self, prompt_id: str, node_id: str, error: str) -> None:
|
|
"""Record an error that occurred during node execution"""
|
|
if prompt_id not in self.active_profiles:
|
|
logger.warning(f"Attempted to record error for non-existent workflow: {prompt_id}")
|
|
return
|
|
if node_id not in self.active_profiles[prompt_id]['nodes']:
|
|
logger.warning(f"Attempted to record error for non-existent node: {node_id}")
|
|
return
|
|
|
|
logger.error(f"Node error - prompt: {prompt_id}, node: {node_id}, error: {error}")
|
|
self.active_profiles[prompt_id]['nodes'][node_id]['error'] = str(error)
|
|
|
|
def get_latest_profile(self) -> Optional[Dict]:
|
|
"""Get the most recent workflow profile"""
|
|
if not self.history:
|
|
logger.debug("No profiles in history")
|
|
return None
|
|
logger.debug("Returning latest profile")
|
|
return self.history[-1]
|
|
|
|
def get_stats(self) -> Dict:
|
|
"""Get all profiling stats including current and historical data"""
|
|
stats = {
|
|
'current': self.active_profiles,
|
|
'latest': self.get_latest_profile(),
|
|
'node_averages': dict(self.node_averages),
|
|
'workflow_averages': self.workflow_averages,
|
|
'history': self.history[-10:] # Return last 10 profiles
|
|
}
|
|
return stats
|
|
|
|
def get_archives(self) -> List[Dict]:
|
|
"""Get list of archived history files"""
|
|
archives = []
|
|
archive_dir = os.path.join(self.data_dir, "archives")
|
|
os.makedirs(archive_dir, exist_ok=True)
|
|
|
|
try:
|
|
for filename in os.listdir(archive_dir):
|
|
if filename.endswith('.json'):
|
|
path = os.path.join(archive_dir, filename)
|
|
stat = os.stat(path)
|
|
archives.append({
|
|
'filename': filename,
|
|
'size': stat.st_size,
|
|
'created': stat.st_ctime,
|
|
'modified': stat.st_mtime
|
|
})
|
|
return sorted(archives, key=lambda x: x['created'], reverse=True)
|
|
except Exception as e:
|
|
logger.error(f"Failed to get archives: {e}")
|
|
return []
|
|
|
|
def archive_history(self) -> Optional[str]:
|
|
"""Archive current history to a file"""
|
|
if not self.history:
|
|
logger.warning("No history to archive")
|
|
return None
|
|
|
|
archive_dir = os.path.join(self.data_dir, "archives")
|
|
os.makedirs(archive_dir, exist_ok=True)
|
|
|
|
try:
|
|
timestamp = int(time.time())
|
|
filename = f"profiling_history_{timestamp}.json"
|
|
path = os.path.join(archive_dir, filename)
|
|
|
|
with open(path, 'w') as f:
|
|
json.dump(self.history, f, indent=2)
|
|
|
|
# Clear current history and force synchronous save
|
|
self.history = []
|
|
self._sync_save_history() # Force sync save for critical operation
|
|
|
|
logger.debug(f"Created archive: {filename}")
|
|
return filename
|
|
except Exception as e:
|
|
logger.error(f"Failed to create archive: {e}")
|
|
return None
|
|
|
|
def load_archive(self, filename: str) -> bool:
|
|
"""Load history from an archive file and delete it after loading"""
|
|
path = os.path.join(self.data_dir, "archives", filename)
|
|
if not os.path.exists(path):
|
|
logger.error(f"Archive not found: {filename}")
|
|
return False
|
|
|
|
try:
|
|
# Auto-archive current history if it exists
|
|
if self.history:
|
|
self.archive_history()
|
|
logger.debug("Auto-archived current history before loading new archive")
|
|
|
|
# Load the archive
|
|
with open(path, 'r') as f:
|
|
archived_history = json.load(f)
|
|
|
|
if not isinstance(archived_history, list):
|
|
logger.error(f"Invalid archive format: {filename}")
|
|
return False
|
|
|
|
# Update current history
|
|
self.history = archived_history
|
|
self._sync_save_history() # Force sync save for critical operation
|
|
|
|
# Delete the archive file since it's now loaded
|
|
os.remove(path)
|
|
logger.debug(f"Loaded and deleted archive: {filename}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to load archive: {e}")
|
|
return False
|
|
|
|
def delete_archive(self, filename: str) -> bool:
|
|
"""Delete an archived history file"""
|
|
path = os.path.join(self.data_dir, "archives", filename)
|
|
if not os.path.exists(path):
|
|
logger.error(f"Archive not found: {filename}")
|
|
return False
|
|
|
|
try:
|
|
os.remove(path)
|
|
logger.debug(f"Deleted archive: {filename}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete archive: {e}")
|
|
return False
|
|
|
|
def _get_tensor_sizes(self, data):
|
|
"""Get sizes of tensors in the data"""
|
|
if data is None:
|
|
return {}
|
|
|
|
sizes = {}
|
|
if isinstance(data, list):
|
|
for i, value in enumerate(data):
|
|
if isinstance(value, torch.Tensor):
|
|
sizes[f"output_{i}"] = list(value.shape)
|
|
elif isinstance(value, (list, tuple)):
|
|
sizes[f"output_{i}"] = [len(value)]
|
|
return sizes
|
|
|
|
def _benchmark_reset_stats(self, iterations=1000):
|
|
"""Benchmark the overhead of reset_peak_memory_stats"""
|
|
start = time.perf_counter_ns()
|
|
for _ in range(iterations):
|
|
torch.cuda.reset_peak_memory_stats()
|
|
end = time.perf_counter_ns()
|
|
avg_ns = (end - start) / iterations
|
|
logger.debug(f"Average reset_peak_memory_stats time: {avg_ns:.2f} nanoseconds")
|
|
return avg_ns |