Repository: numba/numba

Issue:
Numba isn't working with multiprocessing
(From mailing list)

This program:

```
from numba import jit
from multiprocessing import Pool
from numpy import array
pool = Pool()

@jit
def wastetime(i):
  results = []
  for j in range(6**6):
    results.append(i*j)
  return array(results).sum()

results = [pool.apply_async(wastetime,(i),) for i in range(6**6)]
for r in results:
  print r.get()
```

dies with the following output:

```
Traceback (most recent call last):
  File "test.py", line 15, in <module>
    print r.get()
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/multiprocessing/pool.py", line 558, in get
    raise self._value
TypeError: CPUOverloaded object argument after * must be a sequence, not int
```

(Tested on OS X.)



Discussion:

Resolving patch:
diff --git a/numba/lowering.py b/numba/lowering.py
--- a/numba/lowering.py
+++ b/numba/lowering.py
@@ -8,7 +8,7 @@
 
 
 from numba import (_dynfunc, errcode, ir, types, cgutils, utils, config,
-                   cffi_support, typing)
+                   cffi_support, typing, six)
 
 
 class LoweringError(Exception):
@@ -30,6 +30,7 @@ def default_mangler(name, argtypes):
 # A dummy module for dynamically-generated functions
 _dynamic_modname = '<dynamic>'
 _dynamic_module = ModuleType(_dynamic_modname)
+_dynamic_module.__builtins__ = six.moves.builtins
 
 
 class FunctionDescriptor(object):
diff --git a/numba/serialize.py b/numba/serialize.py
--- a/numba/serialize.py
+++ b/numba/serialize.py
@@ -33,6 +33,8 @@ def __reduce__(self):
 
 
 def _rebuild_module(name):
+    if name is None:
+        raise ImportError("cannot import None")
     __import__(name)
     return sys.modules[name]
 
@@ -51,6 +53,9 @@ def _get_function_globals_for_reduction(func):
         # Make modules picklable by name
         if isinstance(v, ModuleType):
             globs[k] = _ModuleRef(k)
+    # Remember the module name so that the function gets a proper __module__
+    # when rebuilding.  This is used to recreate the environment.
+    globs['__name__'] = func.__module__
     return globs
 
 def _reduce_function(func):
@@ -87,6 +92,13 @@ def _rebuild_function(code_reduced, globals, name, cell_values):
     else:
         cells = ()
     code = _rebuild_code(*code_reduced)
+    modname = globals['__name__']
+    try:
+        _rebuild_module(modname)
+    except ImportError:
+        # If the module can't be found, avoid passing it (it would produce
+        # errors when lowering).
+        del globals['__name__']
     return FunctionType(code, globals, name, (), cells)
 
 def _rebuild_code(marshal_version, bytecode_magic, marshalled):


Test patch:
diff --git a/numba/tests/serialize_usecases.py b/numba/tests/serialize_usecases.py
new file mode 100644
--- /dev/null
+++ b/numba/tests/serialize_usecases.py
@@ -0,0 +1,87 @@
+"""
+Separate module with function samples for serialization tests,
+to avoid issues with __main__.
+"""
+
+import math
+
+from numba import jit, types
+
+
+@jit((types.int32, types.int32))
+def add_with_sig(a, b):
+    return a + b
+
+@jit
+def add_without_sig(a, b):
+    return a + b
+
+@jit(nopython=True)
+def add_nopython(a, b):
+    return a + b
+
+@jit(nopython=True)
+def add_nopython_fail(a, b):
+    print(a.__class__)
+    return a + b
+
+def closure(a):
+    @jit(nopython=True)
+    def inner(b, c):
+        return a + b + c
+    return inner
+
+K = 3.0
+
+from math import sqrt
+
+def closure_with_globals(x, **jit_args):
+    @jit(**jit_args)
+    def inner(y):
+        # Exercise a builtin function and a module-level constant
+        k = max(K, K + 1)
+        # Exercise two functions from another module, one accessed with
+        # dotted notation, one imported explicitly.
+        return math.hypot(x, y) + sqrt(k)
+    return inner
+
+@jit(nopython=True)
+def other_function(x, y):
+    return math.hypot(x, y)
+
+@jit(forceobj=True)
+def get_global_objmode(x):
+    return K * x
+
+def closure_calling_other_function(x):
+    @jit(nopython=True)
+    def inner(y, z):
+        return other_function(x, y) + z
+    return inner
+
+def closure_calling_other_closure(x):
+    @jit(nopython=True)
+    def other_inner(y):
+        return math.hypot(x, y)
+
+    @jit(nopython=True)
+    def inner(y):
+        return other_inner(y) + x
+    return inner
+
+
+# A dynamic function calling a builtin function
+def _get_dyn_func(**jit_args):
+    code = """
+        def dyn_func(x):
+            res = 0
+            for i in range(x):
+                res += i
+            return res
+        """
+    ns = {}
+    exec(code.strip(), ns)
+    return jit(**jit_args)(ns['dyn_func'])
+
+dyn_func = _get_dyn_func(nopython=True)
+dyn_func_objmode = _get_dyn_func(forceobj=True)
diff --git a/numba/tests/test_serialize.py b/numba/tests/test_serialize.py
--- a/numba/tests/test_serialize.py
+++ b/numba/tests/test_serialize.py
@@ -1,73 +1,14 @@
 from __future__ import print_function, absolute_import, division
 
-import math
 import pickle
 import subprocess
 import sys
 
 from numba import unittest_support as unittest
-from numba import jit, types
 from numba.targets import registry
 from numba.typeinfer import TypingError
 from .support import TestCase
-
-
-@jit((types.int32, types.int32))
-def add_with_sig(a, b):
-    return a + b
-
-@jit
-def add_without_sig(a, b):
-    return a + b
-
-@jit(nopython=True)
-def add_nopython(a, b):
-    return a + b
-
-@jit(nopython=True)
-def add_nopython_fail(a, b):
-    print(a.__class__)
-    return a + b
-
-def closure(a):
-    @jit(nopython=True)
-    def inner(b, c):
-        return a + b + c
-    return inner
-
-K = 3.0
-
-from math import sqrt
-
-def closure_with_globals(x):
-    @jit(nopython=True)
-    def inner(y):
-        # Exercise a builtin function and a module-level constant
-        k = max(K, K + 1)
-        # Exercise two functions from another module, one accessed with
-        # dotted notation, one imported explicitly.
-        return math.hypot(x, y) + sqrt(k)
-    return inner
-
-@jit(nopython=True)
-def other_function(x, y):
-    return math.hypot(x, y)
-
-def closure_calling_other_function(x):
-    @jit(nopython=True)
-    def inner(y, z):
-        return other_function(x, y) + z
-    return inner
-
-def closure_calling_other_closure(x):
-    @jit(nopython=True)
-    def other_inner(y):
-        return math.hypot(x, y)
-
-    @jit(nopython=True)
-    def inner(y):
-        return other_inner(y) + x
-    return inner
+from .serialize_usecases import *
 
 
 class TestDispatcherPickling(TestCase):
@@ -116,14 +57,23 @@ def test_call_nopython_fail(self):
         # Compilation fails
         self.run_with_protocols(self.check_call, add_nopython_fail, TypingError, (1, 2))
 
+    def test_call_objmode_with_global(self):
+        self.run_with_protocols(self.check_call, get_global_objmode, 7.5, (2.5,))
+
     def test_call_closure(self):
         inner = closure(1)
         self.run_with_protocols(self.check_call, inner, 6, (2, 3))
 
-    def test_call_closure_with_globals(self):
-        inner = closure_with_globals(3.0)
+    def check_call_closure_with_globals(self, **jit_args):
+        inner = closure_with_globals(3.0, **jit_args)
         self.run_with_protocols(self.check_call, inner, 7.0, (4.0,))
 
+    def test_call_closure_with_globals_nopython(self):
+        self.check_call_closure_with_globals(nopython=True)
+
+    def test_call_closure_with_globals_objmode(self):
+        self.check_call_closure_with_globals(forceobj=True)
+
     def test_call_closure_calling_other_function(self):
         inner = closure_calling_other_function(3.0)
         self.run_with_protocols(self.check_call, inner, 11.0, (4.0, 6.0))
@@ -132,6 +82,14 @@ def test_call_closure_calling_other_closure(self):
         inner = closure_calling_other_closure(3.0)
         self.run_with_protocols(self.check_call, inner, 8.0, (4.0,))
 
+    def test_call_dyn_func(self):
+        # Check serializing a dynamically-created function
+        self.run_with_protocols(self.check_call, dyn_func, 6, (4,))
+
+    def test_call_dyn_func_objmode(self):
+        # Same with an object mode function
+        self.run_with_protocols(self.check_call, dyn_func_objmode, 6, (4,))
+
     def test_other_process(self):
         """
         Check that reconstructing doesn't depend on resources already