From 30abe004b9062c08c3ad67141628026fb7d08cfc Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 2 Oct 2025 20:43:02 +0000
Subject: [PATCH 1/2] Initial plan


From 1f112a160be7a3b553e5273d0573a062e2401dab Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 2 Oct 2025 20:51:02 +0000
Subject: [PATCH 2/2] Replace pynvml with nvidia-ml-py in dependencies and
 error messages

Co-authored-by: AlexeyKozhevin <19620584+AlexeyKozhevin@users.noreply.github.com>
---
 nbtools/core.py | 4 ++--
 setup.py        | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/nbtools/core.py b/nbtools/core.py
index 0a3af42..90faa6f 100755
--- a/nbtools/core.py
+++ b/nbtools/core.py
@@ -155,7 +155,7 @@ def get_available_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False,
     try:
         import pynvml
     except ImportError as exception:
-        raise ImportError('Install Python interface for nvidia_smi') from exception
+        raise ImportError('Install nvidia-ml-py package') from exception
 
     try:
         error_message = None
@@ -220,7 +220,7 @@ def get_gpu_free_memory(index, ratio=True):
     try:
         import pynvml
     except ImportError as exception:
-        raise ImportError('Install Python interface for nvidia_smi') from exception
+        raise ImportError('Install nvidia-ml-py package') from exception
 
     pynvml.nvmlInit()
     pynvml.nvmlDeviceGetCount()
diff --git a/setup.py b/setup.py
index 224f0e7..d0782f3 100644
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,7 @@
     zip_safe=False,
     platforms='any',
     install_requires=[
-        'pynvml>=11.5.0',
+        'nvidia-ml-py>=12.0.0',
         'blessed>=1.17',
         'psutil>=5.6',
         'requests>=2.24',