Repository: numba/numba

Issue:
KeyError when caching with multiple files
First of all, thank you for maintaining and developing numba.  It's really awesome and a delight to use.

Using numba v0.42.0, I encounter a KeyError when trying to cache numba functions if I have multiple numba functions in a single file and import one, but not all of them into another file.  I have been able to reproduce this error on both OSX (High Sierra) and Linux (Ubuntu v16.04) in both python2.7 and python3.6.

This error does not arise when using numba v0.35.0.

A (very) minimal reproducible example is below.

file1.py:
<pre><code>
from file2 import function2
</pre></code>

file2.py:
<pre><code>
from numba import njit


@njit('float64(float64)', cache=True)
def function1(x):
    return x


@njit('float64(float64)', cache=True)
def function2(x):
    return x
</code></pre>


Then, running:
<pre><code>
> python file1.py
> python file2.py
</pre></code>
on the command line

results in:
<pre><code>
  File "file2.py", line 4, in <module>
    @njit('float64(float64)', cache=True)
  File "/path/to/anaconda/lib/python3.6/site-packages/numba/decorators.py", line 184, in wrapper
    disp.disable_compile()
  File "/path/to/anaconda/lib/python3.6/contextlib.py", line 89, in __exit__
    next(self.gen)
  File "/path/to/anaconda/lib/python3.6/site-packages/numba/typeinfer.py", line 729, in register_dispatcher
    del _temporary_dispatcher_map[name]
KeyError: 'function1'
</code></pre>




Discussion:
Thanks for the report (edited the code to add in a missing `:`). I can reproduce locally, seems like it's because the dispatcher registration sequence is entered twice but the map is not reference counted so the `del` deletes the sole reference when in actual fact two are live.
Following up on @stuartarchibald 's comment, would it then make sense to add a reference count map, e.g. `_temporary_dispatcher_ref_count_map`, which keeps track of the number of dispatcher registrations and only deletes once the corresponding reference count is zero? Doing this locally seems to fix the originally posted issue

Resolving patch:
diff --git a/numba/typeinfer.py b/numba/typeinfer.py
--- a/numba/typeinfer.py
+++ b/numba/typeinfer.py
@@ -19,7 +19,7 @@
 import contextlib
 import itertools
 from pprint import pprint
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 
 from numba import ir, types, utils, config, typing
 from .errors import (TypingError, UntypedAttributeError, new_error_context,
@@ -770,6 +770,9 @@ def __setitem__(self, name, value):
 
 # A temporary mapping of {function name: dispatcher object}
 _temporary_dispatcher_map = {}
+# A temporary mapping of {function name: dispatcher object reference count}
+# Reference: https://github.com/numba/numba/issues/3658
+_temporary_dispatcher_map_ref_count = defaultdict(int)
 
 
 @contextlib.contextmanager
@@ -784,10 +787,13 @@ def register_dispatcher(disp):
     assert callable(disp.py_func)
     name = disp.py_func.__name__
     _temporary_dispatcher_map[name] = disp
+    _temporary_dispatcher_map_ref_count[name] += 1
     try:
         yield
     finally:
-        del _temporary_dispatcher_map[name]
+        _temporary_dispatcher_map_ref_count[name] -= 1
+        if not _temporary_dispatcher_map_ref_count[name]:
+            del _temporary_dispatcher_map[name]
 
 
 typeinfer_extensions = {}


Test patch:
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -1641,6 +1641,56 @@ def cache_file_collision_tester(q, tempdir, modname_bar1, modname_bar2):
     q.put(r2)
 
 
+class TestCacheMultipleFilesWithSignature(unittest.TestCase):
+    # Regression test for https://github.com/numba/numba/issues/3658
+
+    _numba_parallel_test_ = False
+
+    source_text_file1 = """
+from file2 import function2
+"""
+    source_text_file2 = """
+from numba import njit
+
+@njit('float64(float64)', cache=True)
+def function1(x):
+    return x
+
+@njit('float64(float64)', cache=True)
+def function2(x):
+    return x
+"""
+
+    def setUp(self):
+        self.tempdir = temp_directory('test_cache_file_loc')
+
+        self.file1 = os.path.join(self.tempdir, 'file1.py')
+        with open(self.file1, 'w') as fout:
+            print(self.source_text_file1, file=fout)
+
+        self.file2 = os.path.join(self.tempdir, 'file2.py')
+        with open(self.file2, 'w') as fout:
+            print(self.source_text_file2, file=fout)
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+    def test_caching_mutliple_files_with_signature(self):
+        # Execute file1.py
+        popen = subprocess.Popen([sys.executable, self.file1],
+                                 stdout=subprocess.PIPE,
+                                 stderr=subprocess.PIPE)
+        out, err = popen.communicate()
+        self.assertEqual(popen.returncode, 0)
+
+        # Execute file2.py
+        popen = subprocess.Popen([sys.executable, self.file2],
+                                 stdout=subprocess.PIPE,
+                                 stderr=subprocess.PIPE)
+        out, err = popen.communicate()
+        self.assertEqual(popen.returncode, 0)
+
+
 class TestDispatcherFunctionBoundaries(TestCase):
     def test_pass_dispatcher_as_arg(self):
         # Test that a Dispatcher object can be pass as argument