commit aef8b63ee236de4e9d3804a73c13720fdd30015c Author: Shan2L Date: Wed Jan 28 03:42:17 2026 +0000 feat: introduce excution time record and memory peak record. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9bec2fb --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +outputs/* + +__pycache__ \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..76f8c1b --- /dev/null +++ b/README.md @@ -0,0 +1,34 @@ + + + + _____________ _____ _____ __ + / ___/ __/ __/ / ___/__ _______ / __/ | /| / / + / /___\ \/ _/ / /__/ _ \/ __/ -_)\ \ | |/ |/ / + _____\___/___/___/__\___/\___/_/ _\__/___/_|__/|__/___ __ + /_ __/ _/ |/ / __/ _ \ / _ \/ / / / / / ___/ _/ |/ / + / / _/ // /|_/ / _// , _/ / ___/ /__/ /_/ / (_ // // / + /_/ /___/_/ /_/___/_/|_| /_/ /____/\____/\___/___/_/|_/ + + +# Introduction + +ComfyUI-Simple-Profiler is a pure backend custom node/plugin, which is used to record execution time for each node in a workflow. + +It can also record the maximum device usage used during a give workflow. + +All collected statistical information will be output in a .json or/and .csv file to facilatating thereafter query from web API. + +The profiling information can be fetched after the execution via address '/exec_timer/stat', more detailed use example see: [example.py](./example.py) + + +# How to Use + +1. Clone this repo to the `Path_to_ComfyUI/custom_nodes` + +2. Restart your ComfyUI server. + +3. Run a workflow from UI interface or API request. + +4. The stat data of corresponding workflow will be output to the terminal of your server and become persistent in a .csv/.json file under outputs directory. + + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..767d803 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ + +import importlib + +importlib.import_module('.exec_timer', __name__) + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/example.py b/example.py new file mode 100644 index 0000000..d1efaac --- /dev/null +++ b/example.py @@ -0,0 +1,63 @@ +import time +import requests +import json +import uuid + +COMFY_URL = "http://localhost:8188" +client_id = str(uuid.uuid4()) + +with open("text_to_image.json", "r") as f: + prompt = json.load(f) + +response = requests.post( + f"{COMFY_URL}/prompt", json={"prompt": prompt, "client_id": client_id}, timeout=10 +) + +result = response.json() + +if "prompt_id" in result: + prompt_id = result["prompt_id"] + print(f"The task has been submit prompt ID: {prompt_id}") + + history_url = f"{COMFY_URL}/history/{prompt_id}" + while True: + history_resp = requests.get(history_url) + history_data = history_resp.json() + + if prompt_id in history_data: + outputs = history_data[prompt_id]["outputs"] + print(f"The task has finished, output message: {outputs}") + break + time.sleep(1) + + stats_response = requests.post( + f"{COMFY_URL}/exec-timer/stat", + json={ + "prompt_id": prompt_id, + }, + timeout=10, + ) + + stats_response.raise_for_status() + stats_dict = stats_response.json() + + assert stats_dict["prompt_id"] == prompt_id + resp_msg = stats_dict["msg"] + resp_code = stats_dict["code"] + node_stats = stats_dict["data"] + + TABLE_HEADER = ["GLOBAL_STEP", "CLASS TYPE", "Execution Time(ms)", "Unique ID"] + print( + f"{TABLE_HEADER[0]:<30} {TABLE_HEADER[1]:<30} {TABLE_HEADER[2]:<30} {TABLE_HEADER[3]:<30}" + ) + + for key, value in node_stats.items(): + if key == "peak_memory_usage" or key == "total_time": + print(f"{key:<30} {value:<30}") + else: + print( + f"{key:<30} {value['class_type']:<30} {value['exec_time_ms']:<30} {value['unique_id']:<30}" + ) + +else: + print(f"Failed to submit task, {result.get('error', 'Unknown error')}") diff --git a/exec_timer.py b/exec_timer.py new file mode 100644 index 0000000..128ea6b --- /dev/null +++ b/exec_timer.py @@ -0,0 +1,375 @@ +import time +import server +import execution +import inspect +import os +import json +from aiohttp import web +import csv +import torch + +API_PATH = "/exec-timer/stat" +ENABLE_CSV_GEN = os.environ.get("EXEC_TIMER_ENABLE_CSV_GEN", True) +TIME_PRECISION = int(os.environ.get("EXEC_TIMER_TIME_PRECISION", "2")) +OUTPUT_PATH = os.path.join(os.path.dirname(__file__), "outputs") + +API_ALLOW_ORIGIN = "*" + +NODE_STATS = {} +WORKFLOW_INFO = { + "start_time": None, + "end_time": None, + "prompt_id": None, + "is_running": False, + "peak_memory_usage": None, +} + +GLOBAL_STEP = 1 + + +def init_stats(prompt_id: str): + global NODE_STATS, WORKFLOW_INFO, GLOBAL_STEP + NODE_STATS = {} + WORKFLOW_INFO = { + "start_time": time.perf_counter(), + "end_time": None, + "prompt_id": prompt_id, + "is_running": True, + "peak_memory_usage": None, + } + GLOBAL_STEP = 1 + torch.cuda.reset_max_memory_allocated() + print( + f"\n[Exec_Timer] Workflow begins, prompt_id: {prompt_id}, initialize status..." + ) + + +def reset_stats(): + global NODE_STATS, WORKFLOW_INFO, GLOBAL_STEP + NODE_STATS = {} + WORKFLOW_INFO = { + "start_time": None, + "end_time": None, + "prompt_id": None, + "is_running": False, + "peak_memory_usage": None, + } + GLOBAL_STEP = 1 + torch.cuda.reset_max_memory_allocated() + + +def get_accurate_time(raw_time: float, precision: int): + if precision < 0: + raise ValueError("The precision can't be smaller than zero.") + + if precision == 0: + return int(raw_time) + else: + return round(raw_time, TIME_PRECISION) + + +def generate_csv_file(prompt_id: str): + global WORKFLOW_INFO, NODE_STATS, OUTPUT_PATH + try: + prompt_id = WORKFLOW_INFO["prompt_id"] + csv_filepath = os.path.join(OUTPUT_PATH, f"{prompt_id}.csv") + with open(csv_filepath, "w") as file: + csv_writer = csv.writer(file, delimiter=" ") + csv_writer.writerow( + ["GLOBAL_STEP", "CLASS TYPE", "Execution Time(ms)", "Unique ID"] + ) + for global_step, node_stat in NODE_STATS.items(): + csv_writer.writerow( + [ + global_step, + node_stat["class_type"], + node_stat["exec_time_ms"], + node_stat["unique_id"], + ] + ) + + workflow_execution_time_diff = ( + WORKFLOW_INFO["end_time"] - WORKFLOW_INFO["start_time"] + ) * 1000 + workflow_execution_time_ms = get_accurate_time( + workflow_execution_time_diff, TIME_PRECISION + ) + csv_writer.writerow( + ["Total Execution Time(ms)", workflow_execution_time_ms] + ) + csv_writer.writerow( + ["Peak Memory Usage(Bytes)", WORKFLOW_INFO["peak_memory_usage"]] + ) + + print(f"The csv file has been generated: {csv_filepath}") + except Exception as e: + print(f"[Exec_Timer] Failed to save csv: {str(e)}") + + +def generate_json_file(prompt_id: str): + global NODE_STATS, WORKFLOW_INFO, OUTPUT_PATH + + node_stats_with_mem = NODE_STATS.copy() + workflow_execution_time_diff = ( + WORKFLOW_INFO["end_time"] - WORKFLOW_INFO["start_time"] + ) * 1000 + workflow_execution_time_ms = get_accurate_time( + workflow_execution_time_diff, TIME_PRECISION + ) + node_stats_with_mem["peak_memory_usage"] = WORKFLOW_INFO["peak_memory_usage"] + node_stats_with_mem["total_time"] = workflow_execution_time_ms + + try: + prompt_id = WORKFLOW_INFO["prompt_id"] + json_filepath = os.path.join(OUTPUT_PATH, f"{prompt_id}.json") + with open(json_filepath, "w") as file: + json.dump(node_stats_with_mem, file) + + print(f"The json file has been generated: {json_filepath}") + except Exception as e: + print(f"[Exec_Timer] Failed to save json: {str(e)}") + + +async def send_workflow_timer_record(request): + global NODE_EXEC_LOG, WORKFLOW_INFO, OUTPUT_PATH + post_data = await request.json() + prompt_id = post_data.get("prompt_id") + if prompt_id == WORKFLOW_INFO["prompt_id"] and WORKFLOW_INFO["is_running"]: + return_code = 200 + return_msg = "Succeed to find the target prompt, but it is still running." + node_stats = {} + else: + try: + json_filepath = os.path.join(OUTPUT_PATH, f"{prompt_id}.json") + with open(json_filepath, "r") as file: + node_stats = json.load(file) + return_code = 200 + return_msg = "Succeed to get the workflow execution infomation." + except Exception as e: + print(f"Failed to open {json_filepath}: {str(e)}") + return_code = 500 + return_msg = f"Failed to get status info for {prompt_id}" + node_stats = {} + + response_data = { + "code": return_code, + "msg": return_msg, + "prompt_id": prompt_id, + "data": node_stats, + "timestamp": int(time.time() * 1000), + } + headers = { + "Access-Control-Allow-Origin": API_ALLOW_ORIGIN, + "Access-Control-Allow-Methods": "GET, POST, OPTIONS", + "Access-Control-Allow-Headers": "Content-Type", + } + return web.json_response(response_data, headers=headers) + + +def register_api_route(): + try: + server.PromptServer.instance.app.add_routes( + [ + # web.get(API_PATH, get_workflow_timer_record), + web.post(API_PATH, send_workflow_timer_record), + # web.options(API_PATH, get_workflow_timer_record) + ] + ) + print( + f"[Exec_Timer] API registered | access address: http://0.0.0.0:8188{API_PATH}" + ) + except Exception as e: + print(f"[Exec_Timer] Failed to register API: {str(e)}") + + +origin_execute = execution.execute + +# async +if inspect.iscoroutinefunction(origin_execute): + + async def hooked_execute( + server_inst, + dynprompt, + caches, + current_item, + extra_data, + executed, + prompt_id, + execution_list, + pending_subgraph_results, + pending_async_nodes, + *args, + **kwargs, + ): + global NODE_STATS, WORKFLOW_INFO, GLOBAL_STEP + unique_id = current_item + + if not WORKFLOW_INFO["is_running"]: + init_stats(prompt_id) + + try: + node_info = dynprompt.get_node(unique_id) + class_type = node_info.get("class_type", "Unknown Class") + display_id = dynprompt.get_display_node_id(unique_id) + except Exception: + class_type = "Unknown Class" + display_id = "Unknown Id" + + node_start_time = time.perf_counter() + + result = await origin_execute( + server_inst, + dynprompt, + caches, + current_item, + extra_data, + executed, + prompt_id, + execution_list, + pending_subgraph_results, + pending_async_nodes, + *args, + **kwargs, + ) + + # Caculate the execution time + node_end_time = time.perf_counter() + exec_time_ms = (node_end_time - node_start_time) * 1000 + exec_time_ms = get_accurate_time(exec_time_ms, TIME_PRECISION) + + NODE_STATS[str(GLOBAL_STEP)] = { + "display_id": display_id, + "class_type": class_type, + "start_time": node_start_time, + "exec_time_ms": exec_time_ms, + "unique_id": unique_id, + } + GLOBAL_STEP += 1 + + print( + f"[Exec_Timer] Node Exec Finished | Class Type{class_type:25s} | Front ID: {display_id:5s} | Execution Time: {exec_time_ms} ms | Unique_id: {unique_id}" + ) + + return result +else: + + def hooked_execute( + server_inst, + dynprompt, + caches, + current_item, + extra_data, + executed, + prompt_id, + execution_list, + pending_subgraph_results, + *args, + **kwargs, + ): + global NODE_STATS, WORKFLOW_INFO, GLOBAL_STEP + unique_id = current_item + + if not WORKFLOW_INFO["is_running"]: + init_stats(prompt_id) + + try: + node_info = dynprompt.get_node(unique_id) + class_type = node_info.get("class_type", "Unknown Class") + display_id = dynprompt.get_display_node_id(unique_id) + except Exception: + class_type = "Unknown Class" + display_id = "Unknown ID" + + node_start_time = time.perf_counter() + + result = origin_execute( + server_inst, + dynprompt, + caches, + current_item, + extra_data, + executed, + prompt_id, + execution_list, + pending_subgraph_results, + *args, + **kwargs, + ) + + node_end_time = time.perf_counter() + exec_time_ms = (node_end_time - node_start_time) * 1000 + exec_time_ms = get_accurate_time(exec_time_ms, TIME_PRECISION) + + # 存储统计数据 + NODE_STATS[str(GLOBAL_STEP)] = { + "display_id": display_id, + "class_type": class_type, + "start_time": node_start_time, + "exec_time_ms": exec_time_ms, + "unique_id": unique_id, + } + GLOBAL_STEP += 1 + + print( + f"[Exec_Timer] Node Exec Finished | Class Type{class_type:25s} | Front ID: {display_id:5s} | Execution Time: {exec_time_ms} ms | Unique_id: {unique_id}" + ) + + return result + + +execution.execute = hooked_execute + + +origin_send_sync = server.PromptServer.send_sync + + +def hooked_send_sync(self, event, data, sid=None): + global WORKFLOW_INFO, ENABLE_CSV_GEN + origin_send_sync(self, event, data, sid) + + # The signal that indicate the workflow ends: event = 'executing' && data.node = None + if ( + event == "executing" + and data + and data.get("node") is None + and WORKFLOW_INFO["is_running"] + ): + WORKFLOW_INFO["end_time"] = time.perf_counter() + workflow_execution_time_tmsp = ( + WORKFLOW_INFO["end_time"] - WORKFLOW_INFO["start_time"] + ) * 1000 + workflow_execution_time_ms = round(workflow_execution_time_tmsp, TIME_PRECISION) + print( + f"The total excution time of the workflow {WORKFLOW_INFO['prompt_id']} is {workflow_execution_time_ms}" + ) + + WORKFLOW_INFO["peak_memory_usage"] = torch.cuda.max_memory_allocated() + print( + f"[Exec-Timer] The max memory usage is : {WORKFLOW_INFO['peak_memory_usage']} MB" + ) + prompt_id = WORKFLOW_INFO["prompt_id"] + generate_json_file(prompt_id) + if ENABLE_CSV_GEN: + generate_csv_file(prompt_id) + reset_stats() + + +server.PromptServer.send_sync = hooked_send_sync + +register_api_route() + +print( + """ + _____________ _____ _____ __ + / ___/ __/ __/ / ___/__ _______ / __/ | /| / / + / /___\ \/ _/ / /__/ _ \/ __/ -_)\ \ | |/ |/ / + _____\___/___/___/__\___/\___/_/ _\__/___/_|__/|__/___ __ + /_ __/ _/ |/ / __/ _ \ / _ \/ / / / / / ___/ _/ |/ / + / / _/ // /|_/ / _// , _/ / ___/ /__/ /_/ / (_ // // / + /_/ /___/_/ /_/___/_/|_| /_/ /____/\____/\___/___/_/|_/ + """ +) + +print( + f"[Exec_Timer] The pure backend plugin Exec_Timer load successfully | Generate CSV File:{'True' if ENABLE_CSV_GEN else 'False'} | Time Precision: {TIME_PRECISION}" +)