cudaPackages: updated convention for gpu/runtime checks

Runtime tests (derivations asking for a relaxed sandbox) are now expected at p.gpuCheck, p.gpuChecks.<name>, or at p.tests.<name>.gpuCheck.
NixOS · Jun 26, 2024 · 79a7186 · 79a7186
1 parent 7d667a0
commit 79a7186
Show file tree

Hide file tree

Showing 11 changed files with 79 additions and 81 deletions.
diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
@@ -85,7 +85,6 @@ in
             opengl.paths = config.hardware.opengl.extraPackages ++ [
               config.hardware.opengl.package
               pkgs.addOpenGLRunpath.driverLink
-              "/dev/video*"
               "/dev/dri"
             ];
           }

diff --git a/pkgs/applications/misc/blender/default.nix b/pkgs/applications/misc/blender/default.nix
@@ -7,6 +7,7 @@
   SDL,
   addOpenGLRunpath,
   alembic,
+  blender,
   boost,
   brotli,
   callPackage,
@@ -372,9 +373,21 @@ stdenv.mkDerivation (finalAttrs: {
             --render-frame 1
         done
       '';
-
+      tester-cudaAvailable = cudaPackages.writeGpuTestPython { } ''
+        import subprocess
+        subprocess.run([${
+          lib.concatMapStringsSep ", " (x: ''"${x}"'') [
+            (lib.getExe (blender.override { cudaSupport = true; }))
+            "--background"
+            "-noaudio"
+            "--python-exit-code"
+            "1"
+            "--python"
+            "${./test-cuda.py}"
+          ]
+        }], check=True)  # noqa: E501
+      '';
     };
-    gpuChecks = callPackage ./gpu-checks.nix { };
   };
 
   meta = {
@@ -383,7 +396,8 @@ stdenv.mkDerivation (finalAttrs: {
     # They comment two licenses: GPLv2 and Blender License, but they
     # say: "We've decided to cancel the BL offering for an indefinite period."
     # OptiX, enabled with cudaSupport, is non-free.
-    license = with lib.licenses; [ gpl2Plus ] ++ lib.optional cudaSupport unfree;
+    license = with lib.licenses; [ gpl2Plus ] ++ lib.optional cudaSupport (unfree // { shortName = "NVidia OptiX EULA"; });
+
     platforms = [
       "aarch64-linux"
       "x86_64-darwin"

diff --git a/pkgs/applications/misc/blender/gpu-checks.nix b/pkgs/applications/misc/blender/gpu-checks.nix
diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -16,7 +16,6 @@ let
     cudatoolkit
     flags
     libcublas
-    setupCudaHook
     ;
   inherit (lib) getDev getLib getOutput;
   fs = lib.fileset;
@@ -59,20 +58,19 @@ backendStdenv.mkDerivation {
     (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)
   ];
 
-  passthru.gpuChecks.withCuda = saxpy.overrideAttrs (
-    _: {
-      requiredSystemFeatures = ["cuda"];
-      doInstallCheck = true;
-      postInstallCheck = ''
-        $out/bin/saxpy
-      '';
-    }
-  );
+  passthru.gpuCheck = saxpy.overrideAttrs (_: {
+    requiredSystemFeatures = [ "cuda" ];
+    doInstallCheck = true;
+    postInstallCheck = ''
+      $out/bin/${saxpy.meta.mainProgram or (lib.getName saxpy)}
+    '';
+  });
 
   meta = rec {
     description = "Simple (Single-precision AX Plus Y) FindCUDAToolkit.cmake example for testing cross-compilation";
     license = lib.licenses.mit;
     maintainers = lib.teams.cuda.members;
+    mainProgram = "saxpy";
     platforms = lib.platforms.unix;
     badPlatforms = lib.optionals (flags.isJetsonBuild && cudaOlder "11.4") platforms;
   };

diff --git a/pkgs/development/cuda-modules/write-gpu-python-test.nix b/pkgs/development/cuda-modules/write-gpu-python-test.nix
@@ -0,0 +1,29 @@
+{
+  lib,
+  writers,
+  runCommand,
+}:
+{
+  feature ? "cuda",
+  name ? feature,
+  libraries ? [ ],
+}:
+content:
+
+let
+  tester = writers.writePython3Bin "tester-${name}" { inherit libraries; } content;
+  tester' = tester.overrideAttrs (oldAttrs: {
+    passthru.gpuCheck =
+      runCommand "test-${name}"
+        {
+          nativeBuildInputs = [ tester' ];
+          requiredSystemFeatures = [ feature ];
+        }
+        ''
+          set -e
+          ${tester.meta.mainProgram or (lib.getName tester')}
+          touch $out
+        '';
+  });
+in
+tester'
diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
@@ -1,14 +1,15 @@
 {
   lib,
   buildPythonPackage,
-  callPackage,
+  cudaPackages,
   fetchFromGitHub,
   substituteAll,
   pythonOlder,
   addOpenGLRunpath,
   setuptools,
   pytestCheckHook,
   versioneer,
+  pynvml,
 }:
 
 buildPythonPackage rec {
@@ -51,7 +52,12 @@ buildPythonPackage rec {
   # OSError: /run/opengl-driver/lib/libnvidia-ml.so.1: cannot open shared object file: No such file or directory
   doCheck = false;
 
-  passthru.gpuChecks.nvmlInit = callPackage ./test-gpu.nix { };
+  passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ pynvml ]; } ''
+    import pynvml
+    from pynvml.smi import nvidia_smi  # noqa: F401
+
+    print(f"{pynvml.nvmlInit()=}")
+  '';
 
   meta = with lib; {
     description = "Python bindings for the NVIDIA Management Library";

diff --git a/pkgs/development/python-modules/pynvml/test-gpu.nix b/pkgs/development/python-modules/pynvml/test-gpu.nix
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
@@ -648,7 +648,7 @@ buildPythonPackage rec {
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-    gpuChecks = callPackage ./gpu-checks.nix { };
+    tests = callPackage ./tests.nix { };
   };
 
   meta = {

diff --git a/pkgs/development/python-modules/torch/gpu-checks.nix b/pkgs/development/python-modules/torch/gpu-checks.nix
@@ -1,8 +1,8 @@
 {
   lib,
-  callPackage,
   torchWithCuda,
   torchWithRocm,
+  callPackage,
 }:
 
 let
@@ -11,38 +11,28 @@ let
       feature,
       versionAttr,
       torch,
-      runCommandNoCC,
-      writers,
+      cudaPackages,
     }:
-    let
-      name = "${torch.name}-${feature}-check";
-      unwrapped = writers.writePython3Bin "${name}-unwrapped" { libraries = [ torch ]; } ''
+    cudaPackages.writeGpuPythonTest
+      {
+        inherit feature;
+        libraries = [ torch ];
+        name = "${feature}Available";
+      }
+      ''
         import torch
         message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
         assert torch.cuda.is_available() and torch.version.${versionAttr}, message
         print(message)
       '';
-    in
-    runCommandNoCC name
-      {
-        nativeBuildInputs = [ unwrapped ];
-        requiredSystemFeatures = [ feature ];
-        passthru = {
-          inherit unwrapped;
-        };
-      }
-      ''
-        ${name}-unwrapped
-        touch $out
-      '';
 in
 {
-  cudaAvailable = callPackage accelAvailable {
+  tester-cudaAvailable = callPackage accelAvailable {
     feature = "cuda";
     versionAttr = "cuda";
     torch = torchWithCuda;
   };
-  rocmAvailable = callPackage accelAvailable {
+  tester-rocmAvailable = callPackage accelAvailable {
     feature = "rocm";
     versionAttr = "hip";
     torch = torchWithRocm;

diff --git a/pkgs/development/python-modules/torch/tests.nix b/pkgs/development/python-modules/torch/tests.nix
@@ -0,0 +1,3 @@
+{ callPackage }:
+
+callPackage ./gpu-checks.nix { }
diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix
@@ -77,6 +77,8 @@ let
     saxpy = final.callPackage ../development/cuda-modules/saxpy { };
     nccl = final.callPackage ../development/cuda-modules/nccl { };
     nccl-tests = final.callPackage ../development/cuda-modules/nccl-tests { };
+
+    writeGpuTestPython = final.callPackage ../development/cuda-modules/write-gpu-python-test.nix { };
   });
 
   mkVersionedPackageName =