140 lines
4.4 KiB
Python
140 lines
4.4 KiB
Python
|
import multiprocessing as mp
|
||
|
|
import os
|
||
|
|
|
||
|
|
from numba import cuda
|
||
|
|
from numba.cuda.cudadrv.driver import CudaAPIError, driver
|
||
|
|
from numba.cuda.cudadrv.error import CudaSupportError
|
||
|
|
from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
|
||
|
|
|
||
|
|
|
||
|
|
# A mock of cuInit that always raises a CudaAPIError
|
||
|
|
def cuInit_raising(arg):
|
||
|
|
raise CudaAPIError(999, 'CUDA_ERROR_UNKNOWN')
|
||
|
|
|
||
|
|
|
||
|
|
# Test code to run in a child that patches driver.cuInit to a variant that
|
||
|
|
# always raises. We can't use mock.patch.object here because driver.cuInit is
|
||
|
|
# not assigned until we attempt to initialize - mock.patch.object cannot locate
|
||
|
|
# the non-existent original method, and so fails. Instead we patch
|
||
|
|
# driver.cuInit with our raising version prior to any attempt to initialize.
|
||
|
|
def cuInit_raising_test(result_queue):
|
||
|
|
driver.cuInit = cuInit_raising
|
||
|
|
|
||
|
|
success = False
|
||
|
|
msg = None
|
||
|
|
|
||
|
|
try:
|
||
|
|
# A CUDA operation that forces initialization of the device
|
||
|
|
cuda.device_array(1)
|
||
|
|
except CudaSupportError as e:
|
||
|
|
success = True
|
||
|
|
msg = e.msg
|
||
|
|
|
||
|
|
result_queue.put((success, msg))
|
||
|
|
|
||
|
|
|
||
|
|
# Similar to cuInit_raising_test above, but for testing that the string
|
||
|
|
# returned by cuda_error() is as expected.
|
||
|
|
def initialization_error_test(result_queue):
|
||
|
|
driver.cuInit = cuInit_raising
|
||
|
|
|
||
|
|
success = False
|
||
|
|
msg = None
|
||
|
|
|
||
|
|
try:
|
||
|
|
# A CUDA operation that forces initialization of the device
|
||
|
|
cuda.device_array(1)
|
||
|
|
except CudaSupportError:
|
||
|
|
success = True
|
||
|
|
|
||
|
|
msg = cuda.cuda_error()
|
||
|
|
result_queue.put((success, msg))
|
||
|
|
|
||
|
|
|
||
|
|
# For testing the path where Driver.__init__() catches a CudaSupportError
|
||
|
|
def cuda_disabled_test(result_queue):
|
||
|
|
success = False
|
||
|
|
msg = None
|
||
|
|
|
||
|
|
try:
|
||
|
|
# A CUDA operation that forces initialization of the device
|
||
|
|
cuda.device_array(1)
|
||
|
|
except CudaSupportError as e:
|
||
|
|
success = True
|
||
|
|
msg = e.msg
|
||
|
|
|
||
|
|
result_queue.put((success, msg))
|
||
|
|
|
||
|
|
|
||
|
|
# Similar to cuda_disabled_test, but checks cuda.cuda_error() instead of the
|
||
|
|
# exception raised on initialization
|
||
|
|
def cuda_disabled_error_test(result_queue):
|
||
|
|
success = False
|
||
|
|
msg = None
|
||
|
|
|
||
|
|
try:
|
||
|
|
# A CUDA operation that forces initialization of the device
|
||
|
|
cuda.device_array(1)
|
||
|
|
except CudaSupportError:
|
||
|
|
success = True
|
||
|
|
|
||
|
|
msg = cuda.cuda_error()
|
||
|
|
result_queue.put((success, msg))
|
||
|
|
|
||
|
|
|
||
|
|
@skip_on_cudasim('CUDA Simulator does not initialize driver')
|
||
|
|
class TestInit(CUDATestCase):
|
||
|
|
def _test_init_failure(self, target, expected):
|
||
|
|
# Run the initialization failure test in a separate subprocess
|
||
|
|
ctx = mp.get_context('spawn')
|
||
|
|
result_queue = ctx.Queue()
|
||
|
|
proc = ctx.Process(target=target, args=(result_queue,))
|
||
|
|
proc.start()
|
||
|
|
proc.join(30) # should complete within 30s
|
||
|
|
success, msg = result_queue.get()
|
||
|
|
|
||
|
|
# Ensure the child process raised an exception during initialization
|
||
|
|
# before checking the message
|
||
|
|
if not success:
|
||
|
|
self.fail('CudaSupportError not raised')
|
||
|
|
|
||
|
|
self.assertIn(expected, msg)
|
||
|
|
|
||
|
|
def test_init_failure_raising(self):
|
||
|
|
expected = 'Error at driver init: CUDA_ERROR_UNKNOWN (999)'
|
||
|
|
self._test_init_failure(cuInit_raising_test, expected)
|
||
|
|
|
||
|
|
def test_init_failure_error(self):
|
||
|
|
expected = 'CUDA_ERROR_UNKNOWN (999)'
|
||
|
|
self._test_init_failure(initialization_error_test, expected)
|
||
|
|
|
||
|
|
def _test_cuda_disabled(self, target):
|
||
|
|
# Uses _test_init_failure to launch the test in a separate subprocess
|
||
|
|
# with CUDA disabled.
|
||
|
|
cuda_disabled = os.environ.get('NUMBA_DISABLE_CUDA')
|
||
|
|
os.environ['NUMBA_DISABLE_CUDA'] = "1"
|
||
|
|
try:
|
||
|
|
expected = 'CUDA is disabled due to setting NUMBA_DISABLE_CUDA=1'
|
||
|
|
self._test_init_failure(cuda_disabled_test, expected)
|
||
|
|
finally:
|
||
|
|
if cuda_disabled is not None:
|
||
|
|
os.environ['NUMBA_DISABLE_CUDA'] = cuda_disabled
|
||
|
|
else:
|
||
|
|
os.environ.pop('NUMBA_DISABLE_CUDA')
|
||
|
|
|
||
|
|
def test_cuda_disabled_raising(self):
|
||
|
|
self._test_cuda_disabled(cuda_disabled_test)
|
||
|
|
|
||
|
|
def test_cuda_disabled_error(self):
|
||
|
|
self._test_cuda_disabled(cuda_disabled_error_test)
|
||
|
|
|
||
|
|
def test_init_success(self):
|
||
|
|
# Here we assume that initialization is successful (because many bad
|
||
|
|
# things will happen with the test suite if it is not) and check that
|
||
|
|
# there is no error recorded.
|
||
|
|
self.assertIsNone(cuda.cuda_error())
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|