mrq/bitsandbytes-rocm
1
0
Fork 1
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

factored cuda_setup.main out into smaller modules and functions

Browse Source
This commit is contained in:
Titus von Koeller 2022-08-02 21:26:50 -07:00
parent 3809236428
commit 59a615b386
9 changed files with 384 additions and 240 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
bitsandbytes/__init__.py
View File

@ -22,3 +22,5 @@ __pdoc__ = {
"optim.optimizer.Optimizer8bit": False,
"optim.optimizer.MockArgs": False,
}
PACKAGE_GITHUB_URL = "https://github.com/TimDettmers/bitsandbytes"

97
bitsandbytes/__main__.py
View File

@ -1,3 +1,96 @@
from bitsandbytes.debug_cli import cli
# from bitsandbytes.debug_cli import cli
cli()
# cli()
import os
import sys
import torch
HEADER_WIDTH = 60
def print_header(
txt: str, width: int = HEADER_WIDTH, filler: str = "+"
) -> None:
txt = f" {txt} " if txt else ""
print(txt.center(width, filler))
def print_debug_info() -> None:
print(
"\nAbove we output some debug information. Please provide this info when "
f"creating an issue via {PACKAGE_GITHUB_URL}/issues/new/choose ...\n"
)
print_header("")
print_header("DEBUG INFORMATION")
print_header("")
print()
from . import COMPILED_WITH_CUDA, PACKAGE_GITHUB_URL
from .cuda_setup.main import get_compute_capabilities
from .cuda_setup.env_vars import to_be_ignored
from .utils import print_stderr
print_header("POTENTIALLY LIBRARY-PATH-LIKE ENV VARS")
for k, v in os.environ.items():
if "/" in v and not to_be_ignored(k, v):
print(f"'{k}': '{v}'")
print_header("")
print(
"\nWARNING: Please be sure to sanitize sensible info from any such env vars!\n"
)
print_header("OTHER")
print(f"{COMPILED_WITH_CUDA = }")
print(f"COMPUTE_CAPABILITIES_PER_GPU = {get_compute_capabilities()}")
print_header("")
print_header("DEBUG INFO END")
print_header("")
print(
"""
Running a quick check that:
+ library is importable
+ CUDA function is callable
"""
)
try:
from bitsandbytes.optim import Adam
p = torch.nn.Parameter(torch.rand(10, 10).cuda())
a = torch.rand(10, 10).cuda()
p1 = p.data.sum().item()
adam = Adam([p])
out = a * p
loss = out.sum()
loss.backward()
adam.step()
p2 = p.data.sum().item()
assert p1 != p2
print("SUCCESS!")
print("Installation was successful!")
sys.exit(0)
except ImportError:
print()
print_stderr(
f"WARNING: {__package__} is currently running as CPU-only!\n"
"Therefore, 8-bit optimizers and GPU quantization are unavailable.\n\n"
f"If you think that this is so erroneously,\nplease report an issue!"
)
print_debug_info()
sys.exit(0)
except Exception as e:
print(e)
print_debug_info()
sys.exit(1)

21
bitsandbytes/cextension.py
View File

@ -1,8 +1,8 @@
import ctypes as ct
import os
from pathlib import Path
from warnings import warn
from bitsandbytes.cuda_setup.main import evaluate_cuda_setup
from .cuda_setup.main import evaluate_cuda_setup
class CUDALibrary_Singleton(object):
@ -12,18 +12,17 @@ class CUDALibrary_Singleton(object):
raise RuntimeError("Call get_instance() instead")
def initialize(self):
self.context = {}
binary_name = evaluate_cuda_setup()
if not os.path.exists(os.path.dirname(__file__) + f"/{binary_name}"):
package_dir = Path(__file__).parent
binary_path = package_dir / binary_name
if not binary_path.exists():
print(f"TODO: compile library for specific version: {binary_name}")
print("defaulting to libbitsandbytes.so")
self.lib = ct.cdll.LoadLibrary(
os.path.dirname(__file__) + "/libbitsandbytes.so"
)
legacy_binary_name = "libbitsandbytes.so"
print(f"Defaulting to {legacy_binary_name}...")
self.lib = ct.cdll.LoadLibrary(package_dir / legacy_binary_name)
else:
self.lib = ct.cdll.LoadLibrary(
os.path.dirname(__file__) + f"/{binary_name}"
)
self.lib = ct.cdll.LoadLibrary(package_dir / binary_name)
@classmethod
def get_instance(cls):

51
bitsandbytes/cuda_setup/env_vars.py Normal file
View File

@ -0,0 +1,51 @@
import os
from typing import Dict
def to_be_ignored(env_var: str, value: str) -> bool:
ignorable = {
"PWD", # PWD: this is how the shell keeps track of the current working dir
"OLDPWD",
"SSH_AUTH_SOCK", # SSH stuff, therefore unrelated
"SSH_TTY",
"HOME", # Linux shell default
"TMUX", # Terminal Multiplexer
"XDG_DATA_DIRS", # XDG: Desktop environment stuff
"XDG_RUNTIME_DIR",
"MAIL", # something related to emails
"SHELL", # binary for currently invoked shell
"DBUS_SESSION_BUS_ADDRESS", # hardware related
"PATH", # this is for finding binaries, not libraries
"LESSOPEN", # related to the `less` command
"LESSCLOSE",
"_", # current Python interpreter
}
return env_var in ignorable
def might_contain_a_path(candidate: str) -> bool:
return "/" in candidate
def is_active_conda_env(env_var: str) -> bool:
return "CONDA_PREFIX" == env_var
def is_other_conda_env_var(env_var: str) -> bool:
return "CONDA" in env_var
def is_relevant_candidate_env_var(env_var: str, value: str) -> bool:
return is_active_conda_env(env_var) or (
might_contain_a_path(value) and not
is_other_conda_env_var(env_var) and not
to_be_ignored(env_var, value)
)
def get_potentially_lib_path_containing_env_vars() -> Dict[str, str]:
return {
env_var: value
for env_var, value in os.environ.items()
if is_relevant_candidate_env_var(env_var, value)
}

142
bitsandbytes/cuda_setup/main.py
View File

@ -8,8 +8,6 @@ extract factors the build is dependent on:
- 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:
@ -19,11 +17,10 @@ evaluation:
"""
import ctypes
import os
from pathlib import Path
from typing import Set, Union
from ..utils import print_err, warn_of_missing_prerequisite, execute_and_return
from ..utils import execute_and_return
from .paths import determine_cuda_runtime_lib_path
def check_cuda_result(cuda, result_val):
@ -34,26 +31,23 @@ def check_cuda_result(cuda, result_val):
raise Exception(f"CUDA exception! ERROR: {error_str}")
# taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
def get_compute_capability():
# 1. find libcuda.so library (GPU driver) (/usr/lib)
# init_device -> init variables -> call function by reference
# 2. call extern C function to determine CC
# (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html)
# 3. Check for CUDA errors
# https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api
def get_compute_capabilities():
"""
1. find libcuda.so library (GPU driver) (/usr/lib)
init_device -> init variables -> call function by reference
2. call extern C function to determine CC
(https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html)
3. Check for CUDA errors
https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api
# bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
"""
# 1. find libcuda.so library (GPU driver) (/usr/lib)
libnames = ("libcuda.so",)
for libname in libnames:
try:
cuda = ctypes.CDLL(libname)
except OSError:
continue
else:
break
else:
raise OSError("could not load any of: " + " ".join(libnames))
try:
cuda = ctypes.CDLL("libcuda.so")
except OSError:
# TODO: shouldn't we error or at least warn here?
return None
nGpus = ctypes.c_int()
cc_major = ctypes.c_int()
@ -70,104 +64,64 @@ def get_compute_capability():
check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i))
ref_major = ctypes.byref(cc_major)
ref_minor = ctypes.byref(cc_minor)
# 2. call extern C function to determine CC
check_cuda_result(cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device))
# 2. call extern C function to determine CC
check_cuda_result(
cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device)
)
ccs.append(f"{cc_major.value}.{cc_minor.value}")
# TODO: handle different compute capabilities; for now, take the max
ccs.sort()
max_cc = ccs[-1]
return max_cc
return ccs.sort()
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 resolve_env_variable(env_var):
'''Searches a given envirionmental library or path for the CUDA runtime library (libcudart.so)'''
paths: Set[Path] = tokenize_paths(env_var)
non_existent_directories: Set[Path] = {
path for path in 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 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) == 0: return None # this is not en error, since other envs can contain CUDA
else: return next(iter(cuda_runtime_libs)) # for now just return the first
def get_cuda_runtime_lib_path() -> Union[Path, None]:
'''Searches conda installation and environmental paths for a cuda installations.'''
cuda_runtime_libs = []
# CONDA_PREFIX/lib is the default location for a default conda
# install of pytorch. This location takes priortiy over all
# other defined variables
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]
# if CONDA_PREFIX does not have the library, search the environment
# (in particualr LD_LIBRARY PATH)
for var in os.environ:
cuda_runtime_libs.append(resolve_env_variable(var))
if len(cuda_runtime_libs) < 1:
err_msg = (
f"Did not find {CUDA_RUNTIME_LIB} files: {cuda_runtime_libs}.."
)
raise FileNotFoundError(err_msg)
return cuda_runtime_libs.pop()
# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error
def get_compute_capability():
"""
Extracts the highest compute capbility from all available GPUs, as compute
capabilities are downwards compatible. If no GPUs are detected, it returns
None.
"""
if ccs := get_compute_capabilities() is not None:
# TODO: handle different compute capabilities; for now, take the max
return ccs[-1]
return None
def evaluate_cuda_setup():
cuda_path = get_cuda_runtime_lib_path()
print(f'CUDA SETUP: CUDA path found: {cuda_path}')
cuda_path = determine_cuda_runtime_lib_path()
print(f"CUDA SETUP: CUDA path found: {cuda_path}")
cc = get_compute_capability()
binary_name = "libbitsandbytes_cpu.so"
# FIXME: has_gpu is still unused
if not (has_gpu := bool(cc)):
print(
"WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..."
)
return binary_name
# 7.5 is the minimum CC vor cublaslt
has_cublaslt = cc in ["7.5", "8.0", "8.6"]
# TODO:
# (1) CUDA missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible)
# (2) Multiple CUDA versions installed
# FIXME: cuda_home is still unused
cuda_home = str(Path(cuda_path).parent.parent)
# we use ls -l instead of nvcc to determine the cuda version
# since most installations will have the libcudart.so installed, but not the compiler
ls_output, err = execute_and_return(f"ls -l {cuda_path}")
major, minor, revision = ls_output.split(' ')[-1].replace('libcudart.so.', '').split('.')
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}{("" if has_cublaslt else "_nocublaslt")}.so'
def get_binary_name():
"if not has_cublaslt (CC < 7.5), then we have to choose _nocublaslt.so"
bin_base_name = "libbitsandbytes_cuda"
if has_cublaslt:
return f"{bin_base_name}{cuda_version_string}.so"
else:
return f"{bin_base_name}_nocublaslt.so"
return binary_name

126
bitsandbytes/cuda_setup/paths.py Normal file
View File

@ -0,0 +1,126 @@
from pathlib import Path
from typing import Set, Union
from warnings import warn
from ..utils import print_stderr
from .env_vars import get_potentially_lib_path_containing_env_vars
CUDA_RUNTIME_LIB: str = "libcudart.so"
def purge_unwanted_semicolon(tentative_path: Path) -> Path:
"""
Special function to handle the following exception:
__LMOD_REF_COUNT_PATH=/sw/cuda/11.6.2/bin:2;/mmfs1/home/dettmers/git/sched/bin:1;/mmfs1/home/dettmers/data/anaconda3/bin:1;/mmfs1/home/dettmers/data/anaconda3/condabin:1;/mmfs1/home/dettmers/.local/bin:1;/mmfs1/home/dettmers/bin:1;/usr/local/bin:1;/usr/bin:1;/usr/local/sbin:1;/usr/sbin:1;/mmfs1/home/dettmers/.fzf/bin:1;/mmfs1/home/dettmers/data/local/cuda-11.4/bin:1
"""
# if ';' in str(tentative_path):
# path_as_str, _ = str(tentative_path).split(';')
pass
def extract_candidate_paths(paths_list_candidate: str) -> Set[Path]:
return {Path(ld_path) for ld_path in paths_list_candidate.split(":") if ld_path}
def remove_non_existent_dirs(candidate_paths: Set[Path]) -> Set[Path]:
non_existent_directories: Set[Path] = {
path for path in candidate_paths if not path.exists()
}
if non_existent_directories:
print_stderr(
"WARNING: The following directories listed in your path were found to "
f"be non-existent: {non_existent_directories}"
)
return candidate_paths - non_existent_directories
def get_cuda_runtime_lib_paths(candidate_paths: Set[Path]) -> Set[Path]:
return {
path / CUDA_RUNTIME_LIB
for path in candidate_paths
if (path / CUDA_RUNTIME_LIB).is_file()
}
def resolve_paths_list(paths_list_candidate: str) -> Set[Path]:
"""
Searches a given environmental var for the CUDA runtime library,
i.e. `libcudart.so`.
"""
return remove_non_existent_dirs(extract_candidate_paths(paths_list_candidate))
def find_cuda_lib_in(paths_list_candidate: str) -> Set[Path]:
return get_cuda_runtime_lib_paths(
resolve_paths_list(paths_list_candidate)
)
def warn_in_case_of_duplicates(results_paths: Set[Path]) -> None:
if len(results_paths) > 1:
warning_msg = (
f"Found duplicate {CUDA_RUNTIME_LIB} files: {results_paths}.. "
"We'll flip a coin and try one of these, in order to fail forward.\n"
"Either way, this might cause trouble in the future:\n"
"If you get `CUDA error: invalid device function` errors, the above "
"might be the cause and the solution is to make sure only one "
f"{CUDA_RUNTIME_LIB} in the paths that we search based on your env."
)
warn(warning_msg)
def determine_cuda_runtime_lib_path() -> Union[Path, None]:
"""
Searches for a cuda installations, in the following order of priority:
1. active conda env
2. LD_LIBRARY_PATH
3. any other env vars, while ignoring those that
- are known to be unrelated (see `bnb.cuda_setup.env_vars.to_be_ignored`)
- don't contain the path separator `/`
If multiple libraries are found in part 3, we optimistically try one,
while giving a warning message.
"""
candidate_env_vars = get_potentially_lib_path_containing_env_vars()
if "CONDA_PREFIX" in candidate_env_vars:
conda_libs_path = Path(candidate_env_vars["CONDA_PREFIX"]) / "lib"
conda_cuda_libs = find_cuda_lib_in(str(conda_libs_path))
warn_in_case_of_duplicates(conda_cuda_libs)
if conda_cuda_libs:
return next(iter(conda_cuda_libs))
warn(
f'{candidate_env_vars["CONDA_PREFIX"]} did not contain '
f'{CUDA_RUNTIME_LIB} as expected! Searching further paths...'
)
if "LD_LIBRARY_PATH" in candidate_env_vars:
lib_ld_cuda_libs = find_cuda_lib_in(candidate_env_vars["LD_LIBRARY_PATH"])
if lib_ld_cuda_libs:
return next(iter(lib_ld_cuda_libs))
warn_in_case_of_duplicates(lib_ld_cuda_libs)
warn(
f'{candidate_env_vars["LD_LIBRARY_PATH"]} did not contain '
f'{CUDA_RUNTIME_LIB} as expected! Searching further paths...'
)
remaining_candidate_env_vars = {
env_var: value for env_var, value in candidate_env_vars.items()
if env_var not in {"CONDA_PREFIX", "LD_LIBRARY_PATH"}
}
cuda_runtime_libs = set()
for env_var, value in remaining_candidate_env_vars:
cuda_runtime_libs.update(find_cuda_lib_in(value))
warn_in_case_of_duplicates(cuda_runtime_libs)
return next(iter(cuda_runtime_libs)) if cuda_runtime_libs else set()

8
bitsandbytes/utils.py
View File

@ -1,9 +1,9 @@
import sys
import shlex
import subprocess
import sys
from typing import Tuple
def execute_and_return(command_string: str) -> Tuple[str, str]:
def _decode(subprocess_err_out_tuple):
return tuple(
@ -24,9 +24,9 @@ def execute_and_return(command_string: str) -> Tuple[str, str]:
return std_out, std_err
def print_err(s: str) -> None:
def print_stderr(s: str) -> None:
print(s, file=sys.stderr)
def warn_of_missing_prerequisite(s: str) -> None:
print_err("WARNING, missing pre-requisite: " + s)
print_stderr("WARNING, missing pre-requisite: " + s)

112
quicktest.py
View File

@ -1,112 +0,0 @@
from itertools import product
import torch
import bitsandbytes as bnb
import bitsandbytes.functional as F
def test_igemmlt(dim1, dim2, dim3, dim4, dims, ldb):
k = 25
for i in range(k):
if dims == 2:
A = torch.randint(-128, 127, size=(dim1, dim3), device="cuda").to(
torch.int8
)
elif dims == 3:
A = torch.randint(
-128, 127, size=(dim1, dim2, dim3), device="cuda"
).to(torch.int8)
B = torch.randint(-128, 127, size=(dim4, dim3), device="cuda").to(
torch.int8
)
C1 = torch.matmul(A.float(), B.t().float())
A2, SA = F.transform(A, "col32")
B2, SB = F.transform(B, "colx")
if dims == 2:
C2, SC = F.transform(
torch.zeros(
A.shape[0], B.shape[0], dtype=torch.int32, device="cuda"
),
"col32",
)
else:
C2, SC = F.transform(
torch.zeros(
A.shape[0],
A.shape[1],
B.shape[0],
dtype=torch.int32,
device="cuda",
),
"col32",
)
F.igemmlt(A2, B2, C2, SA, SB, SC)
C3, S = F.transform(C2, "row", state=SC)
# torch.testing.assert_allclose(C1, C3.float())
# print(C1)
# print(C2)
# print(C3)
allclose = torch.allclose(C1, C3.float())
if allclose:
print(C1)
print(C2)
print(C3)
## transposed
# A = torch.randint(-128, 127, size=(dim4, dim3), device='cuda').to(torch.int8)
# if dims == 2:
# B = torch.randint(-128, 127, size=(dim1, dim3), device='cuda').to(torch.int8)
# C1 = torch.matmul(A.float(), B.float().t())
# elif dims == 3:
# B = torch.randint(-128, 127, size=(dim1, dim2, dim3), device='cuda').to(torch.int8)
# C1 = torch.matmul(B.float(), A.t().float())
# C1 = C1.permute([2, 0, 1])
# A2, SA = F.transform(A, 'col32')
# B2, SB = F.transform(B, 'colx')
# if dims == 2:
# C2, SC = F.transform(torch.zeros(A.shape[0], B.shape[0], dtype=torch.int32, device='cuda'), 'col32')
# else:
# C2 = torch.zeros(A.shape[0], B.shape[0], B.shape[1], dtype=torch.int32, device='cuda')
# state = (C2.shape, 'row', A.shape[0])
# C2, SC = F.transform(C2, 'col32', state=state)
# F.igemmlt(A2, B2, C2, SA, SB, SC)
# C3, S = F.transform(C2, 'row', state=SC, ld=[0])
# torch.testing.assert_allclose(C1, C3.float())
## weight update
# if dims == 3:
# A = torch.randint(-128, 127, size=(dim1, dim2, dim3), device='cuda').to(torch.int8)
# B = torch.randint(-128, 127, size=(dim1, dim2, dim4), device='cuda').to(torch.int8)
# C1 = torch.matmul(B.view(-1, B.shape[-1]).t().float(), A.view(-1, A.shape[-1]).float())
# A2, SA = F.transform(A.view(-1, A.shape[-1]).t().contiguous(), 'colx')
# B2, SB = F.transform(B.view(-1, B.shape[-1]).t().contiguous(), 'col32')
# C2 = torch.zeros(B.shape[-1], A.shape[-1], dtype=torch.int32, device='cuda')
# C2, SC = F.transform(C2, 'col32')
# F.igemmlt(B2, A2, C2, SB, SA, SC)
# C3, S = F.transform(C2, 'row', state=SC)
# torch.testing.assert_allclose(C1, C3.float())
dims = (2, 3)
ldb = [0]
n = 2
dim1 = torch.randint(1, 256, size=(n,)).tolist()
dim2 = torch.randint(32, 512, size=(n,)).tolist()
dim3 = torch.randint(32, 1024, size=(n,)).tolist()
dim4 = torch.randint(32, 1024, size=(n,)).tolist()
values = list(product(dim1, dim2, dim3, dim4, dims, ldb))
for ldb in range(32, 4096, 32):
# for ldb in [None]:
val = test_igemmlt(2, 2, 2, 2, 2, ldb)
if val:
print(val, ldb)
else:
print("nope", ldb)
# for val in values:
# test_igemmlt(*val)

65
tests/test_cuda_setup_evaluator.py
View File

@ -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}')
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':
if version == "" and "LD_LIBRARY_PATH":
ld_path = os.environ["LD_LIBRARY_PATH"]
paths = ld_path.split(":")
version = ""