Repository: numba/numba
Pull request: Add reference counting to register_dispatcher
This PR adds a reference counter for the dispatcher registration mapping in `numba/typeinfer.py` (i.e. `_temporary_dispatcher_map`) to avoid deleting a dispatcher when it has multiple references. Note I added the failing example from the originally posted issue, which involves creating a couple of temporary modules to execute, as a regression test. Any suggestions on how this might be improved is welcome. 

Fixes #3658

Problem statement:
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>

Hints from issue 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

Commit message: Use tearDown methods to remove temp dirs

Patch to review:
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
index 2bd6181fd7b..074abd7ffb4 100644
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -140,7 +140,7 @@ def check_access_is_preventable():
             ret = True
     finally:
         os.chmod(test_dir, 0o775)
-        shutil.rmtree(test_dir)
+        shutil.rmtree(tempdir)
     return ret
 
 _access_preventable = check_access_is_preventable()
@@ -934,6 +934,7 @@ def setUp(self):
     def tearDown(self):
         sys.modules.pop(self.modname, None)
         sys.path.remove(self.tempdir)
+        shutil.rmtree(self.tempdir)
 
     def import_module(self):
         # Import a fresh version of the test module.  All jitted functions
@@ -1310,9 +1311,9 @@ def test_non_creatable_pycache(self):
         # Make it impossible to create the __pycache__ directory
         old_perms = os.stat(self.tempdir).st_mode
         os.chmod(self.tempdir, 0o500)
-        self.addCleanup(os.chmod, self.tempdir, old_perms)
 
         self._test_pycache_fallback()
+        os.chmod(self.tempdir, old_perms)
 
     @skip_bad_access
     @unittest.skipIf(os.name == "nt",
@@ -1323,9 +1324,9 @@ def test_non_writable_pycache(self):
         os.mkdir(pycache)
         old_perms = os.stat(pycache).st_mode
         os.chmod(pycache, 0o500)
-        self.addCleanup(os.chmod, pycache, old_perms)
 
         self._test_pycache_fallback()
+        os.chmod(pycache, old_perms)
 
     def test_ipython(self):
         # Test caching in an IPython session
@@ -1548,6 +1549,7 @@ def tearDown(self):
         sys.modules.pop(self.modname_bar1, None)
         sys.modules.pop(self.modname_bar2, None)
         sys.path.remove(self.tempdir)
+        shutil.rmtree(self.tempdir)
 
     def import_bar1(self):
         return import_dynamic(self.modname_bar1).bar
@@ -1641,6 +1643,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
diff --git a/numba/typeinfer.py b/numba/typeinfer.py
index 82609c12996..944ab2f6912 100644
--- 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 = {}


Code review comments:
diff_hunk: @@ -934,6 +934,7 @@ def setUp(self):
     def tearDown(self):
         sys.modules.pop(self.modname, None)
         sys.path.remove(self.tempdir)
+        shutil.rmtree(self.tempdir)
path: numba/tests/test_dispatcher.py
text: @user1:
I think that because this class is a base class of other test classes it is now not safe for concurrent execution. The `temp_directory` call in the `setUp` is safe for concurrent use by virtue of ignoring `error.EEXIST` on `os.mkdir` calls. However, if tests were interleaved such that a `tearDown` occurred whilst another test that relied on the cache root directory were running, then the cache root directory would be unconditionally removed, this potentially causing problems.

Three potential fixes:
1. Add `_numba_parallel_test_ = False` as a member variable to this base class, this will serialize execution of all such tests (alternatively inherit from `numba.tests.support.SerialMixin` which has the same effect). This is probably less desirable as a fix as there's already a lot of execution time wasted in tests that cannot run concurrently.
2. Revoke these changes and make just the necessary changes to do removal of the directory added in the original patches and defer longer term fixes to https://github.com/numba/numba/issues/4274.
3. Something else?! Suggestions welcomed.

@author:
Thanks for pointing that out @user1. I prefer option 2 for this PR. As you mentioned, we can revisit any issues surrounding `temp_directory` following discussion in #4274