GPU环境下Python+Numba编程报错:RuntimeError: Numba cannot operate on non-primary CUDA context 55d3358f3aa0

先贴出报错信息：

Traceback (most recent call last):
  File "/home/liuz/mxxx/tools/pooling_gpu.py", line 181, in <module>
    d_arr = cuda.to_device(arr)
  File "/home/liuz/anaconda3/envs/xxx/lib/python3.7/site-packages/numba/cuda/cudadrv/devices.py", line 223, in _require_cuda_context
    with _runtime.ensure_context():
  File "/home/liuz/anaconda3/envs/xxx/lib/python3.7/contextlib.py", line 112, in __enter__
    return next(self.gen)
  File "/home/liuz/anaconda3/envs/xxx/lib/python3.7/site-packages/numba/cuda/cudadrv/devices.py", line 123, in ensure_context
    newctx = self.get_or_create_context(None)
  File "/home/liuz/anaconda3/envs/xxx/lib/python3.7/site-packages/numba/cuda/cudadrv/devices.py", line 138, in get_or_create_context
    return self._get_or_create_context_uncached(devnum)
  File "/home/liuz/anaconda3/envs/xxx/lib/python3.7/site-packages/numba/cuda/cudadrv/devices.py", line 161, in _get_or_create_context_uncached
    raise RuntimeError(msg.format(ac.context_handle.value))
RuntimeError: Numba cannot operate on non-primary CUDA context 55d3358f3aa0

Process finished with exit code 1

原因：未指定GPU

服务器上有多个GPU，而Numba貌似是只能在某个GPU上编译运行。未指定GPU会导致Context初始化失败，出现以上报错信息。

解决方案

1、numba.cuda.select_device(0) 这里0是GPU的id，表示使用第一块GPU。在main函数内添加这句就可以了。
2、使用with环境，例子如下：

import numpy as np
from numba import cuda


if __name__ == "__main__":
	with cuda.gpus[0]:	# id:0，自己指定任意GPU id
		arr = np.arange(1000)
	    d_arr = cuda.to_device(arr)

注意：使用os.environ并不能解决该问题，即
os.environ[“CUDA_VISIBLE_DEVICES”] = ‘0’

Reference

https://numba.readthedocs.io/en/stable/cuda-reference/host.html
https://numba.pydata.org/numba-doc/latest/index.html