lol

python3Packages.{torchaudio,torchvision,vllm}: use torch.stdenv

authored by

Pavol Rusnak and committed by
Gaetan Lepage
93d1e251 c98edc6b

+14 -12
+4
pkgs/development/python-modules/torch/default.nix
··· 247 247 echo 248 248 echo "# Update using: unroll-src [version]" 249 249 ''; 250 + 251 + stdenv' = if cudaSupport then cudaPackages.backendStdenv else stdenv; 250 252 in 251 253 buildPythonPackage rec { 252 254 pname = "torch"; 253 255 # Don't forget to update torch-bin to the same version. 254 256 version = "2.6.0"; 255 257 pyproject = true; 258 + 259 + stdenv = stdenv'; 256 260 257 261 outputs = [ 258 262 "out" # output standard python package
+2
pkgs/development/python-modules/torchaudio/default.nix
··· 79 79 version = "2.6.0"; 80 80 pyproject = true; 81 81 82 + stdenv = torch.stdenv; 83 + 82 84 src = fetchFromGitHub { 83 85 owner = "pytorch"; 84 86 repo = "audio";
+2 -5
pkgs/development/python-modules/torchvision/default.nix
··· 27 27 28 28 let 29 29 inherit (torch) cudaCapabilities cudaPackages cudaSupport; 30 - inherit (cudaPackages) backendStdenv; 31 30 32 31 pname = "torchvision"; 33 32 version = "0.21.0"; 34 33 in 35 34 buildPythonPackage { 36 35 inherit pname version; 36 + 37 + stdenv = torch.stdenv; 37 38 38 39 src = fetchFromGitHub { 39 40 owner = "pytorch"; ··· 75 76 export TORCHVISION_INCLUDE="${libjpeg_turbo.dev}/include/" 76 77 export TORCHVISION_LIBRARY="${libjpeg_turbo}/lib/" 77 78 '' 78 - # NOTE: We essentially override the compilers provided by stdenv because we don't have a hook 79 - # for cudaPackages to swap in compilers supported by NVCC. 80 79 + lib.optionalString cudaSupport '' 81 - export CC=${backendStdenv.cc}/bin/cc 82 - export CXX=${backendStdenv.cc}/bin/c++ 83 80 export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}" 84 81 export FORCE_CUDA=1 85 82 '';
+6 -7
pkgs/development/python-modules/vllm/default.nix
··· 58 58 depyf, 59 59 opencv-python-headless, 60 60 61 - config, 62 - 63 - cudaSupport ? config.cudaSupport, 61 + cudaSupport ? torch.cudaSupport, 64 62 cudaPackages ? { }, 65 - rocmSupport ? config.rocmSupport, 63 + rocmSupport ? torch.rocmSupport, 66 64 rocmPackages ? { }, 67 65 gpuTargets ? [ ], 68 - }@args: 66 + }: 69 67 70 68 let 71 69 inherit (lib) ··· 117 115 118 116 cpuSupport = !cudaSupport && !rocmSupport; 119 117 120 - # https://github.com/pytorch/pytorch/blob/v2.4.0/torch/utils/cpp_extension.py#L1953 118 + # https://github.com/pytorch/pytorch/blob/v2.6.0/torch/utils/cpp_extension.py#L2046-L2048 121 119 supportedTorchCudaCapabilities = 122 120 let 123 121 real = [ ··· 138 136 "8.9" 139 137 "9.0" 140 138 "9.0a" 139 + "10.0" 141 140 ]; 142 141 ptx = lists.map (x: "${x}+PTX") real; 143 142 in ··· 200 199 version = "0.7.3"; 201 200 pyproject = true; 202 201 203 - stdenv = if cudaSupport then cudaPackages.backendStdenv else args.stdenv; 202 + stdenv = torch.stdenv; 204 203 205 204 src = fetchFromGitHub { 206 205 owner = "vllm-project";