Unverified Commit 4b755b0f authored by Break Yang's avatar Break Yang Committed by GitHub
Browse files

python3Packages.torch: add libnvrtc to torch's runpath (#315694) 



* python3Packages.torch: add libnvrtc to torch's runpath

Co-authored-by: default avatarSomeone <else@someonex.net>
parent 9d428a52

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

+17 −0
Original line number Diff line number Diff line
@@ -97,6 +97,23 @@ buildPythonPackage { 
    addAutoPatchelfSearchPath "$out/${python.sitePackages}/torch/lib"

  '';



  # See https://github.com/NixOS/nixpkgs/issues/296179

  #

  # This is a quick hack to add `libnvrtc` to the runpath so that torch can find

  # it when it is needed at runtime.

  extraRunpaths = lib.optionals stdenv.hostPlatform.isLinux [ "${lib.getLib cudaPackages.cuda_nvrtc}/lib" ];

  postPhases = lib.optionals stdenv.isLinux [ "postPatchelfPhase" ];

  postPatchelfPhase = ''

    while IFS= read -r -d $'\0' elf ; do

      for extra in $extraRunpaths ; do

        echo patchelf "$elf" --add-rpath "$extra" >&2

        patchelf "$elf" --add-rpath "$extra"

      done

    done < <(

      find "''${!outputLib}" "$out" -type f -iname '*.so' -print0

    )

  '';



  # 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;

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

+17 −0
Original line number Diff line number Diff line
@@ -612,6 +612,23 @@ buildPythonPackage rec { 
      install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libshm.dylib

    '';



  # See https://github.com/NixOS/nixpkgs/issues/296179

  #

  # This is a quick hack to add `libnvrtc` to the runpath so that torch can find

  # it when it is needed at runtime.

  extraRunpaths = lib.optionals cudaSupport [ "${lib.getLib cudaPackages.cuda_nvrtc}/lib" ];

  postPhases = lib.optionals stdenv.isLinux [ "postPatchelfPhase" ];

  postPatchelfPhase = ''

    while IFS= read -r -d $'\0' elf ; do

      for extra in $extraRunpaths ; do

        echo patchelf "$elf" --add-rpath "$extra" >&2

        patchelf "$elf" --add-rpath "$extra"

      done

    done < <(

      find "''${!outputLib}" "$out" -type f -iname '*.so' -print0

    )

  '';



  # Builds in 2+h with 2 cores, and ~15m with a big-parallel builder.

  requiredSystemFeatures = [ "big-parallel" ];