diff --git a/pkgs/development/python-modules/torch/source/default.nix b/pkgs/development/python-modules/torch/source/default.nix index b7504df075d1..d78c4ddba654 100644 --- a/pkgs/development/python-modules/torch/source/default.nix +++ b/pkgs/development/python-modules/torch/source/default.nix @@ -277,7 +277,7 @@ let # From here on, `stdenv` shall be `stdenv'`. stdenv = stdenv'; in -buildPythonPackage.override { inherit stdenv; } rec { +buildPythonPackage.override { inherit stdenv; } (finalAttrs: { pname = "torch"; # Don't forget to update torch-bin to the same version. version = "2.9.1"; @@ -293,11 +293,13 @@ buildPythonPackage.override { inherit stdenv; } rec { src = callPackage ./src.nix { inherit - version fetchFromGitHub fetchFromGitLab runCommand ; + inherit (finalAttrs) + version + ; }; patches = [ @@ -420,35 +422,76 @@ buildPythonPackage.override { inherit stdenv; } rec { # causes possible redefinition of _FORTIFY_SOURCE hardeningDisable = [ "fortify3" ]; - BUILD_NAMEDTENSOR = setBool true; - BUILD_DOCS = setBool buildDocs; + env = { + BUILD_NAMEDTENSOR = setBool true; + BUILD_DOCS = setBool buildDocs; - # We only do an imports check, so do not build tests either. - BUILD_TEST = setBool false; + # We only do an imports check, so do not build tests either. + BUILD_TEST = setBool false; - # ninja hook doesn't automatically turn on ninja - # because pytorch setup.py is responsible for this - CMAKE_GENERATOR = "Ninja"; + # ninja hook doesn't automatically turn on ninja + # because pytorch setup.py is responsible for this + CMAKE_GENERATOR = "Ninja"; - # Unlike MKL, oneDNN (née MKLDNN) is FOSS, so we enable support for - # it by default. PyTorch currently uses its own vendored version - # of oneDNN through Intel iDeep. - USE_MKLDNN = setBool mklDnnSupport; - USE_MKLDNN_CBLAS = setBool mklDnnSupport; + # Unlike MKL, oneDNN (née MKLDNN) is FOSS, so we enable support for + # it by default. PyTorch currently uses its own vendored version + # of oneDNN through Intel iDeep. + USE_MKLDNN = setBool mklDnnSupport; + USE_MKLDNN_CBLAS = setBool mklDnnSupport; - # Avoid using pybind11 from git submodule - # Also avoids pytorch exporting the headers of pybind11 - USE_SYSTEM_PYBIND11 = true; + # Avoid using pybind11 from git submodule + # Also avoids pytorch exporting the headers of pybind11 + USE_SYSTEM_PYBIND11 = true; - # Multicore CPU convnet support - USE_NNPACK = 1; + # Multicore CPU convnet support + USE_NNPACK = 1; - # Explicitly enable MPS for Darwin - USE_MPS = setBool stdenv.hostPlatform.isDarwin; + # Explicitly enable MPS for Darwin + USE_MPS = setBool stdenv.hostPlatform.isDarwin; - # building torch.distributed on Darwin is disabled by default - # https://pytorch.org/docs/stable/distributed.html#torch.distributed.is_available - USE_DISTRIBUTED = setBool true; + # building torch.distributed on Darwin is disabled by default + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.is_available + USE_DISTRIBUTED = setBool true; + + # Override the (weirdly) wrong version set by default. See + # https://github.com/NixOS/nixpkgs/pull/52437#issuecomment-449718038 + # https://github.com/pytorch/pytorch/blob/v1.0.0/setup.py#L267 + PYTORCH_BUILD_VERSION = finalAttrs.version; + PYTORCH_BUILD_NUMBER = 0; + + # 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 = finalAttrs.env.USE_NCCL; + USE_STATIC_NCCL = finalAttrs.env.USE_NCCL; + + # Set the correct Python library path, broken since + # https://github.com/pytorch/pytorch/commit/3d617333e + PYTHON_LIB_REL_PATH = "${placeholder "out"}/${python.sitePackages}"; + # disable warnings as errors as they break the build on every compiler + # bump, among other things. + # Also of interest: pytorch ignores CXXFLAGS uses CFLAGS for both C and C++: + # https://github.com/pytorch/pytorch/blob/v1.11.0/setup.py#L17 + NIX_CFLAGS_COMPILE = toString ( + [ + "-Wno-error" + ] + # fix build aarch64-linux build failure with GCC14 + ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ + "-Wno-error=incompatible-pointer-types" + ] + ); + USE_VULKAN = setBool vulkanSupport; + } + // lib.optionalAttrs vulkanSupport { + VULKAN_SDK = shaderc.bin; + } + // lib.optionalAttrs rocmSupport { + AOTRITON_INSTALLED_PREFIX = "${rocmPackages.aotriton}"; + # Broken HIP flag setup, fails to compile due to not finding rocthrust + # Only supports gfx942 so let's turn it off for now + USE_FBGEMM_GENAI = setBool false; + }; cmakeFlags = [ (lib.cmakeFeature "PYTHON_SIX_SOURCE_DIR" "${six.src}") @@ -482,48 +525,6 @@ buildPythonPackage.override { inherit stdenv; } rec { done ''; - # Override the (weirdly) wrong version set by default. See - # https://github.com/NixOS/nixpkgs/pull/52437#issuecomment-449718038 - # https://github.com/pytorch/pytorch/blob/v1.0.0/setup.py#L267 - PYTORCH_BUILD_VERSION = version; - PYTORCH_BUILD_NUMBER = 0; - - # 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; - - # Set the correct Python library path, broken since - # https://github.com/pytorch/pytorch/commit/3d617333e - PYTHON_LIB_REL_PATH = "${placeholder "out"}/${python.sitePackages}"; - - env = { - # disable warnings as errors as they break the build on every compiler - # bump, among other things. - # Also of interest: pytorch ignores CXXFLAGS uses CFLAGS for both C and C++: - # https://github.com/pytorch/pytorch/blob/v1.11.0/setup.py#L17 - NIX_CFLAGS_COMPILE = toString ( - [ - "-Wno-error" - ] - # fix build aarch64-linux build failure with GCC14 - ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ - "-Wno-error=incompatible-pointer-types" - ] - ); - USE_VULKAN = setBool vulkanSupport; - } - // lib.optionalAttrs vulkanSupport { - VULKAN_SDK = shaderc.bin; - } - // lib.optionalAttrs rocmSupport { - AOTRITON_INSTALLED_PREFIX = "${rocmPackages.aotriton}"; - # Broken HIP flag setup, fails to compile due to not finding rocthrust - # Only supports gfx942 so let's turn it off for now - USE_FBGEMM_GENAI = setBool false; - }; - nativeBuildInputs = [ cmake which @@ -650,7 +651,7 @@ buildPythonPackage.override { inherit stdenv; } rec { # ^^^^^^^^^^^^ NOTE: while test_dataloader does return errors, these are acceptable errors and do not interfere with the build # tensorboard has acceptable failures for pytorch 1.3.x due to dependencies on tensorboard-plugins - (optionalString (majorMinor version == "1.3") "tensorboard") + (optionalString (majorMinor finalAttrs.version == "1.3") "tensorboard") ]) "runHook postCheck" ]; @@ -692,7 +693,7 @@ buildPythonPackage.override { inherit stdenv; } rec { --replace-fail "\''${_IMPORT_PREFIX}/lib64" "$lib/lib" substituteInPlace $dev/share/cmake/ATen/ATenConfig.cmake \ - --replace-fail "/build/${src.name}/torch/include" "$dev/include" + --replace-fail "/build/${finalAttrs.src.name}/torch/include" "$dev/include" ''; postFixup = '' @@ -751,7 +752,7 @@ buildPythonPackage.override { inherit stdenv; } rec { }; meta = { - changelog = "https://github.com/pytorch/pytorch/releases/tag/v${version}"; + changelog = "https://github.com/pytorch/pytorch/releases/tag/v${finalAttrs.version}"; # keep PyTorch in the description so the package can be found under that name on search.nixos.org description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration"; homepage = "https://pytorch.org/"; @@ -767,4 +768,4 @@ buildPythonPackage.override { inherit stdenv; } rec { lib.platforms.linux ++ lib.optionals (!cudaSupport && !rocmSupport) lib.platforms.darwin; broken = builtins.any trivial.id (builtins.attrValues brokenConditions); }; -} +})