ai-content-maker/.venv/Lib/site-packages/numba/cuda/tests/cudadrv/test_init.py

140 lines
4.4 KiB
Python
Raw Normal View History

2024-05-03 04:18:51 +03:00
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()