""" The ``numba.core.event`` module provides a simple event system for applications to register callbacks to listen to specific compiler events. The following events are built in: - ``"numba:compile"`` is broadcast when a dispatcher is compiling. Events of this kind have ``data`` defined to be a ``dict`` with the following key-values: - ``"dispatcher"``: the dispatcher object that is compiling. - ``"args"``: the argument types. - ``"return_type"``: the return type. - ``"numba:compiler_lock"`` is broadcast when the internal compiler-lock is acquired. This is mostly used internally to measure time spent with the lock acquired. - ``"numba:llvm_lock"`` is broadcast when the internal LLVM-lock is acquired. This is used internally to measure time spent with the lock acquired. - ``"numba:run_pass"`` is broadcast when a compiler pass is running. - ``"name"``: pass name. - ``"qualname"``: qualified name of the function being compiled. - ``"module"``: module name of the function being compiled. - ``"flags"``: compilation flags. - ``"args"``: argument types. - ``"return_type"`` return type. Applications can register callbacks that are listening for specific events using ``register(kind: str, listener: Listener)``, where ``listener`` is an instance of ``Listener`` that defines custom actions on occurrence of the specific event. """ import os import json import atexit import abc import enum import time import threading from timeit import default_timer as timer from contextlib import contextmanager, ExitStack from collections import defaultdict from numba.core import config class EventStatus(enum.Enum): """Status of an event. """ START = enum.auto() END = enum.auto() # Builtin event kinds. _builtin_kinds = frozenset([ "numba:compiler_lock", "numba:compile", "numba:llvm_lock", "numba:run_pass", ]) def _guard_kind(kind): """Guard to ensure that an event kind is valid. All event kinds with a "numba:" prefix must be defined in the pre-defined ``numba.core.event._builtin_kinds``. Custom event kinds are allowed by not using the above prefix. Parameters ---------- kind : str Return ------ res : str """ if kind.startswith("numba:") and kind not in _builtin_kinds: msg = (f"{kind} is not a valid event kind, " "it starts with the reserved prefix 'numba:'") raise ValueError(msg) return kind class Event: """An event. Parameters ---------- kind : str status : EventStatus data : any; optional Additional data for the event. exc_details : 3-tuple; optional Same 3-tuple for ``__exit__``. """ def __init__(self, kind, status, data=None, exc_details=None): self._kind = _guard_kind(kind) self._status = status self._data = data self._exc_details = (None if exc_details is None or exc_details[0] is None else exc_details) @property def kind(self): """Event kind Returns ------- res : str """ return self._kind @property def status(self): """Event status Returns ------- res : EventStatus """ return self._status @property def data(self): """Event data Returns ------- res : object """ return self._data @property def is_start(self): """Is it a *START* event? Returns ------- res : bool """ return self._status == EventStatus.START @property def is_end(self): """Is it an *END* event? Returns ------- res : bool """ return self._status == EventStatus.END @property def is_failed(self): """Is the event carrying an exception? This is used for *END* event. This method will never return ``True`` in a *START* event. Returns ------- res : bool """ return self._exc_details is None def __str__(self): data = (f"{type(self.data).__qualname__}" if self.data is not None else "None") return f"Event({self._kind}, {self._status}, data: {data})" __repr__ = __str__ _registered = defaultdict(list) def register(kind, listener): """Register a listener for a given event kind. Parameters ---------- kind : str listener : Listener """ assert isinstance(listener, Listener) kind = _guard_kind(kind) _registered[kind].append(listener) def unregister(kind, listener): """Unregister a listener for a given event kind. Parameters ---------- kind : str listener : Listener """ assert isinstance(listener, Listener) kind = _guard_kind(kind) lst = _registered[kind] lst.remove(listener) def broadcast(event): """Broadcast an event to all registered listeners. Parameters ---------- event : Event """ for listener in _registered[event.kind]: listener.notify(event) class Listener(abc.ABC): """Base class for all event listeners. """ @abc.abstractmethod def on_start(self, event): """Called when there is a *START* event. Parameters ---------- event : Event """ pass @abc.abstractmethod def on_end(self, event): """Called when there is a *END* event. Parameters ---------- event : Event """ pass def notify(self, event): """Notify this Listener with the given Event. Parameters ---------- event : Event """ if event.is_start: self.on_start(event) elif event.is_end: self.on_end(event) else: raise AssertionError("unreachable") class TimingListener(Listener): """A listener that measures the total time spent between *START* and *END* events during the time this listener is active. """ def __init__(self): self._depth = 0 def on_start(self, event): if self._depth == 0: self._ts = timer() self._depth += 1 def on_end(self, event): self._depth -= 1 if self._depth == 0: last = getattr(self, "_duration", 0) self._duration = (timer() - self._ts) + last @property def done(self): """Returns a ``bool`` indicating whether a measurement has been made. When this returns ``False``, the matching event has never fired. If and only if this returns ``True``, ``.duration`` can be read without error. """ return hasattr(self, "_duration") @property def duration(self): """Returns the measured duration. This may raise ``AttributeError``. Users can use ``.done`` to check that a measurement has been made. """ return self._duration class RecordingListener(Listener): """A listener that records all events and stores them in the ``.buffer`` attribute as a list of 2-tuple ``(float, Event)``, where the first element is the time the event occurred as returned by ``time.time()`` and the second element is the event. """ def __init__(self): self.buffer = [] def on_start(self, event): self.buffer.append((time.time(), event)) def on_end(self, event): self.buffer.append((time.time(), event)) @contextmanager def install_listener(kind, listener): """Install a listener for event "kind" temporarily within the duration of the context. Returns ------- res : Listener The *listener* provided. Examples -------- >>> with install_listener("numba:compile", listener): >>> some_code() # listener will be active here. >>> other_code() # listener will be unregistered by this point. """ register(kind, listener) try: yield listener finally: unregister(kind, listener) @contextmanager def install_timer(kind, callback): """Install a TimingListener temporarily to measure the duration of an event. If the context completes successfully, the *callback* function is executed. The *callback* function is expected to take a float argument for the duration in seconds. Returns ------- res : TimingListener Examples -------- This is equivalent to: >>> with install_listener(kind, TimingListener()) as res: >>> ... """ tl = TimingListener() with install_listener(kind, tl): yield tl if tl.done: callback(tl.duration) @contextmanager def install_recorder(kind): """Install a RecordingListener temporarily to record all events. Once the context is closed, users can use ``RecordingListener.buffer`` to access the recorded events. Returns ------- res : RecordingListener Examples -------- This is equivalent to: >>> with install_listener(kind, RecordingListener()) as res: >>> ... """ rl = RecordingListener() with install_listener(kind, rl): yield rl def start_event(kind, data=None): """Trigger the start of an event of *kind* with *data*. Parameters ---------- kind : str Event kind. data : any; optional Extra event data. """ evt = Event(kind=kind, status=EventStatus.START, data=data) broadcast(evt) def end_event(kind, data=None, exc_details=None): """Trigger the end of an event of *kind*, *exc_details*. Parameters ---------- kind : str Event kind. data : any; optional Extra event data. exc_details : 3-tuple; optional Same 3-tuple for ``__exit__``. Or, ``None`` if no error. """ evt = Event( kind=kind, status=EventStatus.END, data=data, exc_details=exc_details, ) broadcast(evt) @contextmanager def trigger_event(kind, data=None): """A context manager to trigger the start and end events of *kind* with *data*. The start event is triggered when entering the context. The end event is triggered when exiting the context. Parameters ---------- kind : str Event kind. data : any; optional Extra event data. """ with ExitStack() as scope: @scope.push def on_exit(*exc_details): end_event(kind, data=data, exc_details=exc_details) start_event(kind, data=data) yield def _prepare_chrome_trace_data(listener: RecordingListener): """Prepare events in `listener` for serializing as chrome trace data. """ # The spec for the trace event format can be found at: # https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/edit # noqa # This code only uses the JSON Array Format for simplicity. pid = os.getpid() tid = threading.get_native_id() evs = [] for ts, rec in listener.buffer: data = rec.data cat = str(rec.kind) ts_scaled = ts * 1_000_000 # scale to microseconds ph = 'B' if rec.is_start else 'E' name = data['name'] args = data ev = dict( cat=cat, pid=pid, tid=tid, ts=ts_scaled, ph=ph, name=name, args=args, ) evs.append(ev) return evs def _setup_chrome_trace_exit_handler(): """Setup a RecordingListener and an exit handler to write the captured events to file. """ listener = RecordingListener() register("numba:run_pass", listener) filename = config.CHROME_TRACE @atexit.register def _write_chrome_trace(): # The following output file is not multi-process safe. evs = _prepare_chrome_trace_data(listener) with open(filename, "w") as out: json.dump(evs, out) if config.CHROME_TRACE: _setup_chrome_trace_exit_handler()