libtorch-bin: fix passthru.tests for CUDA

- Provide the CUDA toolkit and CuDNN dependencies.
- Expose libcuda.so stub to run the test in the sandbox.
This commit is contained in:
Daniël de Kok 2021-02-15 17:19:00 +01:00
parent c4482a0c2a
commit d5b676eb34
2 changed files with 44 additions and 7 deletions

View file

@ -8,9 +8,16 @@
, fixDarwinDylibNames
, cudaSupport
, cudatoolkit_10_2
, cudnn_cudatoolkit_10_2
}:
let
# The binary libtorch distribution statically links the CUDA
# toolkit. This means that we do not need to provide CUDA to
# this derivation. However, we should ensure on version bumps
# that the CUDA toolkit for `passthru.tests` is still
# up-to-date.
version = "1.7.1";
device = if cudaSupport then "cuda" else "cpu";
srcs = import ./binary-hashes.nix version;
@ -50,8 +57,7 @@ in stdenv.mkDerivation {
'';
postFixup = let
libPaths = [ stdenv.cc.cc.lib ];
rpath = lib.makeLibraryPath libPaths;
rpath = lib.makeLibraryPath [ stdenv.cc.cc.lib ];
in lib.optionalString stdenv.isLinux ''
find $out/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
echo "setting rpath for $lib..."
@ -101,7 +107,11 @@ in stdenv.mkDerivation {
outputs = [ "out" "dev" ];
passthru.tests.cmake = callPackage ./test { };
passthru.tests.cmake = callPackage ./test {
inherit cudaSupport;
cudatoolkit = cudatoolkit_10_2;
cudnn = cudnn_cudatoolkit_10_2;
};
meta = with lib; {
description = "C++ API of the PyTorch machine learning framework";

View file

@ -1,6 +1,28 @@
{ stdenv, cmake, libtorch-bin, symlinkJoin }:
{ lib
, stdenv
, cmake
, libtorch-bin
, linkFarm
, symlinkJoin
stdenv.mkDerivation {
, cudaSupport
, cudatoolkit
, cudnn
}:
let
cudatoolkit_joined = symlinkJoin {
name = "${cudatoolkit.name}-unsplit";
paths = [ cudatoolkit.out cudatoolkit.lib ];
};
# We do not have access to /run/opengl-driver/lib in the sandbox,
# so use a stub instead.
cudaStub = linkFarm "cuda-stub" [{
name = "libcuda.so.1";
path = "${cudatoolkit}/lib/stubs/libcuda.so";
}];
in stdenv.mkDerivation {
pname = "libtorch-test";
version = libtorch-bin.version;
@ -8,7 +30,11 @@ stdenv.mkDerivation {
nativeBuildInputs = [ cmake ];
buildInputs = [ libtorch-bin ];
buildInputs = [ libtorch-bin ] ++
lib.optionals cudaSupport [ cudnn ];
cmakeFlags = lib.optionals cudaSupport
[ "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit_joined}" ];
doCheck = true;
@ -17,6 +43,7 @@ stdenv.mkDerivation {
'';
checkPhase = ''
./test
LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
./test
'';
}