Merge pull request numba#6141 from sklam/fix/objmode_cache_segfault

Fix numba#6130 objmode cache segfault

numba/_helperlib.c

@@ -275,6 +275,12 @@ numba_recreate_record(void *pdata, int size, PyObject *dtype) {
     PyObject *record = NULL;
     PyArray_Descr *descr = NULL;
 
+    if (dtype == NULL) {
+        PyErr_Format(PyExc_RuntimeError,
+            "In 'numba_recreate_record', 'dtype' is NULL");
+        return NULL;
+    }
+
     numpy = PyImport_ImportModuleNoBlock("numpy");
     if (!numpy) goto CLEANUP;
 

numba/core/compiler.py

@@ -91,6 +91,7 @@ class Flags(utils.ConfigOptions):
              # List of functions to call to initialize on unserialization
              # (i.e cache load).
              "reload_init",
+             "referenced_envs",
              ]
 
 
@@ -158,6 +159,7 @@ def _rebuild(cls, target_context, libdata, fndesc, env,
                  call_helper=None,
                  metadata=None,  # Do not store, arbitrary & potentially large!
                  reload_init=reload_init,
+                 referenced_envs=referenced_envs,
                  )
 
         # Load Environments

numba/core/pythonapi.py

@@ -131,10 +131,28 @@ def read_const(self, index):
         """
         Look up constant number *index* inside the environment body.
         A borrowed reference is returned.
+
+        The returned LLVM value may have NULL value at runtime which indicates
+        an error at runtime.
         """
         assert index < len(self.env.consts)
 
-        return self.pyapi.list_getitem(self.env_body.consts, index)
+        builder = self.pyapi.builder
+        consts = self.env_body.consts
+        ret = cgutils.alloca_once(builder, self.pyapi.pyobj, zfill=True)
+        with builder.if_else(cgutils.is_not_null(builder, consts)) as \
+                (br_not_null, br_null):
+            with br_not_null:
+                getitem = self.pyapi.list_getitem(consts, index)
+                builder.store(getitem, ret)
+            with br_null:
+                # This can happen when the Environment is accidentally released
+                # and has subsequently been garbage collected.
+                self.pyapi.err_set_string(
+                    "PyExc_RuntimeError",
+                    "`env.consts` is NULL in `read_const`",
+                )
+        return builder.load(ret)
 
 
 _IteratorLoop = namedtuple('_IteratorLoop', ('value', 'do_break'))

numba/core/runtime/_nrt_python.c

@@ -295,6 +295,12 @@ NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
     npy_intp *shape, *strides;
     int flags = 0;
 
+    if (descr == NULL) {
+        PyErr_Format(PyExc_RuntimeError,
+                     "In 'NRT_adapt_ndarray_to_python', 'descr' is NULL");
+        return NULL;
+    }
+
     if (!NUMBA_PyArray_DescrCheck(descr)) {
         PyErr_Format(PyExc_TypeError,
                      "expected dtype object, got '%.200s'",

numba/tests/support.py

@@ -19,6 +19,7 @@
 import ctypes
 import multiprocessing as mp
 import warnings
+import traceback
 from contextlib import contextmanager
 
 import numpy as np
@@ -779,6 +780,26 @@ def run_in_new_process_caching(func, cache_dir_prefix=__name__, verbose=True):
     The childprocess's stdout and stderr will be captured and redirected to
     the current process's stdout and stderr.
 
+    Returns
+    -------
+    ret : dict
+        exitcode: 0 for success. 1 for exception-raised.
+        stdout: str
+        stderr: str
+    """
+    cache_dir = temp_directory(cache_dir_prefix)
+    return run_in_new_process_in_cache_dir(func, cache_dir, verbose=verbose)
+
+
+def run_in_new_process_in_cache_dir(func, cache_dir, verbose=True):
+    """Spawn a new process to run `func` with a temporary cache directory.
+
+    The childprocess's stdout and stderr will be captured and redirected to
+    the current process's stdout and stderr.
+
+    Similar to ``run_in_new_process_caching()`` but the ``cache_dir`` is a
+    directory path instead of a name prefix for the directory path.
+
     Returns
     -------
     ret : dict
@@ -788,7 +809,6 @@ def run_in_new_process_caching(func, cache_dir_prefix=__name__, verbose=True):
     """
     ctx = mp.get_context('spawn')
     qout = ctx.Queue()
-    cache_dir = temp_directory(cache_dir_prefix)
     with override_env_config('NUMBA_CACHE_DIR', cache_dir):
         proc = ctx.Process(target=_remote_runner, args=[func, qout])
         proc.start()

numba/tests/test_extending.py

@@ -20,6 +20,7 @@
     captured_stdout,
     temp_directory,
     override_config,
+    run_in_new_process_in_cache_dir,
 )
 from numba.core.errors import LoweringError
 import unittest
@@ -1748,12 +1749,56 @@ def testcase(x):
             got = testcase_cached(123)
             self.assertEqual(got, expect)
 
-            self.assertEqual(
-                testcase_cached._cache_hits[testcase.signatures[0]], 1,
-            )
-            self.assertEqual(
-                testcase_cached._cache_misses[testcase.signatures[0]], 0,
+    @classmethod
+    def check_objmode_cache_ndarray(cls):
+        def do_this(a, b):
+            return np.sum(a + b)
+
+        def do_something(a, b):
+            return np.sum(a + b)
+
+        @overload(do_something)
+        def overload_do_something(a, b):
+            def _do_something_impl(a, b):
+                with objmode(y='float64'):
+                    y = do_this(a, b)
+                return y
+            return _do_something_impl
+
+        @njit(cache=True)
+        def test_caching():
+            a = np.arange(20)
+            b = np.arange(20)
+            return do_something(a, b)
+
+        got = test_caching()
+        expect = test_caching.py_func()
+
+        # Check result
+        if got != expect:
+            raise AssertionError("incorrect result")
+        return test_caching
+
+    @classmethod
+    def check_objmode_cache_ndarray_check_cache(cls):
+        disp = cls.check_objmode_cache_ndarray()
+        if len(disp.stats.cache_misses) != 0:
+            raise AssertionError('unexpected cache miss')
+        if len(disp.stats.cache_hits) <= 0:
+            raise AssertionError("unexpected missing cache hit")
+
+    def test_check_objmode_cache_ndarray(self):
+        # See issue #6130.
+        # Env is missing after cache load.
+        cache_dir = temp_directory(self.__class__.__name__)
+        with override_config("CACHE_DIR", cache_dir):
+            # Test in local process to populate the cache.
+            self.check_objmode_cache_ndarray()
+            # Run in new process to use the cache in a fresh process.
+            res = run_in_new_process_in_cache_dir(
+                self.check_objmode_cache_ndarray_check_cache, cache_dir
             )
+        self.assertEqual(res['exitcode'], 0)
 
 
 class TestMisc(TestCase):