Ray 和 numba jit 不兼容

问题描述

我想结合 Ray 和 numba 进行并行计算。但是,Ray 无法使用 numba 正常工作。

如果 ray 和 numba 内核在同一个 notebook 中,则它工作正常,即

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()))

但是,如果我将以下内容放入 f.py 文件中,则会失败:

# 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.)

如果我运行,我会得到一个错误

import ray
from numba import jit

ray.init(ignore_reinit_error=True)

from f import f1

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

这是错误

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,ip=210.45.78.32)
  File "python/ray/_raylet.pyx",line 473,in ray._raylet.execute_task
  File "python/ray/_raylet.pyx",line 476,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,line 91,in _compile_cached
    retval = self._compile_core(args,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",**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",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

这很奇怪,因为什么都没发生。

解决方法

该问题在以下参考文献中得到解答:

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