From 680c9522600fdb4f3be97c777545e10180ec72e1 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Mon, 3 Feb 2025 17:37:51 -0500 Subject: [PATCH] feat: add stand alone cProfile utility additionally includes a utility to combine separate prof files into one. --- utilities/ComfyUI_ExecutionTracker.py | 126 ++++++++++++++++++++++++++ utilities/combine_profiles.py | 47 ++++++++++ 2 files changed, 173 insertions(+) create mode 100644 utilities/ComfyUI_ExecutionTracker.py create mode 100644 utilities/combine_profiles.py diff --git a/utilities/ComfyUI_ExecutionTracker.py b/utilities/ComfyUI_ExecutionTracker.py new file mode 100644 index 0000000..b31af20 --- /dev/null +++ b/utilities/ComfyUI_ExecutionTracker.py @@ -0,0 +1,126 @@ +""" +Dedicated execution tracking for ComfyUI using cProfile + +Drop this into your custom nodes folder and restart ComfyUI. +You may or may not need to disable ProfilerX to use this, as they may conflict with one another. + +""" +import os +import sys +import cProfile +import logging +import threading +from pathlib import Path + +# Add ComfyUI root to path +comfy_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) +if comfy_path not in sys.path: + sys.path.append(comfy_path) + +import execution + +logger = logging.getLogger('ComfyUI-ExecutionTracker') +logger.setLevel(logging.ERROR) + +class ExecutionTracker: + """Tracks ComfyUI's internal method execution using cProfile""" + _instance = None + _lock = threading.Lock() + ENABLED = True + + def __init__(self): + self.data_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data") + os.makedirs(self.data_dir, exist_ok=True) + self.profiler = None + self._current_prompt_id = None + + @classmethod + def get_instance(cls): + """Get singleton instance""" + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def start_execution(self, prompt_id: str): + """Start profiling a new execution""" + if not self.ENABLED: + return + + with self._lock: + self._current_prompt_id = prompt_id + self.profiler = cProfile.Profile() + self.profiler.enable() + + def end_execution(self): + """End profiling current execution""" + if not self.ENABLED or not self.profiler: + return + + with self._lock: + self.profiler.disable() + # Save as .prof file that snakeviz can read + output_file = os.path.join(self.data_dir, f"execution_{self._current_prompt_id}.prof") + self.profiler.dump_stats(output_file) + self.profiler = None + self._current_prompt_id = None + +# Store original PromptExecutor.execute to wrap it with profiling +original_execute = execution.PromptExecutor.execute + +def execute_with_profiling(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): + """Profile workflow execution""" + tracker = ExecutionTracker.get_instance() + tracker.start_execution(prompt_id) + try: + return original_execute(self, prompt, prompt_id, extra_data, execute_outputs) + finally: + tracker.end_execution() + +def inject_tracking(): + """Inject profiling hook into PromptExecutor.execute""" + logger.info("Verifying ComfyUI components...") + try: + if not hasattr(execution, 'PromptExecutor'): + raise ImportError("Required ComfyUI components not found - incompatible ComfyUI version?") + + # Only need to wrap the main execute method + execution.PromptExecutor.execute = execute_with_profiling + + logger.info("✓ Profiling hook injected successfully") + return True + + except Exception as e: + logger.error(f"❌ Failed to inject profiling hook: {str(e)}") + logger.error("Execution tracking will be disabled for this session") + return False + +# Dummy node for ComfyUI +class ExecutionTrackerNode: + """Dummy node to satisfy ComfyUI's node requirements""" + + @classmethod + def INPUT_TYPES(s): + return {"required": {}} + + RETURN_TYPES = () + FUNCTION = "noop" + OUTPUT_NODE = True + CATEGORY = "profiling" + + def noop(self): + return {} + +# Node class mappings for ComfyUI +NODE_CLASS_MAPPINGS = { + "ExecutionTracker": ExecutionTrackerNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ExecutionTracker": "Execution Tracker" +} + +# Inject tracking on import if enabled +if ExecutionTracker.ENABLED: + inject_tracking() diff --git a/utilities/combine_profiles.py b/utilities/combine_profiles.py new file mode 100644 index 0000000..8fd407c --- /dev/null +++ b/utilities/combine_profiles.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +"""Utility script to combine multiple .prof files into one""" +import os +import glob +import pstats +import argparse +from pathlib import Path + +def combine_profiles(data_dir, output_file, pattern="execution_*.prof"): + """Combine all matching .prof files in data_dir into one""" + combined_stats = pstats.Stats() + + # Find all matching .prof files + prof_files = glob.glob(os.path.join(data_dir, pattern)) + if not prof_files: + print(f"No .prof files found matching pattern: {pattern}") + return None + + try: + # Add each file's stats to the combined stats + for prof_file in prof_files: + print(f"Adding {prof_file}") + combined_stats.add(prof_file) + + # Save combined stats + output_path = os.path.join(data_dir, output_file) + combined_stats.dump_stats(output_path) + print(f"Successfully combined {len(prof_files)} profiles into {output_file}") + return output_path + + except Exception as e: + print(f"Error combining profiles: {e}") + return None + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Combine multiple .prof files into one") + parser.add_argument("--data-dir", default="data", help="Directory containing .prof files") + parser.add_argument("--output", default="combined_executions.prof", help="Output filename") + parser.add_argument("--pattern", default="execution_*.prof", help="Glob pattern for files to combine") + + args = parser.parse_args() + + # Resolve data directory relative to script location + script_dir = os.path.dirname(os.path.realpath(__file__)) + data_dir = os.path.join(script_dir, args.data_dir) + + combine_profiles(data_dir, args.output, args.pattern) \ No newline at end of file