diff --git a/pkgs/development/python-modules/torch/bin.nix b/pkgs/development/python-modules/torch/bin.nix index ef77103c44f1..f9d5cd97c183 100644 --- a/pkgs/development/python-modules/torch/bin.nix +++ b/pkgs/development/python-modules/torch/bin.nix @@ -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; diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index ed1d61697da4..d5d7e823bed7 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -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" ];