diff options
| -rw-r--r-- | bitsandbytes/cuda_setup.py | 54 | ||||
| -rw-r--r-- | bitsandbytes/utils.py | 3 | ||||
| -rw-r--r-- | tests/test_cuda_setup_evaluator.py | 34 |
3 files changed, 52 insertions, 39 deletions
diff --git a/bitsandbytes/cuda_setup.py b/bitsandbytes/cuda_setup.py index 05d4c7f..59e90e4 100644 --- a/bitsandbytes/cuda_setup.py +++ b/bitsandbytes/cuda_setup.py @@ -19,11 +19,11 @@ evaluation: """ import ctypes -from os import environ as env +import os from pathlib import Path from typing import Set, Union -from .utils import print_err, warn_of_missing_prerequisite +from .utils import print_err, warn_of_missing_prerequisite, execute_and_return def check_cuda_result(cuda, result_val): @@ -88,22 +88,11 @@ def tokenize_paths(paths: str) -> Set[Path]: return {Path(ld_path) for ld_path in paths.split(":") if ld_path} -def get_cuda_runtime_lib_path( - # TODO: replace this with logic for all paths in env vars - LD_LIBRARY_PATH: Union[str, None] = env.get("LD_LIBRARY_PATH") -) -> Union[Path, None]: - """# TODO: add doc-string""" - - if not LD_LIBRARY_PATH: - warn_of_missing_prerequisite( - "LD_LIBRARY_PATH is completely missing from environment!" - ) - return None - - ld_library_paths: Set[Path] = tokenize_paths(LD_LIBRARY_PATH) +def resolve_env_variable(env_var): + paths: Set[Path] = tokenize_paths(env_var) non_existent_directories: Set[Path] = { - path for path in ld_library_paths if not path.exists() + path for path in paths if not path.exists() } if non_existent_directories: @@ -114,7 +103,7 @@ def get_cuda_runtime_lib_path( cuda_runtime_libs: Set[Path] = { path / CUDA_RUNTIME_LIB - for path in ld_library_paths + for path in paths if (path / CUDA_RUNTIME_LIB).is_file() } - non_existent_directories @@ -123,19 +112,35 @@ def get_cuda_runtime_lib_path( f"Found duplicate {CUDA_RUNTIME_LIB} files: {cuda_runtime_libs}.." ) raise FileNotFoundError(err_msg) + elif len(cuda_runtime_libs) == 0: return None + else: return next(iter(cuda_runtime_libs)) # for now just return the first + +def get_cuda_runtime_lib_path() -> Union[Path, None]: + """# TODO: add doc-string""" + + cuda_runtime_libs = [] + if 'CONDA_PREFIX' in os.environ: + lib_conda_path = f'{os.environ["CONDA_PREFIX"]}/lib/' + print(lib_conda_path) + cuda_runtime_libs.append(resolve_env_variable(lib_conda_path)) + + if len(cuda_runtime_libs) == 1: return cuda_runtime_libs[0] + + for var in os.environ: + cuda_runtime_libs.append(resolve_env_variable(var)) - elif len(cuda_runtime_libs) < 1: + if len(cuda_runtime_libs) < 1: err_msg = ( f"Did not find {CUDA_RUNTIME_LIB} files: {cuda_runtime_libs}.." ) raise FileNotFoundError(err_msg) - single_cuda_runtime_lib_dir = next(iter(cuda_runtime_libs)) - return single_cuda_runtime_lib_dir + return cuda_runtime_libs.pop() def evaluate_cuda_setup(): cuda_path = get_cuda_runtime_lib_path() + print(f'CUDA SETUP: CUDA path found: {cuda_path}') cc = get_compute_capability() binary_name = "libbitsandbytes_cpu.so" @@ -152,13 +157,10 @@ def evaluate_cuda_setup(): # (2) Multiple CUDA versions installed cuda_home = str(Path(cuda_path).parent.parent) - ls_output, err = execute_and_return(f"{cuda_home}/bin/nvcc --version") - cuda_version = ( - ls_output.split("\n")[3].split(",")[-1].strip().lower().replace("v", "") - ) - major, minor, revision = cuda_version.split(".") + ls_output, err = execute_and_return(f"ls -l {cuda_path}") + major, minor, revision = ls_output.split(' ')[-1].replace('libcudart.so.', '').split('.') cuda_version_string = f"{major}{minor}" - binary_name = f'libbitsandbytes_cuda{cuda_version_string}_{("cublaslt" if has_cublaslt else "")}.so' + binary_name = f'libbitsandbytes_cuda{cuda_version_string}{("" if has_cublaslt else "_nocublaslt")}.so' return binary_name diff --git a/bitsandbytes/utils.py b/bitsandbytes/utils.py index 8a9fc0e..e1d9460 100644 --- a/bitsandbytes/utils.py +++ b/bitsandbytes/utils.py @@ -2,6 +2,7 @@ import sys import shlex import subprocess +from typing import Tuple def execute_and_return(command_string: str) -> Tuple[str, str]: def _decode(subprocess_err_out_tuple): @@ -19,7 +20,7 @@ def execute_and_return(command_string: str) -> Tuple[str, str]: ).communicate() ) - std_out, std_err = execute_and_return_decoded_std_streams() + std_out, std_err = execute_and_return_decoded_std_streams(command_string) return std_out, std_err diff --git a/tests/test_cuda_setup_evaluator.py b/tests/test_cuda_setup_evaluator.py index 5a58be4..5da190d 100644 --- a/tests/test_cuda_setup_evaluator.py +++ b/tests/test_cuda_setup_evaluator.py @@ -1,7 +1,8 @@ import os -from typing import List, NamedTuple - import pytest +import bitsandbytes as bnb + +from typing import List, NamedTuple from bitsandbytes.cuda_setup import ( CUDA_RUNTIME_LIB, @@ -91,16 +92,25 @@ def test_get_cuda_runtime_lib_path__non_existent_dir(capsys, tmp_path): def test_full_system(): ## this only tests the cuda version and not compute capability - ld_path = os.environ["LD_LIBRARY_PATH"] - paths = ld_path.split(":") - version = "" - for p in paths: - if "cuda" in p: - idx = p.rfind("cuda-") - version = p[idx + 5 : idx + 5 + 4].replace("/", "") - version = float(version) - break - + 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 = "" + for p in paths: + if "cuda" in p: + idx = p.rfind("cuda-") + version = p[idx + 5 : idx + 5 + 4].replace("/", "") + version = float(version) + break + + assert version > 0 binary_name = evaluate_cuda_setup() binary_name = binary_name.replace("libbitsandbytes_cuda", "") assert binary_name.startswith(str(version).replace(".", "")) |