Merge pull request #256230 from SomeoneSerge/feat/gpu-tests-py

GPU access in the sandbox

nixos/modules/module-list.nix

@@ -243,6 +243,7 @@
   ./programs/nh.nix
   ./programs/nix-index.nix
   ./programs/nix-ld.nix
+  ./programs/nix-required-mounts.nix
   ./programs/nm-applet.nix
   ./programs/nncp.nix
   ./programs/noisetorch.nix

nixos/modules/programs/nix-required-mounts.nix

@@ -0,0 +1,118 @@
+{
+  config,
+  lib,
+  pkgs,
+  ...
+}:
+
+let
+  cfg = config.programs.nix-required-mounts;
+  package = pkgs.nix-required-mounts;
+
+  Mount =
+    with lib;
+    types.submodule {
+      options.host = mkOption {
+        type = types.str;
+        description = "Host path to mount";
+      };
+      options.guest = mkOption {
+        type = types.str;
+        description = "Location in the sandbox to mount the host path at";
+      };
+    };
+  Pattern =
+    with lib.types;
+    types.submodule (
+      { config, name, ... }:
+      {
+        options.onFeatures = lib.mkOption {
+          type = listOf types.str;
+          description = "Which requiredSystemFeatures should trigger relaxation of the sandbox";
+          default = [ name ];
+        };
+        options.paths = lib.mkOption {
+          type = listOf (oneOf [
+            path
+            Mount
+          ]);
+          description = "A list of glob patterns, indicating which paths to expose to the sandbox";
+        };
+        options.unsafeFollowSymlinks = lib.mkEnableOption ''
+          Instructs the hook to mount the symlink targets as well, when any of
+          the `paths` contain symlinks. This may not work correctly with glob
+          patterns.
+        '';
+      }
+    );
+
+  driverPaths = [
+    pkgs.addOpenGLRunpath.driverLink
+
+    # mesa:
+    config.hardware.opengl.package
+
+    # nvidia_x11, etc:
+  ] ++ config.hardware.opengl.extraPackages; # nvidia_x11
+
+  defaults = {
+    nvidia-gpu.onFeatures = package.allowedPatterns.nvidia-gpu.onFeatures;
+    nvidia-gpu.paths = package.allowedPatterns.nvidia-gpu.paths ++ driverPaths;
+    nvidia-gpu.unsafeFollowSymlinks = false;
+  };
+in
+{
+  meta.maintainers = with lib.maintainers; [ SomeoneSerge ];
+  options.programs.nix-required-mounts = {
+    enable = lib.mkEnableOption "Expose extra paths to the sandbox depending on derivations' requiredSystemFeatures";
+    presets.nvidia-gpu.enable = lib.mkEnableOption ''
+      Declare the support for derivations that require an Nvidia GPU to be
+      available, e.g. derivations with `requiredSystemFeatures = [ "cuda" ]`.
+      This mounts the corresponding userspace drivers and device nodes in the
+      sandbox, but only for derivations that request these special features.
+
+      You may extend or override the exposed paths via the
+      `programs.nix-required-mounts.allowedPatterns.nvidia-gpu.paths` option.
+    '';
+    allowedPatterns =
+      with lib.types;
+      lib.mkOption rec {
+        type = attrsOf Pattern;
+        description = "The hook config, describing which paths to mount for which system features";
+        default = { };
+        defaultText = lib.literalExpression ''
+          {
+            opengl.paths = config.hardware.opengl.extraPackages ++ [
+              config.hardware.opengl.package
+              pkgs.addOpenGLRunpath.driverLink
+              "/dev/dri"
+            ];
+          }
+        '';
+        example.require-ipfs.paths = [ "/ipfs" ];
+        example.require-ipfs.onFeatures = [ "ifps" ];
+      };
+    extraWrapperArgs = lib.mkOption {
+      type = with lib.types; listOf str;
+      default = [ ];
+      description = "List of extra arguments (such as `--add-flags -v`) to pass to the hook's wrapper";
+    };
+    package = lib.mkOption {
+      type = lib.types.package;
+      default = package.override { inherit (cfg) allowedPatterns extraWrapperArgs; };
+      description = "The final package with the final config applied";
+      internal = true;
+    };
+  };
+  config = lib.mkIf cfg.enable (
+    lib.mkMerge [
+      { nix.settings.pre-build-hook = lib.getExe cfg.package; }
+      (lib.mkIf cfg.presets.nvidia-gpu.enable {
+        nix.settings.system-features = cfg.allowedPatterns.nvidia-gpu.onFeatures;
+        programs.nix-required-mounts.allowedPatterns = {
+          inherit (defaults) nvidia-gpu;
+        };
+      })
+    ]
+  );
+}

nixos/tests/all-tests.nix

@@ -651,6 +651,7 @@ in {
   nix-config = handleTest ./nix-config.nix {};
   nix-ld = handleTest ./nix-ld.nix {};
   nix-misc = handleTest ./nix/misc.nix {};
+  nix-required-mounts = runTest ./nix-required-mounts;
   nix-serve = handleTest ./nix-serve.nix {};
   nix-serve-ssh = handleTest ./nix-serve-ssh.nix {};
   nixops = handleTest ./nixops/default.nix {};

nixos/tests/nix-required-mounts/default.nix

@@ -0,0 +1,58 @@
+{ pkgs, ... }:
+
+let
+  inherit (pkgs) lib;
+in
+
+{
+  name = "nix-required-mounts";
+  meta.maintainers = with lib.maintainers; [ SomeoneSerge ];
+  nodes.machine =
+    { config, pkgs, ... }:
+    {
+      virtualisation.writableStore = true;
+      system.extraDependencies = [ (pkgs.runCommand "deps" { } "mkdir $out").inputDerivation ];
+      nix.nixPath = [ "nixpkgs=${../../..}" ];
+      nix.settings.substituters = lib.mkForce [ ];
+      nix.settings.system-features = [ "supported-feature" ];
+      nix.settings.experimental-features = [ "nix-command" ];
+      programs.nix-required-mounts.enable = true;
+      programs.nix-required-mounts.allowedPatterns.supported-feature = {
+        onFeatures = [ "supported-feature" ];
+        paths = [
+          "/supported-feature-files"
+          {
+            host = "/usr/lib/imaginary-fhs-drivers";
+            guest = "/run/opengl-driver/lib";
+          }
+        ];
+        unsafeFollowSymlinks = true;
+      };
+      users.users.person.isNormalUser = true;
+      systemd.tmpfiles.rules = [
+        "d /supported-feature-files 0755 person users -"
+        "f /usr/lib/libcuda.so 0444 root root - fakeContent"
+        "L /usr/lib/imaginary-fhs-drivers/libcuda.so 0444 root root - /usr/lib/libcuda.so"
+      ];
+    };
+  testScript = ''
+    import shlex
+
+    def person_do(cmd, succeed=True):
+        cmd = shlex.quote(cmd)
+        cmd = f"su person -l -c {cmd} &>/dev/console"
+
+        if succeed:
+            return machine.succeed(cmd)
+        else:
+            return machine.fail(cmd)
+
+    start_all()
+
+    person_do("nix-build ${./ensure-path-not-present.nix} --argstr feature supported-feature")
+    person_do("nix-build ${./test-require-feature.nix} --argstr feature supported-feature")
+    person_do("nix-build ${./test-require-feature.nix} --argstr feature unsupported-feature", succeed=False)
+    person_do("nix-build ${./test-structured-attrs.nix} --argstr feature supported-feature")
+    person_do("nix-build ${./test-structured-attrs-empty.nix}")
+  '';
+}

nixos/tests/nix-required-mounts/ensure-path-not-present.nix

@@ -0,0 +1,13 @@
+{
+  pkgs ? import <nixpkgs> { },
+  feature,
+}:
+
+pkgs.runCommandNoCC "${feature}-not-present" { } ''
+  if [[ -e /${feature}-files ]]; then
+    echo "No ${feature} in requiredSystemFeatures, but /${feature}-files was mounted anyway"
+    exit 1
+  else
+    touch $out
+  fi
+''

nixos/tests/nix-required-mounts/test-require-feature.nix

@@ -0,0 +1,26 @@
+{
+  pkgs ? import <nixpkgs> { },
+  feature,
+}:
+
+pkgs.runCommandNoCC "${feature}-present" { requiredSystemFeatures = [ feature ]; } ''
+  if [[ ! -e /${feature}-files ]]; then
+    echo "The host declares ${feature} support, but doesn't expose /${feature}-files" >&2
+    exit 1
+  fi
+  libcudaLocation=/run/opengl-driver/lib/libcuda.so
+  if [[ -e "$libcudaLocation" || -h "$libcudaLocation" ]] ; then
+    true # we're good
+  else
+    echo "The host declares ${feature} support, but it the hook fails to handle the hostPath != guestPath cases" >&2
+    exit 1
+  fi
+  if cat "$libcudaLocation" | xargs test fakeContent = ; then
+    true # we're good
+  else
+    echo "The host declares ${feature} support, but it seems to fail to follow symlinks" >&2
+    echo "The content of /run/opengl-driver/lib/libcuda.so is: $(cat /run/opengl-driver/lib/libcuda.so)" >&2
+    exit 1
+  fi
+  touch $out
+''

nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix

@@ -0,0 +1,8 @@
+{
+  pkgs ? import <nixpkgs> { },
+}:
+
+pkgs.runCommandNoCC "nix-required-mounts-structured-attrs-no-features" { __structuredAttrs = true; }
+  ''
+    touch $out
+  ''

nixos/tests/nix-required-mounts/test-structured-attrs.nix

@@ -0,0 +1,18 @@
+{
+  pkgs ? import <nixpkgs> { },
+  feature,
+}:
+
+pkgs.runCommandNoCC "${feature}-present-structured"
+  {
+    __structuredAttrs = true;
+    requiredSystemFeatures = [ feature ];
+  }
+  ''
+    if [[ -e /${feature}-files ]]; then
+      touch $out
+    else
+      echo "The host declares ${feature} support, but doesn't expose /${feature}-files" >&2
+      echo "Do we fail to parse __structuredAttrs=true derivations?" >&2
+    fi
+  ''

pkgs/applications/misc/blender/default.nix

@@ -7,6 +7,7 @@
   SDL,
   addOpenGLRunpath,
   alembic,
+  blender,
   boost,
   brotli,
   callPackage,
@@ -372,6 +373,20 @@ 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
+      '';
     };
   };
 
@@ -381,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"

pkgs/applications/misc/blender/test-cuda.py

@@ -0,0 +1,8 @@
+import bpy
+
+preferences = bpy.context.preferences.addons["cycles"].preferences
+devices = preferences.get_devices_for_type("CUDA")
+ids = [d.id for d in devices]
+
+assert any("CUDA" in i for i in ids), f"CUDA not present in {ids}"
+print("CUDA is available")

pkgs/by-name/ni/nix-required-mounts/closure.nix

@@ -0,0 +1,37 @@
+# Use exportReferencesGraph to capture the possible dependencies of the
+# drivers (e.g. libc linked through DT_RUNPATH) and ensure they are mounted
+# in the sandbox as well. In practice, things seemed to have worked without
+# this as well, but we go with the safe option until we understand why.
+
+{
+  lib,
+  runCommand,
+  python3Packages,
+  allowedPatterns,
+}:
+runCommand "allowed-patterns.json"
+  {
+    nativeBuildInputs = [ python3Packages.python ];
+    exportReferencesGraph = builtins.concatMap (
+      name:
+      builtins.concatMap (
+        path:
+        let
+          prefix = "${builtins.storeDir}/";
+          # Has to start with a letter: https://github.com/NixOS/nix/blob/516e7ddc41f39ff939b5d5b5dc71e590f24890d4/src/libstore/build/local-derivation-goal.cc#L568
+          exportName = ''references-${lib.strings.removePrefix prefix "${path}"}'';
+          isStorePath = lib.isStorePath path && (lib.hasPrefix prefix "${path}");
+        in
+        lib.optionals isStorePath [
+          exportName
+          path
+        ]
+      ) allowedPatterns.${name}.paths
+    ) (builtins.attrNames allowedPatterns);
+    env.storeDir = "${builtins.storeDir}/";
+    shallowConfig = builtins.toJSON allowedPatterns;
+    passAsFile = [ "shallowConfig" ];
+  }
+  ''
+    python ${./scripts/nix_required_mounts_closure.py}
+  ''