486 lines
12 KiB
Python
486 lines
12 KiB
Python
|
"""
|
||
|
|
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()
|