From 79a7186f1ce8d94b0c136a7cc7c3e2e31facc794 Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Wed, 26 Jun 2024 00:29:42 +0000 Subject: [PATCH] cudaPackages: updated convention for gpu/runtime checks Runtime tests (derivations asking for a relaxed sandbox) are now expected at p.gpuCheck, p.gpuChecks., or at p.tests..gpuCheck. --- .../modules/programs/nix-required-mounts.nix | 1 - pkgs/applications/misc/blender/default.nix | 20 ++++++++++-- pkgs/applications/misc/blender/gpu-checks.nix | 26 --------------- .../cuda-modules/saxpy/default.nix | 18 +++++------ .../cuda-modules/write-gpu-python-test.nix | 29 +++++++++++++++++ .../python-modules/pynvml/default.nix | 10 ++++-- .../python-modules/pynvml/test-gpu.nix | 17 ---------- .../python-modules/torch/default.nix | 2 +- .../python-modules/torch/gpu-checks.nix | 32 +++++++------------ .../python-modules/torch/tests.nix | 3 ++ pkgs/top-level/cuda-packages.nix | 2 ++ 11 files changed, 79 insertions(+), 81 deletions(-) delete mode 100644 pkgs/applications/misc/blender/gpu-checks.nix create mode 100644 pkgs/development/cuda-modules/write-gpu-python-test.nix delete mode 100644 pkgs/development/python-modules/pynvml/test-gpu.nix create mode 100644 pkgs/development/python-modules/torch/tests.nix diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix index c339dd1cfddd..5d25958a7698 100644 --- 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 index 38f3e226d0ff..3d044abaad6d 100644 --- 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 deleted file mode 100644 index bfbaf25b989a..000000000000 --- a/pkgs/applications/misc/blender/gpu-checks.nix +++ /dev/null @@ -1,26 +0,0 @@ -{ - bash, - blender, - callPackage, - lib, - runCommand, - writeScriptBin, -}: - -let - blenderWithCuda = blender.override { cudaSupport = true; }; - name = "${blenderWithCuda.name}-check-cuda"; - unwrapped = writeScriptBin "${name}-unwrapped" '' - #!${lib.getExe bash} - ${lib.getExe blenderWithCuda} --background -noaudio --python-exit-code 1 --python ${./test-cuda.py} - ''; -in -{ - cudaAvailable = runCommand name { - nativeBuildInputs = [ unwrapped ]; - requiredSystemFeatures = [ "cuda" ]; - passthru = { - inherit unwrapped; - }; - } "${name}-unwrapped && touch $out"; -} diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix index 3da3dc08f0eb..5eb0a235ace8 100644 --- 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 new file mode 100644 index 000000000000..5f0d5c6b8fe6 --- /dev/null +++ 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 index 79624e5298a6..762771c66a2b 100644 --- a/pkgs/development/python-modules/pynvml/default.nix +++ b/pkgs/development/python-modules/pynvml/default.nix @@ -1,7 +1,7 @@ { lib, buildPythonPackage, - callPackage, + cudaPackages, fetchFromGitHub, substituteAll, pythonOlder, @@ -9,6 +9,7 @@ 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 deleted file mode 100644 index 6ab4290a2bba..000000000000 --- a/pkgs/development/python-modules/pynvml/test-gpu.nix +++ /dev/null @@ -1,17 +0,0 @@ -{ runCommandNoCC, python }: - -runCommandNoCC "pynvml-gpu-test" - { - nativeBuildInputs = [ (python.withPackages (ps: [ ps.pynvml ])) ]; - requiredSystemFeatures = [ "cuda" ]; - } - '' - python3 << EOF - import pynvml - from pynvml.smi import nvidia_smi - - pynvml.nvmlInit() - EOF - - touch $out - '' diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index 11a3b3df42ce..9597a047bdb4 100644 --- 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 index 371b83f1b778..d01fffe45cb0 100644 --- 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 new file mode 100644 index 000000000000..5a46d0886868 --- /dev/null +++ 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 index d34a37294ae0..7f01f4310c9e 100644 --- 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 =