From 4d8cc5512d84fb5a2085e5d1328c316d407fd93b Mon Sep 17 00:00:00 2001
From: Ben Darwin <bcdarwin@gmail.com>
Date: Mon, 20 Jan 2025 20:18:24 -0500
Subject: [PATCH 1/3] python312Packages.pynvml: 11.5.3 -> 12.0.0

---
 ...01-locate-libnvidia-ml.so.1-on-NixOS.patch | 17 ---------
 .../python-modules/pynvml/default.nix         | 37 ++++++-------------
 2 files changed, 11 insertions(+), 43 deletions(-)
 delete 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
deleted file mode 100644
index 48727de2516a..000000000000
--- a/pkgs/development/python-modules/pynvml/0001-locate-libnvidia-ml.so.1-on-NixOS.patch
+++ /dev/null
@@ -1,17 +0,0 @@
-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):
diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
index ea8c12ad5528..7de3a041ef1e 100644
--- a/pkgs/development/python-modules/pynvml/default.nix
+++ b/pkgs/development/python-modules/pynvml/default.nix
@@ -3,63 +3,48 @@
   buildPythonPackage,
   cudaPackages,
   fetchFromGitHub,
-  substituteAll,
-  addDriverRunpath,
   setuptools,
   pytestCheckHook,
-  versioneer,
+  nvidia-ml-py,
   pynvml,
 }:
 
 buildPythonPackage rec {
   pname = "pynvml";
-  version = "11.5.3";
+  version = "12.0.0";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "gpuopenanalytics";
     repo = "pynvml";
     tag = version;
-    hash = "sha256-8NkYBRpcW3dvxVc6z17TMRPqA0YK/J/CdjuqgdcTdy8=";
+    hash = "sha256-bfES6QqL9DO7rmQ3btXxVzW0KlL0eHyUvYJVeijymBk=";
   };
 
-  patches = [
-    (substituteAll {
-      src = ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch;
-      inherit (addDriverRunpath) driverLink;
-    })
-  ];
-
-  # unvendor versioneer
-  postPatch = ''
-    rm versioneer.py
-  '';
-
   build-system = [
     setuptools
-    versioneer
   ];
 
+  dependencies = [ nvidia-ml-py ];
+
   pythonImportsCheck = [
-    "pynvml"
-    "pynvml.smi"
+    "pynvml_utils"
   ];
 
   nativeCheckInputs = [ pytestCheckHook ];
 
-  # OSError: /run/opengl-driver/lib/libnvidia-ml.so.1: cannot open shared object file: No such file or directory
   doCheck = false;
 
   passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ pynvml ]; } ''
-    import pynvml
-    from pynvml.smi import nvidia_smi  # noqa: F401
-
-    print(f"{pynvml.nvmlInit()=}")
+    from pynvml_utils import nvidia_smi  # noqa: F401
+    nvsmi = nvidia_smi.getInstance()
+    print(nvsmi.DeviceQuery('memory.free, memory.total'))
   '';
 
   meta = {
-    description = "Python bindings for the NVIDIA Management Library";
+    description = "Unofficial Python bindings for the NVIDIA Management Library";
     homepage = "https://github.com/gpuopenanalytics/pynvml";
+    changelog = "https://github.com/gpuopenanalytics/pynvml?tab=readme-ov-file#release-notes";
     license = lib.licenses.bsd3;
     maintainers = [ lib.maintainers.bcdarwin ];
   };

From 2f3b52c51c45fea5a7218d0cbe0ad513002f0f03 Mon Sep 17 00:00:00 2001
From: Ben Darwin <bcdarwin@gmail.com>
Date: Mon, 20 Jan 2025 20:26:56 -0500
Subject: [PATCH 2/3] python312Packages.scalene: drop obsolete pynvml
 dependency

---
 pkgs/development/python-modules/scalene/default.nix | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pkgs/development/python-modules/scalene/default.nix b/pkgs/development/python-modules/scalene/default.nix
index b2fc07bbcebd..b03c0f041e2f 100644
--- a/pkgs/development/python-modules/scalene/default.nix
+++ b/pkgs/development/python-modules/scalene/default.nix
@@ -11,7 +11,6 @@
   nvidia-ml-py,
   psutil,
   pydantic,
-  pynvml,
   pytestCheckHook,
   pythonOlder,
   rich,
@@ -75,7 +74,6 @@ buildPythonPackage rec {
     numpy
     psutil
     pydantic
-    pynvml
     rich
   ] ++ lib.optionals stdenv.hostPlatform.isLinux [ nvidia-ml-py ];
 

From 423e6480e29937b814336dbdc0c340229414827e Mon Sep 17 00:00:00 2001
From: Ben Darwin <bcdarwin@gmail.com>
Date: Mon, 20 Jan 2025 20:29:10 -0500
Subject: [PATCH 3/3] python312Packages.vllm: remove obsolete pynvml dependency

---
 pkgs/development/python-modules/vllm/default.nix | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pkgs/development/python-modules/vllm/default.nix b/pkgs/development/python-modules/vllm/default.nix
index 0d554e5a5491..4231983b0c5b 100644
--- a/pkgs/development/python-modules/vllm/default.nix
+++ b/pkgs/development/python-modules/vllm/default.nix
@@ -25,7 +25,6 @@
   uvicorn,
   pydantic,
   aioprometheus,
-  pynvml,
   openai,
   pyzmq,
   tiktoken,
@@ -149,7 +148,6 @@ buildPythonPackage rec {
     ++ aioprometheus.optional-dependencies.starlette
     ++ lib.optionals cudaSupport [
       cupy
-      pynvml
     ];
 
   dontUseCmakeConfigure = true;