Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
"""
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.
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 abc
import enum
import time
from timeit import default_timer as timer
from contextlib import contextmanager, ExitStack
from collections import defaultdict
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",
])
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