Hooks for compilation start/end

Numba Community Support
1

I’m wondering if there are places to register callbacks before and after compilation? I’d like to “pause” timing tests during jit compilation, as compilation time is not relevant. (However, I don’t want to go to eager compilation everywhere.)

2

Aha – Event API — Numba 0.57.1+0.g04e81073b.dirty documentation

Perhaps this is what I’m looking for? Any examples on using it?

3

I’m not familiar with the hooks, but I get the total compilation time for a function by summing the pipeline steps, inspired by the @stuartarchibald here :slight_smile:

4

Hi @shaunc

No guarantee, but I think you can use the event API as follows:

from time import perf_counter
from dataclasses import dataclass
from numba.core import event
from numba.core.registry import CPUDispatcher


@dataclass
class Timing:
    tic: float | None = None 
    toc: float | None = None 
        
    @property
    def elapsed(self) -> float | None:
        if not None in (self.tic, self.toc):
            return self.toc - self.tic 
        
          
class MyCompileTimeListener(event.Listener):
    def __init__(self) -> None:
        self.timings: dict[CPUDispatcher, Timing] = {}
        
    def on_start(self, event: event.Event) -> None:
        self.timings[event.data["dispatcher"]] = Timing(perf_counter())

    def on_end(self, event: event.Event) -> None:
        self.timings[event.data["dispatcher"]].toc = perf_counter()
        
    def get_elapsed(self, disp: CPUDispatcher) -> float:
        return self.timings[disp].elapsed


my_listener = MyCompileTimeListener()
event.register(kind="numba:compile", listener=my_listener)

That gives me very similar timings compared to the solution from @nelson2005:

# An expensive to compile function
@nb.njit
def foo(a):
    return np.linalg.eig(a)

a = np.zeros((10, 10))
foo(a)

# Using the solution proposed by nelson2005
tot_time = 0
for name, time in foo.overloads[foo.signatures[0]].metadata['pipeline_times']['nopython'].items():
       tot_time += (time.init + time.run + time.finalize)
print(f"Elapsed [s]: {tot_time:5f}")
# Elapsed [s]: 3.226300

# Using the listener
print(f"Elapsed [s]: {my_listener.get_elapsed(foo):5f}")
# Elapsed [s]: 3.251469

Edit:
I have seen that there are also many practical context managers already implemented. For example, simple timing related callbacks can easily be implemented as such:

with event.install_timer(kind="numba:compile", callback=lambda t: print(t)):
    foo(a)

Or there is a ready to use context manager that records time and event:

with event.install_recorder(kind="numba:compile") as recorder:
    foo(a)

    tic = recorder.buffer[-2][0]
    toc = recorder.buffer[-1][0]
    print(f"Elapsed [s]: {toc - tic:5f}")

I’m not quite sure what your goal is, but the event API looks very interesting to me. I wish I looked at this earlier. :slight_smile:

1 Like
5

Thanks!

My goal was to be able to “subtract” compilation time from pytest timeout times. Currently I’m stymied on the pytest-timeout side: Question/feature request: way to pause timing · Issue #150 · pytest-dev/pytest-timeout · GitHub but could roll my own there I guess.