From b422dafc896070297e0d5a038a52f58bc3c11eea Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 19 Sep 2023 12:28:30 +0300
Subject: [PATCH 01/28] nix-required-mounts: init

---
 nixos/modules/module-list.nix                 |   1 +
 .../modules/programs/nix-required-mounts.nix  |  85 +++++++++++++
 pkgs/by-name/ni/nix-required-mounts/main.py   | 118 ++++++++++++++++++
 .../ni/nix-required-mounts/package.nix        |  57 +++++++++
 .../ni/nix-required-mounts/pyproject.toml     |   2 +
 5 files changed, 263 insertions(+)
 create mode 100644 nixos/modules/programs/nix-required-mounts.nix
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/main.py
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/package.nix
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/pyproject.toml

diff --git a/nixos/modules/module-list.nix b/nixos/modules/module-list.nix
index a008c3c5bdea..b7747743f5a6 100644
--- a/nixos/modules/module-list.nix
+++ b/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
diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
new file mode 100644
index 000000000000..72554dbc6cc5
--- /dev/null
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -0,0 +1,85 @@
+{ config, lib, pkgs, ... }:
+
+let
+  cfg = config.programs.nix-required-mounts;
+  hook =
+    pkgs.nix-required-mounts.override { inherit (cfg) allowedPatterns; };
+
+  patternType = with lib.types; submodule ({ config, name, ... }: {
+    options.onFeatures = lib.mkOption {
+      type = listOf str;
+      description = "Which requiredSystemFeatures should trigger relaxation of the sandbox";
+      default = [ name ];
+    };
+    options.paths = lib.mkOption {
+      type = listOf path;
+      description = "A list of glob patterns, indicating which paths to expose to the sandbox";
+    };
+  });
+
+  defaults = {
+    opengl.onFeatures = [ "opengl" ];
+    opengl.paths = [
+      "/dev/video*"
+      "/dev/dri"
+
+      pkgs.addOpenGLRunpath.driverLink
+      # /run/opengl-driver/lib only contains symlinks
+      config.hardware.opengl.package
+    ] ++ config.hardware.opengl.extraPackages;
+    cuda.onFeatures = [ "cuda" ];
+    cuda.paths = defaults.opengl.paths ++ [ "/dev/nvidia*" ];
+  };
+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.opengl.enable = lib.mkOption {
+      type = lib.types.bool;
+      default = config.hardware.opengl.enable;
+      defaultText = lib.literalExpression "hardware.opengl.enable";
+      description = ''
+        Expose OpenGL drivers to derivations marked with requiredSystemFeatures = [ "opengl" ]
+      '';
+    };
+    presets.cuda.enable = lib.mkEnableOption ''
+      Expose CUDA drivers and GPUs to derivations marked with requiredSystemFeatures = [ "cuda" ]
+    '';
+    allowedPatterns = with lib.types;
+      lib.mkOption rec {
+        type = attrsOf patternType;
+        description = "The hook config, describing which paths to mount for which system features";
+        default = { inherit (defaults) opengl; };
+        defaultText = lib.literalExpression ''
+          {
+            opengl.paths = config.hardware.opengl.extraPackages ++ [
+              config.hardware.opengl.package
+              pkgs.addOpenGLRunpath.driverLink
+              "/dev/video*"
+              "/dev/dri"
+            ];
+          }
+        '';
+        example.require-ipfs = [ "/ipfs" ];
+      };
+  };
+  config = lib.mkIf cfg.enable (lib.mkMerge [
+    {
+      nix.settings.pre-build-hook = lib.getExe hook;
+    }
+    (lib.mkIf cfg.presets.opengl.enable {
+      nix.settings.system-features = [ "opengl" ];
+      programs.nix-required-mounts.allowedPatterns = {
+        inherit (defaults) opengl;
+      };
+    })
+    (lib.mkIf cfg.presets.cuda.enable {
+      nix.settings.system-features = [ "cuda" ];
+      programs.nix-required-mounts.allowedPatterns = {
+        inherit (defaults) cuda;
+      };
+    })
+  ]);
+}
diff --git a/pkgs/by-name/ni/nix-required-mounts/main.py b/pkgs/by-name/ni/nix-required-mounts/main.py
new file mode 100644
index 000000000000..50b670785054
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/main.py
@@ -0,0 +1,118 @@
+import glob
+import json
+import subprocess
+import textwrap
+from argparse import ArgumentParser
+from itertools import chain
+from pathlib import Path
+from sys import stderr
+from typing import List
+
+CONFIG = {
+    "nixExe": "nix",
+    "allowedPatterns": {},
+}
+
+parser = ArgumentParser("pre-build-hook")
+parser.add_argument("derivation_path")
+parser.add_argument("sandbox_path", nargs="?")
+parser.add_argument(
+    "--issue-command",
+    choices=("always", "conditional", "never"),
+    default="conditional",
+    help="Whether to print extra-sandbox-paths",
+)
+parser.add_argument(
+    "--issue-stop",
+    choices=("always", "conditional", "never"),
+    default="conditional",
+    help="Whether to print the final empty line",
+)
+
+
+def symlink_parents(p: Path) -> List[Path]:
+    out = []
+    while p.is_symlink() and p not in out:
+        p = p.readlink()
+        out.append(p)
+    return out
+
+
+def entrypoint():
+    if __name__ != "__main__":
+        return
+
+    args = parser.parse_args()
+    drv_path = args.derivation_path
+
+    if not Path(drv_path).exists():
+        print(
+            f"[E] {drv_path} doesn't exist."
+            " This may happen with the remote builds."
+            " Exiting the hook",
+            file=stderr,
+        )
+
+    proc = subprocess.run(
+        [
+            CONFIG["nixExe"],
+            "show-derivation",
+            drv_path,
+        ],
+        capture_output=True,
+    )
+    try:
+        drv = json.loads(proc.stdout)
+    except json.JSONDecodeError:
+        print(
+            "[E] Couldn't parse the output of"
+            "`nix show-derivation`"
+            f". Expected JSON, observed: {proc.stdout}",
+            file=stderr,
+        )
+        print(
+            textwrap.indent(proc.stdout.decode("utf8"), prefix=" " * 4),
+            file=stderr,
+        )
+        print("[I] Exiting the nix-required-binds hook", file=stderr)
+        return
+    [canon_drv_path] = drv.keys()
+
+    allowed_patterns = CONFIG["allowedPatterns"]
+    known_features = set(
+        chain.from_iterable(
+            pattern["onFeatures"] for pattern in allowed_patterns.values()
+        )
+    )
+
+    drv_env = drv[canon_drv_path].get("env", {})
+    features = drv_env.get("requiredSystemFeatures", [])
+    if isinstance(features, str):
+        features = features.split()
+
+    features = list(filter(known_features.__contains__, features))
+
+    patterns = list(chain.from_iterable(allowed_patterns[f]["paths"] for f in features))  # noqa: E501
+
+    roots = sorted(
+        set(Path(path) for pattern in patterns for path in glob.glob(pattern))
+    )
+
+    # the pre-build-hook command
+    if args.issue_command == "always" or (
+        args.issue_command == "conditional" and roots
+    ):
+        print("extra-sandbox-paths")
+
+    # arguments, one per line
+    for p in roots:
+        guest_path, host_path = p, p
+        print(f"{guest_path}={host_path}")
+
+    # terminated by an empty line
+    something_to_terminate = args.issue_stop == "conditional" and roots
+    if args.issue_stop == "always" or something_to_terminate:
+        print()
+
+
+entrypoint()
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
new file mode 100644
index 000000000000..cbac41addbc4
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -0,0 +1,57 @@
+{ addOpenGLRunpath
+, cmake
+, allowedPatterns ? rec {
+    opengl.onFeatures = [ "opengl" ];
+    opengl.paths = [
+      addOpenGLRunpath.driverLink
+      "/dev/video*"
+      "/dev/dri"
+    ];
+    cuda.onFeatures = [ "cuda" ];
+    cuda.paths = opengl.paths ++ [
+      "/dev/nvidia*"
+    ];
+  }
+, buildPackages
+, formats
+, lib
+, nix
+, python3
+, runCommand
+}:
+
+
+let
+  confPath = (formats.pythonVars { }).generate "config.py" {
+    CONFIG = {
+      inherit allowedPatterns;
+      nixExe = lib.getExe nix;
+    };
+  };
+  pname = "nix-required-mounts";
+in
+
+runCommand pname
+{
+  inherit confPath;
+  meta.mainProgram = pname;
+} ''
+  ${lib.getExe buildPackages.python3.pkgs.flake8} ${./main.py}
+
+  cat > main.py << EOF
+  #!${lib.getExe python3}
+
+  $(cat ${./main.py})
+  EOF
+
+  sed -ie '
+    /^entrypoint()$/ {
+      x ;
+      r ${confPath}
+      }' main.py
+
+  echo "entrypoint()" >> main.py
+
+  mkdir -p $out/bin
+  install main.py $out/bin/${pname}
+''
diff --git a/pkgs/by-name/ni/nix-required-mounts/pyproject.toml b/pkgs/by-name/ni/nix-required-mounts/pyproject.toml
new file mode 100644
index 000000000000..a8f43fefdf14
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/pyproject.toml
@@ -0,0 +1,2 @@
+[tool.black]
+line-length = 79

From 340b41815df042fd5e8153b3f8c7a495bbdaba5c Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 17 Oct 2023 15:25:17 +0300
Subject: [PATCH 02/28] nixosTests.nix-required-mounts: init

---
 nixos/tests/all-tests.nix                     |  1 +
 nixos/tests/nix-required-mounts/default.nix   | 44 +++++++++++++++++++
 .../ensure-path-not-present.nix               | 13 ++++++
 .../test-require-feature.nix                  | 12 +++++
 4 files changed, 70 insertions(+)
 create mode 100644 nixos/tests/nix-required-mounts/default.nix
 create mode 100644 nixos/tests/nix-required-mounts/ensure-path-not-present.nix
 create mode 100644 nixos/tests/nix-required-mounts/test-require-feature.nix

diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index 6b5ee429f9a2..a33247ba7dbc 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -649,6 +649,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 {};
diff --git a/nixos/tests/nix-required-mounts/default.nix b/nixos/tests/nix-required-mounts/default.nix
new file mode 100644
index 000000000000..ee6f7db5ee98
--- /dev/null
+++ b/nixos/tests/nix-required-mounts/default.nix
@@ -0,0 +1,44 @@
+{ 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" ];
+    };
+    users.users.person.isNormalUser = true;
+    virtualisation.fileSystems."/supported-feature-files".fsType = "tmpfs";
+  };
+  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)
+  '';
+}
diff --git a/nixos/tests/nix-required-mounts/ensure-path-not-present.nix b/nixos/tests/nix-required-mounts/ensure-path-not-present.nix
new file mode 100644
index 000000000000..871f336ee9bd
--- /dev/null
+++ b/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
+''
+
diff --git a/nixos/tests/nix-required-mounts/test-require-feature.nix b/nixos/tests/nix-required-mounts/test-require-feature.nix
new file mode 100644
index 000000000000..ddfd068b87fc
--- /dev/null
+++ b/nixos/tests/nix-required-mounts/test-require-feature.nix
@@ -0,0 +1,12 @@
+{ pkgs ? import <nixpkgs> { }, feature }:
+
+pkgs.runCommandNoCC "${feature}-present"
+{
+  requiredSystemFeatures = [ feature ];
+} ''
+  if [[ -e /${feature}-files ]]; then
+    touch $out
+  else
+    echo "The host declares ${feature} support, but doesn't expose /${feature}-files" >&2
+  fi
+''

From b9299696abf06a56a5c00f24b343ded3fbed7932 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 16 Oct 2023 21:59:59 +0300
Subject: [PATCH 03/28] blender: add CUDA discovery test

---
 pkgs/applications/misc/blender/default.nix  | 10 ++++++++++
 pkgs/applications/misc/blender/test-cuda.py |  8 ++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 pkgs/applications/misc/blender/test-cuda.py

diff --git a/pkgs/applications/misc/blender/default.nix b/pkgs/applications/misc/blender/default.nix
index e54cae9e56a2..357bbbe0e6f6 100644
--- a/pkgs/applications/misc/blender/default.nix
+++ b/pkgs/applications/misc/blender/default.nix
@@ -372,6 +372,16 @@ stdenv.mkDerivation (finalAttrs: {
             --render-frame 1
         done
       '';
+
+      cudaAvailable = runCommand
+      "blender-cuda-available"
+      {
+        nativeBuildInputs = [ finalAttrs.finalPackage ];
+        requiredSystemFeatures = [ "cuda" ];
+      }
+      ''
+        blender --background -noaudio --python-exit-code 1 --python ${./test-cuda.py} && touch $out
+      '';
     };
   };
 
diff --git a/pkgs/applications/misc/blender/test-cuda.py b/pkgs/applications/misc/blender/test-cuda.py
new file mode 100644
index 000000000000..8a3ec5734759
--- /dev/null
+++ b/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")

From 7ed2cba5e8882c93516f71e3de2bf8a828ccdb06 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Wed, 20 Sep 2023 07:02:35 +0300
Subject: [PATCH 04/28] python3Packages.pynvml: add a gpu test

---
 .../python-modules/pynvml/default.nix         |  4 ++++
 .../python-modules/pynvml/test-gpu.nix        | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 pkgs/development/python-modules/pynvml/test-gpu.nix

diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
index a115cd723998..e76913a63df4 100644
--- a/pkgs/development/python-modules/pynvml/default.nix
+++ b/pkgs/development/python-modules/pynvml/default.nix
@@ -1,6 +1,7 @@
 {
   lib,
   buildPythonPackage,
+  callPackage,
   fetchFromGitHub,
   substituteAll,
   pythonOlder,
@@ -50,6 +51,9 @@ buildPythonPackage rec {
   # OSError: /run/opengl-driver/lib/libnvidia-ml.so.1: cannot open shared object file: No such file or directory
   doCheck = false;
 
+  passthru.tests.nvmlInit = callPackage ./test-gpu.nix { };
+
+
   meta = with lib; {
     description = "Python bindings for the NVIDIA Management Library";
     homepage = "https://github.com/gpuopenanalytics/pynvml";
diff --git a/pkgs/development/python-modules/pynvml/test-gpu.nix b/pkgs/development/python-modules/pynvml/test-gpu.nix
new file mode 100644
index 000000000000..c316d0b5094b
--- /dev/null
+++ b/pkgs/development/python-modules/pynvml/test-gpu.nix
@@ -0,0 +1,23 @@
+{ 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
+''
+

From f22b9da6b8e66b9b325cff5bab5f22b818f099b5 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Wed, 20 Sep 2023 07:14:07 +0300
Subject: [PATCH 05/28] python3Packages.torch{,-bin}: test
 torch.cuda.is_available()

---
 pkgs/development/python-modules/torch/bin.nix |  4 ++++
 .../python-modules/torch/default.nix          | 11 +++++++++-
 .../python-modules/torch/test-cuda.nix        | 21 +++++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 pkgs/development/python-modules/torch/test-cuda.nix

diff --git a/pkgs/development/python-modules/torch/bin.nix b/pkgs/development/python-modules/torch/bin.nix
index f9d5cd97c183..37170ea9adf4 100644
--- a/pkgs/development/python-modules/torch/bin.nix
+++ b/pkgs/development/python-modules/torch/bin.nix
@@ -8,6 +8,7 @@
   pythonAtLeast,
   pythonOlder,
   addOpenGLRunpath,
+  callPackage,
   cudaPackages,
   future,
   numpy,
@@ -15,6 +16,7 @@
   pyyaml,
   requests,
   setuptools,
+  torch-bin,
   typing-extensions,
   sympy,
   jinja2,
@@ -119,6 +121,8 @@ buildPythonPackage {
 
   pythonImportsCheck = [ "torch" ];
 
+  passthru.tests.cudaAvailable = callPackage ./test-cuda.nix { torch = torch-bin; };
+
   meta = {
     description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration";
     homepage = "https://pytorch.org/";
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index d5d7e823bed7..671e76dfe02d 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -24,6 +24,10 @@
   mpi,
   buildDocs ? false,
 
+  # tests.cudaAvailable:
+  callPackage,
+  torch,
+
   # Native build inputs
   cmake,
   symlinkJoin,
@@ -639,11 +643,16 @@ buildPythonPackage rec {
       rocmSupport
       rocmPackages
       ;
+    cudaCapabilities = if cudaSupport then supportedCudaCapabilities else [ ];
     # At least for 1.10.2 `torch.fft` is unavailable unless BLAS provider is MKL. This attribute allows for easy detection of its availability.
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-    cudaCapabilities = if cudaSupport then supportedCudaCapabilities else [ ];
+  } // lib.optionalAttrs cudaSupport {
+
+    tests = lib.optionalAttrs cudaSupport {
+      cudaAvailable = callPackage ./test-cuda.nix { inherit torch; };
+    };
   };
 
   meta = {
diff --git a/pkgs/development/python-modules/torch/test-cuda.nix b/pkgs/development/python-modules/torch/test-cuda.nix
new file mode 100644
index 000000000000..22d73eeb9167
--- /dev/null
+++ b/pkgs/development/python-modules/torch/test-cuda.nix
@@ -0,0 +1,21 @@
+{ runCommandNoCC
+, python
+, torch
+}:
+
+runCommandNoCC "${torch.name}-gpu-test"
+{
+  nativeBuildInputs = [
+    (python.withPackages (_: [ torch ]))
+  ];
+  requiredSystemFeatures = [
+    "cuda"
+  ];
+} ''
+  python3 << EOF
+  import torch
+  assert torch.cuda.is_available(), f"{torch.cuda.is_available()=}"
+  EOF
+
+  touch $out
+''

From 6859a2dabc356bc0d575e98c46a215fd83d31fb7 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Wed, 18 Oct 2023 12:44:59 +0300
Subject: [PATCH 06/28] nix-required-mounts: use wrappers instead of statically
 embedding config into the script

---
 pkgs/by-name/ni/nix-required-mounts/main.py   | 36 ++++++++++++-------
 .../ni/nix-required-mounts/package.nix        | 34 +++++++-----------
 2 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/main.py b/pkgs/by-name/ni/nix-required-mounts/main.py
index 50b670785054..e8cbc70d4c52 100644
--- a/pkgs/by-name/ni/nix-required-mounts/main.py
+++ b/pkgs/by-name/ni/nix-required-mounts/main.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 import glob
 import json
 import subprocess
@@ -6,16 +8,23 @@ from argparse import ArgumentParser
 from itertools import chain
 from pathlib import Path
 from sys import stderr
-from typing import List
+from typing import Dict, List, TypedDict
+
+
+class Pattern(TypedDict):
+    onFeatures: List[str]
+    paths: List[str]  # List of glob patterns
+
+
+class HookConfig(TypedDict):
+    nixExe: str
+    allowedPatterns: Dict[str, Pattern]
 
-CONFIG = {
-    "nixExe": "nix",
-    "allowedPatterns": {},
-}
 
 parser = ArgumentParser("pre-build-hook")
 parser.add_argument("derivation_path")
 parser.add_argument("sandbox_path", nargs="?")
+parser.add_argument("--config", type=Path)
 parser.add_argument(
     "--issue-command",
     choices=("always", "conditional", "never"),
@@ -39,12 +48,12 @@ def symlink_parents(p: Path) -> List[Path]:
 
 
 def entrypoint():
-    if __name__ != "__main__":
-        return
-
     args = parser.parse_args()
     drv_path = args.derivation_path
 
+    with open(args.config, "r") as f:
+        config = json.load(f)
+
     if not Path(drv_path).exists():
         print(
             f"[E] {drv_path} doesn't exist."
@@ -55,7 +64,7 @@ def entrypoint():
 
     proc = subprocess.run(
         [
-            CONFIG["nixExe"],
+            config["nixExe"],
             "show-derivation",
             drv_path,
         ],
@@ -78,7 +87,7 @@ def entrypoint():
         return
     [canon_drv_path] = drv.keys()
 
-    allowed_patterns = CONFIG["allowedPatterns"]
+    allowed_patterns = config["allowedPatterns"]
     known_features = set(
         chain.from_iterable(
             pattern["onFeatures"] for pattern in allowed_patterns.values()
@@ -92,7 +101,9 @@ def entrypoint():
 
     features = list(filter(known_features.__contains__, features))
 
-    patterns = list(chain.from_iterable(allowed_patterns[f]["paths"] for f in features))  # noqa: E501
+    patterns = list(
+        chain.from_iterable(allowed_patterns[f]["paths"] for f in features)
+    )  # noqa: E501
 
     roots = sorted(
         set(Path(path) for pattern in patterns for path in glob.glob(pattern))
@@ -115,4 +126,5 @@ def entrypoint():
         print()
 
 
-entrypoint()
+if __name__ == "__main__":
+    entrypoint()
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index cbac41addbc4..c8cae067b24d 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -16,42 +16,32 @@
 , formats
 , lib
 , nix
-, python3
+, python3Packages
+, makeWrapper
 , runCommand
 }:
 
 
 let
-  confPath = (formats.pythonVars { }).generate "config.py" {
-    CONFIG = {
-      inherit allowedPatterns;
-      nixExe = lib.getExe nix;
-    };
+  confPath = (formats.json { }).generate "config.py" {
+    inherit allowedPatterns;
+    nixExe = lib.getExe nix;
   };
   pname = "nix-required-mounts";
 in
 
 runCommand pname
 {
-  inherit confPath;
+  nativeBuildInputs = [
+    makeWrapper
+    python3Packages.wrapPython
+  ];
   meta.mainProgram = pname;
 } ''
   ${lib.getExe buildPackages.python3.pkgs.flake8} ${./main.py}
 
-  cat > main.py << EOF
-  #!${lib.getExe python3}
-
-  $(cat ${./main.py})
-  EOF
-
-  sed -ie '
-    /^entrypoint()$/ {
-      x ;
-      r ${confPath}
-      }' main.py
-
-  echo "entrypoint()" >> main.py
-
   mkdir -p $out/bin
-  install main.py $out/bin/${pname}
+  install ${./main.py} $out/bin/${pname}
+  wrapProgram $out/bin/${pname} --add-flags "--config ${confPath}"
+  wrapPythonPrograms
 ''

From 6662b099419d568a7d682e5b851f7f51b34cd589 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Thu, 19 Oct 2023 11:42:26 +0300
Subject: [PATCH 07/28] nix-required-mounts: handle __structuredAttrs

---
 nixos/tests/nix-required-mounts/default.nix   |  5 ++++-
 .../test-structured-attrs-empty.nix           | 10 ++++++++++
 .../test-structured-attrs.nix                 | 15 +++++++++++++++
 pkgs/by-name/ni/nix-required-mounts/main.py   | 19 ++++++++++++-------
 4 files changed, 41 insertions(+), 8 deletions(-)
 create mode 100644 nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix
 create mode 100644 nixos/tests/nix-required-mounts/test-structured-attrs.nix

diff --git a/nixos/tests/nix-required-mounts/default.nix b/nixos/tests/nix-required-mounts/default.nix
index ee6f7db5ee98..4550e6ac50a2 100644
--- a/nixos/tests/nix-required-mounts/default.nix
+++ b/nixos/tests/nix-required-mounts/default.nix
@@ -14,7 +14,8 @@ in
     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" ];
+    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" ];
@@ -40,5 +41,7 @@ in
     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}")
   '';
 }
diff --git a/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix b/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix
new file mode 100644
index 000000000000..d788c6773c8e
--- /dev/null
+++ b/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix
@@ -0,0 +1,10 @@
+{ pkgs ? import <nixpkgs> { } }:
+
+pkgs.runCommandNoCC "nix-required-mounts-structured-attrs-no-features"
+{
+  __structuredAttrs = true;
+} ''
+  touch $out
+''
+
+
diff --git a/nixos/tests/nix-required-mounts/test-structured-attrs.nix b/nixos/tests/nix-required-mounts/test-structured-attrs.nix
new file mode 100644
index 000000000000..fecd2c32eec0
--- /dev/null
+++ b/nixos/tests/nix-required-mounts/test-structured-attrs.nix
@@ -0,0 +1,15 @@
+{ 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
+''
+
diff --git a/pkgs/by-name/ni/nix-required-mounts/main.py b/pkgs/by-name/ni/nix-required-mounts/main.py
index e8cbc70d4c52..8e32483e8a4b 100644
--- a/pkgs/by-name/ni/nix-required-mounts/main.py
+++ b/pkgs/by-name/ni/nix-required-mounts/main.py
@@ -47,6 +47,13 @@ def symlink_parents(p: Path) -> List[Path]:
     return out
 
 
+def get_strings(drv_env: dict, name: str) -> List[str]:
+    if "__json" in drv_env:
+        return list(json.loads(drv_env["__json"]).get(name, []))
+    else:
+        return drv_env.get(name, "").split()
+
+
 def entrypoint():
     args = parser.parse_args()
     drv_path = args.derivation_path
@@ -71,7 +78,7 @@ def entrypoint():
         capture_output=True,
     )
     try:
-        drv = json.loads(proc.stdout)
+        parsed_drv = json.loads(proc.stdout)
     except json.JSONDecodeError:
         print(
             "[E] Couldn't parse the output of"
@@ -85,7 +92,7 @@ def entrypoint():
         )
         print("[I] Exiting the nix-required-binds hook", file=stderr)
         return
-    [canon_drv_path] = drv.keys()
+    [canon_drv_path] = parsed_drv.keys()
 
     allowed_patterns = config["allowedPatterns"]
     known_features = set(
@@ -94,11 +101,9 @@ def entrypoint():
         )
     )
 
-    drv_env = drv[canon_drv_path].get("env", {})
-    features = drv_env.get("requiredSystemFeatures", [])
-    if isinstance(features, str):
-        features = features.split()
-
+    parsed_drv = parsed_drv[canon_drv_path]
+    drv_env = parsed_drv.get("env", {})
+    features = get_strings(drv_env, "requiredSystemFeatures")
     features = list(filter(known_features.__contains__, features))
 
     patterns = list(

From 50d4382114e51386fda72c1983ed5c2c62c53ade Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 6 Nov 2023 16:44:14 +0000
Subject: [PATCH 08/28] programs.nix-required-mounts: inherit defaults from the
 package

---
 .../modules/programs/nix-required-mounts.nix  | 66 ++++++++++---------
 .../ni/nix-required-mounts/package.nix        |  3 +
 2 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index 72554dbc6cc5..c2a81ff03bff 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -2,33 +2,39 @@
 
 let
   cfg = config.programs.nix-required-mounts;
-  hook =
-    pkgs.nix-required-mounts.override { inherit (cfg) allowedPatterns; };
+  package = pkgs.nix-required-mounts;
+  overridenPackage = package.override { inherit (cfg) allowedPatterns; };
 
-  patternType = with lib.types; submodule ({ config, name, ... }: {
-    options.onFeatures = lib.mkOption {
-      type = listOf str;
-      description = "Which requiredSystemFeatures should trigger relaxation of the sandbox";
-      default = [ name ];
-    };
-    options.paths = lib.mkOption {
-      type = listOf path;
-      description = "A list of glob patterns, indicating which paths to expose to the sandbox";
-    };
-  });
+  Pattern = with lib.types;
+    submodule ({ config, name, ... }: {
+      options.onFeatures = lib.mkOption {
+        type = listOf str;
+        description =
+          "Which requiredSystemFeatures should trigger relaxation of the sandbox";
+        default = [ name ];
+      };
+      options.paths = lib.mkOption {
+        type = listOf path;
+        description =
+          "A list of glob patterns, indicating which paths to expose to the sandbox";
+      };
+    });
+
+  driverPaths = [
+    # symlinks in /run/opengl-driver/lib:
+    pkgs.addOpenGLRunpath.driverLink
+
+    # mesa:
+    config.hardware.opengl.package
+
+    # nvidia_x11, etc:
+  ] ++ config.hardware.opengl.extraPackages; # nvidia_x11
 
   defaults = {
-    opengl.onFeatures = [ "opengl" ];
-    opengl.paths = [
-      "/dev/video*"
-      "/dev/dri"
-
-      pkgs.addOpenGLRunpath.driverLink
-      # /run/opengl-driver/lib only contains symlinks
-      config.hardware.opengl.package
-    ] ++ config.hardware.opengl.extraPackages;
-    cuda.onFeatures = [ "cuda" ];
-    cuda.paths = defaults.opengl.paths ++ [ "/dev/nvidia*" ];
+    opengl.onFeatures = package.allowedPatterns.opengl.onFeatures;
+    opengl.paths = package.allowedPatterns.opengl.paths ++ driverPaths;
+    cuda.onFeatures = package.allowedPatterns.cuda.onFeatures;
+    cuda.paths = package.allowedPatterns.cuda.paths ++ driverPaths;
   };
 in
 {
@@ -49,8 +55,9 @@ in
     '';
     allowedPatterns = with lib.types;
       lib.mkOption rec {
-        type = attrsOf patternType;
-        description = "The hook config, describing which paths to mount for which system features";
+        type = attrsOf Pattern;
+        description =
+          "The hook config, describing which paths to mount for which system features";
         default = { inherit (defaults) opengl; };
         defaultText = lib.literalExpression ''
           {
@@ -62,13 +69,12 @@ in
             ];
           }
         '';
-        example.require-ipfs = [ "/ipfs" ];
+        example.require-ipfs.paths = [ "/ipfs" ];
+        example.require-ipfs.onFeatures = [ "ifps" ];
       };
   };
   config = lib.mkIf cfg.enable (lib.mkMerge [
-    {
-      nix.settings.pre-build-hook = lib.getExe hook;
-    }
+    { nix.settings.pre-build-hook = lib.getExe overridenPackage; }
     (lib.mkIf cfg.presets.opengl.enable {
       nix.settings.system-features = [ "opengl" ];
       programs.nix-required-mounts.allowedPatterns = {
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index c8cae067b24d..58db69487a03 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -36,6 +36,9 @@ runCommand pname
     makeWrapper
     python3Packages.wrapPython
   ];
+  passthru = {
+    inherit allowedPatterns;
+  };
   meta.mainProgram = pname;
 } ''
   ${lib.getExe buildPackages.python3.pkgs.flake8} ${./main.py}

From 3d84ab0b09773f0a4edef19897ce81e46508fac9 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 6 Nov 2023 16:44:36 +0000
Subject: [PATCH 09/28] nix-required-mounts: expose the VM test in passthru

---
 pkgs/by-name/ni/nix-required-mounts/package.nix | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index 58db69487a03..dbb35356b7af 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -16,6 +16,7 @@
 , formats
 , lib
 , nix
+, nixosTests
 , python3Packages
 , makeWrapper
 , runCommand
@@ -38,6 +39,9 @@ runCommand pname
   ];
   passthru = {
     inherit allowedPatterns;
+    tests = {
+      inherit (nixosTests) nix-required-mounts;
+    };
   };
   meta.mainProgram = pname;
 } ''

From 3a0d777486191cbdd2ffc7e31f25156da94c4831 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 6 Nov 2023 16:49:54 +0000
Subject: [PATCH 10/28] nix-required-mounts: link the issue about unavailable
 .drvs

---
 pkgs/by-name/ni/nix-required-mounts/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/main.py b/pkgs/by-name/ni/nix-required-mounts/main.py
index 8e32483e8a4b..1d263dcd1028 100644
--- a/pkgs/by-name/ni/nix-required-mounts/main.py
+++ b/pkgs/by-name/ni/nix-required-mounts/main.py
@@ -64,7 +64,7 @@ def entrypoint():
     if not Path(drv_path).exists():
         print(
             f"[E] {drv_path} doesn't exist."
-            " This may happen with the remote builds."
+            " Cf. https://github.com/NixOS/nix/issues/9272"
             " Exiting the hook",
             file=stderr,
         )

From 6a0f2aedc14d38cc2ad184f6c9a448729de8f8a5 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 6 Nov 2023 17:21:47 +0000
Subject: [PATCH 11/28] nix-required-mounts: fix: add missing metadata

---
 .../{main.py => nix_required_mounts.py}       |  0
 .../ni/nix-required-mounts/package.nix        | 35 ++++++++++++-------
 .../ni/nix-required-mounts/pyproject.toml     | 18 ++++++++++
 3 files changed, 41 insertions(+), 12 deletions(-)
 rename pkgs/by-name/ni/nix-required-mounts/{main.py => nix_required_mounts.py} (100%)

diff --git a/pkgs/by-name/ni/nix-required-mounts/main.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
similarity index 100%
rename from pkgs/by-name/ni/nix-required-mounts/main.py
rename to pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index dbb35356b7af..c1411a6d2a80 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -28,27 +28,38 @@ let
     inherit allowedPatterns;
     nixExe = lib.getExe nix;
   };
-  pname = "nix-required-mounts";
+  attrs = builtins.fromTOML (builtins.readFile ./pyproject.toml);
+  pname = attrs.project.name;
+  inherit (attrs.project) version;
 in
 
-runCommand pname
+python3Packages.buildPythonApplication
 {
+  inherit pname version;
+  pyproject = true;
+
+  src = lib.cleanSource ./.;
+
   nativeBuildInputs = [
     makeWrapper
-    python3Packages.wrapPython
+    python3Packages.setuptools
   ];
+
+  postFixup = ''
+    wrapProgram $out/bin/${pname} --add-flags "--config ${confPath}"
+  '';
+
   passthru = {
     inherit allowedPatterns;
     tests = {
       inherit (nixosTests) nix-required-mounts;
     };
   };
-  meta.mainProgram = pname;
-} ''
-  ${lib.getExe buildPackages.python3.pkgs.flake8} ${./main.py}
-
-  mkdir -p $out/bin
-  install ${./main.py} $out/bin/${pname}
-  wrapProgram $out/bin/${pname} --add-flags "--config ${confPath}"
-  wrapPythonPrograms
-''
+  meta = {
+    inherit (attrs.project) description;
+    homepage = attrs.project.urls.Homepage;
+    license = lib.licenses.mit;
+    mainProgram = attrs.project.name;
+    maintainers = with lib.maintainers; [ SomeoneSerge ];
+  };
+}
diff --git a/pkgs/by-name/ni/nix-required-mounts/pyproject.toml b/pkgs/by-name/ni/nix-required-mounts/pyproject.toml
index a8f43fefdf14..bb754e08ab1d 100644
--- a/pkgs/by-name/ni/nix-required-mounts/pyproject.toml
+++ b/pkgs/by-name/ni/nix-required-mounts/pyproject.toml
@@ -1,2 +1,20 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = [ "setuptools" ]
+
+[project]
+name = "nix-required-mounts"
+version = "0.0.1"
+description = """
+A --pre-build-hook for Nix, \
+that allows to expose extra paths in the build sandbox \
+based on derivations' requiredSystemFeatrues"""
+
+[project.urls]
+Homepage = "https://github.com/NixOS/nixpkgs/tree/master/pkgs/by-name/ni/nix-required-mounts"
+
+[project.scripts]
+nix-required-mounts = "nix_required_mounts:entrypoint"
+
 [tool.black]
 line-length = 79

From 7418e4fefd622eb1d1709353aa842e136896d70f Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 11 Nov 2023 02:40:44 +0000
Subject: [PATCH 12/28] programs.nix-required-mounts: presets.cuda ->
 nvidia-gpu

This hopefully clarifies that the preset configures the hook to expose
"nvidia devices", which includse both the userspace driver and the
device nodes.

The derivations  still declare requiredSystemFeatures = [ "cuda" ] to
explicitly indicate they need to use the CUDA functionality and expect a
libcuda.so and a CUDA-capable device. Ideally, we'd also include the
specific CUDA architectures (sm_86, etc) in feature names.

Derivations that use a co-processor but do not care about the vendor or
even the particular interface may ask for the more generic "opengl",
"vulkan", or "gpu" features. It is then responsibility of the host
declaring the support for this feature to ensure the drivers and
hardware are appropriately set up.
---
 .../modules/programs/nix-required-mounts.nix  | 38 +++++++------------
 .../ni/nix-required-mounts/package.nix        | 15 +++++---
 2 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index c2a81ff03bff..c05c2016a6be 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -31,10 +31,8 @@ let
   ] ++ config.hardware.opengl.extraPackages; # nvidia_x11
 
   defaults = {
-    opengl.onFeatures = package.allowedPatterns.opengl.onFeatures;
-    opengl.paths = package.allowedPatterns.opengl.paths ++ driverPaths;
-    cuda.onFeatures = package.allowedPatterns.cuda.onFeatures;
-    cuda.paths = package.allowedPatterns.cuda.paths ++ driverPaths;
+    nvidia-gpu.onFeatures = package.allowedPatterns.nvidia-gpu.onFeatures;
+    nvidia-gpu.paths = package.allowedPatterns.nvidia-gpu.paths ++ driverPaths;
   };
 in
 {
@@ -42,23 +40,21 @@ in
   options.programs.nix-required-mounts = {
     enable = lib.mkEnableOption
       "Expose extra paths to the sandbox depending on derivations' requiredSystemFeatures";
-    presets.opengl.enable = lib.mkOption {
-      type = lib.types.bool;
-      default = config.hardware.opengl.enable;
-      defaultText = lib.literalExpression "hardware.opengl.enable";
-      description = ''
-        Expose OpenGL drivers to derivations marked with requiredSystemFeatures = [ "opengl" ]
-      '';
-    };
-    presets.cuda.enable = lib.mkEnableOption ''
-      Expose CUDA drivers and GPUs to derivations marked with requiredSystemFeatures = [ "cuda" ]
+    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 = { inherit (defaults) opengl; };
+        default = { };
         defaultText = lib.literalExpression ''
           {
             opengl.paths = config.hardware.opengl.extraPackages ++ [
@@ -75,16 +71,10 @@ in
   };
   config = lib.mkIf cfg.enable (lib.mkMerge [
     { nix.settings.pre-build-hook = lib.getExe overridenPackage; }
-    (lib.mkIf cfg.presets.opengl.enable {
-      nix.settings.system-features = [ "opengl" ];
+    (lib.mkIf cfg.presets.nvidia-gpu.enable {
+      nix.settings.system-features = cfg.allowedPatterns.nvidia-gpu.onFeatures;
       programs.nix-required-mounts.allowedPatterns = {
-        inherit (defaults) opengl;
-      };
-    })
-    (lib.mkIf cfg.presets.cuda.enable {
-      nix.settings.system-features = [ "cuda" ];
-      programs.nix-required-mounts.allowedPatterns = {
-        inherit (defaults) cuda;
+        inherit (defaults) nvidia-gpu;
       };
     })
   ]);
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index c1411a6d2a80..42fa3ff1d90d 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -1,15 +1,18 @@
 { addOpenGLRunpath
 , cmake
 , allowedPatterns ? rec {
-    opengl.onFeatures = [ "opengl" ];
-    opengl.paths = [
+    # This config is just an example.
+    # When the hook observes either of the following requiredSystemFeatures:
+    nvidia-gpu.onFeatures = [ "gpu" "opengl" "vulkan" "cuda" ];
+    # It exposes these paths in the sandbox:
+    nvidia-gpu.paths = [
+      # Note that mounting /run/opengl-driver/lib actually isn't sufficient,
+      # because it's populated with symlinks. One most also mount their
+      # targets, which is what the NixOS module additionaly does.
       addOpenGLRunpath.driverLink
-      "/dev/video*"
       "/dev/dri"
-    ];
-    cuda.onFeatures = [ "cuda" ];
-    cuda.paths = opengl.paths ++ [
       "/dev/nvidia*"
+      "/dev/video*"
     ];
   }
 , buildPackages

From 5560f6a5141d9bcc6c3e0ac2e30c714874c1e076 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 11 Nov 2023 14:02:54 +0000
Subject: [PATCH 13/28] nix-required-mounts: guest and host paths may differ

---
 .../modules/programs/nix-required-mounts.nix  | 13 ++++++--
 nixos/tests/nix-required-mounts/default.nix   | 13 ++++++--
 .../test-require-feature.nix                  | 10 +++++--
 .../nix_required_mounts.py                    | 30 ++++++++++++++-----
 .../ni/nix-required-mounts/package.nix        |  2 +-
 5 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index c05c2016a6be..611c065815e1 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -5,16 +5,23 @@ let
   package = pkgs.nix-required-mounts;
   overridenPackage = package.override { inherit (cfg) allowedPatterns; };
 
+  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;
-    submodule ({ config, name, ... }: {
+    types.submodule ({ config, name, ... }: {
       options.onFeatures = lib.mkOption {
-        type = listOf str;
+        type = listOf types.str;
         description =
           "Which requiredSystemFeatures should trigger relaxation of the sandbox";
         default = [ name ];
       };
       options.paths = lib.mkOption {
-        type = listOf path;
+        type = listOf (oneOf [ path Mount ]);
         description =
           "A list of glob patterns, indicating which paths to expose to the sandbox";
       };
diff --git a/nixos/tests/nix-required-mounts/default.nix b/nixos/tests/nix-required-mounts/default.nix
index 4550e6ac50a2..38f94bf6fd98 100644
--- a/nixos/tests/nix-required-mounts/default.nix
+++ b/nixos/tests/nix-required-mounts/default.nix
@@ -19,10 +19,19 @@ in
     programs.nix-required-mounts.enable = true;
     programs.nix-required-mounts.allowedPatterns.supported-feature = {
       onFeatures = [ "supported-feature" ];
-      paths = [ "/supported-feature-files" ];
+      paths = [
+        "/supported-feature-files"
+        {
+          host = "/usr/lib/imaginary-fhs-drivers";
+          guest = "/run/opengl-driver/lib";
+        }
+      ];
     };
     users.users.person.isNormalUser = true;
-    virtualisation.fileSystems."/supported-feature-files".fsType = "tmpfs";
+    systemd.tmpfiles.rules = [
+      "d /supported-feature-files 0755 person users -"
+      "f /usr/lib/imaginary-fhs-drivers/libcuda.so 0444 root root -"
+    ];
   };
   testScript = ''
     import shlex
diff --git a/nixos/tests/nix-required-mounts/test-require-feature.nix b/nixos/tests/nix-required-mounts/test-require-feature.nix
index ddfd068b87fc..061b59e1628a 100644
--- a/nixos/tests/nix-required-mounts/test-require-feature.nix
+++ b/nixos/tests/nix-required-mounts/test-require-feature.nix
@@ -4,9 +4,13 @@ pkgs.runCommandNoCC "${feature}-present"
 {
   requiredSystemFeatures = [ feature ];
 } ''
-  if [[ -e /${feature}-files ]]; then
-    touch $out
-  else
+  if [[ ! -e /${feature}-files ]]; then
     echo "The host declares ${feature} support, but doesn't expose /${feature}-files" >&2
+    exit 1
   fi
+  if [[ ! -f /run/opengl-driver/lib/libcuda.so ]] ; then
+    echo "The host declares ${feature} support, but it the hook fails to handle the hostPath != guestPath cases" >&2
+    exit 1
+  fi
+  touch $out
 ''
diff --git a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
index 1d263dcd1028..5edf61ff115b 100644
--- a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
+++ b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
@@ -8,12 +8,20 @@ from argparse import ArgumentParser
 from itertools import chain
 from pathlib import Path
 from sys import stderr
-from typing import Dict, List, TypedDict
+from typing import Dict, List, Tuple, TypeAlias, TypedDict
+
+Glob: TypeAlias = str
+PathString: TypeAlias = str
+
+
+class Mount(TypedDict):
+    host: PathString
+    guest: PathString
 
 
 class Pattern(TypedDict):
     onFeatures: List[str]
-    paths: List[str]  # List of glob patterns
+    paths: List[Glob | Mount]
 
 
 class HookConfig(TypedDict):
@@ -106,12 +114,21 @@ def entrypoint():
     features = get_strings(drv_env, "requiredSystemFeatures")
     features = list(filter(known_features.__contains__, features))
 
-    patterns = list(
+    patterns: List[PathString | Mount] = list(
         chain.from_iterable(allowed_patterns[f]["paths"] for f in features)
     )  # noqa: E501
 
-    roots = sorted(
-        set(Path(path) for pattern in patterns for path in glob.glob(pattern))
+    # TODO: Would it make sense to preserve the original order instead?
+    roots: List[Tuple[PathString, PathString]] = sorted(
+        set(
+            mnt
+            for pattern in patterns
+            for mnt in (
+                ((path, path) for path in glob.glob(pattern))
+                if isinstance(pattern, PathString)
+                else [(pattern["guest"], pattern["host"])]
+            )
+        )
     )
 
     # the pre-build-hook command
@@ -121,8 +138,7 @@ def entrypoint():
         print("extra-sandbox-paths")
 
     # arguments, one per line
-    for p in roots:
-        guest_path, host_path = p, p
+    for guest_path, host_path in roots:
         print(f"{guest_path}={host_path}")
 
     # terminated by an empty line
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index 42fa3ff1d90d..a17962183de1 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -3,7 +3,7 @@
 , allowedPatterns ? rec {
     # This config is just an example.
     # When the hook observes either of the following requiredSystemFeatures:
-    nvidia-gpu.onFeatures = [ "gpu" "opengl" "vulkan" "cuda" ];
+    nvidia-gpu.onFeatures = [ "gpu" "nvidia-gpu" "opengl" "cuda" ];
     # It exposes these paths in the sandbox:
     nvidia-gpu.paths = [
       # Note that mounting /run/opengl-driver/lib actually isn't sufficient,

From 3cf5bcfe49c74ba429b87e04cf652833e7e1b260 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 11 Nov 2023 19:43:04 +0000
Subject: [PATCH 14/28] nix-required-mounts: restore the followSymlinks option

This way pkgs.nix-required-mounts is "correct" even before
we override it in the NixOS module
---
 .../modules/programs/nix-required-mounts.nix  | 22 +++++--
 .../nix_required_mounts.py                    | 57 ++++++++++++++-----
 .../ni/nix-required-mounts/package.nix        |  4 +-
 3 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index 611c065815e1..1c79e532a036 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -5,13 +5,17 @@ let
   package = pkgs.nix-required-mounts;
   overridenPackage = package.override { inherit (cfg) allowedPatterns; };
 
-  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";
+  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 {
@@ -25,6 +29,11 @@ let
         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 = [
@@ -40,6 +49,7 @@ let
   defaults = {
     nvidia-gpu.onFeatures = package.allowedPatterns.nvidia-gpu.onFeatures;
     nvidia-gpu.paths = package.allowedPatterns.nvidia-gpu.paths ++ driverPaths;
+    nvidia-gpu.unsafeFollowSymlinks = false;
   };
 in
 {
diff --git a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
index 5edf61ff115b..50f0f80f0bf2 100644
--- a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
+++ b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
@@ -5,10 +5,11 @@ import json
 import subprocess
 import textwrap
 from argparse import ArgumentParser
+from collections import deque
 from itertools import chain
 from pathlib import Path
 from sys import stderr
-from typing import Dict, List, Tuple, TypeAlias, TypedDict
+from typing import Deque, Dict, List, Set, Tuple, TypeAlias, TypedDict
 
 Glob: TypeAlias = str
 PathString: TypeAlias = str
@@ -22,6 +23,7 @@ class Mount(TypedDict):
 class Pattern(TypedDict):
     onFeatures: List[str]
     paths: List[Glob | Mount]
+    unsafeFollowSymlinks: bool
 
 
 class HookConfig(TypedDict):
@@ -50,7 +52,11 @@ parser.add_argument(
 def symlink_parents(p: Path) -> List[Path]:
     out = []
     while p.is_symlink() and p not in out:
-        p = p.readlink()
+        parent = p.readlink()
+        if parent.is_relative_to("."):
+            p = p / parent
+        else:
+            p = parent
         out.append(p)
     return out
 
@@ -111,38 +117,59 @@ def entrypoint():
 
     parsed_drv = parsed_drv[canon_drv_path]
     drv_env = parsed_drv.get("env", {})
-    features = get_strings(drv_env, "requiredSystemFeatures")
-    features = list(filter(known_features.__contains__, features))
+    required_features = get_strings(drv_env, "requiredSystemFeatures")
+    required_features = list(filter(known_features.__contains__, required_features))
 
-    patterns: List[PathString | Mount] = list(
-        chain.from_iterable(allowed_patterns[f]["paths"] for f in features)
+    patterns: List[Tuple[PathString | Mount, bool]] = list(
+        (path, pattern["unsafeFollowSymlinks"])
+        for pattern in allowed_patterns.values()
+        for path in pattern["paths"]
+        if any(feature in required_features for feature in pattern["onFeatures"])
     )  # noqa: E501
 
-    # TODO: Would it make sense to preserve the original order instead?
-    roots: List[Tuple[PathString, PathString]] = sorted(
-        set(
+    queue: Deque[Tuple[PathString, PathString, bool]] = deque(
+        (
             mnt
-            for pattern in patterns
+            for (pattern, follow_symlinks) in patterns
             for mnt in (
-                ((path, path) for path in glob.glob(pattern))
+                ((path, path, follow_symlinks) for path in glob.glob(pattern))
                 if isinstance(pattern, PathString)
-                else [(pattern["guest"], pattern["host"])]
+                else [(pattern["guest"], pattern["host"], follow_symlinks)]
             )
         )
     )
 
+    unique_mounts: Set[Tuple[PathString, PathString]] = set()
+    mounts: List[Tuple[PathString, PathString]] = []
+
+    while queue:
+        guest_path, host_path, follow_symlinks = queue.popleft()
+        if (guest_path, host_path) not in unique_mounts:
+            mounts.append((guest_path, host_path))
+            unique_mounts.add((guest_path, host_path))
+
+        if not follow_symlinks:
+            continue
+
+        for parent in symlink_parents(Path(host_path)):
+            parent_str = parent.absolute().as_posix()
+            queue.append((parent_str, parent_str, follow_symlinks))
+
     # the pre-build-hook command
     if args.issue_command == "always" or (
-        args.issue_command == "conditional" and roots
+        args.issue_command == "conditional" and mounts
     ):
         print("extra-sandbox-paths")
+        print_paths = True
+    else:
+        print_paths = False
 
     # arguments, one per line
-    for guest_path, host_path in roots:
+    for guest_path, host_path in mounts if print_paths else []:
         print(f"{guest_path}={host_path}")
 
     # terminated by an empty line
-    something_to_terminate = args.issue_stop == "conditional" and roots
+    something_to_terminate = args.issue_stop == "conditional" and mounts
     if args.issue_stop == "always" or something_to_terminate:
         print()
 
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index a17962183de1..8e2611552d05 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -6,14 +6,12 @@
     nvidia-gpu.onFeatures = [ "gpu" "nvidia-gpu" "opengl" "cuda" ];
     # It exposes these paths in the sandbox:
     nvidia-gpu.paths = [
-      # Note that mounting /run/opengl-driver/lib actually isn't sufficient,
-      # because it's populated with symlinks. One most also mount their
-      # targets, which is what the NixOS module additionaly does.
       addOpenGLRunpath.driverLink
       "/dev/dri"
       "/dev/nvidia*"
       "/dev/video*"
     ];
+    nvidia-gpu.unsafeFollowSymlinks = true;
   }
 , buildPackages
 , formats

From 55f54cc2c33e52ae8ac57bced83046b7d9dbc6cd Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 11 Nov 2023 20:52:31 +0000
Subject: [PATCH 15/28] nix-required-mounts: restore (optional) symlink support

---
 nixos/tests/nix-required-mounts/default.nix   |  4 ++-
 .../test-require-feature.nix                  | 12 ++++++++-
 .../nix_required_mounts.py                    | 25 ++++++++++++-------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/nixos/tests/nix-required-mounts/default.nix b/nixos/tests/nix-required-mounts/default.nix
index 38f94bf6fd98..49bf3b0268a6 100644
--- a/nixos/tests/nix-required-mounts/default.nix
+++ b/nixos/tests/nix-required-mounts/default.nix
@@ -26,11 +26,13 @@ in
           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/imaginary-fhs-drivers/libcuda.so 0444 root root -"
+      "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 = ''
diff --git a/nixos/tests/nix-required-mounts/test-require-feature.nix b/nixos/tests/nix-required-mounts/test-require-feature.nix
index 061b59e1628a..647d9a92d4a3 100644
--- a/nixos/tests/nix-required-mounts/test-require-feature.nix
+++ b/nixos/tests/nix-required-mounts/test-require-feature.nix
@@ -8,9 +8,19 @@ pkgs.runCommandNoCC "${feature}-present"
     echo "The host declares ${feature} support, but doesn't expose /${feature}-files" >&2
     exit 1
   fi
-  if [[ ! -f /run/opengl-driver/lib/libcuda.so ]] ; then
+  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
 ''
diff --git a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
index 50f0f80f0bf2..029740e4b80c 100644
--- a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
+++ b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
@@ -143,17 +143,24 @@ def entrypoint():
     mounts: List[Tuple[PathString, PathString]] = []
 
     while queue:
-        guest_path, host_path, follow_symlinks = queue.popleft()
-        if (guest_path, host_path) not in unique_mounts:
-            mounts.append((guest_path, host_path))
-            unique_mounts.add((guest_path, host_path))
+        guest_path_str, host_path_str, follow_symlinks = queue.popleft()
+        if (guest_path_str, host_path_str) not in unique_mounts:
+            mounts.append((guest_path_str, host_path_str))
+            unique_mounts.add((guest_path_str, host_path_str))
 
         if not follow_symlinks:
             continue
 
-        for parent in symlink_parents(Path(host_path)):
-            parent_str = parent.absolute().as_posix()
-            queue.append((parent_str, parent_str, follow_symlinks))
+        host_path = Path(host_path_str)
+        if not (host_path.is_dir() or host_path.is_symlink()):
+            continue
+
+        # assert host_path_str == guest_path_str, (host_path_str, guest_path_str)
+
+        for child in host_path.iterdir() if host_path.is_dir() else [host_path]:
+            for parent in symlink_parents(child):
+                parent_str = parent.absolute().as_posix()
+                queue.append((parent_str, parent_str, follow_symlinks))
 
     # the pre-build-hook command
     if args.issue_command == "always" or (
@@ -165,8 +172,8 @@ def entrypoint():
         print_paths = False
 
     # arguments, one per line
-    for guest_path, host_path in mounts if print_paths else []:
-        print(f"{guest_path}={host_path}")
+    for guest_path_str, host_path_str in mounts if print_paths else []:
+        print(f"{guest_path_str}={host_path_str}")
 
     # terminated by an empty line
     something_to_terminate = args.issue_stop == "conditional" and mounts

From 075dd8b536c0a0765bdafb4926187349bd162b52 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Mon, 20 Nov 2023 01:17:42 +0000
Subject: [PATCH 16/28] nix-required-mounts: allow overriding the rendered
 config

---
 .../ni/nix-required-mounts/nix_required_mounts.py  | 14 ++++++--------
 pkgs/by-name/ni/nix-required-mounts/package.nix    |  9 ++++-----
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
index 029740e4b80c..388468f7894e 100644
--- a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
+++ b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
@@ -26,15 +26,14 @@ class Pattern(TypedDict):
     unsafeFollowSymlinks: bool
 
 
-class HookConfig(TypedDict):
-    nixExe: str
-    allowedPatterns: Dict[str, Pattern]
+AllowedPatterns: TypeAlias = Dict[str, Pattern]
 
 
 parser = ArgumentParser("pre-build-hook")
 parser.add_argument("derivation_path")
 parser.add_argument("sandbox_path", nargs="?")
-parser.add_argument("--config", type=Path)
+parser.add_argument("--patterns", type=Path, required=True)
+parser.add_argument("--nix-exe", type=Path, required=True)
 parser.add_argument(
     "--issue-command",
     choices=("always", "conditional", "never"),
@@ -72,8 +71,8 @@ def entrypoint():
     args = parser.parse_args()
     drv_path = args.derivation_path
 
-    with open(args.config, "r") as f:
-        config = json.load(f)
+    with open(args.patterns, "r") as f:
+        allowed_patterns = json.load(f)
 
     if not Path(drv_path).exists():
         print(
@@ -85,7 +84,7 @@ def entrypoint():
 
     proc = subprocess.run(
         [
-            config["nixExe"],
+            args.nix_exe,
             "show-derivation",
             drv_path,
         ],
@@ -108,7 +107,6 @@ def entrypoint():
         return
     [canon_drv_path] = parsed_drv.keys()
 
-    allowed_patterns = config["allowedPatterns"]
     known_features = set(
         chain.from_iterable(
             pattern["onFeatures"] for pattern in allowed_patterns.values()
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index 8e2611552d05..3edc40c06245 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -1,5 +1,6 @@
 { addOpenGLRunpath
 , cmake
+, allowedPatternsPath ? (formats.json { }).generate "patterns.json" allowedPatterns
 , allowedPatterns ? rec {
     # This config is just an example.
     # When the hook observes either of the following requiredSystemFeatures:
@@ -25,10 +26,6 @@
 
 
 let
-  confPath = (formats.json { }).generate "config.py" {
-    inherit allowedPatterns;
-    nixExe = lib.getExe nix;
-  };
   attrs = builtins.fromTOML (builtins.readFile ./pyproject.toml);
   pname = attrs.project.name;
   inherit (attrs.project) version;
@@ -47,7 +44,9 @@ python3Packages.buildPythonApplication
   ];
 
   postFixup = ''
-    wrapProgram $out/bin/${pname} --add-flags "--config ${confPath}"
+    wrapProgram $out/bin/${pname} \
+      --add-flags "--patterns ${allowedPatternsPath}" \
+      --add-flags "--nix-exe ${lib.getExe nix}"
   '';
 
   passthru = {

From dd70727622d1484253b8993d166be46915d8ce03 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 21 Nov 2023 00:21:33 +0000
Subject: [PATCH 17/28] nixos/nix-required-mounts: mount the runtime closures

---
 .../modules/programs/nix-required-mounts.nix  |  5 ++-
 .../ni/nix-required-mounts/closure.nix        | 34 ++++++++++++++
 .../ni/nix-required-mounts/package.nix        |  3 +-
 .../scripts/nix_required_mounts_closure.py    | 45 +++++++++++++++++++
 4 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/closure.nix
 create mode 100644 pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index 1c79e532a036..b3c11a51f6fc 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -3,7 +3,9 @@
 let
   cfg = config.programs.nix-required-mounts;
   package = pkgs.nix-required-mounts;
-  overridenPackage = package.override { inherit (cfg) allowedPatterns; };
+  overridenPackage = package.override {
+    inherit (cfg) allowedPatterns;
+  };
 
   Mount = with lib;
     types.submodule {
@@ -37,7 +39,6 @@ let
     });
 
   driverPaths = [
-    # symlinks in /run/opengl-driver/lib:
     pkgs.addOpenGLRunpath.driverLink
 
     # mesa:
diff --git a/pkgs/by-name/ni/nix-required-mounts/closure.nix b/pkgs/by-name/ni/nix-required-mounts/closure.nix
new file mode 100644
index 000000000000..70a00e86f729
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/closure.nix
@@ -0,0 +1,34 @@
+# 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}
+  ''
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index 3edc40c06245..8f5ad450a21c 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -1,6 +1,6 @@
 { addOpenGLRunpath
 , cmake
-, allowedPatternsPath ? (formats.json { }).generate "patterns.json" allowedPatterns
+, allowedPatternsPath ? callPackage ./closure.nix { inherit allowedPatterns; }
 , allowedPatterns ? rec {
     # This config is just an example.
     # When the hook observes either of the following requiredSystemFeatures:
@@ -15,6 +15,7 @@
     nvidia-gpu.unsafeFollowSymlinks = true;
   }
 , buildPackages
+, callPackage
 , formats
 , lib
 , nix
diff --git a/pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py b/pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py
new file mode 100644
index 000000000000..4425e98d0925
--- /dev/null
+++ b/pkgs/by-name/ni/nix-required-mounts/scripts/nix_required_mounts_closure.py
@@ -0,0 +1,45 @@
+import json
+import os
+
+store_dir = os.environ["storeDir"]
+
+with open(os.environ["shallowConfigPath"], "r") as f:
+    config = json.load(f)
+
+cache = {}
+
+
+def read_edges(path: str | dict) -> list[str | dict]:
+    if isinstance(path, dict):
+        return [path]
+    assert isinstance(path, str)
+
+    if not path.startswith(store_dir):
+        return [path]
+    if path in cache:
+        return cache[path]
+
+    name = f"references-{path.removeprefix(store_dir)}"
+
+    assert os.path.exists(name)
+
+    with open(name, "r") as f:
+        return [p.strip() for p in f.readlines() if p.startswith(store_dir)]
+
+
+def host_path(mount: str | dict) -> str:
+    if isinstance(mount, dict):
+        return mount["host"]
+    assert isinstance(mount, str), mount
+    return mount
+
+
+for pattern in config:
+    closure = []
+    for path in config[pattern]["paths"]:
+        closure.append(path)
+        closure.extend(read_edges(path))
+    config[pattern]["paths"] = list({host_path(m): m for m in closure}.values())
+
+with open(os.environ["out"], "w") as f:
+    json.dump(config, f)

From 61001a31c320abc59c363af265273d87491074ba Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 21 Nov 2023 01:59:02 +0000
Subject: [PATCH 18/28] nix-required-mounts: enforce that host paths exist

---
 .../nix_required_mounts.py                    | 38 +++++++++++++------
 1 file changed, 27 insertions(+), 11 deletions(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
index 388468f7894e..72ca086d63b5 100644
--- a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
+++ b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
@@ -67,6 +67,30 @@ def get_strings(drv_env: dict, name: str) -> List[str]:
         return drv_env.get(name, "").split()
 
 
+def validate_mounts(pattern: Pattern) -> List[Tuple[PathString, PathString, bool]]:
+    roots = []
+    for mount in pattern["paths"]:
+        if isinstance(mount, PathString):
+            matches = glob.glob(mount)
+            assert matches, f"Specified host paths do not exist: {mount}"
+
+            roots.extend((m, m, pattern["unsafeFollowSymlinks"]) for m in matches)
+        else:
+            assert isinstance(mount, dict) and "host" in mount, mount
+            assert Path(
+                mount["host"]
+            ).exists(), f"Specified host paths do not exist: {mount['host']}"
+            roots.append(
+                (
+                    mount["guest"],
+                    mount["host"],
+                    pattern["unsafeFollowSymlinks"],
+                )
+            )
+
+    return roots
+
+
 def entrypoint():
     args = parser.parse_args()
     drv_path = args.derivation_path
@@ -118,23 +142,15 @@ def entrypoint():
     required_features = get_strings(drv_env, "requiredSystemFeatures")
     required_features = list(filter(known_features.__contains__, required_features))
 
-    patterns: List[Tuple[PathString | Mount, bool]] = list(
-        (path, pattern["unsafeFollowSymlinks"])
+    patterns: List[Pattern] = list(
+        pattern
         for pattern in allowed_patterns.values()
         for path in pattern["paths"]
         if any(feature in required_features for feature in pattern["onFeatures"])
     )  # noqa: E501
 
     queue: Deque[Tuple[PathString, PathString, bool]] = deque(
-        (
-            mnt
-            for (pattern, follow_symlinks) in patterns
-            for mnt in (
-                ((path, path, follow_symlinks) for path in glob.glob(pattern))
-                if isinstance(pattern, PathString)
-                else [(pattern["guest"], pattern["host"], follow_symlinks)]
-            )
-        )
+        (mnt for pattern in patterns for mnt in validate_mounts(pattern))
     )
 
     unique_mounts: Set[Tuple[PathString, PathString]] = set()

From 6a6b6ac3590abc8020d4fbb332296703551ea867 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 21 Nov 2023 15:37:46 +0000
Subject: [PATCH 19/28] nix-required-mounts: print -> logging

---
 .../nix_required_mounts.py                    | 26 ++++++++++---------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
index 72ca086d63b5..6f05ee913a5a 100644
--- a/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
+++ b/pkgs/by-name/ni/nix-required-mounts/nix_required_mounts.py
@@ -8,8 +8,8 @@ from argparse import ArgumentParser
 from collections import deque
 from itertools import chain
 from pathlib import Path
-from sys import stderr
 from typing import Deque, Dict, List, Set, Tuple, TypeAlias, TypedDict
+import logging
 
 Glob: TypeAlias = str
 PathString: TypeAlias = str
@@ -46,6 +46,7 @@ parser.add_argument(
     default="conditional",
     help="Whether to print the final empty line",
 )
+parser.add_argument("-v", "--verbose", action="count", default=0)
 
 
 def symlink_parents(p: Path) -> List[Path]:
@@ -93,17 +94,22 @@ def validate_mounts(pattern: Pattern) -> List[Tuple[PathString, PathString, bool
 
 def entrypoint():
     args = parser.parse_args()
+
+    VERBOSITY_LEVELS = [logging.ERROR, logging.INFO, logging.DEBUG]
+
+    level_index = min(args.verbose, len(VERBOSITY_LEVELS) - 1)
+    logging.basicConfig(level=VERBOSITY_LEVELS[level_index])
+
     drv_path = args.derivation_path
 
     with open(args.patterns, "r") as f:
         allowed_patterns = json.load(f)
 
     if not Path(drv_path).exists():
-        print(
-            f"[E] {drv_path} doesn't exist."
+        logging.error(
+            f"{drv_path} doesn't exist."
             " Cf. https://github.com/NixOS/nix/issues/9272"
             " Exiting the hook",
-            file=stderr,
         )
 
     proc = subprocess.run(
@@ -117,17 +123,13 @@ def entrypoint():
     try:
         parsed_drv = json.loads(proc.stdout)
     except json.JSONDecodeError:
-        print(
-            "[E] Couldn't parse the output of"
+        logging.error(
+            "Couldn't parse the output of"
             "`nix show-derivation`"
             f". Expected JSON, observed: {proc.stdout}",
-            file=stderr,
         )
-        print(
-            textwrap.indent(proc.stdout.decode("utf8"), prefix=" " * 4),
-            file=stderr,
-        )
-        print("[I] Exiting the nix-required-binds hook", file=stderr)
+        logging.error(textwrap.indent(proc.stdout.decode("utf8"), prefix=" " * 4))
+        logging.info("Exiting the nix-required-binds hook")
         return
     [canon_drv_path] = parsed_drv.keys()
 

From 927b15ed6db1e7da9f35e0f4a3ae85293cac5454 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 21 Nov 2023 16:03:58 +0000
Subject: [PATCH 20/28] nixos/nix-required-mounts: allow passing extra
 arguments to the hook

---
 .../modules/programs/nix-required-mounts.nix  | 22 +++++++++++++++----
 .../ni/nix-required-mounts/package.nix        |  6 +++--
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index b3c11a51f6fc..98ab819af55e 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -3,9 +3,6 @@
 let
   cfg = config.programs.nix-required-mounts;
   package = pkgs.nix-required-mounts;
-  overridenPackage = package.override {
-    inherit (cfg) allowedPatterns;
-  };
 
   Mount = with lib;
     types.submodule {
@@ -86,9 +83,26 @@ in
         example.require-ipfs.paths = [ "/ipfs" ];
         example.require-ipfs.onFeatures = [ "ifps" ];
       };
+    extraWrapperArgs = lib.mkOption {
+      type = with lib.types; listOf str;
+      default = [ ];
+      description =
+        lib.mdDoc
+          "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 = lib.mdDoc "The final package with the final config applied";
+      internal = true;
+    };
   };
   config = lib.mkIf cfg.enable (lib.mkMerge [
-    { nix.settings.pre-build-hook = lib.getExe overridenPackage; }
+    { 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 = {
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index 8f5ad450a21c..e674539fb909 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -16,12 +16,13 @@
   }
 , buildPackages
 , callPackage
+, extraWrapperArgs ? [ ]
 , formats
 , lib
+, makeWrapper
 , nix
 , nixosTests
 , python3Packages
-, makeWrapper
 , runCommand
 }:
 
@@ -47,7 +48,8 @@ python3Packages.buildPythonApplication
   postFixup = ''
     wrapProgram $out/bin/${pname} \
       --add-flags "--patterns ${allowedPatternsPath}" \
-      --add-flags "--nix-exe ${lib.getExe nix}"
+      --add-flags "--nix-exe ${lib.getExe nix}" \
+      ${builtins.concatStringsSep " " extraWrapperArgs}
   '';
 
   passthru = {

From 9aa0403154ad65ccbaf7bfa92f8eb55c8073bb0d Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 21 Nov 2023 20:48:09 +0000
Subject: [PATCH 21/28] cudaPackages.saxpy: passthru: test gpu/runtime

---
 pkgs/development/cuda-modules/saxpy/default.nix | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
index 9b7326cd321f..520fa5b68a44 100644
--- a/pkgs/development/cuda-modules/saxpy/default.nix
+++ b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -3,6 +3,7 @@
   cmake,
   cudaPackages,
   lib,
+  saxpy,
 }:
 let
   inherit (cudaPackages)
@@ -58,6 +59,16 @@ backendStdenv.mkDerivation {
     (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)
   ];
 
+  passthru.tests.withCuda = saxpy.overrideAttrs (
+    _: {
+      requiredSystemFeatures = ["cuda"];
+      doInstallCheck = true;
+      postInstallCheck = ''
+        $out/bin/saxpy
+      '';
+    }
+  );
+
   meta = rec {
     description = "Simple (Single-precision AX Plus Y) FindCUDAToolkit.cmake example for testing cross-compilation";
     license = lib.licenses.mit;

From efd64b5f0b0794fd15492e867078be2bb0c0c95c Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Fri, 8 Dec 2023 19:09:41 +0000
Subject: [PATCH 22/28] cudaPackages: move cuda tests from passthru.tests

Otherwise we crash Ofborg
---
 pkgs/applications/misc/blender/default.nix         | 8 ++++----
 pkgs/development/cuda-modules/saxpy/default.nix    | 2 +-
 pkgs/development/python-modules/pynvml/default.nix | 3 +--
 pkgs/development/python-modules/torch/bin.nix      | 2 +-
 pkgs/development/python-modules/torch/default.nix  | 8 ++------
 5 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/pkgs/applications/misc/blender/default.nix b/pkgs/applications/misc/blender/default.nix
index 357bbbe0e6f6..284ee5778101 100644
--- a/pkgs/applications/misc/blender/default.nix
+++ b/pkgs/applications/misc/blender/default.nix
@@ -373,16 +373,16 @@ stdenv.mkDerivation (finalAttrs: {
         done
       '';
 
-      cudaAvailable = runCommand
+    };
+    gpuChecks.cudaAvailable = { blenderWithCuda, runCommand }: runCommand
       "blender-cuda-available"
       {
-        nativeBuildInputs = [ finalAttrs.finalPackage ];
+        nativeBuildInputs = [ blenderWithCuda ];
         requiredSystemFeatures = [ "cuda" ];
       }
       ''
         blender --background -noaudio --python-exit-code 1 --python ${./test-cuda.py} && touch $out
-      '';
-    };
+      '' { };
   };
 
   meta = {
diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
index 520fa5b68a44..3da3dc08f0eb 100644
--- a/pkgs/development/cuda-modules/saxpy/default.nix
+++ b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -59,7 +59,7 @@ backendStdenv.mkDerivation {
     (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)
   ];
 
-  passthru.tests.withCuda = saxpy.overrideAttrs (
+  passthru.gpuChecks.withCuda = saxpy.overrideAttrs (
     _: {
       requiredSystemFeatures = ["cuda"];
       doInstallCheck = true;
diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
index e76913a63df4..79624e5298a6 100644
--- a/pkgs/development/python-modules/pynvml/default.nix
+++ b/pkgs/development/python-modules/pynvml/default.nix
@@ -51,8 +51,7 @@ buildPythonPackage rec {
   # OSError: /run/opengl-driver/lib/libnvidia-ml.so.1: cannot open shared object file: No such file or directory
   doCheck = false;
 
-  passthru.tests.nvmlInit = callPackage ./test-gpu.nix { };
-
+  passthru.gpuChecks.nvmlInit = callPackage ./test-gpu.nix { };
 
   meta = with lib; {
     description = "Python bindings for the NVIDIA Management Library";
diff --git a/pkgs/development/python-modules/torch/bin.nix b/pkgs/development/python-modules/torch/bin.nix
index 37170ea9adf4..e2899c081e08 100644
--- a/pkgs/development/python-modules/torch/bin.nix
+++ b/pkgs/development/python-modules/torch/bin.nix
@@ -121,7 +121,7 @@ buildPythonPackage {
 
   pythonImportsCheck = [ "torch" ];
 
-  passthru.tests.cudaAvailable = callPackage ./test-cuda.nix { torch = torch-bin; };
+  passthru.gpuChecks.cudaAvailable = callPackage ./test-cuda.nix { torch = torch-bin; };
 
   meta = {
     description = "PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration";
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index 671e76dfe02d..62c69905609e 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -26,7 +26,7 @@
 
   # tests.cudaAvailable:
   callPackage,
-  torch,
+  torchWithCuda,
 
   # Native build inputs
   cmake,
@@ -648,11 +648,7 @@ buildPythonPackage rec {
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-  } // lib.optionalAttrs cudaSupport {
-
-    tests = lib.optionalAttrs cudaSupport {
-      cudaAvailable = callPackage ./test-cuda.nix { inherit torch; };
-    };
+    gpuChecks.cudaAvailable = callPackage ./test-cuda.nix { torch = torchWithCuda; };
   };
 
   meta = {

From 39f33456e4443cdeb4b56c0b87a667b6a0ee3699 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 9 Dec 2023 21:52:13 +0000
Subject: [PATCH 23/28] python3Packages.torch.gpuChecks: add rocm

---
 .../python-modules/torch/default.nix          |  2 +-
 .../python-modules/torch/gpu-checks.nix       | 50 +++++++++++++++++++
 .../python-modules/torch/test-cuda.nix        | 21 --------
 3 files changed, 51 insertions(+), 22 deletions(-)
 create mode 100644 pkgs/development/python-modules/torch/gpu-checks.nix
 delete mode 100644 pkgs/development/python-modules/torch/test-cuda.nix

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index 62c69905609e..11a3b3df42ce 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.cudaAvailable = callPackage ./test-cuda.nix { torch = torchWithCuda; };
+    gpuChecks = callPackage ./gpu-checks.nix { };
   };
 
   meta = {
diff --git a/pkgs/development/python-modules/torch/gpu-checks.nix b/pkgs/development/python-modules/torch/gpu-checks.nix
new file mode 100644
index 000000000000..71004d8457b2
--- /dev/null
+++ b/pkgs/development/python-modules/torch/gpu-checks.nix
@@ -0,0 +1,50 @@
+{
+  lib,
+  callPackage,
+  torchWithCuda,
+  torchWithRocm,
+}:
+
+let
+  accelAvailable =
+    {
+      feature,
+      versionAttr,
+      torch,
+      runCommandNoCC,
+      writers,
+    }:
+    let
+      name = "${torch.name}-${feature}-check";
+      unwrapped = writers.writePython3Bin "${name}-unwrapped" {libraries = [torch];} ''
+        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 {
+    feature = "cuda";
+    versionAttr = "cuda";
+    torch = torchWithCuda;
+  };
+  rocmAvailable = callPackage accelAvailable {
+    feature = "rocm";
+    versionAttr = "hip";
+    torch = torchWithRocm;
+  };
+}
diff --git a/pkgs/development/python-modules/torch/test-cuda.nix b/pkgs/development/python-modules/torch/test-cuda.nix
deleted file mode 100644
index 22d73eeb9167..000000000000
--- a/pkgs/development/python-modules/torch/test-cuda.nix
+++ /dev/null
@@ -1,21 +0,0 @@
-{ runCommandNoCC
-, python
-, torch
-}:
-
-runCommandNoCC "${torch.name}-gpu-test"
-{
-  nativeBuildInputs = [
-    (python.withPackages (_: [ torch ]))
-  ];
-  requiredSystemFeatures = [
-    "cuda"
-  ];
-} ''
-  python3 << EOF
-  import torch
-  assert torch.cuda.is_available(), f"{torch.cuda.is_available()=}"
-  EOF
-
-  touch $out
-''

From da430f48723a3c728b8cb5ea630aba600105e568 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 9 Dec 2023 23:07:01 +0000
Subject: [PATCH 24/28] blender.gpuChecks: add unwrapped
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

An unwrapped check for `nix run`-ing on the host platform,
instead of `nix build`-ing in the sandbox

E.g.:

```
❯ nix run -f ./. --arg config '{ cudaSupport = true; cudaCapabilities = [ "8.6" ]; cudaEnableForwardCompat = false; allowUnfree = true; }' -L blender.gpuChecks.cudaAvailable.unwrapped
Blender 4.0.1
Read prefs: "/home/user/.config/blender/4.0/config/userpref.blend"
CUDA is available
Blender quit
❯ nix build -f ./. --arg config '{ cudaSupport = true; cudaCapabilities = [ "8.6" ]; cudaEnableForwardCompat = false; allowUnfree = true; }' -L blender.gpuChecks
blender> Blender 4.0.1
blender> could not get a list of mounted file-systems
blender> CUDA is available
blender> Blender quit
```
---
 pkgs/applications/misc/blender/default.nix    | 10 +------
 pkgs/applications/misc/blender/gpu-checks.nix | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 9 deletions(-)
 create mode 100644 pkgs/applications/misc/blender/gpu-checks.nix

diff --git a/pkgs/applications/misc/blender/default.nix b/pkgs/applications/misc/blender/default.nix
index 284ee5778101..38f3e226d0ff 100644
--- a/pkgs/applications/misc/blender/default.nix
+++ b/pkgs/applications/misc/blender/default.nix
@@ -374,15 +374,7 @@ stdenv.mkDerivation (finalAttrs: {
       '';
 
     };
-    gpuChecks.cudaAvailable = { blenderWithCuda, runCommand }: runCommand
-      "blender-cuda-available"
-      {
-        nativeBuildInputs = [ blenderWithCuda ];
-        requiredSystemFeatures = [ "cuda" ];
-      }
-      ''
-        blender --background -noaudio --python-exit-code 1 --python ${./test-cuda.py} && touch $out
-      '' { };
+    gpuChecks = callPackage ./gpu-checks.nix { };
   };
 
   meta = {
diff --git a/pkgs/applications/misc/blender/gpu-checks.nix b/pkgs/applications/misc/blender/gpu-checks.nix
new file mode 100644
index 000000000000..144cdeb968c4
--- /dev/null
+++ b/pkgs/applications/misc/blender/gpu-checks.nix
@@ -0,0 +1,29 @@
+{
+  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";
+}

From ff430d16995ca3cc16cb971c6aa5eb37cae43872 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Thu, 20 Jun 2024 21:18:22 +0000
Subject: [PATCH 25/28] nix-required-mounts: cuda: /dev/video* may not exist
 and aren't relevant

---
 pkgs/by-name/ni/nix-required-mounts/package.nix | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index e674539fb909..eded427635dc 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -10,7 +10,6 @@
       addOpenGLRunpath.driverLink
       "/dev/dri"
       "/dev/nvidia*"
-      "/dev/video*"
     ];
     nvidia-gpu.unsafeFollowSymlinks = true;
   }

From ebeb6b9d1df9351a76af366aa0821fec78d70e63 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Thu, 20 Jun 2024 21:33:51 +0000
Subject: [PATCH 26/28] nix-required-mounts: nixfmt

---
 .../modules/programs/nix-required-mounts.nix  | 94 ++++++++++---------
 nixos/tests/nix-required-mounts/default.nix   | 54 +++++------
 .../ensure-path-not-present.nix               | 10 +-
 .../test-require-feature.nix                  | 10 +-
 .../test-structured-attrs-empty.nix           | 14 ++-
 .../test-structured-attrs.nix                 | 29 +++---
 pkgs/applications/misc/blender/gpu-checks.nix | 19 ++--
 .../ni/nix-required-mounts/closure.nix        | 51 +++++-----
 .../ni/nix-required-mounts/package.nix        | 42 +++++----
 .../python-modules/pynvml/test-gpu.nix        | 32 +++----
 .../python-modules/torch/gpu-checks.nix       |  6 +-
 11 files changed, 183 insertions(+), 178 deletions(-)

diff --git a/nixos/modules/programs/nix-required-mounts.nix b/nixos/modules/programs/nix-required-mounts.nix
index 98ab819af55e..c339dd1cfddd 100644
--- a/nixos/modules/programs/nix-required-mounts.nix
+++ b/nixos/modules/programs/nix-required-mounts.nix
@@ -1,10 +1,16 @@
-{ config, lib, pkgs, ... }:
+{
+  config,
+  lib,
+  pkgs,
+  ...
+}:
 
 let
   cfg = config.programs.nix-required-mounts;
   package = pkgs.nix-required-mounts;
 
-  Mount = with lib;
+  Mount =
+    with lib;
     types.submodule {
       options.host = mkOption {
         type = types.str;
@@ -15,25 +21,30 @@ let
         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.
-      '';
-    });
+  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
@@ -53,8 +64,7 @@ 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";
+    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" ]`.
@@ -64,11 +74,11 @@ in
       You may extend or override the exposed paths via the
       `programs.nix-required-mounts.allowedPatterns.nvidia-gpu.paths` option.
     '';
-    allowedPatterns = with lib.types;
+    allowedPatterns =
+      with lib.types;
       lib.mkOption rec {
         type = attrsOf Pattern;
-        description =
-          "The hook config, describing which paths to mount for which system features";
+        description = "The hook config, describing which paths to mount for which system features";
         default = { };
         defaultText = lib.literalExpression ''
           {
@@ -86,28 +96,24 @@ in
     extraWrapperArgs = lib.mkOption {
       type = with lib.types; listOf str;
       default = [ ];
-      description =
-        lib.mdDoc
-          "List of extra arguments (such as `--add-flags -v`) to pass to the hook's wrapper";
+      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 = lib.mdDoc "The final package with the final config applied";
+      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;
-      };
-    })
-  ]);
+  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;
+        };
+      })
+    ]
+  );
 }
diff --git a/nixos/tests/nix-required-mounts/default.nix b/nixos/tests/nix-required-mounts/default.nix
index 49bf3b0268a6..60f894ce0bcc 100644
--- a/nixos/tests/nix-required-mounts/default.nix
+++ b/nixos/tests/nix-required-mounts/default.nix
@@ -1,6 +1,4 @@
-{ pkgs
-, ...
-}:
+{ pkgs, ... }:
 
 let
   inherit (pkgs) lib;
@@ -9,32 +7,34 @@ 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";
-        }
+  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"
       ];
-      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
 
diff --git a/nixos/tests/nix-required-mounts/ensure-path-not-present.nix b/nixos/tests/nix-required-mounts/ensure-path-not-present.nix
index 871f336ee9bd..270c268fcbd9 100644
--- a/nixos/tests/nix-required-mounts/ensure-path-not-present.nix
+++ b/nixos/tests/nix-required-mounts/ensure-path-not-present.nix
@@ -1,8 +1,9 @@
-{ pkgs ? import <nixpkgs> { }, feature }:
-
-pkgs.runCommandNoCC "${feature}-not-present"
 {
-} ''
+  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
@@ -10,4 +11,3 @@ pkgs.runCommandNoCC "${feature}-not-present"
     touch $out
   fi
 ''
-
diff --git a/nixos/tests/nix-required-mounts/test-require-feature.nix b/nixos/tests/nix-required-mounts/test-require-feature.nix
index 647d9a92d4a3..447fd49a300a 100644
--- a/nixos/tests/nix-required-mounts/test-require-feature.nix
+++ b/nixos/tests/nix-required-mounts/test-require-feature.nix
@@ -1,9 +1,9 @@
-{ pkgs ? import <nixpkgs> { }, feature }:
-
-pkgs.runCommandNoCC "${feature}-present"
 {
-  requiredSystemFeatures = [ feature ];
-} ''
+  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
diff --git a/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix b/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix
index d788c6773c8e..86f275330936 100644
--- a/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix
+++ b/nixos/tests/nix-required-mounts/test-structured-attrs-empty.nix
@@ -1,10 +1,8 @@
-{ pkgs ? import <nixpkgs> { } }:
-
-pkgs.runCommandNoCC "nix-required-mounts-structured-attrs-no-features"
 {
-  __structuredAttrs = true;
-} ''
-  touch $out
-''
-
+  pkgs ? import <nixpkgs> { },
+}:
 
+pkgs.runCommandNoCC "nix-required-mounts-structured-attrs-no-features" { __structuredAttrs = true; }
+  ''
+    touch $out
+  ''
diff --git a/nixos/tests/nix-required-mounts/test-structured-attrs.nix b/nixos/tests/nix-required-mounts/test-structured-attrs.nix
index fecd2c32eec0..874910eee7bb 100644
--- a/nixos/tests/nix-required-mounts/test-structured-attrs.nix
+++ b/nixos/tests/nix-required-mounts/test-structured-attrs.nix
@@ -1,15 +1,18 @@
-{ pkgs ? import <nixpkgs> { }, feature }:
+{
+  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
-''
-
+  {
+    __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
+  ''
diff --git a/pkgs/applications/misc/blender/gpu-checks.nix b/pkgs/applications/misc/blender/gpu-checks.nix
index 144cdeb968c4..bfbaf25b989a 100644
--- a/pkgs/applications/misc/blender/gpu-checks.nix
+++ b/pkgs/applications/misc/blender/gpu-checks.nix
@@ -8,7 +8,7 @@
 }:
 
 let
-  blenderWithCuda = blender.override {cudaSupport = true;};
+  blenderWithCuda = blender.override { cudaSupport = true; };
   name = "${blenderWithCuda.name}-check-cuda";
   unwrapped = writeScriptBin "${name}-unwrapped" ''
     #!${lib.getExe bash}
@@ -16,14 +16,11 @@ let
   '';
 in
 {
-  cudaAvailable =
-    runCommand name
-      {
-        nativeBuildInputs = [unwrapped];
-        requiredSystemFeatures = ["cuda"];
-        passthru = {
-          inherit unwrapped;
-        };
-      }
-      "${name}-unwrapped && touch $out";
+  cudaAvailable = runCommand name {
+    nativeBuildInputs = [ unwrapped ];
+    requiredSystemFeatures = [ "cuda" ];
+    passthru = {
+      inherit unwrapped;
+    };
+  } "${name}-unwrapped && touch $out";
 }
diff --git a/pkgs/by-name/ni/nix-required-mounts/closure.nix b/pkgs/by-name/ni/nix-required-mounts/closure.nix
index 70a00e86f729..3e361114bc4c 100644
--- a/pkgs/by-name/ni/nix-required-mounts/closure.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/closure.nix
@@ -3,32 +3,35 @@
 # 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
+{
+  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" ];
-}
+  {
+    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}
   ''
diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index eded427635dc..a7b9c3093e3f 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -1,10 +1,16 @@
-{ addOpenGLRunpath
-, cmake
-, allowedPatternsPath ? callPackage ./closure.nix { inherit allowedPatterns; }
-, allowedPatterns ? rec {
+{
+  addOpenGLRunpath,
+  cmake,
+  allowedPatternsPath ? callPackage ./closure.nix { inherit allowedPatterns; },
+  allowedPatterns ? rec {
     # This config is just an example.
     # When the hook observes either of the following requiredSystemFeatures:
-    nvidia-gpu.onFeatures = [ "gpu" "nvidia-gpu" "opengl" "cuda" ];
+    nvidia-gpu.onFeatures = [
+      "gpu"
+      "nvidia-gpu"
+      "opengl"
+      "cuda"
+    ];
     # It exposes these paths in the sandbox:
     nvidia-gpu.paths = [
       addOpenGLRunpath.driverLink
@@ -12,28 +18,26 @@
       "/dev/nvidia*"
     ];
     nvidia-gpu.unsafeFollowSymlinks = true;
-  }
-, buildPackages
-, callPackage
-, extraWrapperArgs ? [ ]
-, formats
-, lib
-, makeWrapper
-, nix
-, nixosTests
-, python3Packages
-, runCommand
+  },
+  buildPackages,
+  callPackage,
+  extraWrapperArgs ? [ ],
+  formats,
+  lib,
+  makeWrapper,
+  nix,
+  nixosTests,
+  python3Packages,
+  runCommand,
 }:
 
-
 let
   attrs = builtins.fromTOML (builtins.readFile ./pyproject.toml);
   pname = attrs.project.name;
   inherit (attrs.project) version;
 in
 
-python3Packages.buildPythonApplication
-{
+python3Packages.buildPythonApplication {
   inherit pname version;
   pyproject = true;
 
diff --git a/pkgs/development/python-modules/pynvml/test-gpu.nix b/pkgs/development/python-modules/pynvml/test-gpu.nix
index c316d0b5094b..6ab4290a2bba 100644
--- a/pkgs/development/python-modules/pynvml/test-gpu.nix
+++ b/pkgs/development/python-modules/pynvml/test-gpu.nix
@@ -1,23 +1,17 @@
-{ runCommandNoCC
-, python
-}:
+{ runCommandNoCC, python }:
 
 runCommandNoCC "pynvml-gpu-test"
-{
-  nativeBuildInputs = [
-    (python.withPackages (ps: [ ps.pynvml ]))
-  ];
-  requiredSystemFeatures = [
-    "cuda"
-  ];
-} ''
-  python3 << EOF
-  import pynvml
-  from pynvml.smi import nvidia_smi
+  {
+    nativeBuildInputs = [ (python.withPackages (ps: [ ps.pynvml ])) ];
+    requiredSystemFeatures = [ "cuda" ];
+  }
+  ''
+    python3 << EOF
+    import pynvml
+    from pynvml.smi import nvidia_smi
 
-  pynvml.nvmlInit()
-  EOF
-
-  touch $out
-''
+    pynvml.nvmlInit()
+    EOF
 
+    touch $out
+  ''
diff --git a/pkgs/development/python-modules/torch/gpu-checks.nix b/pkgs/development/python-modules/torch/gpu-checks.nix
index 71004d8457b2..371b83f1b778 100644
--- a/pkgs/development/python-modules/torch/gpu-checks.nix
+++ b/pkgs/development/python-modules/torch/gpu-checks.nix
@@ -16,7 +16,7 @@ let
     }:
     let
       name = "${torch.name}-${feature}-check";
-      unwrapped = writers.writePython3Bin "${name}-unwrapped" {libraries = [torch];} ''
+      unwrapped = writers.writePython3Bin "${name}-unwrapped" { libraries = [ torch ]; } ''
         import torch
         message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
         assert torch.cuda.is_available() and torch.version.${versionAttr}, message
@@ -25,8 +25,8 @@ let
     in
     runCommandNoCC name
       {
-        nativeBuildInputs = [unwrapped];
-        requiredSystemFeatures = [feature];
+        nativeBuildInputs = [ unwrapped ];
+        requiredSystemFeatures = [ feature ];
         passthru = {
           inherit unwrapped;
         };

From 7d667a0996d718b2384ac280ff2b4c9ad22e2beb Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Thu, 20 Jun 2024 22:16:36 +0000
Subject: [PATCH 27/28] nix-required-mounts: refactor: drop unused arguments

---
 pkgs/by-name/ni/nix-required-mounts/package.nix | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/pkgs/by-name/ni/nix-required-mounts/package.nix b/pkgs/by-name/ni/nix-required-mounts/package.nix
index a7b9c3093e3f..197e0812a8ec 100644
--- a/pkgs/by-name/ni/nix-required-mounts/package.nix
+++ b/pkgs/by-name/ni/nix-required-mounts/package.nix
@@ -1,6 +1,5 @@
 {
   addOpenGLRunpath,
-  cmake,
   allowedPatternsPath ? callPackage ./closure.nix { inherit allowedPatterns; },
   allowedPatterns ? rec {
     # This config is just an example.
@@ -19,16 +18,13 @@
     ];
     nvidia-gpu.unsafeFollowSymlinks = true;
   },
-  buildPackages,
   callPackage,
   extraWrapperArgs ? [ ],
-  formats,
   lib,
   makeWrapper,
   nix,
   nixosTests,
   python3Packages,
-  runCommand,
 }:
 
 let

From 79a7186f1ce8d94b0c136a7cc7c3e2e31facc794 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Wed, 26 Jun 2024 00:29:42 +0000
Subject: [PATCH 28/28] 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.
---
 .../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 =