Benchmarks: Add Feature - Add nvml package to provide python interfaces of nvidia. (#91)

.azure-pipelines/cuda-unit-test.yml

@@ -16,7 +16,7 @@ steps:
       echo "##vso[task.prependpath]$HOME/.local/bin"
     displayName: Export path
   - script: |
-      python3 -m pip install .[test,torch]
+      python3 -m pip install .[test,nvidia,torch]
       make postinstall
     displayName: Install dependencies
   - script: |

dockerfile/cuda11.1.1.dockerfile

@@ -88,6 +88,6 @@ WORKDIR ${SB_HOME}
 ADD . .
 
 RUN cd ${SB_HOME} && \
-    python3 -m pip install .[torch] && \
+    python3 -m pip install .[nvidia,torch] && \
     make cppbuild && \
     make thirdparty
\ No newline at end of file

setup.py

@@ -158,6 +158,7 @@ def run(self):
             'torchvision>=0.8.0',
             'transformers>=4.3.3',
         ],
+        'nvidia': ['py3nvml>=0.2.6']
     },
     include_package_data=True,
     entry_points={

superbench/common/utils/__init__.py

@@ -7,4 +7,6 @@
 from superbench.common.utils.file_handler import create_output_dir, get_sb_config
 from superbench.common.utils.lazy_import import LazyImport
 
-__all__ = ['SuperBenchLogger', 'logger', 'create_output_dir', 'get_sb_config', 'LazyImport']
+nv_helper = LazyImport('superbench.common.utils.nvidia_helper')
+
+__all__ = ['SuperBenchLogger', 'logger', 'create_output_dir', 'get_sb_config', 'LazyImport', 'nv_helper']

superbench/common/utils/nvidia_helper.py

+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Nvidia Utility."""
+
+import py3nvml.py3nvml as nvml
+
+
+def get_device_compute_capability():
+    """Get the compute capability of device.
+
+    Return:
+        cap (float): the compute capability of device, None means no device found.
+    """
+    nvml.nvmlInit()
+    device_count = nvml.nvmlDeviceGetCount()
+    if device_count == 0:
+        return None
+
+    handle = nvml.nvmlDeviceGetHandleByIndex(0)
+    cap = nvml.nvmlDeviceGetCudaComputeCapability(handle)
+    return cap