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

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

Co-authored-by: Someone <else@someonex.net>
This commit is contained in:
Break Yang
2024-06-12 16:49:46 +00:00
committed by GitHub
co-authored by Someone
parent 9d428a5276
commit 4b755b0f08
2 changed files with 34 additions and 0 deletions
@@ -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;
@@ -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" ];