Spaces:
Running
Running
| """ | |
| Lightweight, hierarchical timers for profiling sections of code. | |
| Example: | |
| @timed | |
| def foo(t): | |
| time.sleep(t) | |
| def main(): | |
| for i in range(3): | |
| foo(i + 1) | |
| with hierarchical_timer("context"): | |
| foo(1) | |
| print(get_timer_tree()) | |
| This would produce a timer tree like | |
| (root) | |
| "foo" | |
| "context" | |
| "foo" | |
| The total time and counts are tracked for each block of code; in this example "foo" and "context.foo" are considered | |
| distinct blocks, and are tracked separately. | |
| The decorator and contextmanager are equivalent; the context manager may be more useful if you want more control | |
| over the timer name, or are splitting up multiple sections of a large function. | |
| """ | |
| import math | |
| import sys | |
| import time | |
| import threading | |
| from contextlib import contextmanager | |
| from typing import Any, Callable, Dict, Generator, Optional, TypeVar | |
| TIMER_FORMAT_VERSION = "0.1.0" | |
| class TimerNode: | |
| """ | |
| Represents the time spent in a block of code. | |
| """ | |
| __slots__ = ["children", "total", "count", "is_parallel"] | |
| def __init__(self): | |
| # Note that since dictionary keys are the node names, we don't explicitly store the name on the TimerNode. | |
| self.children: Dict[str, TimerNode] = {} | |
| self.total: float = 0.0 | |
| self.count: int = 0 | |
| self.is_parallel = False | |
| def get_child(self, name: str) -> "TimerNode": | |
| """ | |
| Get the child node corresponding to the name (and create if it doesn't already exist). | |
| """ | |
| child = self.children.get(name) | |
| if child is None: | |
| child = TimerNode() | |
| self.children[name] = child | |
| return child | |
| def add_time(self, elapsed: float) -> None: | |
| """ | |
| Accumulate the time spent in the node (and increment the count). | |
| """ | |
| self.total += elapsed | |
| self.count += 1 | |
| def merge( | |
| self, other: "TimerNode", root_name: str = None, is_parallel: bool = True | |
| ) -> None: | |
| """ | |
| Add the other node to this node, then do the same recursively on its children. | |
| :param other: The other node to merge | |
| :param root_name: Optional name of the root node being merged. | |
| :param is_parallel: Whether or not the code block was executed in parallel. | |
| :return: | |
| """ | |
| if root_name: | |
| node = self.get_child(root_name) | |
| else: | |
| node = self | |
| node.total += other.total | |
| node.count += other.count | |
| node.is_parallel |= is_parallel | |
| for other_child_name, other_child_node in other.children.items(): | |
| child = node.get_child(other_child_name) | |
| child.merge(other_child_node, is_parallel=is_parallel) | |
| class GaugeNode: | |
| """ | |
| Tracks the most recent value of a metric. This is analogous to gauges in statsd. | |
| """ | |
| __slots__ = ["value", "min_value", "max_value", "count", "_timestamp"] | |
| def __init__(self, value: float): | |
| self.value = value | |
| self.min_value = value | |
| self.max_value = value | |
| self.count = 1 | |
| # Internal timestamp so we can determine priority. | |
| self._timestamp = time.time() | |
| def update(self, new_value: float) -> None: | |
| self.min_value = min(self.min_value, new_value) | |
| self.max_value = max(self.max_value, new_value) | |
| self.value = new_value | |
| self.count += 1 | |
| self._timestamp = time.time() | |
| def merge(self, other: "GaugeNode") -> None: | |
| if self._timestamp < other._timestamp: | |
| # Keep the "later" value | |
| self.value = other.value | |
| self._timestamp = other._timestamp | |
| self.min_value = min(self.min_value, other.min_value) | |
| self.max_value = max(self.max_value, other.max_value) | |
| self.count += other.count | |
| def as_dict(self) -> Dict[str, float]: | |
| return { | |
| "value": self.value, | |
| "min": self.min_value, | |
| "max": self.max_value, | |
| "count": self.count, | |
| } | |
| class TimerStack: | |
| """ | |
| Tracks all the time spent. Users shouldn't use this directly, they should use the contextmanager below to make | |
| sure that pushes and pops are already matched. | |
| """ | |
| __slots__ = ["root", "stack", "start_time", "gauges", "metadata"] | |
| def __init__(self): | |
| self.root = TimerNode() | |
| self.stack = [self.root] | |
| self.start_time = time.perf_counter() | |
| self.gauges: Dict[str, GaugeNode] = {} | |
| self.metadata: Dict[str, str] = {} | |
| self._add_default_metadata() | |
| def reset(self): | |
| self.root = TimerNode() | |
| self.stack = [self.root] | |
| self.start_time = time.perf_counter() | |
| self.gauges: Dict[str, GaugeNode] = {} | |
| self.metadata: Dict[str, str] = {} | |
| self._add_default_metadata() | |
| def push(self, name: str) -> TimerNode: | |
| """ | |
| Called when entering a new block of code that is timed (e.g. with a contextmanager). | |
| """ | |
| current_node: TimerNode = self.stack[-1] | |
| next_node = current_node.get_child(name) | |
| self.stack.append(next_node) | |
| return next_node | |
| def pop(self) -> None: | |
| """ | |
| Called when exiting a new block of code that is timed (e.g. with a contextmanager). | |
| """ | |
| self.stack.pop() | |
| def get_root(self) -> TimerNode: | |
| """ | |
| Update the total time and count of the root name, and return it. | |
| """ | |
| root = self.root | |
| root.total = time.perf_counter() - self.start_time | |
| root.count = 1 | |
| return root | |
| def get_timing_tree(self, node: TimerNode = None) -> Dict[str, Any]: | |
| """ | |
| Recursively build a tree of timings, suitable for output/archiving. | |
| """ | |
| res: Dict[str, Any] = {} | |
| if node is None: | |
| # Special case the root - total is time since it was created, and count is 1 | |
| node = self.get_root() | |
| res["name"] = "root" | |
| # Only output gauges at top level | |
| if self.gauges: | |
| res["gauges"] = self._get_gauges() | |
| if self.metadata: | |
| self.metadata["end_time_seconds"] = str(int(time.time())) | |
| res["metadata"] = self.metadata | |
| res["total"] = node.total | |
| res["count"] = node.count | |
| if node.is_parallel: | |
| # Note when the block ran in parallel, so that it's less confusing that a timer is less that its children. | |
| res["is_parallel"] = True | |
| child_total = 0.0 | |
| child_dict = {} | |
| for child_name, child_node in node.children.items(): | |
| child_res: Dict[str, Any] = self.get_timing_tree(child_node) | |
| child_dict[child_name] = child_res | |
| child_total += child_res["total"] | |
| # "self" time is total time minus all time spent on children | |
| res["self"] = max(0.0, node.total - child_total) | |
| if child_dict: | |
| res["children"] = child_dict | |
| return res | |
| def set_gauge(self, name: str, value: float) -> None: | |
| if math.isnan(value): | |
| return | |
| gauge_node = self.gauges.get(name) | |
| if gauge_node: | |
| gauge_node.update(value) | |
| else: | |
| self.gauges[name] = GaugeNode(value) | |
| def add_metadata(self, key: str, value: str) -> None: | |
| self.metadata[key] = value | |
| def _get_gauges(self) -> Dict[str, Dict[str, float]]: | |
| gauges = {} | |
| for gauge_name, gauge_node in self.gauges.items(): | |
| gauges[gauge_name] = gauge_node.as_dict() | |
| return gauges | |
| def _add_default_metadata(self): | |
| self.metadata["timer_format_version"] = TIMER_FORMAT_VERSION | |
| self.metadata["start_time_seconds"] = str(int(time.time())) | |
| self.metadata["python_version"] = sys.version | |
| self.metadata["command_line_arguments"] = " ".join(sys.argv) | |
| # Maintain a separate "global" timer per thread, so that they don't accidentally conflict with each other. | |
| _thread_timer_stacks: Dict[int, TimerStack] = {} | |
| def _get_thread_timer() -> TimerStack: | |
| ident = threading.get_ident() | |
| if ident not in _thread_timer_stacks: | |
| timer_stack = TimerStack() | |
| _thread_timer_stacks[ident] = timer_stack | |
| return _thread_timer_stacks[ident] | |
| def get_timer_stack_for_thread(t: threading.Thread) -> Optional[TimerStack]: | |
| if t.ident is None: | |
| # Thread hasn't started, shouldn't ever happen | |
| return None | |
| return _thread_timer_stacks.get(t.ident) | |
| def hierarchical_timer(name: str, timer_stack: TimerStack = None) -> Generator: | |
| """ | |
| Creates a scoped timer around a block of code. This time spent will automatically be incremented when | |
| the context manager exits. | |
| """ | |
| timer_stack = timer_stack or _get_thread_timer() | |
| timer_node = timer_stack.push(name) | |
| start_time = time.perf_counter() | |
| try: | |
| # The wrapped code block will run here. | |
| yield timer_node | |
| finally: | |
| # This will trigger either when the context manager exits, or an exception is raised. | |
| # We'll accumulate the time, and the exception (if any) gets raised automatically. | |
| elapsed = time.perf_counter() - start_time | |
| timer_node.add_time(elapsed) | |
| timer_stack.pop() | |
| # This is used to ensure the signature of the decorated function is preserved | |
| # See also https://github.com/python/mypy/issues/3157 | |
| FuncT = TypeVar("FuncT", bound=Callable[..., Any]) | |
| def timed(func: FuncT) -> FuncT: | |
| """ | |
| Decorator for timing a function or method. The name of the timer will be the qualified name of the function. | |
| Usage: | |
| @timed | |
| def my_func(x, y): | |
| return x + y | |
| Note that because this doesn't take arguments, the global timer stack is always used. | |
| """ | |
| def wrapped(*args, **kwargs): | |
| with hierarchical_timer(func.__qualname__): | |
| return func(*args, **kwargs) | |
| return wrapped # type: ignore | |
| def set_gauge(name: str, value: float, timer_stack: TimerStack = None) -> None: | |
| """ | |
| Updates the value of the gauge (or creates it if it hasn't been set before). | |
| """ | |
| timer_stack = timer_stack or _get_thread_timer() | |
| timer_stack.set_gauge(name, value) | |
| def merge_gauges(gauges: Dict[str, GaugeNode], timer_stack: TimerStack = None) -> None: | |
| """ | |
| Merge the gauges from another TimerStack with the provided one (or the | |
| current thread's stack if none is provided). | |
| :param gauges: | |
| :param timer_stack: | |
| :return: | |
| """ | |
| timer_stack = timer_stack or _get_thread_timer() | |
| for n, g in gauges.items(): | |
| if n in timer_stack.gauges: | |
| timer_stack.gauges[n].merge(g) | |
| else: | |
| timer_stack.gauges[n] = g | |
| def add_metadata(key: str, value: str, timer_stack: TimerStack = None) -> None: | |
| timer_stack = timer_stack or _get_thread_timer() | |
| timer_stack.add_metadata(key, value) | |
| def get_timer_tree(timer_stack: TimerStack = None) -> Dict[str, Any]: | |
| """ | |
| Return the tree of timings from the TimerStack as a dictionary (or the | |
| current thread's stack if none is provided) | |
| """ | |
| timer_stack = timer_stack or _get_thread_timer() | |
| return timer_stack.get_timing_tree() | |
| def get_timer_root(timer_stack: TimerStack = None) -> TimerNode: | |
| """ | |
| Get the root TimerNode of the timer_stack (or the current thread's | |
| TimerStack if not specified) | |
| """ | |
| timer_stack = timer_stack or _get_thread_timer() | |
| return timer_stack.get_root() | |
| def reset_timers(timer_stack: TimerStack = None) -> None: | |
| """ | |
| Reset the timer_stack (or the current thread's TimerStack if not specified) | |
| """ | |
| timer_stack = timer_stack or _get_thread_timer() | |
| timer_stack.reset() | |