pkgs/development/python-modules/torch/bin.nix

   1 { lib, stdenv
   2 , buildPythonPackage
   3 , fetchurl
   4 , python
   5 , pythonAtLeast
   6 , pythonOlder
   7 , addOpenGLRunpath
   8 , cudaPackages
   9 , future
  10 , numpy
  11 , autoPatchelfHook
  12 , patchelf
  13 , pyyaml
  14 , requests
  15 , setuptools
  16 , typing-extensions
  17 , sympy
  18 , jinja2
  19 , networkx
  20 , filelock
  21 , openai-triton
  22 }:
  23
  24 let
  25   pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
  26   srcs = import ./binary-hashes.nix version;
  27   unsupported = throw "Unsupported system";
  28   version = "2.0.1";
  29 in buildPythonPackage {
  30   inherit version;
  31
  32   pname = "torch";
  33   # Don't forget to update torch to the same version.
  34
  35   format = "wheel";
  36
  37   disabled = (pythonOlder "3.8") || (pythonAtLeast "3.12");
  38
  39   src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
  40
  41   nativeBuildInputs = lib.optionals stdenv.isLinux [
  42     addOpenGLRunpath
  43     autoPatchelfHook
  44     cudaPackages.autoAddOpenGLRunpathHook
  45   ];
  46
  47   buildInputs = lib.optionals stdenv.isLinux (with cudaPackages; [
  48     # $out/${sitePackages}/nvfuser/_C*.so wants libnvToolsExt.so.1 but torch/lib only ships
  49     # libnvToolsExt-$hash.so.1
  50     cuda_nvtx
  51   ]);
  52
  53   autoPatchelfIgnoreMissingDeps = lib.optionals stdenv.isLinux [
  54     # This is the hardware-dependent userspace driver that comes from
  55     # nvidia_x11 package. It must be deployed at runtime in
  56     # /run/opengl-driver/lib or pointed at by LD_LIBRARY_PATH variable, rather
  57     # than pinned in runpath
  58     "libcuda.so.1"
  59   ];
  60
  61   propagatedBuildInputs = [
  62     future
  63     numpy
  64     pyyaml
  65     requests
  66     setuptools
  67     typing-extensions
  68     sympy
  69     jinja2
  70     networkx
  71     filelock
  72   ] ++ lib.optionals stdenv.isx86_64 [
  73     openai-triton
  74   ];
  75
  76   postInstall = ''
  77     # ONNX conversion
  78     rm -rf $out/bin
  79   '';
  80
  81   postFixup = lib.optionalString stdenv.isLinux ''
  82     addAutoPatchelfSearchPath "$out/${python.sitePackages}/torch/lib"
  83
  84     patchelf $out/${python.sitePackages}/torch/lib/libcudnn.so.8 --add-needed libcudnn_cnn_infer.so.8
  85
  86     pushd $out/${python.sitePackages}/torch/lib || exit 1
  87       for LIBNVRTC in ./libnvrtc*
  88       do
  89         case "$LIBNVRTC" in
  90           ./libnvrtc-builtins*) true;;
  91           ./libnvrtc*) patchelf "$LIBNVRTC" --add-needed libnvrtc-builtins* ;;
  92         esac
  93       done
  94     popd || exit 1
  95   '';
  96
  97   # The wheel-binary is not stripped to avoid the error of `ImportError: libtorch_cuda_cpp.so: ELF load command address/offset not properly aligned.`.
  98   dontStrip = true;
  99
 100   pythonImportsCheck = [ "torch" ];
 101
 102   meta = with lib; {
 103     description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration";
 104     homepage = "https://pytorch.org/";
 105     changelog = "https://github.com/pytorch/pytorch/releases/tag/v${version}";
 106     # Includes CUDA and Intel MKL, but redistributions of the binary are not limited.
 107     # https://docs.nvidia.com/cuda/eula/index.html
 108     # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
 109     # torch's license is BSD3.
 110     # torch-bin includes CUDA and MKL binaries, therefore unfreeRedistributable is set.
 111     license = with licenses; [ bsd3 issl unfreeRedistributable ];
 112     sourceProvenance = with sourceTypes; [ binaryNativeCode ];
 113     platforms = [ "aarch64-darwin" "aarch64-linux" "x86_64-darwin" "x86_64-linux" ];
 114     hydraPlatforms = []; # output size 3.2G on 1.11.0
 115     maintainers = with maintainers; [ junjihashimoto ];
 116   };
 117 }