From 92c03649e397c2a0b1f9e12e472f9573fbd5ca64 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 16 Oct 2024 11:56:48 +0200
Subject: [PATCH] python312Packages.nvidia-ml-py: fix and improve test

---
 .../python-modules/nvidia-ml-py/default.nix      | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/nvidia-ml-py/default.nix b/pkgs/development/python-modules/nvidia-ml-py/default.nix
index 95604dcb20d7..08beac0daf50 100644
--- a/pkgs/development/python-modules/nvidia-ml-py/default.nix
+++ b/pkgs/development/python-modules/nvidia-ml-py/default.nix
@@ -40,10 +40,20 @@ buildPythonPackage rec {
   passthru.tests.tester-nvmlInit =
     cudaPackages.writeGpuTestPython { libraries = [ nvidia-ml-py ]; }
       ''
-        import pynvml
-        from pynvml.smi import nvidia_smi  # noqa: F401
+        from pynvml import (
+          nvmlInit,
+          nvmlSystemGetDriverVersion,
+          nvmlDeviceGetCount,
+          nvmlDeviceGetHandleByIndex,
+          nvmlDeviceGetName,
+        )
 
-        print(f"{pynvml.nvmlInit()=}")
+        nvmlInit()
+        print(f"Driver Version: {nvmlSystemGetDriverVersion()}")
+
+        for i in range(nvmlDeviceGetCount()):
+            handle = nvmlDeviceGetHandleByIndex(i)
+            print(f"Device {i} : {nvmlDeviceGetName(handle)}")
       '';
 
   meta = {