tangled
alpha
login
or
join now
pyrox.dev
/
nixpkgs
0
fork
atom
lol
0
fork
atom
overview
issues
pulls
pipelines
tensorrt: dont break eval for unrelated packages
Someone Serge
2 years ago
3ee37e43
5bda2ec6
+32
-13
3 changed files
expand all
collapse all
unified
split
pkgs
development
libraries
science
math
tensorrt
extension.nix
generic.nix
top-level
python-packages.nix
+22
-6
pkgs/development/libraries/science/math/tensorrt/extension.nix
···
17
17
isSupported = fileData: elem cudaVersion fileData.supportedCudaVersions;
18
18
# Return the first file that is supported. In practice there should only ever be one anyway.
19
19
supportedFile = files: findFirst isSupported null files;
20
20
-
# Supported versions with versions as keys and file as value
21
21
-
supportedVersions = filterAttrs (version: file: file !=null ) (mapAttrs (version: files: supportedFile files) tensorRTVersions);
20
20
+
22
21
# Compute versioned attribute name to be used in this package set
23
22
computeName = version: "tensorrt_${toUnderscore version}";
23
23
+
24
24
+
# Supported versions with versions as keys and file as value
25
25
+
supportedVersions = lib.recursiveUpdate
26
26
+
{
27
27
+
tensorrt = {
28
28
+
enable = false;
29
29
+
fileVersionCuda = null;
30
30
+
fileVersionCudnn = null;
31
31
+
fullVersion = "0.0.0";
32
32
+
sha256 = null;
33
33
+
tarball = null;
34
34
+
supportedCudaVersions = [ ];
35
35
+
};
36
36
+
}
37
37
+
(mapAttrs' (version: attrs: nameValuePair (computeName version) attrs)
38
38
+
(filterAttrs (version: file: file != null) (mapAttrs (version: files: supportedFile files) tensorRTVersions)));
39
39
+
24
40
# Add all supported builds as attributes
25
25
-
allBuilds = mapAttrs' (version: file: nameValuePair (computeName version) (buildTensorRTPackage (removeAttrs file ["fileVersionCuda"]))) supportedVersions;
41
41
+
allBuilds = mapAttrs (name: file: buildTensorRTPackage (removeAttrs file ["fileVersionCuda"])) supportedVersions;
42
42
+
26
43
# Set the default attributes, e.g. tensorrt = tensorrt_8_4;
27
27
-
defaultBuild = { "tensorrt" = if allBuilds ? ${computeName tensorRTDefaultVersion}
28
28
-
then allBuilds.${computeName tensorRTDefaultVersion}
29
29
-
else throw "tensorrt-${tensorRTDefaultVersion} does not support your cuda version ${cudaVersion}"; };
44
44
+
defaultName = computeName tensorRTDefaultVersion;
45
45
+
defaultBuild = lib.optionalAttrs (allBuilds ? ${defaultName}) { tensorrt = allBuilds.${computeName tensorRTDefaultVersion}; };
30
46
in {
31
47
inherit buildTensorRTPackage;
32
48
} // allBuilds // defaultBuild;
+9
-6
pkgs/development/libraries/science/math/tensorrt/generic.nix
···
8
8
, cudnn
9
9
}:
10
10
11
11
-
{ fullVersion
11
11
+
{ enable ? true
12
12
+
, fullVersion
12
13
, fileVersionCudnn ? null
13
14
, tarball
14
15
, sha256
15
16
, supportedCudaVersions ? [ ]
16
17
}:
17
18
18
18
-
assert fileVersionCudnn == null || lib.assertMsg (lib.strings.versionAtLeast cudnn.version fileVersionCudnn)
19
19
+
assert !enable || fileVersionCudnn == null || lib.assertMsg (lib.strings.versionAtLeast cudnn.version fileVersionCudnn)
19
20
"This version of TensorRT requires at least cuDNN ${fileVersionCudnn} (current version is ${cudnn.version})";
20
21
21
22
backendStdenv.mkDerivation rec {
22
23
pname = "cudatoolkit-${cudatoolkit.majorVersion}-tensorrt";
23
24
version = fullVersion;
24
24
-
src = requireFile rec {
25
25
+
src = if !enable then null else
26
26
+
requireFile rec {
25
27
name = tarball;
26
28
inherit sha256;
27
29
message = ''
···
38
40
39
41
outputs = [ "out" "dev" ];
40
42
41
41
-
nativeBuildInputs = [
43
43
+
nativeBuildInputs = lib.optionals enable [
42
44
autoPatchelfHook
43
45
autoAddOpenGLRunpathHook
44
46
];
45
47
46
48
# Used by autoPatchelfHook
47
47
-
buildInputs = [
49
49
+
buildInputs = lib.optionals enable [
48
50
backendStdenv.cc.cc.lib # libstdc++
49
51
cudatoolkit
50
52
cudnn
···
75
77
'';
76
78
77
79
passthru.stdenv = backendStdenv;
80
80
+
passthru.enable = enable;
78
81
79
82
meta = with lib; {
80
83
# Check that the cudatoolkit version satisfies our min/max constraints (both
···
82
85
# official version constraints (as recorded in default.nix). In some cases
83
86
# you _may_ be able to smudge version constraints, just know that you're
84
87
# embarking into unknown and unsupported territory when doing so.
85
85
-
broken = !(elem cudaVersion supportedCudaVersions);
88
88
+
broken = !enable || !(elem cudaVersion supportedCudaVersions);
86
89
description = "TensorRT: a high-performance deep learning interface";
87
90
homepage = "https://developer.nvidia.com/tensorrt";
88
91
license = licenses.unfree;
+1
-1
pkgs/top-level/python-packages.nix
···
13956
13956
13957
13957
tensorly = callPackage ../development/python-modules/tensorly { };
13958
13958
13959
13959
-
tensorrt = callPackage ../development/python-modules/tensorrt { };
13959
13959
+
tensorrt = callPackage ../development/python-modules/tensorrt { cudaPackages = pkgs.cudaPackages_11; };
13960
13960
13961
13961
tensorstore = callPackage ../development/python-modules/tensorstore { };
13962
13962