mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
import time
|
|
from collections import OrderedDict, deque
|
|
|
|
|
|
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
|
|
|
|
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 print(self):
|
|
print(f"\nTimer '{self.name}':")
|
|
# 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)
|
|
print(f" - {avg_time:.4f}s avg - {timer_name}, num = {len(timings)}")
|
|
|
|
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)
|