push sheeet

pkgs/development/python-modules/nvidia-ml-py/0001-locate-libnvidia-ml.so.1-on-NixOS.patch

@@ -0,0 +1,17 @@
+diff --git a/pynvml.py b/pynvml.py
+index 9a424de..669afe0 100644
+--- a/pynvml.py
++++ b/pynvml.py
+@@ -1676,7 +1676,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/nvidia-ml-py/default.nix

@@ -0,0 +1,64 @@
+{
+  lib,
+  fetchPypi,
+  buildPythonPackage,
+  replaceVars,
+  addDriverRunpath,
+  setuptools,
+  cudaPackages,
+  nvidia-ml-py,
+}:
+
+buildPythonPackage rec {
+  pname = "nvidia-ml-py";
+  version = "13.580.82";
+
+  pyproject = true;
+
+  src = fetchPypi {
+    pname = "nvidia_ml_py";
+    inherit version;
+    hash = "sha256-DAKIBdxToOKmmF6oAYiBl3Zawu+PHJ4pp78NNhal78c=";
+  };
+
+  patches = [
+    (replaceVars ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch {
+      inherit (addDriverRunpath) driverLink;
+    })
+  ];
+
+  build-system = [
+    setuptools
+  ];
+
+  # no tests
+  doCheck = false;
+
+  pythonImportsCheck = [ "pynvml" ];
+
+  passthru.tests.tester-nvmlInit =
+    cudaPackages.writeGpuTestPython { libraries = [ nvidia-ml-py ]; }
+      ''
+        from pynvml import (
+          nvmlInit,
+          nvmlSystemGetDriverVersion,
+          nvmlDeviceGetCount,
+          nvmlDeviceGetHandleByIndex,
+          nvmlDeviceGetName,
+        )
+
+        nvmlInit()
+        print(f"Driver Version: {nvmlSystemGetDriverVersion()}")
+
+        for i in range(nvmlDeviceGetCount()):
+            handle = nvmlDeviceGetHandleByIndex(i)
+            print(f"Device {i} : {nvmlDeviceGetName(handle)}")
+      '';
+
+  meta = {
+    description = "Python Bindings for the NVIDIA Management Library";
+    homepage = "https://pypi.org/project/nvidia-ml-py";
+    license = lib.licenses.bsd3;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}