Repository: numba/numba

Issue:
Extension types loaded via entry points are registered too late
If an extension project uses [entry points](https://numba.pydata.org/numba-doc/latest/extending/entrypoints.html) and the first function to be JIT-compiled takes that project's custom types as arguments, Numba will not recognize the custom types (error message is "non-precise type pyobject").

If the custom types are used in the second function to be compiled (i.e. start by compiling a dummy, `numba.njit(lambda x: x + 1)(123)`), then the custom types are correctly recognized (no "non-precise type pyobject" error).

If the custom types are registered with

```python
@numba.extending.typeof_impl.register(custom_python_type)
def typeof(val, c):
    return custom_numba_type()
```

before the entry point gets triggered (i.e. before compiling a dummy), then the custom types are correctly recognized.

I see that the extension project used as an entry point example, [scipy-numba](https://github.com/numba/numba-scipy), doesn't define any custom types (only functions), so it hasn't encountered this issue.

I can still put most of my type and lowering definitions behind the entry point, since `numba.extending.typeof_impl.register` only has to reference the custom types inside a function definition, which is delayed. However, the register decorator itself requires Numba to be loaded, so I can't escape the Numba load-time in processes where Numba is not actually used. (I'm avoiding any time involved in constructing my custom types, but it's more the LLVM-load that I was trying to avoid.)

Since this involves setting up entry points, I don't see how to make a minimal reproducer that's really "minimal," but I'm willing to be a guinea pig for proposed fixes.


Discussion:
Thanks for reporting this, without investigating I can guess that this is as described just from knowing where the entry points are loaded in Numba. I think it perhaps may be helpful to discuss some use cases and then try and work out how to accommodate them?
You could say that Python lets you have three levels of dependency on another project:

   1. Strict load-time dependency: when your project is imported, it explicitly imports the other.
   2. Conditional load-time dependency: when your project is imported, it imports the other if the other exists on the user's system.
   3. Runtime dependency: your project might import the other if certain functions are called or certain submodules are imported by the user.

A project that has connectors to many others shouldn't depend on them all strictly, since that would require users to install everything before doing anything, even though they might only use a few of them. Conditional load-time dependencies are a little better because they don't require the user to install things, but they do have a startup time cost. (For example, Pandas adds 1.5 seconds to startup time on a Chromebook—the slowest plausible user system.) Runtime dependencies are best because then users only pay for what they use: both in installation effort and in startup time cost.

Decorating a function with `numba.extending.typeof_impl.register` can't be a runtime dependency because I have to `import numba` to use it. I had thought that the entry point mechanism would allow it to move from conditional load-time to runtime, but this type registration is the only step that prevents that.

Numba adds 1.0 seconds to startup time on a Chromebook. It's not terrible, but I just thought that this was the reason for using the entry points.
Is there any chance of this being addressed by 0.50?

The issue is that even with an entry point, things are resolved in this order:

   1. User defines a function decorated with `@nb.njit`.
   2. User creates an Awkward Array.
   3. User passes the Awkward Array to the JIT-compiled function.
   4. Numba identifies the arguments to the function and holds that information somewhere.
   5. While loading Numba machinery, entry points get triggered, including Awkward's, which registers Awkward Arrays as Numba extensions.
   6. Attempting to compile the function, Numba uses the old information about the function arguments, which was made before the entry point was triggered, so it doesn't recognize the Awkward Arrays as Numba extensions.
   7. Compilation fails with a type error.

The second time the function is invoked, the Awkward Arrays _are_ recognized as arguments and there is no type error. It's because the entry point had already been triggered before reaching the second step 4, so its information about the arguments isn't stale anymore. This inconsistency looks like a bug.

For the time being, I've resolved the order dependence by explicitly importing Numba in Awkward (if it exists on the system), which means that I really don't need to have an entry point. In fact, I don't see how entry points are useful until this gets fixed.

The reason I'd like to only import Numba if Numba gets used is in part because of the startup time (1.0 seconds on a slow computer), but also because not all users of Awkward use Numba, and the explicit dependence would cause Awkward to be non-importable if Numba is non-importable for some reason. This can happen if the user incorrectly installs Numba—a common occurrence among my users, since they usually install several versions of ROOT (not through pip or conda), which comes with a custom LLVM build, and it can happen that symbols clash. (Incidentally, this also prevents using ROOT and Julia together in any case, not just bad cases.)

I've been advocating Numba to users who wouldn't ordinarily use it, but if they get hung up in the installation and give up with a broken Numba, they'll also have a broken Awkward. I recognize that the real problem is that broken installations are possible, but like unhappy families, each situation is different.

If Numba 0.50 fixes this issue, I'd make it Awkward's minimum version. (Right now, the minimum version is 0.49 because of the API changes.)

I don't know where to look to solve the issue, but I can review the PR.
I've taken a bit of a look at this and can't reproduce with a dummy extension package I've created. The package has a dummy type and a nonsense overload of `np.linalg.norm` for that dummy type. I then have a test script which creates one of these types and passes it into a jit function that calls the norm, seemed to work, so I think I'm missing something. Do you have a way to reproduce? Having that would likely speed up finding a fix. Thanks in advance.
I'm working on a minimal reproducer now. I remember that the issue only affects new types, not new functions or new overloads of existing functions. That's why it worked fine for scumba (numba-scipy now?), because the latter is a collection of functions, not types.
Here is a minimally reproducing example:

[dummy-numba-entry-points.zip](https://github.com/numba/numba/files/4635445/dummy-numba-entry-points.zip)

```
dummy-numba-entry-points
├── dummy
│   ├── define_type.py
│   ├── __init__.py
│   └── register_numba.py
├── setup.py
└── test_dummy.py
```

Enter this directory and run

```bash
pip install .
```

From another directory (to avoid confusion with the `dummy` subdirectory in it), run `test_dummy.py`. For me, the output looks like this:

```
% python dummy-numba-entry-points/test_dummy.py 
decorating pass_through as a numba.njit function
calling that function with an Interval as an argument
registering the dummy module with Numba
defining IntervalType, its boxers, etc.
it failed!
calling that function with an Interval as an argument (again)
the function was successfully executed!
```

(and then clean up with `pip uninstall dummy`). The reason for the pip-install is to have the entry point in setup.py registered with Numba.

What we see in the output is that the new type failed to be passed as an argument into the Numbafied function the first time but succeeded the second time. This is because the registration does get triggered; it's too late for the first function invocation but not too late for the second. A user who doesn't recognize that pattern would be inclined to think that it's an intermittent bug, which looks bad.

At one point, I tried this with overloading functions, rather than types, and the issue only presents itself when overloading types. I think that's because function overloads are encountered in the body of the jitted function, but types have to be recognized when the input arguments are examined. If that input argument examination could be moved a little later or performed a second time if any modules had to be loaded with entry points (potentially introducing new types to handle), then that would fix the issue.

Thanks!
Brilliant, thanks for that, I can reproduce, will take a look.
I think I've worked it out. The dispatcher object call, which is made from C code, calls into `numba.core.dispatcher._DispatcherBase::_compile_for_args`, and this is where the types of the args is computed, immediately after this `numba.core.dispatcher._DispatcherBase::compile` is called, which triggers the extension loading, by which time it's too late.

Resolving patch:
diff --git a/numba/core/dispatcher.py b/numba/core/dispatcher.py
--- a/numba/core/dispatcher.py
+++ b/numba/core/dispatcher.py
@@ -20,6 +20,7 @@
 from numba.core.bytecode import get_code_object
 from numba.core.utils import reraise
 from numba.core.caching import NullCache, FunctionCache
+from numba.core import entrypoints
 
 
 class OmittedArg(object):
@@ -226,6 +227,16 @@ def __init__(self, arg_count, py_func, pysig, can_fallback,
         self._compiling_counter = _CompilingCounter()
         weakref.finalize(self, self._make_finalizer())
 
+    def _compilation_chain_init_hook(self):
+        """
+        This will be called ahead of any part of compilation taking place (this
+        even includes being ahead of working out the types of the arguments).
+        This permits activities such as initialising extension entry points so
+        that the compiler knows about additional externally defined types etc
+        before it does anything.
+        """
+        entrypoints.init_all()
+
     def _reset_overloads(self):
         self._clear()
         self.overloads.clear()
@@ -329,6 +340,9 @@ def _compile_for_args(self, *args, **kws):
         for the given *args* and *kws*, and return the resulting callable.
         """
         assert not kws
+        # call any initialisation required for the compilation chain (e.g.
+        # extension point registration).
+        self._compilation_chain_init_hook()
 
         def error_rewrite(e, issue_type):
             """


Test patch:
diff --git a/numba/tests/test_entrypoints.py b/numba/tests/test_entrypoints.py
--- a/numba/tests/test_entrypoints.py
+++ b/numba/tests/test_entrypoints.py
@@ -1,10 +1,25 @@
 import sys
 import types
 import warnings
-
+import unittest
+import os
+import subprocess
+import threading
 import pkg_resources
 
+from numba import njit
 from numba.tests.support import TestCase
+from numba.testing.main import _TIMEOUT as _RUNNER_TIMEOUT
+
+_TEST_TIMEOUT = _RUNNER_TIMEOUT - 60.
+
+
+class _DummyClass(object):
+    def __init__(self, value):
+        self.value = value
+
+    def __repr__(self):
+        return '_DummyClass(%f, %f)' % self.value
 
 
 class TestEntrypoints(TestCase):
@@ -35,9 +50,9 @@ def init_function():
             dist = "numba"
             entrypoints = pkg_resources.get_entry_map(dist)
             my_entrypoint = pkg_resources.EntryPoint(
-                "init", # name of entry point
-                mod.__name__, # module with entry point object
-                attrs=['init_func'], # name of entry point object
+                "init",  # name of entry point
+                mod.__name__,  # module with entry point object
+                attrs=['init_func'],  # name of entry point object
                 dist=pkg_resources.get_distribution(dist)
             )
             entrypoints.setdefault('numba_extensions',
@@ -84,9 +99,9 @@ def init_function():
             dist = "numba"
             entrypoints = pkg_resources.get_entry_map(dist)
             my_entrypoint = pkg_resources.EntryPoint(
-                "init", # name of entry point
-                mod.__name__, # module with entry point object
-                attrs=['init_func'], # name of entry point object
+                "init",  # name of entry point
+                mod.__name__,  # module with entry point object
+                attrs=['init_func'],  # name of entry point object
                 dist=pkg_resources.get_distribution(dist)
             )
             entrypoints.setdefault('numba_extensions',
@@ -113,3 +128,119 @@ def init_function():
             # remove fake module
             if mod.__name__ in sys.modules:
                 del sys.modules[mod.__name__]
+
+    _EP_MAGIC_TOKEN = 'RUN_ENTRY'
+
+    @unittest.skipIf(os.environ.get('_EP_MAGIC_TOKEN', None) != _EP_MAGIC_TOKEN,
+                     "needs token")
+    def test_entrypoint_handles_type_extensions(self):
+        # loosely based on Pandas test from:
+        #   https://github.com/pandas-dev/pandas/pull/27488
+        import numba
+
+        def init_function():
+            # This init function would normally just call a module init via
+            # import or similar, for the sake of testing, inline registration
+            # of how to handle the global "_DummyClass".
+            class DummyType(numba.types.Type):
+                def __init__(self):
+                    super(DummyType, self).__init__(name='DummyType')
+
+            @numba.extending.typeof_impl.register(_DummyClass)
+            def typer_DummyClass(val, c):
+                return DummyType()
+
+            @numba.extending.register_model(DummyType)
+            class DummyModel(numba.extending.models.StructModel):
+                def __init__(self, dmm, fe_type):
+                    members = [
+                        ('value', numba.types.float64), ]
+                    super(DummyModel, self).__init__(dmm, fe_type, members)
+
+            @numba.extending.unbox(DummyType)
+            def unbox_dummy(typ, obj, c):
+                value_obj = c.pyapi.object_getattr_string(obj, "value")
+                dummy_struct_proxy = numba.core.cgutils.create_struct_proxy(typ)
+                dummy_struct = dummy_struct_proxy(c.context, c.builder)
+                dummy_struct.value = c.pyapi.float_as_double(value_obj)
+                c.pyapi.decref(value_obj)
+                err_flag = c.pyapi.err_occurred()
+                is_error = numba.core.cgutils.is_not_null(c.builder, err_flag)
+                return numba.extending.NativeValue(dummy_struct._getvalue(),
+                                                   is_error=is_error)
+
+            @numba.extending.box(DummyType)
+            def box_dummy(typ, val, c):
+                dummy_struct_proxy = numba.core.cgutils.create_struct_proxy(typ)
+                dummy_struct = dummy_struct_proxy(c.context, c.builder)
+                value_obj = c.pyapi.float_from_double(dummy_struct.value)
+                serialized_clazz = c.pyapi.serialize_object(_DummyClass)
+                class_obj = c.pyapi.unserialize(serialized_clazz)
+                res = c.pyapi.call_function_objargs(class_obj, (value_obj,))
+                c.pyapi.decref(value_obj)
+                c.pyapi.decref(class_obj)
+                return res
+
+        mod = types.ModuleType("_test_numba_init_sequence")
+        mod.init_func = init_function
+
+        try:
+            # will remove this module at the end of the test
+            sys.modules[mod.__name__] = mod
+
+            # We are registering an entry point using the "numba" package
+            # ("distribution" in pkg_resources-speak) itself, though these are
+            # normally registered by other packages.
+            dist = "numba"
+            entrypoints = pkg_resources.get_entry_map(dist)
+            my_entrypoint = pkg_resources.EntryPoint(
+                "init",  # name of entry point
+                mod.__name__,  # module with entry point object
+                attrs=['init_func'],  # name of entry point object
+                dist=pkg_resources.get_distribution(dist)
+            )
+            entrypoints.setdefault('numba_extensions',
+                                   {})['init'] = my_entrypoint
+
+            @njit
+            def foo(x):
+                return x
+
+            ival = _DummyClass(10)
+            foo(ival)
+        finally:
+            # remove fake module
+            if mod.__name__ in sys.modules:
+                del sys.modules[mod.__name__]
+
+    def run_cmd(self, cmdline, env):
+        popen = subprocess.Popen(cmdline,
+                                 stdout=subprocess.PIPE,
+                                 stderr=subprocess.PIPE,
+                                 env=env)
+        # finish in _TEST_TIMEOUT seconds or kill it
+        timeout = threading.Timer(_TEST_TIMEOUT, popen.kill)
+        try:
+            timeout.start()
+            out, err = popen.communicate()
+            if popen.returncode != 0:
+                raise AssertionError(
+                    "process failed with code %s: stderr follows\n%s\n" %
+                    (popen.returncode, err.decode()))
+            return out.decode(), err.decode()
+        finally:
+            timeout.cancel()
+        return None, None
+
+    def test_entrypoint_extension_sequence(self):
+        env_copy = os.environ.copy()
+        env_copy['_EP_MAGIC_TOKEN'] = str(self._EP_MAGIC_TOKEN)
+        themod = self.__module__
+        thecls = type(self).__name__
+        methname = 'test_entrypoint_handles_type_extensions'
+        injected_method = '%s.%s.%s' % (themod, thecls, methname)
+        cmdline = [sys.executable, "-m", "numba.runtests", injected_method]
+        out, err = self.run_cmd(cmdline, env_copy)
+        _DEBUG = False
+        if _DEBUG:
+            print(out, err)