diff --git a/pkgs/development/python-modules/torchaudio/bin.nix b/pkgs/development/python-modules/torchaudio/bin.nix new file mode 100644 index 000000000000..1f6ae80b0329 --- /dev/null +++ b/pkgs/development/python-modules/torchaudio/bin.nix @@ -0,0 +1,52 @@ +{ lib +, stdenv +, buildPythonPackage +, fetchurl +, python +, pytorch-bin +, pythonOlder +, pythonAtLeast +}: + +buildPythonPackage rec { + pname = "torchaudio"; + version = "0.10.0"; + format = "wheel"; + + src = + let pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion; + unsupported = throw "Unsupported system"; + srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported; + in fetchurl srcs; + + disabled = ! (pythonAtLeast "3.7" && pythonOlder "3.10"); + + propagatedBuildInputs = [ + pytorch-bin + ]; + + # The wheel-binary is not stripped to avoid the error of `ImportError: libtorch_cuda_cpp.so: ELF load command address/offset not properly aligned.`. + dontStrip = true; + + pythonImportsCheck = [ "torchaudio" ]; + + postFixup = '' + # Note: after patchelf'ing, libcudart can still not be found. However, this should + # not be an issue, because PyTorch is loaded before torchvision and brings + # in the necessary symbols. + patchelf --set-rpath "${lib.makeLibraryPath [ stdenv.cc.cc.lib ]}:${pytorch-bin}/${python.sitePackages}/torch/lib:" \ + "$out/${python.sitePackages}/torchaudio/_torchaudio.so" + ''; + + meta = with lib; { + description = "PyTorch audio library"; + homepage = "https://pytorch.org/"; + changelog = "https://github.com/pytorch/audio/releases/tag/v${version}"; + # Includes CUDA and Intel MKL, but redistributions of the binary are not limited. + # https://docs.nvidia.com/cuda/eula/index.html + # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html + license = licenses.bsd3; + platforms = platforms.linux; + maintainers = with maintainers; [ junjihashimoto ]; + }; +} diff --git a/pkgs/development/python-modules/torchaudio/binary-hashes.nix b/pkgs/development/python-modules/torchaudio/binary-hashes.nix new file mode 100644 index 000000000000..764b9db426b9 --- /dev/null +++ b/pkgs/development/python-modules/torchaudio/binary-hashes.nix @@ -0,0 +1,26 @@ +# Warning: Need to update at the same time as pytorch-bin +# +# Precompiled wheels can be found at: +# https://download.pytorch.org/whl/torch_stable.html + +# To add a new version, run "prefetch.sh 'new-version'" to paste the generated file as follows. + +version : builtins.getAttr version { + "0.10.0" = { + x86_64-linux-37 = { + name = "torchaudio-0.10.0-cp37-cp37m-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/cu113/torchaudio-0.10.0%2Bcu113-cp37-cp37m-linux_x86_64.whl"; + hash = "sha256-FspXTTODdkO0nPUJcJm8+vLIvckUa8gRfBPBT9LcKPw="; + }; + x86_64-linux-38 = { + name = "torchaudio-0.10.0-cp38-cp38-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/cu113/torchaudio-0.10.0%2Bcu113-cp38-cp38-linux_x86_64.whl"; + hash = "sha256-Mf7QdXBSIIWRfT7ACthEwFA1V2ieid8legbMnRQnzqI="; + }; + x86_64-linux-39 = { + name = "torchaudio-0.10.0-cp39-cp39-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/cu113/torchaudio-0.10.0%2Bcu113-cp39-cp39-linux_x86_64.whl"; + hash = "sha256-LMSGNdmku1iHRy1jCRTTOYcQlRL+Oc9jjZC1nx++skA="; + }; + }; +} diff --git a/pkgs/development/python-modules/torchaudio/prefetch.sh b/pkgs/development/python-modules/torchaudio/prefetch.sh new file mode 100755 index 000000000000..4e62ab4ed359 --- /dev/null +++ b/pkgs/development/python-modules/torchaudio/prefetch.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env nix-shell +#!nix-shell -i bash -p nix-prefetch-scripts + +set -eou pipefail + +version=$1 + +bucket="https://download.pytorch.org/whl/cu113" + +url_and_key_list=( + "x86_64-linux-37 $bucket/torchaudio-${version}%2Bcu113-cp37-cp37m-linux_x86_64.whl torchaudio-${version}-cp37-cp37m-linux_x86_64.whl" + "x86_64-linux-38 $bucket/torchaudio-${version}%2Bcu113-cp38-cp38-linux_x86_64.whl torchaudio-${version}-cp38-cp38-linux_x86_64.whl" + "x86_64-linux-39 $bucket/torchaudio-${version}%2Bcu113-cp39-cp39-linux_x86_64.whl torchaudio-${version}-cp39-cp39-linux_x86_64.whl" +) + +hashfile=binary-hashes-"$version".nix +echo " \"$version\" = {" >> $hashfile + +for url_and_key in "${url_and_key_list[@]}"; do + key=$(echo "$url_and_key" | cut -d' ' -f1) + url=$(echo "$url_and_key" | cut -d' ' -f2) + name=$(echo "$url_and_key" | cut -d' ' -f3) + + echo "prefetching ${url}..." + hash=$(nix hash to-sri --type sha256 `nix-prefetch-url "$url" --name "$name"`) + + echo " $key = {" >> $hashfile + echo " name = \"$name\";" >> $hashfile + echo " url = \"$url\";" >> $hashfile + echo " hash = \"$hash\";" >> $hashfile + echo " };" >> $hashfile + + echo +done + +echo " };" >> $hashfile +echo "done." diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index bf39935714aa..1a1759463623 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -9068,6 +9068,8 @@ in { toposort = callPackage ../development/python-modules/toposort { }; + torchaudio-bin = callPackage ../development/python-modules/torchaudio/bin.nix { }; + torchgpipe = callPackage ../development/python-modules/torchgpipe { }; torchvision = callPackage ../development/python-modules/torchvision { };