How to accelerate import of site-package using numba jit functions?

Numba Support: How do I do ...?
1

I have installed a small python code library named “tasc” as site-package in development mode.
The package has 1673 lines of python code.
If I want to import a module from this package for the first time in a session it takes very long.
I have measured the first import of a typical module called “tasc.bands.py” with the “line_profiler” package.
It took 144 seconds. My laptop is not the newest but this feels rather extreme.
All functions have a signature and are decorated with nopython=True, fastmath=True and cache=True.
Here is a typical example of a function used in the library:

import numpy as np
from numba import jit

@jit(['f8[:](f8[:], i8)',
      'f8[:](f8[:], Omitted(50))'],
     nopython=True, fastmath=True, cache=True)
def sma(x: np.ndarray, length: int = 50) -> np.ndarray:
    """Calculate rolling moving average.

    Args:
        x (np.ndarray): Values as 1D-array of float.
        length (int): Rolling window.

    Returns:
        np.ndarray : Moving average as 1D-array of float.

    Example:
        vals = np.cumsum(np.random.normal(scale=0.01, size=50))
        res = sma(vals, 10)
        print(np.column_stack([vals, res]))
    """
    res = np.empty_like(x)
    res[0] = mavg = x[0]
    for i in range(1, length):
        mavg += ((x[i]-mavg)/(i+1))
        res[i] = mavg
    for i in range(length, res.size):
        mavg += ((x[i]-x[i-length])/length)
        res[i] = mavg
    return res

All functions are running without problems/warnings.
There are no circular import references in the modules.
It seems every time I start a session in the virtual environment Numba compiles all functions
of the package again before these accelerated functions can be called.
Is this the case?
Is this a problem in development mode only?
Can I suppress this behavior and force Numba to use the already cached functions from previous sessions?

My problem sounds similar to this topic “Cache njit function at import or install”.
Can you help me?

Below are my system information, the number of lines of code and the result of line_profiler.

System information:
Operating System: Linux Mint 21.1
RAM: 16 GB
CPU(s): 4 x Intel(R) Core™ i7-5500U CPU @ 2.40GHz

Statistics for package “tasc”:
Language: Python
Files: 11
blank: 536 lines
comment: 1815 lines
code: 1673 lines
Result of line_profiler (1st time: “import tasc.bands”): Total time: 143.807 s

2

I believe I may have found a solution to my specific problem. By setting a system variable, NUMBA_CACHE_DIR, on my local machine and providing it with a specific path, I was able to prevent Numba from recompiling and caching the package functions with each new Python session. This has significantly improved the import times, reducing them from 144 seconds to just 1-2 seconds at the start of a session.

However, I should note that this method is not ideal as it makes the code less portable. A better solution would likely involve ahead-of-time compilation, but this may increase the effort required to write and maintain the code. If anyone in the community has alternative solutions or further suggestions, I would greatly appreciate your input.

3

Indeed, if site-packages isn’t writable, then you will have to set the cache dir to somewhere else writable. Presently there isn’t an alternative, and compiling ahead-of-time (e.g. at installation time) presents a sizeable set of challenges to resolve, such as how to ensure that the compiled code matches the CPU it’s run on - there’s a broad set of situations where the build and installation of the package is distinct from the use of it (e.g. building Linux distro packages, etc.).

1 Like
4

Thank you for your time and effort to respond to my request.
The option to set a NUMBA_CACH_DIR as system variable was very helpful.
It is not an ideal solution but I am happy to make my site-package usable on my local machine at all.
I appreciate your input and the information you provided.
It was helpful understanding caching better.