Spaces:
Running
Running
File size: 2,713 Bytes
fcc02a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import time
from collections import OrderedDict, deque
import sys
import os
# check if is ui process will have IS_AI_TOOLKIT_UI in env
is_ui = os.environ.get("IS_AI_TOOLKIT_UI", "0") == "1"
class Timer:
def __init__(self, name='Timer', max_buffer=10):
self.name = name
self.max_buffer = max_buffer
self.timers = OrderedDict()
self.active_timers = {}
self.current_timer = None # Used for the context manager functionality
self._after_print_hooks = []
def start(self, timer_name):
if timer_name not in self.timers:
self.timers[timer_name] = deque(maxlen=self.max_buffer)
self.active_timers[timer_name] = time.time()
def cancel(self, timer_name):
"""Cancel an active timer."""
if timer_name in self.active_timers:
del self.active_timers[timer_name]
def stop(self, timer_name):
if timer_name not in self.active_timers:
raise ValueError(f"Timer '{timer_name}' was not started!")
elapsed_time = time.time() - self.active_timers[timer_name]
self.timers[timer_name].append(elapsed_time)
# Clean up active timers
del self.active_timers[timer_name]
# Check if this timer's buffer exceeds max_buffer and remove the oldest if it does
if len(self.timers[timer_name]) > self.max_buffer:
self.timers[timer_name].popleft()
def add_after_print_hook(self, hook):
self._after_print_hooks.append(hook)
def print(self):
if not is_ui:
print(f"\nTimer '{self.name}':")
timing_dict = {}
# sort by longest at top
for timer_name, timings in sorted(self.timers.items(), key=lambda x: sum(x[1]), reverse=True):
avg_time = sum(timings) / len(timings)
if not is_ui:
print(f" - {avg_time:.4f}s avg - {timer_name}, num = {len(timings)}")
timing_dict[timer_name] = avg_time
for hook in self._after_print_hooks:
hook(timing_dict)
if not is_ui:
print('')
def reset(self):
self.timers.clear()
self.active_timers.clear()
def __call__(self, timer_name):
"""Enable the use of the Timer class as a context manager."""
self.current_timer = timer_name
self.start(timer_name)
return self
def __enter__(self):
pass
def __exit__(self, exc_type, exc_value, traceback):
if exc_type is None:
# No exceptions, stop the timer normally
self.stop(self.current_timer)
else:
# There was an exception, cancel the timer
self.cancel(self.current_timer)
|