Repository: numba/numba

Issue:
Universal function made with vectorize is not thread safe
I'm using the `numba.vectorize` decorator to make a ufunc.  This program crashes, but I don't think it should:
```
import concurrent.futures
import numba
import numpy as np

@numba.vectorize
def between(value, low, high):
    return low <= value <= high

def task(n):
    arr = np.random.randn(100000)
    between(arr, -0.5, 1)
    print(n, "ok")

with concurrent.futures.ThreadPoolExecutor(10) as pool:
    for result in pool.map(task, range(10)):
        pass
```

The traceback when it crashes:
```
  File "./t.py", line 13, in task
    between(arr, -0.5, 1)
  File "numba/npyufunc/dufunc.py", line 164, in _compile_for_args
    return self._compile_for_argtys(tuple(argtys))
  File "numba/npyufunc/dufunc.py", line 182, in _compile_for_argtys
    self._dispatcher, self.targetoptions, sig)
  File "numba/npyufunc/ufuncbuilder.py", line 111, in _compile_element_wise_function
    cres = nb_func.compile(sig, **targetoptions)
  File "numba/npyufunc/ufuncbuilder.py", line 78, in compile
    self.overloads[cres.signature] = cres
  File "numba/utils.py", line 261, in __setitem__
    raise AssertionError("key already in dictionary: %r" % (key,))
AssertionError: key already in dictionary: (float64, float64, int64) -> bool
```

The expected result can be seen by changing `numba.vectorize` to `np.vectorize` (ordering of the output does not matter):
```
0 ok
1 ok
2 ok
4 ok
3 ok
5 ok
6 ok
8 ok
9 ok
7 ok
```

It seems to me that `vectorize` is running Python code while the GIL is not held--at least during signature lookup as seen in the above traceback.

I am using Python 3.5.1, NumPy 1.11.1, and Numba 0.33.0 on 64-bit Linux.

#1096 may have the same root cause.



Discussion:

Resolving patch:
diff --git a/numba/npyufunc/ufuncbuilder.py b/numba/npyufunc/ufuncbuilder.py
--- a/numba/npyufunc/ufuncbuilder.py
+++ b/numba/npyufunc/ufuncbuilder.py
@@ -3,6 +3,7 @@
 
 import warnings
 import inspect
+from contextlib import contextmanager
 
 import numpy as np
 
@@ -74,9 +75,7 @@ def compile(self, sig, locals={}, **targetoptions):
         # The feature requires a real python function
         flags.unset("enable_looplift")
 
-        cres = self._compile_core(sig, flags, locals)
-        self.overloads[cres.signature] = cres
-        return cres
+        return self._compile_core(sig, flags, locals)
 
     def _compile_core(self, sig, flags, locals):
         """
@@ -85,19 +84,38 @@ def _compile_core(self, sig, flags, locals):
         """
         typingctx = self.targetdescr.typing_context
         targetctx = self.targetdescr.target_context
+
+        @contextmanager
+        def store_overloads_on_success():
+            # use to ensure overloads are stored on success
+            try:
+                yield
+            except:
+                raise
+            else:
+                exists = self.overloads.get(cres.signature)
+                if exists is None:
+                    self.overloads[cres.signature] = cres
+
         # Use cache and compiler in a critical section
         with compiler.lock_compiler:
-            cres = self.cache.load_overload(sig, targetctx)
-            if cres is not None:
-                # Use cached version
+            with store_overloads_on_success():
+                # attempt look up of existing
+                cres = self.cache.load_overload(sig, targetctx)
+                if cres is not None:
+                    return cres
+
+                # Compile
+                args, return_type = sigutils.normalize_signature(sig)
+                cres = compiler.compile_extra(typingctx, targetctx,
+                                              self.py_func, args=args,
+                                              return_type=return_type,
+                                              flags=flags, locals=locals)
+
+                # cache lookup failed before so safe to save
+                self.cache.save_overload(sig, cres)
+
                 return cres
-            # Compile
-            args, return_type = sigutils.normalize_signature(sig)
-            cres = compiler.compile_extra(typingctx, targetctx, self.py_func,
-                                        args=args, return_type=return_type,
-                                        flags=flags, locals=locals)
-            self.cache.save_overload(sig, cres)
-            return cres
 
 
 dispatcher_registry['npyufunc'] = UFuncDispatcher


Test patch:
diff --git a/numba/tests/test_ufuncs.py b/numba/tests/test_ufuncs.py
--- a/numba/tests/test_ufuncs.py
+++ b/numba/tests/test_ufuncs.py
@@ -5,6 +5,7 @@
 import re
 import sys
 import warnings
+import threading
 
 import numpy as np
 
@@ -1806,6 +1807,34 @@ def func(x):
         self.assertRaises(TypingError, compile_isolated, func, [types.float64],
                           return_type=types.float64, flags=self._compile_flags)
 
+class TestUFuncCompilationThreadSafety(TestCase):
+
+    def test_lock(self):
+        """
+        Test that (lazy) compiling from several threads at once doesn't
+        produce errors (see issue #2403).
+        """
+        errors = []
+
+        @vectorize
+        def foo(x):
+            return x + 1
+
+        def wrapper():
+            try:
+                a = np.ones((10,), dtype = np.float64)
+                expected = np.ones((10,), dtype = np.float64) + 1.
+                np.testing.assert_array_equal(foo(a), expected)
+            except BaseException as e:
+                errors.append(e)
+
+        threads = [threading.Thread(target=wrapper) for i in range(16)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        self.assertFalse(errors)
+
 
 if __name__ == '__main__':
     unittest.main()