python3Packages.torch*: move env variables into env for structuredAttrs (#491399)
This commit is contained in:
Gaétan Lepage
@@ -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);
|
||||
};
|
||||
}
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user