Merge pull request #90 from 907Resident/dev

GPU Info

Author: rasbt

README.md

@@ -176,6 +176,9 @@ In line with [NEP 29][nep-29], this project supports:
 
 [[top](#sections)]
 
+#### v. 2.4.0 (May 23, 2023)
+
+- Adds a new `--gpu` flag to print out GPU information (currently limited to NVIDIA devices) ([#90](https://github.com/rasbt/watermark/pull/63), via contribution by [907Resident](https://github.com/907Resident))
 
 
 #### v. 2.3.1 (May 27, 2022)

appveyor.yml

@@ -18,6 +18,7 @@ install:
 
 test_script:
   - set PYTHONPATH=%PYTHONPATH%;%CD%
+  - pip install -e .
   - pytest -sv
 
 notifications:

binder/requirements.txt

@@ -3,5 +3,4 @@ numpy
 scipy
 scikit-learn
 jupyter
-
--e .
+py3nvml

requirements.txt

@@ -0,0 +1,3 @@
+ipython >= 6.0
+importlib-metadata >= 1.4
+py3nvml >= 0.2

setup.cfg

@@ -1,5 +1,5 @@
 [metadata]
-version = 2.3.1
+version = 2.4.0
 license_file = LICENSE
 classifiers =
     Development Status :: 5 - Production/Stable

setup.py

@@ -5,10 +5,20 @@
 #
 # License: BSD 3 clause
 
+from os.path import dirname, join, realpath
 from textwrap import dedent
 
 from setuptools import find_packages, setup
 
+
+PROJECT_ROOT = dirname(realpath(__file__))
+REQUIREMENTS_FILE = join(PROJECT_ROOT, "requirements.txt")
+
+with open(REQUIREMENTS_FILE) as f:
+    install_reqs = f.read().splitlines()
+
+install_reqs.append("setuptools")
+
 # Also see settings in setup.cfg
 setup(
     name="watermark",
@@ -21,10 +31,7 @@
     author_email="[email protected]",
     url="https://github.com/rasbt/watermark",
     packages=find_packages(exclude=[]),
-    install_requires=[
-        "ipython",
-        'importlib-metadata >= 1.4 ; python_version < "3.8"',
-    ],
+    install_requires=install_reqs,
     long_description=dedent(
         """\
         An IPython magic extension for printing date and time stamps, version

watermark/magic.py

@@ -70,6 +70,9 @@ class WaterMark(Magics):
               help='prints the current version of watermark')
     @argument('-iv', '--iversions', action='store_true',
               help='prints the name/version of all imported modules')
+    @argument('--gpu', action='store_true',
+              help='prints GPU information (currently limited to NVIDIA GPUs),'
+                   ' if available')
     @line_magic
     def watermark(self, line):
         """

watermark/tests/test_watermark.py

@@ -1,5 +1,10 @@
 # -*- coding: utf-8 -*-
 
+import sys
+import os
+
+sys.path.append(os.path.join("../watermark"))
+
 import watermark
 
 

watermark/tests/test_watermark_gpu.py

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+import sys
+import os
+
+sys.path.append(os.path.join("../watermark"))
+
+import watermark
+
+def test_gpu_info():
+    a = watermark.watermark(gpu=True)
+    txt = a.split('\n')
+    clean_txt = []
+    for t in txt:
+        t = t.strip()
+        if t:
+            t = t.split(':')[0]
+            clean_txt.append(t.strip())
+    clean_txt = set(clean_txt)
+
+    expected = [
+        'GPU Info',
+    ]
+
+    for i in expected:
+        assert i in clean_txt, print(f'{i} not in {clean_txt}')

watermark/watermark.py

@@ -18,6 +18,8 @@
 import types
 from multiprocessing import cpu_count
 from socket import gethostname
+import platform
+from py3nvml import py3nvml
 
 try:
     import importlib.metadata as importlib_metadata
@@ -30,14 +32,32 @@
 from .version import __version__
 
 
-def watermark(author=None, email=None, github_username=None,
-              website=None, current_date=False, datename=False,
-              current_time=False, iso8601=False, timezone=False,
-              updated=False, custom_time=None, python=False,
-              packages=None, conda=False, hostname=False, machine=False,
-              githash=False, gitrepo=False, gitbranch=False,
-              watermark=False, iversions=False, watermark_self=None,
-              globals_=None):
+def watermark(
+        author=None, 
+        email=None, 
+        github_username=None,
+        website=None, 
+        current_date=False, 
+        datename=False,
+        current_time=False, 
+        iso8601=False, 
+        timezone=False,
+        updated=False, 
+        custom_time=None, 
+        python=False,
+        packages=None, 
+        conda=False, 
+        hostname=False, 
+        machine=False,
+        githash=False, 
+        gitrepo=False, 
+        gitbranch=False,
+        watermark=False, 
+        iversions=False, 
+        gpu=False,
+        watermark_self=None,
+        globals_=None
+):
 
     '''Function to print date/time stamps and various system information.
 
@@ -107,6 +127,9 @@ def watermark(author=None, email=None, github_username=None,
 
     iversions :
         prints the name/version of all imported modules
+    
+    gpu :
+        prints GPU information (currently limited to NVIDIA GPUs), if available
 
     watermark_self :
         instance of the watermark magics class, which is required
@@ -182,6 +205,8 @@ def watermark(author=None, email=None, github_username=None,
                     "to show imported package versions."
                 )
             output.append(_get_all_import_versions(ns))
+        if args['gpu']:
+            output.append(_get_gpu_info())
         if args['watermark']:
             output.append({"Watermark": __version__})
 
@@ -306,3 +331,23 @@ def _get_all_import_versions(vars):
 def _get_conda_env():
     name = os.getenv('CONDA_DEFAULT_ENV', 'n/a')
     return {"conda environment": name}
+
+
+def _get_gpu_info():
+    try:
+        gpu_info = [""]
+        py3nvml.nvmlInit()
+        num_gpus = py3nvml.nvmlDeviceGetCount()
+        for i in range(num_gpus):
+            handle = py3nvml.nvmlDeviceGetHandleByIndex(i)
+            gpu_name = py3nvml.nvmlDeviceGetName(handle)
+            gpu_info.append(f"GPU {i}: {gpu_name}")
+        py3nvml.nvmlShutdown()
+        return {"GPU Info": "\n  ".join(gpu_info)}
+
+    except py3nvml.NVMLError_LibraryNotFound:
+        return {"GPU Info": "NVIDIA drivers do not appear "
+                "to be installed on this machine."}
+    except:
+        return {"GPU Info": "GPU information is not "
+                "available for this machine."}