82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
|
import os
|
||
|
import sys
|
||
|
import subprocess
|
||
|
import threading
|
||
|
from numba import cuda
|
||
|
from numba.cuda.testing import (unittest, CUDATestCase, skip_on_cudasim,
|
||
|
skip_under_cuda_memcheck)
|
||
|
from numba.tests.support import captured_stdout
|
||
|
|
||
|
|
||
|
class TestCudaDetect(CUDATestCase):
|
||
|
def test_cuda_detect(self):
|
||
|
# exercise the code path
|
||
|
with captured_stdout() as out:
|
||
|
cuda.detect()
|
||
|
output = out.getvalue()
|
||
|
self.assertIn('Found', output)
|
||
|
self.assertIn('CUDA devices', output)
|
||
|
|
||
|
|
||
|
@skip_under_cuda_memcheck('Hangs cuda-memcheck')
|
||
|
class TestCUDAFindLibs(CUDATestCase):
|
||
|
|
||
|
def run_cmd(self, cmdline, env):
|
||
|
popen = subprocess.Popen(cmdline,
|
||
|
stdout=subprocess.PIPE,
|
||
|
stderr=subprocess.PIPE,
|
||
|
env=env)
|
||
|
|
||
|
# finish in 5 minutes or kill it
|
||
|
timeout = threading.Timer(5 * 60., popen.kill)
|
||
|
try:
|
||
|
timeout.start()
|
||
|
out, err = popen.communicate()
|
||
|
# the process should exit with an error
|
||
|
return out.decode(), err.decode()
|
||
|
finally:
|
||
|
timeout.cancel()
|
||
|
return None, None
|
||
|
|
||
|
def run_test_in_separate_process(self, envvar, envvar_value):
|
||
|
env_copy = os.environ.copy()
|
||
|
env_copy[envvar] = str(envvar_value)
|
||
|
code = """if 1:
|
||
|
from numba import cuda
|
||
|
@cuda.jit('(int64,)')
|
||
|
def kernel(x):
|
||
|
pass
|
||
|
kernel(1,)
|
||
|
"""
|
||
|
cmdline = [sys.executable, "-c", code]
|
||
|
return self.run_cmd(cmdline, env_copy)
|
||
|
|
||
|
@skip_on_cudasim('Simulator does not hit device library search code path')
|
||
|
@unittest.skipIf(not sys.platform.startswith('linux'), "linux only")
|
||
|
def test_cuda_find_lib_errors(self):
|
||
|
"""
|
||
|
This tests that the find_libs works as expected in the case of an
|
||
|
environment variable being used to set the path.
|
||
|
"""
|
||
|
# one of these is likely to exist on linux, it's also unlikely that
|
||
|
# someone has extracted the contents of libdevice into here!
|
||
|
locs = ['lib', 'lib64']
|
||
|
|
||
|
looking_for = None
|
||
|
for l in locs:
|
||
|
looking_for = os.path.join(os.path.sep, l)
|
||
|
if os.path.exists(looking_for):
|
||
|
break
|
||
|
|
||
|
# This is the testing part, the test will only run if there's a valid
|
||
|
# path in which to look
|
||
|
if looking_for is not None:
|
||
|
out, err = self.run_test_in_separate_process("NUMBA_CUDA_DRIVER",
|
||
|
looking_for)
|
||
|
self.assertTrue(out is not None)
|
||
|
self.assertTrue(err is not None)
|
||
|
|
||
|
|
||
|
if __name__ == '__main__':
|
||
|
unittest.main()
|