205 lines
7.5 KiB
Python
205 lines
7.5 KiB
Python
"""Data persistence for ProfilerX — ported from profiler_core.py"""
|
|
import json
|
|
import os
|
|
import time
|
|
import threading
|
|
import logging
|
|
import concurrent.futures
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Optional
|
|
|
|
logger = logging.getLogger('ComfyUI-ProfilerX')
|
|
|
|
class StorageManager:
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __init__(self):
|
|
self.history: List[Dict] = []
|
|
self.max_history = 10000
|
|
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
|
|
}
|
|
|
|
self.data_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data")
|
|
os.makedirs(self.data_dir, exist_ok=True)
|
|
self.history_file = os.path.join(self.data_dir, "profiling_history.json")
|
|
|
|
# Load existing history
|
|
if os.path.exists(self.history_file):
|
|
try:
|
|
with open(self.history_file, 'r') as f:
|
|
self.history = json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load history: {e}")
|
|
self.history = []
|
|
|
|
self._save_counter = 0
|
|
self._save_batch_size = 5
|
|
self._save_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="ProfilerSave")
|
|
|
|
# --- Averages ---
|
|
|
|
def _update_node_average(self, node_type: str, execution_time: float, vram_used: float, ram_used: float) -> Dict:
|
|
avg = self.node_averages[node_type]
|
|
avg['count'] += 1
|
|
n = avg['count']
|
|
avg['total_time'] = (avg['total_time'] * (n - 1) + execution_time) / n
|
|
avg['vram_usage'] = (avg['vram_usage'] * (n - 1) + vram_used) / n
|
|
avg['ram_usage'] = (avg['ram_usage'] * (n - 1) + ram_used) / n
|
|
return avg
|
|
|
|
def _update_workflow_average(self, execution_time: float, vram_peak: float, ram_peak: float) -> Dict:
|
|
wa = self.workflow_averages
|
|
wa['count'] += 1
|
|
n = wa['count']
|
|
wa['total_time'] = (wa['total_time'] * (n - 1) + execution_time) / n
|
|
wa['vram_peak'] = (wa['vram_peak'] * (n - 1) + vram_peak) / n
|
|
wa['ram_peak'] = (wa['ram_peak'] * (n - 1) + ram_peak) / n
|
|
return wa
|
|
|
|
# --- Save / Load ---
|
|
|
|
def _save_history(self, force=False):
|
|
if not force:
|
|
self._save_counter += 1
|
|
if self._save_counter < self._save_batch_size:
|
|
return
|
|
self._save_counter = 0
|
|
self._save_executor.submit(self._sync_save_history)
|
|
|
|
def _sync_save_history(self):
|
|
try:
|
|
with open(self.history_file, 'w') as f:
|
|
json.dump(self.history[-self.max_history:], f, indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Failed to save history: {e}")
|
|
|
|
# --- Run storage ---
|
|
|
|
def save_run(self, run_data: Dict):
|
|
"""Save a completed run and update averages."""
|
|
with self._lock:
|
|
# Update node averages
|
|
for node in run_data.get('nodes', {}).values():
|
|
if not node.get('cacheHit', False):
|
|
self._update_node_average(
|
|
node.get('nodeType', 'unknown'),
|
|
node.get('endTime', 0) - node.get('startTime', 0),
|
|
node.get('vramAfter', 0) - node.get('vramBefore', 0),
|
|
node.get('ramAfter', 0) - node.get('ramBefore', 0),
|
|
)
|
|
|
|
# Update workflow averages
|
|
exec_time = run_data.get('endTime', 0) - run_data.get('startTime', 0)
|
|
self._update_workflow_average(
|
|
exec_time,
|
|
run_data.get('totalVramPeak', 0),
|
|
run_data.get('totalRamPeak', 0),
|
|
)
|
|
|
|
# Inject averages into run_data for frontend
|
|
run_data['averages'] = {
|
|
'execution_time': self.workflow_averages['total_time'],
|
|
'vram_peak': self.workflow_averages['vram_peak'],
|
|
'ram_peak': self.workflow_averages['ram_peak'],
|
|
'count': self.workflow_averages['count'],
|
|
}
|
|
for node in run_data.get('nodes', {}).values():
|
|
nt = node.get('nodeType', 'unknown')
|
|
if nt in self.node_averages:
|
|
a = self.node_averages[nt]
|
|
node['averages'] = {
|
|
'execution_time': a['total_time'],
|
|
'vram_usage': a['vram_usage'],
|
|
'ram_usage': a['ram_usage'],
|
|
'count': a['count'],
|
|
}
|
|
|
|
self.history.append(run_data)
|
|
if len(self.history) >= self.max_history:
|
|
self.archive_history()
|
|
else:
|
|
self._save_history()
|
|
|
|
def get_stats(self) -> Dict:
|
|
with self._lock:
|
|
return {
|
|
'current': {},
|
|
'latest': self.history[-1] if self.history else None,
|
|
'node_averages': dict(self.node_averages),
|
|
'workflow_averages': self.workflow_averages,
|
|
'history': self.history[-10:],
|
|
}
|
|
|
|
# --- Archives ---
|
|
|
|
def _archive_dir(self):
|
|
d = os.path.join(self.data_dir, "archives")
|
|
os.makedirs(d, exist_ok=True)
|
|
return d
|
|
|
|
def get_archives(self) -> List[Dict]:
|
|
archives = []
|
|
try:
|
|
for fn in os.listdir(self._archive_dir()):
|
|
if fn.endswith('.json'):
|
|
path = os.path.join(self._archive_dir(), fn)
|
|
stat = os.stat(path)
|
|
archives.append({
|
|
'filename': fn, '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]:
|
|
if not self.history:
|
|
return None
|
|
try:
|
|
fn = f"profiling_history_{int(time.time())}.json"
|
|
path = os.path.join(self._archive_dir(), fn)
|
|
with open(path, 'w') as f:
|
|
json.dump(self.history, f, indent=2)
|
|
self.history = []
|
|
self._sync_save_history()
|
|
return fn
|
|
except Exception as e:
|
|
logger.error(f"Failed to create archive: {e}")
|
|
return None
|
|
|
|
def load_archive(self, filename: str) -> bool:
|
|
path = os.path.join(self._archive_dir(), filename)
|
|
if not os.path.exists(path):
|
|
return False
|
|
try:
|
|
if self.history:
|
|
self.archive_history()
|
|
with open(path, 'r') as f:
|
|
data = json.load(f)
|
|
if not isinstance(data, list):
|
|
return False
|
|
self.history = data
|
|
self._sync_save_history()
|
|
os.remove(path)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to load archive: {e}")
|
|
return False
|
|
|
|
def delete_archive(self, filename: str) -> bool:
|
|
path = os.path.join(self._archive_dir(), filename)
|
|
if not os.path.exists(path):
|
|
return False
|
|
try:
|
|
os.remove(path)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete archive: {e}")
|
|
return False
|