diff options
Diffstat (limited to 'infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch')
3 files changed, 58 insertions, 29 deletions
diff --git a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/bin.nix b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/bin.nix index b1b662e95b3e..f073c408f709 100644 --- a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/bin.nix +++ b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/bin.nix @@ -11,17 +11,20 @@ , patchelf , pyyaml , requests +, typing-extensions }: let pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion; platform = if stdenv.isDarwin then "darwin" else "linux"; - srcs = import ./binary-hashes.nix; + srcs = import ./binary-hashes.nix version; unsupported = throw "Unsupported system"; + version = "1.7.0"; in buildPythonPackage { + inherit version; + pname = "pytorch"; # Don't forget to update pytorch to the same version. - version = "1.6.0"; format = "wheel"; @@ -39,6 +42,17 @@ in buildPythonPackage { numpy pyyaml requests + typing-extensions + ]; + + # PyTorch are broken: the dataclasses wheel is required, but ships with + # Python >= 3.7. Our dataclasses derivation is incompatible with >= 3.7. + # + # https://github.com/pytorch/pytorch/issues/46930 + # + # Should be removed with the next PyTorch version. + pipInstallFlags = [ + "--no-deps" ]; postInstall = '' diff --git a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/binary-hashes.nix b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/binary-hashes.nix index 481dc6825c7d..afdde5057dbd 100644 --- a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/binary-hashes.nix +++ b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/binary-hashes.nix @@ -1,10 +1,10 @@ -{ +version: { x86_64-linux-37 = { - url = "https://download.pytorch.org/whl/cu102/torch-1.6.0-cp37-cp37m-linux_x86_64.whl"; - sha256 = "0xhwv68j8gvahfzcp43bqp2x71iwv6zjhkw2f1hb82xps40mrml7"; + url = "https://download.pytorch.org/whl/cu102/torch-${version}-cp37-cp37m-linux_x86_64.whl"; + sha256 = "0cvafwzd44d7i2d0y01lyk75lhxffcf6c94cx0wla0cnhcg9mqlf"; }; x86_64-linux-38 = { - url = "https://download.pytorch.org/whl/cu102/torch-1.6.0-cp38-cp38-linux_x86_64.whl"; - sha256 = "05m2l04wqzw5xvjam6zwvlmc3979cksl3hrdqc2aikrv4hz8fmsk"; + url = "https://download.pytorch.org/whl/cu102/torch-${version}-cp38-cp38-linux_x86_64.whl"; + sha256 = "1vxswkb2gzp87sbb5ycvyvs8aw1szp3v2mwdn4mp47yarlnbx8r6"; }; } diff --git a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/default.nix b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/default.nix index 4635f813c3a6..ff7a0e5107f3 100644 --- a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/default.nix +++ b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch/default.nix @@ -4,8 +4,18 @@ openMPISupport ? false, openmpi ? null, buildDocs ? false, cudaArchList ? null, - numpy, pyyaml, cffi, click, typing, cmake, hypothesis, numactl, psutil, - linkFarm, symlinkJoin, + + # Native build inputs + cmake, util-linux, linkFarm, symlinkJoin, which, + + # Build inputs + numactl, + + # Propagated build inputs + dataclasses, numpy, pyyaml, cffi, click, typing-extensions, + + # Unit tests + hypothesis, psutil, # virtual pkg that consistently instantiates blas across nixpkgs # See https://github.com/NixOS/nixpkgs/pull/83888 @@ -17,7 +27,7 @@ # dependencies for torch.utils.tensorboard pillow, six, future, tensorflow-tensorboard, protobuf, - utillinux, which, isPy3k }: + isPy3k, pythonOlder }: assert !openMPISupport || openmpi != null; @@ -102,7 +112,7 @@ let in buildPythonPackage rec { pname = "pytorch"; # Don't forget to update pytorch-bin to the same version. - version = "1.6.0"; + version = "1.7.0"; disabled = !isPy3k; @@ -117,23 +127,10 @@ in buildPythonPackage rec { repo = "pytorch"; rev = "v${version}"; fetchSubmodules = true; - sha256 = "14hhjsi6fnpaw9m1a3bhvdinsks6fhss6bbcrfk6jgns64abqdaz"; + sha256 = "0hb11m5kvs4nsi5bc9ijnv6k07593p9dw40rcn33s9vspinrcjd4"; }; - patches = lib.optionals stdenv.isAarch64 [ - # GNU aarch64 assembler does not support 4s on neon mov: - # https://github.com/pytorch/pytorch/issues/33124 - # - # Fix from: - # https://github.com/pytorch/pytorch/pull/40584 - # - # This patch can be removed with the next major version (1.7.0). - (fetchpatch { - name = "qnnpack-neon-fix.patch"; - url = "https://github.com/pytorch/pytorch/commit/7676682584d0caf9243bce74ea0a88711ec4a807.diff"; - sha256 = "13spncaqlpsp8qk2850yly7xqwmhhfwznhmzkk8jgpslkbx75vgq"; - }) - ] ++ lib.optionals stdenv.isDarwin [ + patches = lib.optionals stdenv.isDarwin [ # pthreadpool added support for Grand Central Dispatch in April # 2020. However, this relies on functionality (DISPATCH_APPLY_AUTO) # that is available starting with macOS 10.13. However, our current @@ -142,6 +139,13 @@ in buildPythonPackage rec { ./pthreadpool-disable-gcd.diff ]; + # The dataclasses module is included with Python >= 3.7. This should + # be fixed with the next PyTorch release. + postPatch = '' + substituteInPlace setup.py \ + --replace "'dataclasses'" "'dataclasses; python_version < \"3.7\"'" + ''; + preConfigure = lib.optionalString cudaSupport '' export TORCH_CUDA_ARCH_LIST="${lib.strings.concatStringsSep ";" final_cudaArchList}" export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++ @@ -203,7 +207,7 @@ in buildPythonPackage rec { nativeBuildInputs = [ cmake - utillinux + util-linux which ninja ] ++ lib.optionals cudaSupport [ cudatoolkit_joined ]; @@ -217,9 +221,11 @@ in buildPythonPackage rec { click numpy pyyaml + typing-extensions # the following are required for tensorboard support pillow six future tensorflow-tensorboard protobuf - ] ++ lib.optionals openMPISupport [ openmpi ]; + ] ++ lib.optionals openMPISupport [ openmpi ] + ++ lib.optionals (pythonOlder "3.7") [ dataclasses ]; checkInputs = [ hypothesis ninja psutil ]; @@ -248,6 +254,15 @@ in buildPythonPackage rec { cp -r $out/${python.sitePackages}/torch/include $dev/include cp -r $out/${python.sitePackages}/torch/share $dev/share + # Fix up library paths for split outputs + substituteInPlace \ + $dev/share/cmake/Torch/TorchConfig.cmake \ + --replace \''${TORCH_INSTALL_PREFIX}/lib "$lib/lib" + + substituteInPlace \ + $dev/share/cmake/Caffe2/Caffe2Targets-release.cmake \ + --replace \''${_IMPORT_PREFIX}/lib "$lib/lib" + mkdir $lib cp -r $out/${python.sitePackages}/torch/lib $lib/lib ''; @@ -282,6 +297,6 @@ in buildPythonPackage rec { homepage = "https://pytorch.org/"; license = lib.licenses.bsd3; platforms = with lib.platforms; linux ++ lib.optionals (!cudaSupport) darwin; - maintainers = with lib.maintainers; [ teh thoughtpolice tscholak ]; # tscholak esp. for darwin-related builds + maintainers = with lib.maintainers; [ danieldk teh thoughtpolice tscholak ]; # tscholak esp. for darwin-related builds }; } |