Compile full package with nested modules using numba Ahead-of-Time Compilation

Question

I split up my numerical code into multiple modules, which all are part of a package. I would like to completely pre-compile the full package. The documentation just gives information on [how to compile a single module], but not a package with multiple modules inside it.

This the way I am currently doing it to keep the package __init__.py python source and individually compile the contained modules, then importing those into the __init__.py to export them from the package.

Is there a way to compile the complete package?

My current implementation gives me

aotnumba
|-- __init__.py
|-- main.cpython-38-x86_64-linux-gnu.so
`-- utils.cpython-38-x86_64-linux-gnu.so

but what I want is

aotnumba.cpython-38-x86_64-linux-gnu.so

My current implementation works like this:

.
├── aotnumba
│   ├── __init__.py
│   ├── main.py
│   └── utils.py
├── compile.py
├── Dockerfile
└── Makefile

# aotnumba/__init__.py
from .main import main

__all__ = ['main']

# once compiled, I cannot import the compiler anymore
try:
    from .main import compiler as compiler_main
    from .utils import compiler as compiler_utils
    __compilers__ = [compiler_main, compiler_utils]
except ImportError:
    __compilers__ = []

# aotnumba/main.py
from sys import argv
from numba.pycc import CC as Compiler
from .utils import add

compiler = Compiler('main')

@compiler.export('main', '(int32, int32)')
def main(a, b):
    res = add(a, b)
    print(res)

# aotnumba/utils.py
from numba import njit
from numba.pycc import CC as Compiler

compiler = Compiler('utils')

@njit  # if I don't add this, it doesn't compile?
@compiler.export('add', 'int32(int32, int32)')
def add(a, b):
    return a + b

(why do I need to @njit here when I compile?)

Then I compile it like

# compile.py
from distutils.core import setup

import aotnumba

for compiler in aotnumba.__compilers__:
    compiler.verbose = True
    compiler.compile()

setup(name='aotnumba', ext_modules=[compiler.distutils_extension()
                                    for compiler in aotnumba.__compilers__])

To run the whole thing, I use docker, so it's a clean reproducible setup

FROM ubuntu:20.04

WORKDIR /src

RUN apt-get update
RUN apt-get install -qqy --no-install-recommends \
    tree \
    python3-pip \
    build-essential \
    python3-dev

RUN pip3 install pip --upgrade
RUN pip3 install \
    numba==0.56.3 \
    numpy==1.23.4

COPY . .

RUN python3 compile.py build_ext --inplace
RUN rm -rf \
    aotnumba/__pycache__ \
    aotnumba/main.py \
    aotnumba/utils.py \
    build

RUN tree

To test it I run it using

IMAGE = aotnumba:latest
.PHONY: build test

build:
    docker build -t $(IMAGE) -f Dockerfile .

test:
    docker run $(IMAGE) python3 -c 'import aotnumba; aotnumba.main(23, 42)'

That gives me

$ make build test
…
Step 10/10 : RUN tree
 ---> Running in 43d204599bdb
.
|-- Dockerfile
|-- Makefile
|-- aotnumba
|   |-- __init__.py
|   |-- main.cpython-38-x86_64-linux-gnu.so
|   `-- utils.cpython-38-x86_64-linux-gnu.so
`-- compile.py

1 directory, 6 files
Removing intermediate container 43d204599bdb
 ---> 2544ea438d61
Successfully built 2544ea438d61
Successfully tagged aotnumba:latest
docker run aotnumba:latest python3 -c 'import aotnumba; aotnumba.main(23, 42)'
65

I have tried passing nested module identifiers to CC, but then I get an error

ValueError: basename should be a simple module name, not qualified name

from distutils.core import setup
from numba.pycc import CC as Compiler

import aotnumba

compiler_package = Compiler('aotnumba')
compiler_package.export('main', '(int32, int32)')(aotnumba.main)
compiler_package.compile()

compiler_utils = Compiler('aotnumba.utils')
compiler_utils.export('add', 'int32(int32, int32)')(aotnumba.utils.add)
compiler_utils.compile()

setup(name='aotnumba', ext_modules=[compiler_package.distutils_extension(),
                                    compiler_utils.distutils_extension()])

I have tried passing qualified identifiers to CC.export, but then these are taken to be the bare function name that is exported

from distutils.core import setup
from numba.pycc import CC as Compiler

import aotnumba

compiler = Compiler('aotnumba')
compiler.verbose = True
compiler.export('main.main', '(int32, int32)')(aotnumba.main)
compiler.export('utils.add', 'int32(int32, int32)')(aotnumba.utils.add)
compiler.compile()

setup(name='aotnumba', ext_modules=[compiler.distutils_extension()])

>>> import aotnumba
>>> dir(aotnumba)
['__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__', 'main.main', 'utils.add']

Compile full package with nested modules using numba Ahead-of-Time Compilation

0 Answers0