How to use Numba extension in setup.py?

Question

I try to use the Numba for some fast calculations. I got the following issue while creating a package that use a Numba extension.

I did similar things as suggested on the official website. I have the following folder structure:

-test_numba
  -test_numba
    -__init__.py
    -source_module.py
-setup.py

Then I put the following code in source_module.py:

from numba.pycc import CC

cc = CC('my_module')
cc.verbose = True

@cc.export('multf', 'f8(f8, f8)')
@cc.export('multi', 'i4(i4, i4)')
def mult(a, b):
    return a * b

@cc.export('square', 'f8(f8)')
def square(a):
    return a ** 2

if __name__ == "__main__":
    cc.compile()

And complied it. and for init.py, I simply import the functions:

from .my_module import *

and for setup.py:

from setuptools import setup
from test_numba.source_module import cc
from setuptools import Extension

if __name__ == "__main__":
    setup(
        name="test_numba",
        version="0.0.1",
        packages=["test_numba"],
        ext_modules=[cc.distutils_extension()]
    )

I used pip install to installed the package. but the extension is not installed in the correct folder location.

    /Users/xxx/miniconda3/envs/py/lib/python3.7/site-packages/my_module.cpython-37m-darwin.so
    /Users/xxx/miniconda3/envs/py/lib/python3.7/site-packages/test_numba-0.0.1.dist-info/*
    /Users/xxx/miniconda3/envs/py/lib/python3.7/site-packages/test_numba/*

and when I try to import this package from jupyter. I will get the following error: No module named 'test_numba.my_module'

How to correct this?

Answer 1

First of all, I assume your setup.py is actually within your first test_numba directory instead of outside of it as indicated in your diagram above. IOW, I assume it's actually like this:

-test_numba
  -test_numba
    -__init__.py
    -source_module.py
  -setup.py

There's a step that's required which is to do the Ahead-Of-Time (AOT) compilation that you want. To accomplish that, you can add cc.compile() in your setup.py:

from setuptools import setup
from setuptools import Extension

from test_numba.source_module import cc

if __name__ == "__main__":
    cc.compile()  ### <--------  ADD THIS TO DO AOT COMPILATION
    setup(
        name="test_numba",
        version="0.0.1",
        packages=["test_numba"],
        ext_modules=[cc.distutils_extension()]
    )

But to be able to get that far, you have to allow __init__.py to continue if it can't find my_module (because it isn't built yet).

try:
  from .my_module import *
except:
  pass

You can then install it with python3 -m pip install . while in the first test_numba directory.

Here's my successful test (after changing to my home directory to make sure it's not importing locally):

>>> import test_numba as tn
>>> tn.multf(4,5)
20.0

Answer 2

The answer by @Mark H didn't work for me. It also seems strange to call cc.compile as the numba AOT documentation for distutils integration (i.e. setup.py) doesn't mention running cc.compile before setup(..., ext_modules=[cc.distutils_extension()]).

Here's what worked for me:

1. project structure:

-project_name
  -src
    -my_module
      -__init__.py
      -code2compile.py
  -setup.py

2. In the python file containing the functions to be compiled I had to change the name of the output directory of the DLL, which is defined using cc._source_module and not cc.output_dir for cc.distutils_extension() (see the definition).

cc = CC("compiled_module_name")  # name of the DLL/pyd file
cc._source_module = "my_module.code2compile" 

Note that if this wasn't done numba would attempt to copy the DLL to src/src/my_module, which doesn't exist. Instead, the above fix ensures that numba copies the DLL to src/my_module.

Update: I've created an issue on numba's github here as it seems like a bug.