From 27513ed374ff36b7039ae7d594c85f28f49962ac Mon Sep 17 00:00:00 2001
From: David Wilemski <david@wilemski.org>
Date: Sat, 7 Jan 2023 14:45:03 +0000
Subject: [PATCH 1/2] pynvml: Add patch for finding libnvidia-ml.so.1 on NixOS

This fixes loading the shared library used for interacting with Nvidia
hardware. A similar package, nvidia-ml-py, already provides a patch for
the same thing. This commit adds a similar patch for pynvml. We could
not reuse the existing patch for the other package because the files
being patched have different names despite the patch itself being nearly
the same.

Fixes https://github.com/NixOS/nixpkgs/issues/209416
---
 ...0001-locate-libnvidia-ml.so.1-on-NixOS.patch | 17 +++++++++++++++++
 .../python-modules/pynvml/default.nix           |  4 ++++
 2 files changed, 21 insertions(+)
 create mode 100644 pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch

diff --git a/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch b/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch
new file mode 100644
index 000000000000..ada9eecf8db6
--- /dev/null
+++ b/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch
@@ -0,0 +1,17 @@
+diff --git a/pynvml/nvml.py b/pynvml/nvml.py
+index 56d908f..1de0b97 100644
+--- a/pynvml/nvml.py
++++ b/pynvml/nvml.py
+@@ -1475,7 +1475,11 @@ def _LoadNvmlLibrary():
+                             nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll"))
+                     else:
+                         # assume linux
+-                        nvmlLib = CDLL("libnvidia-ml.so.1")
++                        try:
++                            nvmlLib = CDLL("libnvidia-ml.so.1")
++                        except OSError:
++                            # assume NixOS
++                            nvmlLib = CDLL("/run/opengl-driver/lib/libnvidia-ml.so.1")
+                 except OSError as ose:
+                     _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
+                 if (nvmlLib == None):
diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
index 63d445e0923f..d67208490ac7 100644
--- a/pkgs/development/python-modules/pynvml/default.nix
+++ b/pkgs/development/python-modules/pynvml/default.nix
@@ -15,6 +15,10 @@ buildPythonPackage rec {
     sha256 = "b2e4a33b80569d093b513f5804db0c7f40cfc86f15a013ae7a8e99c5e175d5dd";
   };
 
+  patches = [
+    ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch
+  ];
+
   propagatedBuildInputs = [ cudatoolkit ];
 
   doCheck = false;  # no tests in PyPi dist

From 7650f5f86f20fb37ac1fe4349ae72154d1e19638 Mon Sep 17 00:00:00 2001
From: David Wilemski <david@wilemski.org>
Date: Sat, 7 Jan 2023 16:44:45 +0000
Subject: [PATCH 2/2] template patch with addOpenGLRunpath

---
 .../pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch    | 2 +-
 pkgs/development/python-modules/pynvml/default.nix         | 7 ++++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch b/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch
index ada9eecf8db6..48727de2516a 100644
--- a/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch
+++ b/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch
@@ -11,7 +11,7 @@ index 56d908f..1de0b97 100644
 +                            nvmlLib = CDLL("libnvidia-ml.so.1")
 +                        except OSError:
 +                            # assume NixOS
-+                            nvmlLib = CDLL("/run/opengl-driver/lib/libnvidia-ml.so.1")
++                            nvmlLib = CDLL("@driverLink@/lib/libnvidia-ml.so.1")
                  except OSError as ose:
                      _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
                  if (nvmlLib == None):
diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
index d67208490ac7..8f02bcdbe816 100644
--- a/pkgs/development/python-modules/pynvml/default.nix
+++ b/pkgs/development/python-modules/pynvml/default.nix
@@ -1,8 +1,10 @@
 { lib
 , buildPythonPackage
 , fetchPypi
+, substituteAll
 , pythonOlder
 , cudatoolkit
+, addOpenGLRunpath
 }:
 
 buildPythonPackage rec {
@@ -16,7 +18,10 @@ buildPythonPackage rec {
   };
 
   patches = [
-    ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch
+    (substituteAll {
+      src = ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch;
+      inherit (addOpenGLRunpath) driverLink;
+    })
   ];
 
   propagatedBuildInputs = [ cudatoolkit ];