Merge pull request #281576 from yannham/refactor/cuda-setup-hooks-refactor

cudaPackages: generalize and refactor setup hooks

28 files changed, 154 insertions, 95 deletions

diff --git a/doc/languages-frameworks/cuda.section.md b/doc/languages-frameworks/cuda.section.md
index 11c86e375c61..09af824531a2 100644
--- a/doc/languages-frameworks/cuda.section.md
+++ b/doc/languages-frameworks/cuda.section.md
@@ -144,4 +144,4 @@ All new projects should use the CUDA redistributables available in [`cudaPackage
 | Find libraries | `configurePhase` | Missing dependency on a `dev` output | Add the missing dependency | The `dev` output typically contain CMake configuration files |
 | Find libraries | `buildPhase` or `patchelf` | Missing dependency on a `lib` or `static` output | Add the missing dependency | The `lib` or `static` output typically contain the libraries |
 
-In the scenario you are unable to run the resulting binary: this is arguably the most complicated as it could be any combination of the previous reasons. This type of failure typically occurs when a library attempts to load or open a library it depends on that it does not declare in its `DT_NEEDED` section. As a first step, ensure that dependencies are patched with [`cudaPackages.autoAddOpenGLRunpath`](https://search.nixos.org/packages?channel=unstable&type=packages&query=cudaPackages.autoAddOpenGLRunpath). Failing that, try running the application with [`nixGL`](https://github.com/guibou/nixGL) or a similar wrapper tool. If that works, it likely means that the application is attempting to load a library that is not in the `RPATH` or `RUNPATH` of the binary.
+In the scenario you are unable to run the resulting binary: this is arguably the most complicated as it could be any combination of the previous reasons. This type of failure typically occurs when a library attempts to load or open a library it depends on that it does not declare in its `DT_NEEDED` section. As a first step, ensure that dependencies are patched with [`cudaPackages.autoAddDriverRunpath`](https://search.nixos.org/packages?channel=unstable&type=packages&query=cudaPackages.autoAddDriverRunpath). Failing that, try running the application with [`nixGL`](https://github.com/guibou/nixGL) or a similar wrapper tool. If that works, it likely means that the application is attempting to load a library that is not in the `RPATH` or `RUNPATH` of the binary.
diff --git a/pkgs/applications/science/misc/colmap/default.nix b/pkgs/applications/science/misc/colmap/default.nix
index a029f5440b25..64a7952be4cc 100644
--- a/pkgs/applications/science/misc/colmap/default.nix
+++ b/pkgs/applications/science/misc/colmap/default.nix
@@ -37,7 +37,7 @@ mkDerivation rec {
   nativeBuildInputs = [
     cmake
   ] ++ lib.optionals cudaSupport [
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   meta = with lib; {
diff --git a/pkgs/applications/science/molecular-dynamics/lammps/default.nix b/pkgs/applications/science/molecular-dynamics/lammps/default.nix
index db27bb188cf5..4c2a80c2f1b7 100644
--- a/pkgs/applications/science/molecular-dynamics/lammps/default.nix
+++ b/pkgs/applications/science/molecular-dynamics/lammps/default.nix
@@ -63,7 +63,7 @@ stdenv.mkDerivation (finalAttrs: {
     pkg-config
     # Although not always needed, it is needed if cmakeFlags include
     # GPU_API=cuda, and it doesn't users that don't enable the GPU package.
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   passthru = {
diff --git a/pkgs/by-name/ll/llama-cpp/package.nix b/pkgs/by-name/ll/llama-cpp/package.nix
index cb3cd49f09af..3b683d28a26a 100644
--- a/pkgs/by-name/ll/llama-cpp/package.nix
+++ b/pkgs/by-name/ll/llama-cpp/package.nix
@@ -86,10 +86,7 @@ effectiveStdenv.mkDerivation (finalAttrs: {
   nativeBuildInputs = [ cmake ninja pkg-config git ]
     ++ optionals cudaSupport [
     cudaPackages.cuda_nvcc
-
-    # TODO: Replace with autoAddDriverRunpath
-    # once https://github.com/NixOS/nixpkgs/pull/275241 has been merged
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = optionals effectiveStdenv.isDarwin darwinBuildInputs
diff --git a/pkgs/by-name/nv/nvidia-container-toolkit/package.nix b/pkgs/by-name/nv/nvidia-container-toolkit/package.nix
index a584be35e7b2..ba76e4de97e4 100644
--- a/pkgs/by-name/nv/nvidia-container-toolkit/package.nix
+++ b/pkgs/by-name/nv/nvidia-container-toolkit/package.nix
@@ -87,7 +87,7 @@ buildGoModule rec {
   ];
 
   nativeBuildInputs = [
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
     makeWrapper
   ];
 
diff --git a/pkgs/by-name/ta/tabby/package.nix b/pkgs/by-name/ta/tabby/package.nix
index fa52d372ab9a..8aabe52907d1 100644
--- a/pkgs/by-name/ta/tabby/package.nix
+++ b/pkgs/by-name/ta/tabby/package.nix
@@ -139,7 +139,7 @@ rustPlatform.buildRustPackage {
   ] ++ optionals enableCuda [
     # TODO: Replace with autoAddDriverRunpath
     # once https://github.com/NixOS/nixpkgs/pull/275241 has been merged
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = [ openssl ]
diff --git a/pkgs/development/cuda-modules/cuda-samples/generic.nix b/pkgs/development/cuda-modules/cuda-samples/generic.nix
index fb3d7cc99da9..3d1dac015e16 100644
--- a/pkgs/development/cuda-modules/cuda-samples/generic.nix
+++ b/pkgs/development/cuda-modules/cuda-samples/generic.nix
@@ -1,5 +1,5 @@
 {
-  autoAddOpenGLRunpathHook,
+  autoAddDriverRunpath,
   backendStdenv,
   cmake,
   cudatoolkit,
@@ -31,7 +31,7 @@ backendStdenv.mkDerivation (
 
     nativeBuildInputs =
       [
-        autoAddOpenGLRunpathHook
+        autoAddDriverRunpath
         pkg-config
       ]
       # CMake has to run as a native, build-time dependency for libNVVM samples.
diff --git a/pkgs/development/cuda-modules/cudatoolkit/default.nix b/pkgs/development/cuda-modules/cudatoolkit/default.nix
index aca0c7ad0b78..231a153bf7e6 100644
--- a/pkgs/development/cuda-modules/cudatoolkit/default.nix
+++ b/pkgs/development/cuda-modules/cudatoolkit/default.nix
@@ -2,7 +2,7 @@
   cudaVersion,
   runPatches ? [],
   autoPatchelfHook,
-  autoAddOpenGLRunpathHook,
+  autoAddDriverRunpath,
   addOpenGLRunpath,
   alsa-lib,
   curlMinimal,
@@ -76,7 +76,7 @@ backendStdenv.mkDerivation rec {
       rsync
       addOpenGLRunpath
       autoPatchelfHook
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
       markForCudatoolkitRootHook
     ]
     ++ lib.optionals (lib.versionOlder version "11") [libsForQt5.wrapQtAppsHook]
diff --git a/pkgs/development/cuda-modules/generic-builders/manifest.nix b/pkgs/development/cuda-modules/generic-builders/manifest.nix
index 5e837fa36b5e..4f40b7f01dc2 100644
--- a/pkgs/development/cuda-modules/generic-builders/manifest.nix
+++ b/pkgs/development/cuda-modules/generic-builders/manifest.nix
@@ -1,7 +1,7 @@
 {
   # General callPackage-supplied arguments
-  autoAddOpenGLRunpathHook,
-  autoAddCudaCompatRunpathHook,
+  autoAddDriverRunpath,
+  autoAddCudaCompatRunpath,
   autoPatchelfHook,
   backendStdenv,
   fetchurl,
@@ -193,16 +193,16 @@ backendStdenv.mkDerivation (
       # in typically /lib/opengl-driver by adding that
       # directory to the rpath of all ELF binaries.
       # Check e.g. with `patchelf --print-rpath path/to/my/binary
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
       markForCudatoolkitRootHook
     ]
-    # autoAddCudaCompatRunpathHook depends on cuda_compat and would cause
+    # autoAddCudaCompatRunpath depends on cuda_compat and would cause
     # infinite recursion if applied to `cuda_compat` itself (beside the fact
     # that it doesn't make sense in the first place)
     ++ lib.optionals (pname != "cuda_compat" && flags.isJetsonBuild) [
-      # autoAddCudaCompatRunpathHook must appear AFTER autoAddOpenGLRunpathHook.
+      # autoAddCudaCompatRunpath must appear AFTER autoAddDriverRunpath.
       # See its documentation in ./setup-hooks/extension.nix.
-      autoAddCudaCompatRunpathHook
+      autoAddCudaCompatRunpath
     ];
 
     buildInputs =
diff --git a/pkgs/development/cuda-modules/nccl/default.nix b/pkgs/development/cuda-modules/nccl/default.nix
index 25296c21365d..b8bdc69bba4f 100644
--- a/pkgs/development/cuda-modules/nccl/default.nix
+++ b/pkgs/development/cuda-modules/nccl/default.nix
@@ -12,7 +12,7 @@
 }:
 let
   inherit (cudaPackages)
-    autoAddOpenGLRunpathHook
+    autoAddDriverRunpath
     backendStdenv
     cuda_cccl
     cuda_cudart
@@ -44,7 +44,7 @@ backendStdenv.mkDerivation (
     nativeBuildInputs =
       [
         which
-        autoAddOpenGLRunpathHook
+        autoAddDriverRunpath
         python3
       ]
       ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
index 68c60e2d8446..bc299dea006f 100644
--- a/pkgs/development/cuda-modules/saxpy/default.nix
+++ b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -5,7 +5,7 @@
 }:
 let
   inherit (cudaPackages)
-    autoAddOpenGLRunpathHook
+    autoAddDriverRunpath
     backendStdenv
     cuda_cccl
     cuda_cudart
@@ -29,7 +29,7 @@ backendStdenv.mkDerivation {
   nativeBuildInputs =
     [
       cmake
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
     ]
     ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [cudatoolkit]
     ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [cuda_nvcc];
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
index ae25cebaca6f..fc41024f1551 100644
--- a/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-add-cuda-compat-runpath.sh
@@ -3,25 +3,25 @@
 # coming from the cuda_compat package by adding it to the RUNPATH.
 echo "Sourcing auto-add-cuda-compat-runpath-hook"
 
-elfHasDynamicSection() {
-    patchelf --print-rpath "$1" >& /dev/null
-}
+addCudaCompatRunpath() {
+  local libPath
+  local origRpath
+
+  if [[ $# -eq 0 ]]; then
+    echo "addCudaCompatRunpath: no library path provided" >&2
+    exit 1
+  elif [[ $# -gt 1 ]]; then
+    echo "addCudaCompatRunpath: too many arguments" >&2
+    exit 1
+  elif [[ "$1" == "" ]]; then
+    echo "addCudaCompatRunpath: empty library path" >&2
+    exit 1
+  else
+    libPath="$1"
+  fi
 
-autoAddCudaCompatRunpathPhase() (
-  local outputPaths
-  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
-  find "${outputPaths[@]}" -type f -print0  | while IFS= read -rd "" f; do
-    if isELF "$f"; then
-      # patchelf returns an error on statically linked ELF files
-      if elfHasDynamicSection "$f" ; then
-        echo "autoAddCudaCompatRunpathHook: patching $f"
-        local origRpath="$(patchelf --print-rpath "$f")"
-        patchelf --set-rpath "@libcudaPath@:$origRpath" "$f"
-      elif (( "${NIX_DEBUG:-0}" >= 1 )) ; then
-        echo "autoAddCudaCompatRunpathHook: skipping a statically-linked ELF file $f"
-      fi
-    fi
-  done
-)
+  origRpath="$(patchelf --print-rpath "$libPath")"
+  patchelf --set-rpath "@libcudaPath@:$origRpath" "$libPath"
+}
 
-postFixupHooks+=(autoAddCudaCompatRunpathPhase)
+postFixupHooks+=("autoFixElfFiles addCudaCompatRunpath")
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-driver-runpath-hook.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-driver-runpath-hook.sh
new file mode 100644
index 000000000000..ecff2a032d64
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-add-driver-runpath-hook.sh
@@ -0,0 +1,8 @@
+# shellcheck shell=bash
+# Run addDriverRunpath on all dynamically linked ELF files
+echo "Sourcing auto-add-driver-runpath-hook"
+
+if [ -z "${dontUseAutoAddDriverRunpath-}" ]; then
+  echo "Using autoAddDriverRunpath"
+  postFixupHooks+=("autoFixElfFiles addDriverRunpath")
+fi
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh b/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh
deleted file mode 100644
index a6eeef7c7699..000000000000
--- a/pkgs/development/cuda-modules/setup-hooks/auto-add-opengl-runpath-hook.sh
+++ /dev/null
@@ -1,28 +0,0 @@
-# shellcheck shell=bash
-# Run addOpenGLRunpath on all dynamically linked, ELF files
-echo "Sourcing auto-add-opengl-runpath-hook"
-
-elfHasDynamicSection() {
-    patchelf --print-rpath "$1" >& /dev/null
-}
-
-autoAddOpenGLRunpathPhase() (
-  local outputPaths
-  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
-  find "${outputPaths[@]}" -type f -print0  | while IFS= read -rd "" f; do
-    if isELF "$f"; then
-      # patchelf returns an error on statically linked ELF files
-      if elfHasDynamicSection "$f" ; then
-        echo "autoAddOpenGLRunpathHook: patching $f"
-        addOpenGLRunpath "$f"
-      elif (( "${NIX_DEBUG:-0}" >= 1 )) ; then
-        echo "autoAddOpenGLRunpathHook: skipping a statically-linked ELF file $f"
-      fi
-    fi
-  done
-)
-
-if [ -z "${dontUseAutoAddOpenGLRunpath-}" ]; then
-  echo "Using autoAddOpenGLRunpathPhase"
-  postFixupHooks+=(autoAddOpenGLRunpathPhase)
-fi
diff --git a/pkgs/development/cuda-modules/setup-hooks/auto-fix-elf-files.sh b/pkgs/development/cuda-modules/setup-hooks/auto-fix-elf-files.sh
new file mode 100644
index 000000000000..1d57dfb17a66
--- /dev/null
+++ b/pkgs/development/cuda-modules/setup-hooks/auto-fix-elf-files.sh
@@ -0,0 +1,64 @@
+# shellcheck shell=bash
+# List all dynamically linked ELF files in the outputs and apply a generic fix
+# action provided as a parameter (currently used to add the CUDA or the
+# cuda_compat driver to the runpath of binaries)
+echo "Sourcing cuda/fix-elf-files.sh"
+
+# Returns the exit code of patchelf --print-rpath.
+# A return code of 0 (success) means the ELF file has a dynamic section, while
+# a non-zero return code means the ELF file is statically linked (or is not an
+# ELF file).
+elfHasDynamicSection() {
+  local libPath
+
+  if [[ $# -eq 0 ]]; then
+    echo "elfHasDynamicSection: no library path provided" >&2
+    exit 1
+  elif [[ $# -gt 1 ]]; then
+    echo "elfHasDynamicSection: too many arguments" >&2
+    exit 1
+  elif [[ "$1" == "" ]]; then
+    echo "elfHasDynamicSection: empty library path" >&2
+    exit 1
+  else
+    libPath="$1"
+    shift 1
+  fi
+
+  patchelf --print-rpath "$libPath" >& /dev/null
+  return $?
+}
+
+# Run a fix action on all dynamically linked ELF files in the outputs.
+autoFixElfFiles() {
+  local fixAction
+  local outputPaths
+
+  if [[ $# -eq 0 ]]; then
+    echo "autoFixElfFiles: no fix action provided" >&2
+    exit 1
+  elif [[ $# -gt 1 ]]; then
+    echo "autoFixElfFiles: too many arguments" >&2
+    exit 1
+  elif [[ "$1" == "" ]]; then
+    echo "autoFixElfFiles: empty fix action" >&2
+    exit 1
+  else
+    fixAction="$1"
+  fi
+
+  mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "${!o}"; done)
+
+  find "${outputPaths[@]}" -type f -print0  | while IFS= read -rd "" f; do
+    if ! isELF "$f"; then
+      continue
+    elif elfHasDynamicSection "$f"; then
+      # patchelf returns an error on statically linked ELF files, and in
+      # practice fixing actions all involve patchelf
+      echo "autoFixElfFiles: using $fixAction to fix $f" >&2
+      $fixAction "$f"
+    elif (( "${NIX_DEBUG:-0}" >= 1 )); then
+      echo "autoFixElfFiles: skipping a statically-linked ELF file $f"
+    fi
+  done
+}
diff --git a/pkgs/development/cuda-modules/setup-hooks/extension.nix b/pkgs/development/cuda-modules/setup-hooks/extension.nix
index 9e352bd5b3af..ece70da52b02 100644
--- a/pkgs/development/cuda-modules/setup-hooks/extension.nix
+++ b/pkgs/development/cuda-modules/setup-hooks/extension.nix
@@ -1,4 +1,19 @@
 final: _: {
+  # Helper hook used in both autoAddCudaCompatRunpath and
+  # autoAddDriverRunpath that applies a generic patching action to all elf
+  # files with a dynamic linking section.
+  autoFixElfFiles =
+    final.callPackage
+      (
+        {makeSetupHook}:
+         makeSetupHook
+          {
+            name = "auto-fix-elf-files";
+          }
+          ./auto-fix-elf-files.sh
+      )
+      {};
+
   # Internal hook, used by cudatoolkit and cuda redist packages
   # to accommodate automatic CUDAToolkit_ROOT construction
   markForCudatoolkitRootHook =
@@ -32,31 +47,36 @@ final: _: {
       {}
     );
 
-  autoAddOpenGLRunpathHook =
+  autoAddDriverRunpath =
     final.callPackage
       (
-        {addOpenGLRunpath, makeSetupHook}:
+        {addDriverRunpath, autoFixElfFiles, makeSetupHook}:
         makeSetupHook
           {
             name = "auto-add-opengl-runpath-hook";
-            propagatedBuildInputs = [addOpenGLRunpath];
+            propagatedBuildInputs = [addDriverRunpath autoFixElfFiles];
           }
-          ./auto-add-opengl-runpath-hook.sh
+          ./auto-add-driver-runpath-hook.sh
       )
       {};
 
-  # autoAddCudaCompatRunpathHook hook must be added AFTER `setupCudaHook`. Both
+  # Deprecated: an alias kept for compatibility. Consider removing after 24.11
+  autoAddOpenGLRunpathHook = final.autoAddDriverRunpath;
+
+  # autoAddCudaCompatRunpath hook must be added AFTER `setupCudaHook`. Both
   # hooks prepend a path with `libcuda.so` to the `DT_RUNPATH` section of
   # patched elf files, but `cuda_compat` path must take precedence (otherwise,
   # it doesn't have any effect) and thus appear first. Meaning this hook must be
   # executed last.
-  autoAddCudaCompatRunpathHook =
+  autoAddCudaCompatRunpath =
     final.callPackage
       (
-        {makeSetupHook, cuda_compat ? null }:
+        {makeSetupHook, autoFixElfFiles, cuda_compat ? null }:
         makeSetupHook
           {
             name = "auto-add-cuda-compat-runpath-hook";
+            propagatedBuildInputs = [autoFixElfFiles];
+
             substitutions = {
               # Hotfix Ofborg evaluation
               libcudaPath = if final.flags.isJetsonBuild then "${cuda_compat}/compat" else null;
diff --git a/pkgs/development/libraries/ucx/default.nix b/pkgs/development/libraries/ucx/default.nix
index 3b923d8efdd2..8e0772479ccb 100644
--- a/pkgs/development/libraries/ucx/default.nix
+++ b/pkgs/development/libraries/ucx/default.nix
@@ -36,7 +36,7 @@ stdenv.mkDerivation rec {
   ]
   ++ lib.optionals enableCuda [
     cudaPackages.cuda_nvcc
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = [
diff --git a/pkgs/development/libraries/xgboost/default.nix b/pkgs/development/libraries/xgboost/default.nix
index b700dd2581c4..67190d7cac76 100644
--- a/pkgs/development/libraries/xgboost/default.nix
+++ b/pkgs/development/libraries/xgboost/default.nix
@@ -57,7 +57,7 @@ stdenv.mkDerivation rec {
 
   nativeBuildInputs = [ cmake ]
     ++ lib.optionals stdenv.isDarwin [ llvmPackages.openmp ]
-    ++ lib.optionals cudaSupport [ cudaPackages.autoAddOpenGLRunpathHook ]
+    ++ lib.optionals cudaSupport [ cudaPackages.autoAddDriverRunpath ]
     ++ lib.optionals rLibrary [ R ];
 
   buildInputs = [ gtest ] ++ lib.optional cudaSupport cudaPackages.cudatoolkit
diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix
index 199b352b2aa2..e2260aec4b14 100644
--- a/pkgs/development/python-modules/jaxlib/bin.nix
+++ b/pkgs/development/python-modules/jaxlib/bin.nix
@@ -23,7 +23,7 @@
 }:
 
 let
-  inherit (cudaPackagesGoogle) autoAddOpenGLRunpathHook cudaVersion;
+  inherit (cudaPackagesGoogle) autoAddDriverRunpath cudaVersion;
 
   version = "0.4.24";
 
@@ -180,7 +180,7 @@ buildPythonPackage {
   # Prebuilt wheels are dynamically linked against things that nix can't find.
   # Run `autoPatchelfHook` to automagically fix them.
   nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ]
-    ++ lib.optionals cudaSupport [ autoAddOpenGLRunpathHook ];
+    ++ lib.optionals cudaSupport [ autoAddDriverRunpath ];
   # Dynamic link dependencies
   buildInputs = [ stdenv.cc.cc.lib ];
 
diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix
index 657a6e52c084..cbce801888e2 100644
--- a/pkgs/development/python-modules/jaxlib/default.nix
+++ b/pkgs/development/python-modules/jaxlib/default.nix
@@ -51,7 +51,7 @@
 }@inputs:
 
 let
-  inherit (cudaPackagesGoogle) autoAddOpenGLRunpathHook cudaFlags cudaVersion cudnn nccl;
+  inherit (cudaPackagesGoogle) autoAddDriverRunpath cudaFlags cudaVersion cudnn nccl;
 
   pname = "jaxlib";
   version = "0.4.24";
@@ -420,7 +420,7 @@ buildPythonPackage {
     done
   '';
 
-  nativeBuildInputs = lib.optionals cudaSupport [ autoAddOpenGLRunpathHook ];
+  nativeBuildInputs = lib.optionals cudaSupport [ autoAddDriverRunpath ];
 
   propagatedBuildInputs = [
     absl-py
diff --git a/pkgs/development/python-modules/tensorrt/default.nix b/pkgs/development/python-modules/tensorrt/default.nix
index e108b1a773cb..58a3e64d5610 100644
--- a/pkgs/development/python-modules/tensorrt/default.nix
+++ b/pkgs/development/python-modules/tensorrt/default.nix
@@ -22,7 +22,7 @@ buildPythonPackage rec {
   nativeBuildInputs = [
     unzip
     autoPatchelfHook
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   preUnpack = ''
diff --git a/pkgs/development/python-modules/torch/bin.nix b/pkgs/development/python-modules/torch/bin.nix
index 0bb415574e39..bee32b616345 100644
--- a/pkgs/development/python-modules/torch/bin.nix
+++ b/pkgs/development/python-modules/torch/bin.nix
@@ -40,7 +40,7 @@ in buildPythonPackage {
   nativeBuildInputs = lib.optionals stdenv.isLinux [
     addOpenGLRunpath
     autoPatchelfHook
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = lib.optionals stdenv.isLinux (with cudaPackages; [
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index d6c51904bd9d..10eecd1de99b 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -338,7 +338,7 @@ in buildPythonPackage rec {
     pythonRelaxDepsHook
     removeReferencesTo
   ] ++ lib.optionals cudaSupport (with cudaPackages; [
-    autoAddOpenGLRunpathHook
+    autoAddDriverRunpath
     cuda_nvcc
   ])
   ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];
diff --git a/pkgs/os-specific/linux/dcgm/default.nix b/pkgs/os-specific/linux/dcgm/default.nix
index f3ebdf1427eb..cc9e26d2b707 100644
--- a/pkgs/os-specific/linux/dcgm/default.nix
+++ b/pkgs/os-specific/linux/dcgm/default.nix
@@ -105,10 +105,10 @@ in gcc11Stdenv.mkDerivation rec {
   strictDeps = true;
 
   nativeBuildInputs = [
-    # autoAddOpenGLRunpathHook does not actually depend on or incur any dependency
+    # autoAddDriverRunpath does not actually depend on or incur any dependency
     # of cudaPackages. It merely adds an impure, non-Nix PATH to the RPATHs of
     # executables that need to use cuda at runtime.
-    cudaPackages_12.autoAddOpenGLRunpathHook
+    cudaPackages_12.autoAddDriverRunpath
 
     cmake
     git
diff --git a/pkgs/servers/monitoring/prometheus/dcgm-exporter/default.nix b/pkgs/servers/monitoring/prometheus/dcgm-exporter/default.nix
index c995d4036ae0..7c4d2cfae06e 100644
--- a/pkgs/servers/monitoring/prometheus/dcgm-exporter/default.nix
+++ b/pkgs/servers/monitoring/prometheus/dcgm-exporter/default.nix
@@ -48,7 +48,7 @@ buildGoModule rec {
   vendorHash = "sha256-Fjvx15e/psxoqoS6c6GhiQfe7g2aI40EmPR26xLhrzg=";
 
   nativeBuildInputs = [
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   # Tests try to interact with running DCGM service.
diff --git a/pkgs/servers/sunshine/default.nix b/pkgs/servers/sunshine/default.nix
index aa7b079fea19..5185553262c3 100644
--- a/pkgs/servers/sunshine/default.nix
+++ b/pkgs/servers/sunshine/default.nix
@@ -94,7 +94,7 @@ stdenv'.mkDerivation rec {
     autoPatchelfHook
     makeWrapper
   ] ++ lib.optionals cudaSupport [
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = [
diff --git a/pkgs/tools/audio/openai-whisper-cpp/default.nix b/pkgs/tools/audio/openai-whisper-cpp/default.nix
index 191aae946e21..20d2a23d1dd5 100644
--- a/pkgs/tools/audio/openai-whisper-cpp/default.nix
+++ b/pkgs/tools/audio/openai-whisper-cpp/default.nix
@@ -45,9 +45,7 @@ effectiveStdenv.mkDerivation (finalAttrs: {
     ] ++ lib.optionals cudaSupport ( with cudaPackages ;[
       cuda_nvcc
 
-      # TODO: Replace with autoAddDriverRunpath
-      # once https://github.com/NixOS/nixpkgs/pull/275241 has been merged
-      autoAddOpenGLRunpathHook
+      autoAddDriverRunpath
     ]);
 
   buildInputs = [
diff --git a/pkgs/tools/system/btop/default.nix b/pkgs/tools/system/btop/default.nix
index 9b6956fb5780..5961a0fed754 100644
--- a/pkgs/tools/system/btop/default.nix
+++ b/pkgs/tools/system/btop/default.nix
@@ -23,7 +23,7 @@ stdenv.mkDerivation rec {
   };
 
   nativeBuildInputs = [ cmake ] ++ lib.optionals cudaSupport [
-    cudaPackages.autoAddOpenGLRunpathHook
+    cudaPackages.autoAddDriverRunpath
   ];
 
   buildInputs = lib.optionals stdenv.isDarwin [