1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
{ lib
, fetchgit
, buildPythonPackage
, pytorch
, setuptools
, typer
, linuxPackages
, pytest
# CUDA
, cudaSupport ? true
, cudaPackages ? {}
, addOpenGLRunpath ? null
, gcc
}:
let
inherit (linuxPackages) nvidia_x11;
inherit (cudaPackages) cudatoolkit;
cudaVersion =
lib.strings.removeSuffix "0"
(lib.strings.concatStrings
(lib.strings.splitString "." cudatoolkit.version));
libraryPath = lib.strings.makeLibraryPath [
cudatoolkit
cudatoolkit.lib
nvidia_x11
];
in buildPythonPackage rec {
pname = "bitsandbytes";
version = "unstable-2022-12-21";
src = fetchgit {
url = "https://simatime.com/git/ben/bitsandbytes.git";
rev = "31ef751bea48eeee2e0e95aca79df8e59b4c25c4";
sha256 = "sha256-/a2NFVuMSvSDELeXyfDdU9FZEJcDxCLa4VbMWBUCRI0=";
};
# any exe must be wrapped with this, it doesn't effect the build but it does
# show up on the logs so keep it here for documentation purposes
LD_LIBRARY_PATH = libraryPath;
# this allows the build system to find nvcc
CUDA_HOME = "${cudatoolkit}";
preBuild = ''
make cuda11x CUDA_VERSION=${cudaVersion} GPP=${gcc}/bin/g++ -j3
## this python doesn't know where to get libcuda, so explicitly tell it
## this is probably really bad practice, fix this
substituteInPlace bitsandbytes/cuda_setup/main.py \
--replace "libcuda.so" "${nvidia_x11}/lib/libcuda.so"
'';
propagatedBuildInputs = [
(pytorch.override({ inherit cudaSupport;}))
setuptools
] ++ lib.optionals cudaSupport [
typer
cudatoolkit
cudatoolkit.lib
nvidia_x11
];
nativeBuildInputs = lib.optionals cudaSupport [
gcc
addOpenGLRunpath
];
preFixup = lib.optionalString cudaSupport ''
find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib
do
addOpenGLRunpath "$lib"
patchelf \
--set-rpath "${libraryPath}" \
"$lib"
done
'';
checkInputs = [
pytest
];
# disabled because the test suite cannot find any GPUs in the nix sandbox
doCheck = false;
pythonImportsCheck = [
"torch"
# this tries to load and verify the cuda device on import, since this runs
# in the nix sandbox it will fail with a bunch of errors, but curiously the
# build still succeeds...
"bitsandbytes"
];
}
|
