python3Packages.torch: make overriding cudaPackages easier (#540915)

pkgs/development/python-modules/torch/source/default.nix

@@ -14,7 +14,7 @@
   autoAddDriverRunpath,
   effectiveMagma ?
     if cudaSupport then
-      magma-cuda-static
+      magma-cuda-static.override { inherit cudaPackages; }
     else if rocmSupport then
       magma-hip
     else
@@ -81,7 +81,7 @@
   #          (dependencies without cuda support).
   #          Instead we should rely on overlays and nixpkgsFun.
   # (@SomeoneSerge)
-  _tritonEffective ? if cudaSupport then triton-cuda else triton,
+  _tritonEffective ? if cudaSupport then triton-cuda.override { inherit cudaPackages; } else triton,
   triton-cuda,
 
   # Disable MKLDNN on aarch64-darwin, it negatively impacts performance,
@@ -263,6 +263,9 @@ let
     "Magma cudaPackages does not match cudaPackages" =
       cudaSupport
       && (effectiveMagma.cudaPackages.cudaMajorMinorVersion != cudaPackages.cudaMajorMinorVersion);
+    "Triton cudaPackages does not match cudaPackages" =
+      cudaSupport
+      && (_tritonEffective.cudaPackages.cudaMajorMinorVersion != cudaPackages.cudaMajorMinorVersion);
   };
 
   unroll-src = writeShellScript "unroll-src" ''

pkgs/development/python-modules/torchaudio/default.nix

@@ -28,12 +28,12 @@
   torchaudio,
 
   cudaSupport ? torch.cudaSupport,
-  cudaPackages,
   rocmSupport ? torch.rocmSupport,
 }:
 
 let
   inherit (stdenv) hostPlatform;
+  inherit (torch) cudaCapabilities cudaPackages;
 in
 buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
   pname = "torchaudio";
@@ -53,7 +53,7 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
 
     # CUDA
     USE_CUDA = cudaSupport;
-    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" cudaCapabilities}";
     TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_CUCTC_DECODER = 1;
     TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_CUDA = 1;
     TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_MULTIGPU_CUDA = 1;

pkgs/development/python-modules/torchcodec/default.nix

@@ -20,10 +20,12 @@
   torchvision,
 
   cudaSupport ? torch.cudaSupport,
-  cudaPackages,
   rocmSupport ? torch.rocmSupport,
 }:
 
+let
+  inherit (torch) cudaCapabilities cudaPackages;
+in
 buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
   pname = "torchcodec";
   version = "0.14.0";
@@ -100,7 +102,7 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
     ENABLE_CUDA = cudaSupport;
   }
   // lib.optionalAttrs cudaSupport {
-    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" cudaCapabilities}";
   }
   // lib.optionalAttrs rocmSupport {
     ROCM_PATH = torch.rocmtoolkit_joined;

pkgs/development/python-modules/triton/default.nix

@@ -218,6 +218,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } (finalAttrs: {
   ];
 
   passthru = {
+    inherit cudaPackages;
     gpuCheck = effectiveStdenv.mkDerivation {
       pname = "triton-pytest";
       inherit (triton) version src;