diff options
| author | Someone <sergei.kozlukov@aalto.fi> | 2024-01-11 04:32:59 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-01-11 04:32:59 +0000 |
| commit | e529aea84f3b8cc2ddda94de4cecb583adb2c31f (patch) | |
| tree | 8a7d2a32ce30c807a82397b8166a3442445d9c5a /pkgs/development/python-modules | |
| parent | 6f2d80664de9a493429881f1559e32c616695461 (diff) | |
| parent | b2f97e14aee25445d324a53d3374cb3547242ba7 (diff) | |
Merge pull request #274319 from ConnorBaker/feat/cudaPackages-all-packages-eval
tree-wide: cudaPackages attributes should not cause default eval to fail
Diffstat (limited to 'pkgs/development/python-modules')
| -rw-r--r-- | pkgs/development/python-modules/jaxlib/default.nix | 3 | ||||
| -rw-r--r-- | pkgs/development/python-modules/torch/default.nix | 13 |
2 files changed, 10 insertions, 6 deletions
diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix index 27b9e61fbc82..d8dc4d67a594 100644 --- a/pkgs/development/python-modules/jaxlib/default.nix +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -64,7 +64,8 @@ let # aarch64-darwin is broken because of https://github.com/bazelbuild/rules_cc/pull/136 # however even with that fix applied, it doesn't work for everyone: # https://github.com/NixOS/nixpkgs/pull/184395#issuecomment-1207287129 - broken = stdenv.isDarwin; + # NOTE: We always build with NCCL; if it is unsupported, then our build is broken. + broken = stdenv.isDarwin || nccl.meta.unsupported; }; cudatoolkit_joined = symlinkJoin { diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index 8fb227cbd36b..8a499d763a4a 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -7,7 +7,8 @@ magma, magma-hip, magma-cuda-static, - useSystemNccl ? true, + # Use the system NCCL as long as we're targeting CUDA on a supported platform. + useSystemNccl ? (cudaSupport && !cudaPackages.nccl.meta.unsupported), MPISupport ? false, mpi, buildDocs ? false, @@ -273,9 +274,11 @@ in buildPythonPackage rec { PYTORCH_BUILD_VERSION = version; PYTORCH_BUILD_NUMBER = 0; - USE_NCCL = setBool (cudaSupport && cudaPackages ? nccl); - USE_SYSTEM_NCCL = setBool useSystemNccl; # don't build pytorch's third_party NCCL - USE_STATIC_NCCL = setBool useSystemNccl; + # In-tree builds of NCCL are not supported. + # Use NCCL when cudaSupport is enabled and nccl is available. + USE_NCCL = setBool useSystemNccl; + USE_SYSTEM_NCCL = USE_NCCL; + USE_STATIC_NCCL = USE_NCCL; # Suppress a weird warning in mkl-dnn, part of ideep in pytorch # (upstream seems to have fixed this in the wrong place?) @@ -363,7 +366,7 @@ in buildPythonPackage rec { ] ++ lists.optionals (cudaPackages ? cudnn) [ cudnn.dev cudnn.lib - ] ++ lists.optionals (useSystemNccl && cudaPackages ? nccl) [ + ] ++ lists.optionals useSystemNccl [ # Some platforms do not support NCCL (i.e., Jetson) nccl.dev # Provides nccl.h AND a static copy of NCCL! ] ++ lists.optionals (strings.versionOlder cudaVersion "11.8") [ |