pyrox.dev
1{
2 lib,
3 stdenv,
4 python,
5 buildPythonPackage,
6 pythonOlder,
7 pythonAtLeast,
8 fetchurl,
9
10 # nativeBuildInputs
11 addDriverRunpath,
12 autoAddDriverRunpath,
13 autoPatchelfHook,
14
15 # buildInputs
16 cudaPackages,
17
18 # dependencies
19 filelock,
20 future,
21 jinja2,
22 networkx,
23 numpy,
24 pyyaml,
25 requests,
26 setuptools,
27 sympy,
28 typing-extensions,
29 triton,
30
31 callPackage,
32}:
33
34let
35 pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
36 srcs = import ./binary-hashes.nix version;
37 unsupported = throw "Unsupported system";
38 version = "2.7.0";
39in
40buildPythonPackage {
41 inherit version;
42
43 pname = "torch";
44 # Don't forget to update torch to the same version.
45
46 format = "wheel";
47
48 disabled = (pythonOlder "3.9") || (pythonAtLeast "3.14");
49
50 src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
51
52 nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [
53 addDriverRunpath
54 autoAddDriverRunpath
55 autoPatchelfHook
56 ];
57
58 buildInputs = lib.optionals stdenv.hostPlatform.isLinux (
59 with cudaPackages;
60 [
61 # $out/${sitePackages}/nvfuser/_C*.so wants libnvToolsExt.so.1 but torch/lib only ships
62 # libnvToolsExt-$hash.so.1
63 cuda_nvtx
64
65 cuda_cudart
66 cuda_cupti
67 cuda_nvrtc
68 cudnn
69 cusparselt
70 libcublas
71 libcufft
72 libcufile
73 libcurand
74 libcusolver
75 libcusparse
76 nccl
77 ]
78 );
79
80 autoPatchelfIgnoreMissingDeps = lib.optionals stdenv.hostPlatform.isLinux [
81 # This is the hardware-dependent userspace driver that comes from
82 # nvidia_x11 package. It must be deployed at runtime in
83 # /run/opengl-driver/lib or pointed at by LD_LIBRARY_PATH variable, rather
84 # than pinned in runpath
85 "libcuda.so.1"
86 ];
87
88 dependencies = [
89 filelock
90 future
91 jinja2
92 networkx
93 numpy
94 pyyaml
95 requests
96 setuptools
97 sympy
98 typing-extensions
99 ] ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [ triton ];
100
101 postInstall = ''
102 # ONNX conversion
103 rm -rf $out/bin
104 '';
105
106 postFixup = lib.optionalString stdenv.hostPlatform.isLinux ''
107 addAutoPatchelfSearchPath "$out/${python.sitePackages}/torch/lib"
108 '';
109
110 # See https://github.com/NixOS/nixpkgs/issues/296179
111 #
112 # This is a quick hack to add `libnvrtc` to the runpath so that torch can find
113 # it when it is needed at runtime.
114 extraRunpaths = lib.optionals stdenv.hostPlatform.isLinux [
115 "${lib.getLib cudaPackages.cuda_nvrtc}/lib"
116 ];
117 postPhases = lib.optionals stdenv.hostPlatform.isLinux [ "postPatchelfPhase" ];
118 postPatchelfPhase = ''
119 while IFS= read -r -d $'\0' elf ; do
120 for extra in $extraRunpaths ; do
121 echo patchelf "$elf" --add-rpath "$extra" >&2
122 patchelf "$elf" --add-rpath "$extra"
123 done
124 done < <(
125 find "''${!outputLib}" "$out" -type f -iname '*.so' -print0
126 )
127 '';
128
129 # The wheel-binary is not stripped to avoid the error of `ImportError: libtorch_cuda_cpp.so: ELF load command address/offset not properly aligned.`.
130 dontStrip = true;
131
132 pythonImportsCheck = [ "torch" ];
133
134 passthru.tests = callPackage ../tests { };
135
136 meta = {
137 description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration";
138 homepage = "https://pytorch.org/";
139 changelog = "https://github.com/pytorch/pytorch/releases/tag/v${version}";
140 # Includes CUDA and Intel MKL, but redistributions of the binary are not limited.
141 # https://docs.nvidia.com/cuda/eula/index.html
142 # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
143 # torch's license is BSD3.
144 # torch-bin used to vendor CUDA. It still links against CUDA and MKL.
145 license = with lib.licenses; [
146 bsd3
147 issl
148 unfreeRedistributable
149 ];
150 sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
151 platforms = [
152 "aarch64-darwin"
153 "aarch64-linux"
154 "x86_64-linux"
155 ];
156 hydraPlatforms = [ ]; # output size 3.2G on 1.11.0
157 maintainers = with lib.maintainers; [
158 GaetanLepage
159 junjihashimoto
160 ];
161 };
162}