summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--bitsandbytes/__main__.py3
-rw-r--r--bitsandbytes/cuda_setup.py83
-rw-r--r--bitsandbytes/debug_cli.py27
-rw-r--r--bitsandbytes/utils.py7
-rw-r--r--environment.yml14
-rw-r--r--setup.py3
-rw-r--r--tests/test_cuda_setup_evaluator.py66
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
diff --git a/setup.py b/setup.py
index 6275ddd..3292e30 100644
--- a/setup.py
+++ b/setup.py
@@ -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"}
+ )