From 79ae9c4e64cb4f64d54c25b1c487501752b8fa84 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 6 Feb 2024 17:48:54 +1100 Subject: [PATCH] feat(nodes): move profiler/stats cleanup logic to function Harder to miss something going forward. --- .../invocation_processor_default.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/invokeai/app/services/invocation_processor/invocation_processor_default.py b/invokeai/app/services/invocation_processor/invocation_processor_default.py index a689a0a058..54342c0da1 100644 --- a/invokeai/app/services/invocation_processor/invocation_processor_default.py +++ b/invokeai/app/services/invocation_processor/invocation_processor_default.py @@ -54,6 +54,17 @@ class DefaultInvocationProcessor(InvocationProcessorABC): else None ) + def stats_cleanup(graph_execution_state_id: str) -> None: + if profiler: + profile_path = profiler.stop() + stats_path = profile_path.with_suffix(".json") + self.__invoker.services.performance_statistics.dump_stats( + graph_execution_state_id=graph_execution_state_id, output_path=stats_path + ) + with suppress(GESStatsNotFoundError): + self.__invoker.services.performance_statistics.log_stats(graph_execution_state_id) + self.__invoker.services.performance_statistics.reset_stats(graph_execution_state_id) + while not stop_event.is_set(): try: queue_item = self.__invoker.services.queue.get() @@ -156,15 +167,7 @@ class DefaultInvocationProcessor(InvocationProcessorABC): pass except CanceledException: - with suppress(GESStatsNotFoundError): - if profiler: - profile_path = profiler.stop() - stats_path = profile_path.with_suffix(".json") - self.__invoker.services.performance_statistics.dump_stats( - graph_execution_state_id=graph_execution_state.id, output_path=stats_path - ) - self.__invoker.services.performance_statistics.log_stats(graph_execution_state.id) - self.__invoker.services.performance_statistics.reset_stats(graph_execution_state.id) + stats_cleanup(graph_execution_state.id) pass except Exception as e: @@ -226,15 +229,7 @@ class DefaultInvocationProcessor(InvocationProcessorABC): queue_id=queue_item.session_queue_id, graph_execution_state_id=graph_execution_state.id, ) - if profiler: - profile_path = profiler.stop() - stats_path = profile_path.with_suffix(".json") - self.__invoker.services.performance_statistics.dump_stats( - graph_execution_state_id=graph_execution_state.id, output_path=stats_path - ) - with suppress(GESStatsNotFoundError): - self.__invoker.services.performance_statistics.log_stats(graph_execution_state.id) - self.__invoker.services.performance_statistics.reset_stats(graph_execution_state.id) + stats_cleanup(graph_execution_state.id) except KeyboardInterrupt: pass # Log something? KeyboardInterrupt is probably not going to be seen by the processor