diff options
Diffstat (limited to 'tests')
-rw-r--r-- | tests/test_cuda_setup_evaluator.py | 67 |
1 files changed, 49 insertions, 18 deletions
diff --git a/tests/test_cuda_setup_evaluator.py b/tests/test_cuda_setup_evaluator.py index 119e21a..3d34c29 100644 --- a/tests/test_cuda_setup_evaluator.py +++ b/tests/test_cuda_setup_evaluator.py @@ -7,10 +7,38 @@ from typing import List, NamedTuple from bitsandbytes.cuda_setup import ( CUDA_RUNTIME_LIB, evaluate_cuda_setup, - get_cuda_runtime_lib_path, - tokenize_paths, + determine_cuda_runtime_lib_path, + extract_candidate_paths, ) +""" +'LD_LIBRARY_PATH': ':/mnt/D/titus/local/cuda-11.1/lib64/' +'CONDA_EXE': '/mnt/D/titus/miniconda/bin/conda' +'LESSCLOSE': '/usr/bin/lesspipe %s %s' +'OLDPWD': '/mnt/D/titus/src' +'CONDA_PREFIX': '/mnt/D/titus/miniconda/envs/8-bit' +'SSH_AUTH_SOCK': '/mnt/D/titus/.ssh/ssh-agent.tim-uw.sock' +'CONDA_PREFIX_1': '/mnt/D/titus/miniconda' +'PWD': '/mnt/D/titus/src/8-bit' +'HOME': '/mnt/D/titus' +'CONDA_PYTHON_EXE': '/mnt/D/titus/miniconda/bin/python' +'CUDA_HOME': '/mnt/D/titus/local/cuda-11.1/' +'TMUX': '/tmp/tmux-1007/default,59286,1' +'XDG_DATA_DIRS': '/usr/local/share:/usr/share:/var/lib/snapd/desktop' +'SSH_TTY': '/dev/pts/0' +'MAIL': '/var/mail/titus' +'SHELL': '/bin/bash' +'DBUS_SESSION_BUS_ADDRESS': 'unix:path=/run/user/1007/bus' +'XDG_RUNTIME_DIR': '/run/user/1007' +'PATH': '/mnt/D/titus/miniconda/envs/8-bit/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/mnt/D/titus/local/cuda-11.1/bin' +'LESSOPEN': '| /usr/bin/lesspipe %s' +'_': '/mnt/D/titus/miniconda/envs/8-bit/bin/python' +# any that include 'CONDA' that are not 'CONDA_PREFIX' + +# we search for +'CUDA_HOME': '/mnt/D/titus/local/cuda-11.1/' +""" + class InputAndExpectedOutput(NamedTuple): input: str @@ -47,20 +75,20 @@ HAPPY_PATH__LD_LIB_TEST_PATHS: List[InputAndExpectedOutput] = [ @pytest.fixture(params=HAPPY_PATH__LD_LIB_TEST_PATHS) def happy_path_path_string(tmpdir, request): - for path in tokenize_paths(request.param): + for path in extract_candidate_paths(request.param): test_dir.mkdir() if CUDA_RUNTIME_LIB in path: (test_input / CUDA_RUNTIME_LIB).touch() @pytest.mark.parametrize("test_input, expected", HAPPY_PATH__LD_LIB_TEST_PATHS) -def test_get_cuda_runtime_lib_path__happy_path( +def test_determine_cuda_runtime_lib_path__happy_path( tmp_path, test_input: str, expected: str ): - for path in tokenize_paths(test_input): + for path in extract_candidate_paths(test_input): path.mkdir() (path / CUDA_RUNTIME_LIB).touch() - assert get_cuda_runtime_lib_path(test_input) == expected + assert determine_cuda_runtime_lib_path(test_input) == expected UNHAPPY_PATH__LD_LIB_TEST_PATHS = [ @@ -70,21 +98,21 @@ UNHAPPY_PATH__LD_LIB_TEST_PATHS = [ @pytest.mark.parametrize("test_input", UNHAPPY_PATH__LD_LIB_TEST_PATHS) -def test_get_cuda_runtime_lib_path__unhappy_path(tmp_path, test_input: str): +def test_determine_cuda_runtime_lib_path__unhappy_path(tmp_path, test_input: str): test_input = tmp_path / test_input (test_input / CUDA_RUNTIME_LIB).touch() with pytest.raises(FileNotFoundError) as err_info: - get_cuda_runtime_lib_path(test_input) + determine_cuda_runtime_lib_path(test_input) assert all(match in err_info for match in {"duplicate", CUDA_RUNTIME_LIB}) -def test_get_cuda_runtime_lib_path__non_existent_dir(capsys, tmp_path): +def test_determine_cuda_runtime_lib_path__non_existent_dir(capsys, tmp_path): existent_dir = tmp_path / "a/b" existent_dir.mkdir() non_existent_dir = tmp_path / "c/d" # non-existent dir test_input = ":".join([str(existent_dir), str(non_existent_dir)]) - get_cuda_runtime_lib_path(test_input) + determine_cuda_runtime_lib_path(test_input) std_err = capsys.readouterr().err assert all(match in std_err for match in {"WARNING", "non-existent"}) @@ -95,14 +123,17 @@ def test_full_system(): # if CONDA_PREFIX exists, it has priority before all other env variables # but it does not contain the library directly, so we need to look at the a sub-folder - version = '' - if 'CONDA_PREFIX' in os.environ: - ls_output, err = bnb.utils.execute_and_return(f'ls -l {os.environ["CONDA_PREFIX"]}/lib/libcudart.so') - major, minor, revision = ls_output.split(' ')[-1].replace('libcudart.so.', '').split('.') - version = float(f'{major}.{minor}') - - - if version == '' and 'LD_LIBRARY_PATH': + version = "" + if "CONDA_PREFIX" in os.environ: + ls_output, err = bnb.utils.execute_and_return( + f'ls -l {os.environ["CONDA_PREFIX"]}/lib/libcudart.so' + ) + major, minor, revision = ( + ls_output.split(" ")[-1].replace("libcudart.so.", "").split(".") + ) + version = float(f"{major}.{minor}") + + if version == "" and "LD_LIBRARY_PATH": ld_path = os.environ["LD_LIBRARY_PATH"] paths = ld_path.split(":") version = "" |