diff options
-rw-r--r-- | bitsandbytes/__main__.py | 3 | ||||
-rw-r--r-- | bitsandbytes/cuda_setup.py | 83 | ||||
-rw-r--r-- | bitsandbytes/debug_cli.py | 27 | ||||
-rw-r--r-- | bitsandbytes/utils.py | 7 | ||||
-rw-r--r-- | environment.yml | 14 | ||||
-rw-r--r-- | setup.py | 3 | ||||
-rw-r--r-- | tests/test_cuda_setup_evaluator.py | 66 |
7 files changed, 203 insertions, 0 deletions
diff --git a/bitsandbytes/__main__.py b/bitsandbytes/__main__.py new file mode 100644 index 0000000..a91e942 --- /dev/null +++ b/bitsandbytes/__main__.py @@ -0,0 +1,3 @@ +from bitsandbytes.debug_cli import cli + +cli() diff --git a/bitsandbytes/cuda_setup.py b/bitsandbytes/cuda_setup.py new file mode 100644 index 0000000..48423b5 --- /dev/null +++ b/bitsandbytes/cuda_setup.py @@ -0,0 +1,83 @@ +""" +build is dependent on +- compute capability + - dependent on GPU family +- CUDA version +- Software: + - CPU-only: only CPU quantization functions (no optimizer, no matrix multipl) + - CuBLAS-LT: full-build 8-bit optimizer + - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`) + +alle Binaries packagen + +evaluation: + - if paths faulty, return meaningful error + - else: + - determine CUDA version + - determine capabilities + - based on that set the default path +""" + +from os import environ as env +from pathlib import Path +from typing import Set, Union +from .utils import warn_of_missing_prerequisite, print_err + + +CUDA_RUNTIME_LIB: str = "libcudart.so" + +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) + + non_existent_directories: Set[Path] = { + path for path in ld_library_paths + if not path.exists() + } + + if non_existent_directories: + print_err( + "WARNING: The following directories listed your path were found to " + f"be non-existent: {non_existent_directories}" + ) + + cuda_runtime_libs: Set[Path] = { + path / CUDA_RUNTIME_LIB for path in ld_library_paths + if (path / CUDA_RUNTIME_LIB).is_file() + } - non_existent_directories + + if len(cuda_runtime_libs) > 1: + err_msg = f"Found duplicate {CUDA_RUNTIME_LIB} files: {cuda_runtime_libs}.." + raise FileNotFoundError(err_msg) + + elif 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 ld_library_paths + +def evaluate_cuda_setup(): + # - if paths faulty, return meaningful error + # - else: + # - determine CUDA version + # - determine capabilities + # - based on that set the default path + pass diff --git a/bitsandbytes/debug_cli.py b/bitsandbytes/debug_cli.py new file mode 100644 index 0000000..88307a6 --- /dev/null +++ b/bitsandbytes/debug_cli.py @@ -0,0 +1,27 @@ +import typer + + +cli = typer.Typer() + + +@cli.callback() +def callback(): + """ + Awesome Portal Gun + """ + + +@cli.command() +def shoot(): + """ + Shoot the portal gun + """ + typer.echo("Shooting portal gun") + + +@cli.command() +def load(): + """ + Load the portal gun + """ + typer.echo("Loading portal gun") diff --git a/bitsandbytes/utils.py b/bitsandbytes/utils.py new file mode 100644 index 0000000..a9eddf9 --- /dev/null +++ b/bitsandbytes/utils.py @@ -0,0 +1,7 @@ +import sys + +def print_err(s: str) -> None: + print(s, file=sys.stderr) + +def warn_of_missing_prerequisite(s: str) -> None: + print_err('WARNING, missing pre-requisite: ' + s) diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..6bc6f9a --- /dev/null +++ b/environment.yml @@ -0,0 +1,14 @@ +name: 8-bit +channels: + - conda-forge +dependencies: + - python=3.9 + - pytest + - pytorch + - torchaudio + - torchvision + - cudatoolkit=11.1 + - typer + - ca-certificates + - certifi + - openssl @@ -23,6 +23,9 @@ setup( keywords="gpu optimizers optimization 8-bit quantization compression", url="http://packages.python.org/bitsandbytes", packages=find_packages(), + entry_points={ + "console_scripts": ["debug_cuda = bitsandbytes.debug_cli:cli"], + }, package_data={'': ['libbitsandbytes.so']}, long_description=read('README.md'), long_description_content_type='text/markdown', diff --git a/tests/test_cuda_setup_evaluator.py b/tests/test_cuda_setup_evaluator.py new file mode 100644 index 0000000..96ee6c5 --- /dev/null +++ b/tests/test_cuda_setup_evaluator.py @@ -0,0 +1,66 @@ +import pytest + +from typing import List + +from bitsandbytes.cuda_setup import ( + CUDA_RUNTIME_LIB, + get_cuda_runtime_lib_path, + evaluate_cuda_setup, + tokenize_paths, +) + + +HAPPY_PATH__LD_LIB_TEST_PATHS: List[tuple[str,str]] = [ + (f"some/other/dir:dir/with/{CUDA_RUNTIME_LIB}", f"dir/with/{CUDA_RUNTIME_LIB}"), + (f":some/other/dir:dir/with/{CUDA_RUNTIME_LIB}", f"dir/with/{CUDA_RUNTIME_LIB}"), + (f"some/other/dir:dir/with/{CUDA_RUNTIME_LIB}:", f"dir/with/{CUDA_RUNTIME_LIB}"), + (f"some/other/dir::dir/with/{CUDA_RUNTIME_LIB}", f"dir/with/{CUDA_RUNTIME_LIB}"), + (f"dir/with/{CUDA_RUNTIME_LIB}:some/other/dir", f"dir/with/{CUDA_RUNTIME_LIB}"), +] + + +@pytest.mark.parametrize( + "test_input, expected", + HAPPY_PATH__LD_LIB_TEST_PATHS +) +def test_get_cuda_runtime_lib_path__happy_path( + tmp_path, test_input: str, expected: str +): + for path in tokenize_paths(test_input): + assert False == tmp_path / test_input + test_dir.mkdir() + (test_input / CUDA_RUNTIME_LIB).touch() + assert get_cuda_runtime_lib_path(test_input) == expected + + +UNHAPPY_PATH__LD_LIB_TEST_PATHS = [ + f"a/b/c/{CUDA_RUNTIME_LIB}:d/e/f/{CUDA_RUNTIME_LIB}", + f"a/b/c/{CUDA_RUNTIME_LIB}:d/e/f/{CUDA_RUNTIME_LIB}:g/h/j/{CUDA_RUNTIME_LIB}", +] + + +@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): + 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) + 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): + 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) + std_err = capsys.readouterr().err + + assert all( + match in std_err + for match in {"WARNING", "non-existent"} + ) |