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

reset stats on exception

Browse Source
This commit is contained in:
Lincoln Stein 2023-08-01 19:39:42 -04:00
parent fd7b842419
commit 8a4e5f73aa
2 changed files with 21 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
invokeai/app/services/invocation_stats.py
View File

@ -43,14 +43,15 @@ from ..invocations.baseinvocation import BaseInvocation
import invokeai.backend.util.logging as logger
class InvocationStats():
class InvocationStats:
"""Accumulate performance information about a running graph. Collects time spent in each node,
as well as the maximum and current VRAM utilisation for CUDA systems"""
def __init__(self):
self._stats: Dict[str, int] = {}
class StatsContext():
class StatsContext:
def __init__(self, invocation: BaseInvocation, collector):
self.invocation = invocation
self.collector = collector
@ -61,17 +62,18 @@ class InvocationStats():
def __exit__(self, *args):
self.collector.log_time(self.invocation.type, time.time() - self.start_time)
def collect_stats(self,
invocation: BaseInvocation,
graph_execution_state: GraphExecutionState,
) -> StatsContext:
def collect_stats(
self,
invocation: BaseInvocation,
graph_execution_state: GraphExecutionState,
) -> StatsContext:
"""
Return a context object that will capture the statistics.
:param invocation: BaseInvocation object from the current graph.
:param graph_execution_state: GraphExecutionState object from the current session.
"""
if len(graph_execution_state.executed)==0: # new graph is starting
if len(graph_execution_state.executed) == 0: # new graph is starting
self.reset_stats()
self._current_graph_state = graph_execution_state
sc = self.StatsContext(invocation, self)
@ -83,7 +85,6 @@ class InvocationStats():
torch.cuda.reset_peak_memory_stats()
self._stats: Dict[str, List[int, float]] = {}
def log_time(self, invocation_type: str, time_used: float):
"""
Add timing information on execution of a node. Usually
@ -95,7 +96,7 @@ class InvocationStats():
self._stats[invocation_type] = [0, 0.0]
self._stats[invocation_type][0] += 1
self._stats[invocation_type][1] += time_used
def log_stats(self):
"""
Send the statistics to the system logger at the info level.
@ -103,13 +104,12 @@ class InvocationStats():
is complete.
"""
if self._current_graph_state.is_complete():
logger.info('Node Calls Seconds')
logger.info("Node Calls Seconds")
for node_type, (calls, time_used) in self._stats.items():
logger.info(f'{node_type:<20} {calls:>5} {time_used:4.3f}s')
total_time = sum([ticks for _,ticks in self._stats.values()])
logger.info(f'TOTAL: {total_time:4.3f}s')
logger.info(f"{node_type:<20} {calls:>5} {time_used:4.3f}s")
total_time = sum([ticks for _, ticks in self._stats.values()])
logger.info(f"TOTAL: {total_time:4.3f}s")
if torch.cuda.is_available():
logger.info('Max VRAM used for execution: '+'%4.2fG' % (torch.cuda.max_memory_allocated() / 1e9))
logger.info('Current VRAM utilization '+'%4.2fG' % (torch.cuda.memory_allocated() / 1e9))
logger.info("Max VRAM used for execution: " + "%4.2fG" % (torch.cuda.max_memory_allocated() / 1e9))
logger.info("Current VRAM utilization " + "%4.2fG" % (torch.cuda.memory_allocated() / 1e9))

4
invokeai/app/services/processor.py
View File

@ -116,6 +116,7 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
pass
except CanceledException:
statistics.reset_stats()
pass
except Exception as e:
@ -137,7 +138,7 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
error_type=e.__class__.__name__,
error=error,
)
statistics.reset_stats()
pass
# Check queue to see if this is canceled, and skip if so
@ -165,4 +166,3 @@ class DefaultInvocationProcessor(InvocationProcessorABC):
pass # Log something? KeyboardInterrupt is probably not going to be seen by the processor
finally:
self.__threadLimit.release()