Search code examples
pythonnumbaray

Ray and numba jit are not compatible


I want to combine Ray and numba for parallel computing. However, Ray cannot work properly with numba.

If ray and numba kernel are in a same notebook, it is working properly, i.e.,

import ray
from numba import jit

ray.init(ignore_reinit_error=True)

@ray.remote
def f1():

    @jit
    def compute(x,y):
        return x+y
    re = compute(1,2)
    return re
print(ray.get(f1.remote()))

However, if I put the following in the f.py file, it fails:

# the following is in f.py
import ray
@ray.remote
def f1():
    
    from numba import jit, float64

    @jit
    def f(x, y):
        # A somewhat trivial example
        return x + y
    
    return f(1., 2.)

Then I get an error if I run:

import ray
from numba import jit

ray.init(ignore_reinit_error=True)

from f import f1

print(ray.get(f1.remote()))

this is the error:

2021-03-16 16:42:07,067 INFO services.py:1172 -- View the Ray dashboard at http://127.0.0.1:8265
(pid=3673426) /data2/zhangjuenjie/zjj/CRBMG/f.py:9: NumbaWarning: 
(pid=3673426) Compilation is falling back to object mode WITH looplifting enabled because Function f failed at nopython mode lowering due to: can't compile f1.<locals>.f: import of module f failed
(pid=3673426)   @jit
---------------------------------------------------------------------------
RayTaskError(ModuleNotFoundError)         Traceback (most recent call last)
<ipython-input-1-78c9292bc84b> in <module>
      6 from f import f1
      7 
----> 8 print(ray.get(f1.remote()))

~/anaconda3/envs/crbmg/lib/python3.8/site-packages/ray/_private/client_mode_hook.py in wrapper(*args, **kwargs)
     45         if client_mode_enabled and _client_hook_enabled:
     46             return getattr(ray, func.__name__)(*args, **kwargs)
---> 47         return func(*args, **kwargs)
     48 
     49     return wrapper

~/anaconda3/envs/crbmg/lib/python3.8/site-packages/ray/worker.py in get(object_refs, timeout)
   1454                     worker.core_worker.dump_object_store_memory_usage()
   1455                 if isinstance(value, RayTaskError):
-> 1456                     raise value.as_instanceof_cause()
   1457                 else:
   1458                     raise value

RayTaskError(ModuleNotFoundError): ray::f1() (pid=3673426, ip=210.45.78.32)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/weakref.py", line 131, in __getitem__
    o = self.data[key]()
KeyError: '_ZN08NumbaEnv1f2f112$3clocals$3e5f$241E8pyobject8pyobject'

During handling of the above exception, another exception occurred:

ray::f1() (pid=3673426, ip=210.45.78.32)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/funcdesc.py", line 93, in lookup_module
    return sys.modules[self.modname]
KeyError: 'f'

During handling of the above exception, another exception occurred:

ray::f1() (pid=3673426, ip=210.45.78.32)
  File "python/ray/_raylet.pyx", line 473, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 476, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 480, in ray._raylet.execute_task
  File "/data2/zhangjuenjie/zjj/CRBMG/f.py", line 14, in f1
    return f(1., 2.)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/dispatcher.py", line 433, in _compile_for_args
    raise e
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/dispatcher.py", line 366, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/dispatcher.py", line 857, in compile
    cres = self._compiler.compile(args, return_type)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/dispatcher.py", line 77, in compile
    status, retval = self._compile_cached(args, return_type)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/dispatcher.py", line 91, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/dispatcher.py", line 104, in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler.py", line 602, in compile_extra
    return pipeline.compile_extra(func)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler.py", line 352, in compile_extra
    return self._compile_bytecode()
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler.py", line 414, in _compile_bytecode
    return self._compile_core()
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler.py", line 394, in _compile_core
    raise e
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler.py", line 385, in _compile_core
    pm.run(self.state)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler_machinery.py", line 339, in run
    raise patched_exception
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler_machinery.py", line 330, in run
    self._runPass(idx, pass_inst, state)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler_machinery.py", line 289, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/compiler_machinery.py", line 262, in check
    mangled = func(compiler_state)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/object_mode_passes.py", line 120, in run_pass
    lowered = backend_object_mode()
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/object_mode_passes.py", line 115, in backend_object_mode
    return self._py_lowering_stage(state.targetctx,
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/object_mode_passes.py", line 77, in _py_lowering_stage
    lower = pylowering.PyLower(targetctx, library, fndesc, interp)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/lowering.py", line 37, in __init__
    self.env = Environment.from_fndesc(self.fndesc)
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/environment.py", line 23, in from_fndesc
    inst = cls(fndesc.lookup_globals())
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/funcdesc.py", line 81, in lookup_globals
    return self.global_dict or self.lookup_module().__dict__
  File "/data2/zhangjuenjie/anaconda3/envs/crbmg/lib/python3.8/site-packages/numba/core/funcdesc.py", line 95, in lookup_module
    raise ModuleNotFoundError(
ModuleNotFoundError: can't compile f1.<locals>.f: import of module f failed

This is very strange since nothing really happened.


Solution

  • The question is answered in the following ref:

    https://discuss.ray.io/t/ray-is-not-compatible-with-numba-cuda-when-import/1274/5