Merge pull request #209505 from davidwilemski/pynvml-py3

pynvml: Add patch for finding libnvidia-ml.so.1 on NixOS

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("@driverLink@/lib/libnvidia-ml.so.1")
+                 except OSError as ose:
+                     _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
+                 if (nvmlLib == None):

pkgs/development/python-modules/pynvml/default.nix

@@ -1,8 +1,10 @@
 { lib
 , buildPythonPackage
 , fetchPypi
+, substituteAll
 , pythonOlder
 , cudatoolkit
+, addOpenGLRunpath
 }:
 
 buildPythonPackage rec {
@@ -15,6 +17,13 @@ buildPythonPackage rec {
     sha256 = "b2e4a33b80569d093b513f5804db0c7f40cfc86f15a013ae7a8e99c5e175d5dd";
   };
 
+  patches = [
+    (substituteAll {
+      src = ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch;
+      inherit (addOpenGLRunpath) driverLink;
+    })
+  ];
+
   propagatedBuildInputs = [ cudatoolkit ];
 
   doCheck = false;  # no tests in PyPi dist