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.

2 files changed, 44 insertions, 7 deletions

diff --git a/pkgs/development/libraries/science/math/libtorch/bin.nix b/pkgs/development/libraries/science/math/libtorch/bin.nix
index 5e82d50ed9c..a301b0f881e 100644
--- a/pkgs/development/libraries/science/math/libtorch/bin.nix
+++ b/pkgs/development/libraries/science/math/libtorch/bin.nix
@@ -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";
diff --git a/pkgs/development/libraries/science/math/libtorch/test/default.nix b/pkgs/development/libraries/science/math/libtorch/test/default.nix
index e69807871f4..60f9b5ad884 100644
--- a/pkgs/development/libraries/science/math/libtorch/test/default.nix
+++ b/pkgs/development/libraries/science/math/libtorch/test/default.nix
@@ -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
   '';
 }