Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add cProfile for profiling graph execution.

Browse Source
This commit is contained in:
Ryan Dick 2024-01-12 10:58:03 -05:00
parent d4c36da3ee
commit e2684b45af
1 changed files with 26 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
invokeai/app/services/invocation_processor/invocation_processor_default.py
View File

@ -1,3 +1,4 @@
import cProfile
import time import time
import traceback import traceback
from threading import BoundedSemaphore, Event, Thread from threading import BoundedSemaphore, Event, Thread
@ -39,6 +40,9 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
self.__threadLimit.acquire() self.__threadLimit.acquire()
queue_item: Optional[InvocationQueueItem] = None queue_item: Optional[InvocationQueueItem] = None
profiler = None
last_gesid = None
while not stop_event.is_set(): while not stop_event.is_set():
try: try:
queue_item = self.__invoker.services.queue.get() queue_item = self.__invoker.services.queue.get()
@ -49,6 +53,21 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
# do not hammer the queue # do not hammer the queue
time.sleep(0.5) time.sleep(0.5)
continue continue
if last_gesid != queue_item.graph_execution_state_id:
if profiler is not None:
# I'm not sure what would cause us to get here, but if we do, we should restart the profiler for
# the new graph_execution_state_id.
profiler.disable()
logger.info(f"Stopped profiler for {last_gesid}.")
profiler = None
last_gesid = None
profiler = cProfile.Profile()
profiler.enable()
last_gesid = queue_item.graph_execution_state_id
logger.info(f"Started profiling {last_gesid}.")
try: try:
graph_execution_state = self.__invoker.services.graph_execution_manager.get( graph_execution_state = self.__invoker.services.graph_execution_manager.get(
queue_item.graph_execution_state_id queue_item.graph_execution_state_id
@ -201,6 +220,13 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
queue_id=queue_item.session_queue_id, queue_id=queue_item.session_queue_id,
graph_execution_state_id=graph_execution_state.id, graph_execution_state_id=graph_execution_state.id,
) )
if profiler is not None:
profiler.disable()
dump_path = f"{last_gesid}.prof"
profiler.dump_stats(dump_path)
logger.info(f"Saved profile to {dump_path}.")
profiler = None
last_gesid = None
except KeyboardInterrupt: except KeyboardInterrupt:
pass # Log something? KeyboardInterrupt is probably not going to be seen by the processor pass # Log something? KeyboardInterrupt is probably not going to be seen by the processor