Repository: gpuopenanalytics/pynvml (language: python)

Issue:
Update nvml.py to CUDA 11; smi.py DeviceQuery to R460

Update to sync with latest nvidia-ml-py deployment



Discussion / hints:
Thanks @kenhester!

In our offline discussion earlier, we were trying to establish if the smi python bindings were being used in downstream projects like NVDashboard, dask/distributed, and dask-CUDA.  As far as I can tell, these three projects are using the nvml bindings directly (not smi).

Those familiar with nvml usage in dask-CUDA and dask dashboard (@jacobtomlinson , @pentschev, @quasiben, @charlesbluca), pease do correct me if I am mistaken here.

The smi bindings are awesome to have, but if they are not a requirement for these projects, we can treat smi-feature completeness as a "nice-to-have" priority for now.  With that said, it is clear that NVDashboard and Dask-Dashboard should use the same handle-caching approach you are using in smi (it seems that both dashboard projects are currently defining fresh device handles for every query)


Fix patch:
diff --git a/pynvml/nvml.py b/pynvml/nvml.py
index 419929f..2bdbb1a 100644
--- a/pynvml/nvml.py
+++ b/pynvml/nvml.py
@@ -319,7 +319,7 @@
 NVML_TOPOLOGY_SYSTEM = 50
 
 _nvmlGpuP2PCapsIndex_t = c_uint
-NVML_P2P_CAPS_INDEX_READ = 0,
+NVML_P2P_CAPS_INDEX_READ = 0
 NVML_P2P_CAPS_INDEX_WRITE = 1
 NVML_P2P_CAPS_INDEX_NVLINK =2
 NVML_P2P_CAPS_INDEX_ATOMICS = 3


Test patch:
diff --git a/pynvml/tests/test_nvml.py b/pynvml/tests/test_nvml.py
index c9e35a9..529af40 100644
--- a/pynvml/tests/test_nvml.py
+++ b/pynvml/tests/test_nvml.py
@@ -8,47 +8,79 @@
 NVML_PCIE_UTIL_RX_BYTES = pynvml.NVML_PCIE_UTIL_RX_BYTES
 NVML_PCIE_UTIL_COUNT = pynvml.NVML_PCIE_UTIL_COUNT
 
+XFAIL_LEGACY_NVLINK_MSG = "Legacy NVLink test expected to fail."
+
 # Fixture to initialize and finalize nvml
-@pytest.fixture(scope='module')
+@pytest.fixture(scope="module")
 def nvml(request):
     pynvml.nvmlInit()
+
     def nvml_close():
         pynvml.nvmlShutdown()
+
     request.addfinalizer(nvml_close)
 
+
+# Fixture to check driver version
+@pytest.fixture(scope="module")
+def driver(request):
+    return float(pynvml.nvmlSystemGetDriverVersion())
+
+
 # Get GPU count
 @pytest.fixture
 def ngpus(nvml):
     result = pynvml.nvmlDeviceGetCount()
     assert result > 0
-    print('['+str(result)+' GPUs]', end =' ')
+    print("[" + str(result) + " GPUs]", end=" ")
     return result
 
+
 # Get handles using pynvml.nvmlDeviceGetHandleByIndex
 @pytest.fixture
 def handles(ngpus):
-    handles = [ pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(ngpus) ]
+    handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(ngpus)]
     assert len(handles) == ngpus
     return handles
 
+
+# Get GPU MIG count
+@pytest.fixture
+def nmigs(handles):
+    result = pynvml.nvmlDeviceGetMaxMigDeviceCount(handles[0])
+    print("[" + str(result) + " MIGs]", end=" ")
+    return result
+
+
+# Get MIG device handles
+@pytest.fixture
+def mig_handles(nmigs):
+    handles = [pynvml.nvmlDeviceGetMigDeviceHandleByIndex(i) for i in range(nmigs)]
+    assert len(handles) == nmigs
+    return handles
+
+
 @pytest.fixture
 def serials(ngpus, handles):
-    serials = [ pynvml.nvmlDeviceGetSerial(handles[i]) for i in range(ngpus) ]
+    serials = [pynvml.nvmlDeviceGetSerial(handles[i]) for i in range(ngpus)]
     assert len(serials) == ngpus
     return serials
 
+
 @pytest.fixture
 def uuids(ngpus, handles):
-    uuids = [ pynvml.nvmlDeviceGetUUID(handles[i]) for i in range(ngpus) ]
+    uuids = [pynvml.nvmlDeviceGetUUID(handles[i]) for i in range(ngpus)]
     assert len(uuids) == ngpus
     return uuids
 
+
 @pytest.fixture
 def pci_info(ngpus, handles):
-    pci_info = [ pynvml.nvmlDeviceGetPciInfo(handles[i]) for i in range(ngpus) ]
+    pci_info = [pynvml.nvmlDeviceGetPciInfo(handles[i]) for i in range(ngpus)]
     assert len(pci_info) == ngpus
     return pci_info
 
+
 ## ---------------------------- ##
 ## Start Fucntion-pointer Tests ##
 ## ---------------------------- ##
@@ -57,87 +89,109 @@ def pci_info(ngpus, handles):
 def test_nvmlSystemGetNVMLVersion(nvml):
     vsn = 0.0
     vsn = pynvml.nvmlSystemGetNVMLVersion().decode()
-    print('[NVML Version: '+vsn+']', end =' ')
+    print("[NVML Version: " + vsn + "]", end=" ")
     assert vsn > LooseVersion("0.0")
 
-def test_nvmlSystemGetCudaDriverVersion():
-    vsn = 0;
-    vsn = nvml.nvmlSystemGetCudaDriverVersion()
-    assert vsn != 0
+
+def test_nvmlSystemGetCudaDriverVersion(nvml):
+    vsn = 0.0
+    vsn = pynvml.nvmlSystemGetCudaDriverVersion()
+    assert vsn != 0.0
+
 
 # Test pynvml.nvmlSystemGetProcessName
 def test_nvmlSystemGetProcessName(nvml):
     procname = None
     procname = pynvml.nvmlSystemGetProcessName(os.getpid())
-    print('[Process: '+str(procname.decode())+']', end =' ')
+    print("[Process: " + str(procname.decode()) + "]", end=" ")
     assert procname != None
 
+
 # Test pynvml.nvmlSystemGetDriverVersion
 def test_nvmlSystemGetDriverVersion(nvml):
     vsn = 0.0
     vsn = pynvml.nvmlSystemGetDriverVersion().decode()
-    print('[Driver Version: '+vsn+']', end =' ')
-    assert vsn > LooseVersion("0.0") # Developing with 396.44
+    print("[Driver Version: " + vsn + "]", end=" ")
+    assert vsn > LooseVersion("0.0")  # Developing with 396.44
+
 
 ## Unit "Get" Functions (Skipping for now) ##
 
 ## Device "Get" Functions (Note: Some functions are fixtures, above) ##
 
-def test_nvmlDeviceGetAttributes(handles):
-    for handle in handles:
-        att = pynvml.nvmlDeviceGetAttributes(handle)
-        assert att is not None
+
+def test_nvmlDeviceGetAttributes(mig_handles):
+    # nvmlDeviceGetAttributes requires MIG device handle
+
+    if mig_handles:
+        for handle in mig_handles:
+            att = pynvml.nvmlDeviceGetAttributes(handle)
+            assert att is not None
+    # No MIG devices, Skip
+    pytest.skip()
+
 
 # Test pynvml.nvmlDeviceGetHandleBySerial
 def test_nvmlDeviceGetHandleBySerial(ngpus, serials):
-    handles = [ pynvml.nvmlDeviceGetHandleBySerial(serials[i])
-                for i in range(ngpus) ]
+    handles = [pynvml.nvmlDeviceGetHandleBySerial(serials[i]) for i in range(ngpus)]
     assert len(handles) == ngpus
 
+
 # Test pynvml.nvmlDeviceGetHandleByUUID
 def test_nvmlDeviceGetHandleByUUID(ngpus, uuids):
-    handles = [ pynvml.nvmlDeviceGetHandleByUUID(uuids[i])
-                for i in range(ngpus) ]
+    handles = [pynvml.nvmlDeviceGetHandleByUUID(uuids[i]) for i in range(ngpus)]
     assert len(handles) == ngpus
 
+
 # Test pynvml.nvmlDeviceGetHandleByPciBusId
 def test_nvmlDeviceGetHandleByPciBusId(ngpus, pci_info):
-    handles = [ pynvml.nvmlDeviceGetHandleByPciBusId(pci_info[i].busId)
-                for i in range(ngpus) ]
+    handles = [
+        pynvml.nvmlDeviceGetHandleByPciBusId(pci_info[i].busId) for i in range(ngpus)
+    ]
     assert len(handles) == ngpus
 
-@pytest.mark.parametrize("scope", [
-    pynvml.NVML_AFFINITY_SCOPE_NODE,
-    pynvml.NVML_AFFINITY_SCOPE_SOCKET])
+
+@pytest.mark.parametrize(
+    "scope", [pynvml.NVML_AFFINITY_SCOPE_NODE, pynvml.NVML_AFFINITY_SCOPE_SOCKET]
+)
 def test_nvmlDeviceGetMemoryAffinity(handles, scope):
     size = 1024
     for handle in handles:
         nodeSet = pynvml.nvmlDeviceGetMemoryAffinity(handle, size, scope)
         assert nodeSet is not None
 
-@pytest.mark.parametrize("scope", [
-    pynvml.NVML_AFFINITY_SCOPE_NODE,
-    pynvml.NVML_AFFINITY_SCOPE_SOCKET])
+
+@pytest.mark.parametrize(
+    "scope", [pynvml.NVML_AFFINITY_SCOPE_NODE, pynvml.NVML_AFFINITY_SCOPE_SOCKET]
+)
 def test_nvmlDeviceGetCpuAffinityWithinScope(handles, scope):
     size = 1024
     for handle in handles:
         cpuSet = pynvml.nvmlDeviceGetCpuAffinityWithinScope(handle, size, scope)
         assert cpuSet is not None
 
-@pytest.mark.parametrize("index", [
-    pynvml.NVML_P2P_CAPS_INDEX_READ,
-    pynvml.NVML_P2P_CAPS_INDEX_WRITE,
-    pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
-    pynvml.NVML_P2P_CAPS_INDEX_ATOMICS,
-    pynvml.NVML_P2P_CAPS_INDEX_PROP,
-    pynvml.NVML_P2P_CAPS_INDEX_UNKNOWN
-])
+
+@pytest.mark.parametrize(
+    "index",
+    [
+        pynvml.NVML_P2P_CAPS_INDEX_READ,
+        pynvml.NVML_P2P_CAPS_INDEX_WRITE,
+        pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
+        pynvml.NVML_P2P_CAPS_INDEX_ATOMICS,
+        pynvml.NVML_P2P_CAPS_INDEX_PROP,
+        pynvml.NVML_P2P_CAPS_INDEX_UNKNOWN,
+    ],
+)
 def test_nvmlDeviceGetP2PStatus(handles, index):
     for h1 in handles:
         for h2 in handles:
             if h1 is not h2:
                 status = pynvml.nvmlDeviceGetP2PStatus(h1, h2, index)
-                assert pynvml.NVML_P2P_STATUS_OK <= status <= pynvml.NVML_P2P_STATUS_UNKNOWN
+                assert (
+                    pynvml.NVML_P2P_STATUS_OK
+                    <= status
+                    <= pynvml.NVML_P2P_STATUS_UNKNOWN
+                )
 
 
 # [Skipping] pynvml.nvmlDeviceGetName
@@ -178,21 +232,21 @@ def test_nvmlDeviceGetPowerUsage(ngpus, handles):
         power_mWatts = pynvml.nvmlDeviceGetPowerUsage(handles[i])
         assert power_mWatts >= 0.0
 
+
 # Test pynvml.nvmlDeviceGetTotalEnergyConsumption
 def test_nvmlDeviceGetTotalEnergyConsumption(ngpus, handles):
     for i in range(ngpus):
-        energy_mJoules1 = pynvml.nvmlDeviceGetTotalEnergyConsumption(
-            handles[i])
+        energy_mJoules1 = pynvml.nvmlDeviceGetTotalEnergyConsumption(handles[i])
         for j in range(10):  # idle for 150 ms
             time.sleep(0.015)  # and check for increase every 15 ms
-            energy_mJoules2 = pynvml.nvmlDeviceGetTotalEnergyConsumption(
-                handles[i])
+            energy_mJoules2 = pynvml.nvmlDeviceGetTotalEnergyConsumption(handles[i])
             assert energy_mJoules2 >= energy_mJoules1
             if energy_mJoules2 > energy_mJoules1:
                 break
         else:
             assert False, "energy did not increase across 150 ms interval"
 
+
 # [Skipping] pynvml.nvmlDeviceGetGpuOperationMode
 # [Skipping] pynvml.nvmlDeviceGetCurrentGpuOperationMode
 # [Skipping] pynvml.nvmlDeviceGetPendingGpuOperationMode
@@ -200,9 +254,9 @@ def test_nvmlDeviceGetTotalEnergyConsumption(ngpus, handles):
 # Test pynvml.nvmlDeviceGetMemoryInfo
 def test_nvmlDeviceGetMemoryInfo(ngpus, handles):
     for i in range(ngpus):
-        meminfo = pynvml.nvmlDeviceGetMemoryInfo( handles[i] )
-        assert ((meminfo.used <= meminfo.total)
-                and (meminfo.free <= meminfo.total))
+        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handles[i])
+        assert (meminfo.used <= meminfo.total) and (meminfo.free <= meminfo.total)
+
 
 # [Skipping] pynvml.nvmlDeviceGetBAR1MemoryInfo
 # [Skipping] pynvml.nvmlDeviceGetComputeMode
@@ -216,10 +270,11 @@ def test_nvmlDeviceGetMemoryInfo(ngpus, handles):
 # Test pynvml.nvmlDeviceGetUtilizationRates
 def test_nvmlDeviceGetUtilizationRates(ngpus, handles):
     for i in range(ngpus):
-        urate = pynvml.nvmlDeviceGetUtilizationRates( handles[i] )
+        urate = pynvml.nvmlDeviceGetUtilizationRates(handles[i])
         assert urate.gpu >= 0
         assert urate.memory >= 0
 
+
 # [Skipping] pynvml.nvmlDeviceGetEncoderUtilization
 # [Skipping] pynvml.nvmlDeviceGetDecoderUtilization
 # [Skipping] pynvml.nvmlDeviceGetPcieReplayCounter
@@ -273,19 +328,17 @@ def test_nvmlDeviceGetUtilizationRates(ngpus, handles):
 def test_nvmlDeviceGetPcieThroughput(ngpus, handles):
     for i in range(ngpus):
         tx_bytes_tp = pynvml.nvmlDeviceGetPcieThroughput(
-            handles[i],
-            NVML_PCIE_UTIL_TX_BYTES
+            handles[i], NVML_PCIE_UTIL_TX_BYTES
         )
         assert tx_bytes_tp >= 0
         rx_bytes_tp = pynvml.nvmlDeviceGetPcieThroughput(
             handles[i], NVML_PCIE_UTIL_RX_BYTES
         )
         assert rx_bytes_tp >= 0
-        count_tp = pynvml.nvmlDeviceGetPcieThroughput(
-            handles[i], NVML_PCIE_UTIL_COUNT
-        )
+        count_tp = pynvml.nvmlDeviceGetPcieThroughput(handles[i], NVML_PCIE_UTIL_COUNT)
         assert count_tp >= 0
 
+
 # [Skipping] pynvml.nvmlSystemGetTopologyGpuSet
 # [Skipping] pynvml.nvmlDeviceGetTopologyNearestGpus
 # [Skipping] pynvml.nvmlDeviceGetTopologyCommonAncestor
@@ -293,7 +346,11 @@ def test_nvmlDeviceGetPcieThroughput(ngpus, handles):
 # Test pynvml.nvmlDeviceGetNvLinkVersion
 # Test pynvml.nvmlDeviceGetNvLinkState
 # Test pynvml.nvmlDeviceGetNvLinkRemotePciInfo
-def test_nvml_nvlink_properties(ngpus, handles):
+def test_nvml_nvlink_properties(ngpus, handles, driver):
+
+    if driver > 450.0:
+        pytest.xfail(XFAIL_LEGACY_NVLINK_MSG)
+
     for i in range(ngpus):
         for j in range(pynvml.NVML_NVLINK_MAX_LINKS):
             version = pynvml.nvmlDeviceGetNvLinkVersion(handles[i], j)
@@ -303,66 +360,91 @@ def test_nvml_nvlink_properties(ngpus, handles):
             pci_info = pynvml.nvmlDeviceGetNvLinkRemotePciInfo(handles[i], j)
             assert isinstance(pci_info, pynvml.c_nvmlPciInfo_t)
 
+
 # Test pynvml.nvmlDeviceGetNvLinkCapability
-@pytest.mark.parametrize('cap_type', [
-    pynvml.NVML_NVLINK_CAP_P2P_SUPPORTED,  # P2P over NVLink is supported
-    pynvml.NVML_NVLINK_CAP_SYSMEM_ACCESS,  # Access to system memory is supported
-    pynvml.NVML_NVLINK_CAP_P2P_ATOMICS,    # P2P atomics are supported
-    pynvml.NVML_NVLINK_CAP_SYSMEM_ATOMICS, # System memory atomics are supported
-    pynvml.NVML_NVLINK_CAP_SLI_BRIDGE,     # SLI is supported over this link
-    pynvml.NVML_NVLINK_CAP_VALID])         # Link is supported on this device
+@pytest.mark.parametrize(
+    "cap_type",
+    [
+        pynvml.NVML_NVLINK_CAP_P2P_SUPPORTED,  # P2P over NVLink is supported
+        pynvml.NVML_NVLINK_CAP_SYSMEM_ACCESS,  # Access to system memory is supported
+        pynvml.NVML_NVLINK_CAP_P2P_ATOMICS,  # P2P atomics are supported
+        pynvml.NVML_NVLINK_CAP_SYSMEM_ATOMICS,  # System memory atomics are supported
+        pynvml.NVML_NVLINK_CAP_SLI_BRIDGE,  # SLI is supported over this link
+        pynvml.NVML_NVLINK_CAP_VALID,
+    ],
+)  # Link is supported on this device
 def test_nvml_nvlink_capability(ngpus, handles, cap_type):
     for i in range(ngpus):
         for j in range(pynvml.NVML_NVLINK_MAX_LINKS):
             cap = pynvml.nvmlDeviceGetNvLinkCapability(handles[i], j, cap_type)
             assert cap >= 0
 
+
 # Test pynvml.nvmlDeviceResetNvLinkUtilizationCounter
 # Test pynvml.nvmlDeviceSetNvLinkUtilizationControl
 # Test pynvml.nvmlDeviceGetNvLinkUtilizationCounter
 # Test pynvml.nvmlDeviceGetNvLinkUtilizationControl
 # Test pynvml.nvmlDeviceFreezeNvLinkUtilizationCounter
-@pytest.mark.parametrize('counter', [0, 1])
-@pytest.mark.parametrize('control', [0, 1, 2])
-def test_nvml_nvlink_counters(ngpus, handles, counter, control):
+@pytest.mark.parametrize("counter", [0, 1])
+@pytest.mark.parametrize("control", [0, 1, 2])
+def test_nvml_nvlink_counters(ngpus, handles, counter, control, driver):
+
+    if driver > 450.0:
+        pytest.xfail(XFAIL_LEGACY_NVLINK_MSG)
+
     reset = 0
     for i in range(ngpus):
         for j in range(pynvml.NVML_NVLINK_MAX_LINKS):
-            assert pynvml.nvmlDeviceResetNvLinkUtilizationCounter(
-                handles[i], j, counter
-            ) == pynvml.NVML_SUCCESS
+            assert (
+                pynvml.nvmlDeviceResetNvLinkUtilizationCounter(handles[i], j, counter)
+                == pynvml.NVML_SUCCESS
+            )
             pynvml.nvmlDeviceSetNvLinkUtilizationControl(
                 handles[i], j, counter, control, reset
             )
             countdict = pynvml.nvmlDeviceGetNvLinkUtilizationCounter(
                 handles[i], j, counter
             )
-            ctl = pynvml.nvmlDeviceGetNvLinkUtilizationControl(
-                handles[i], j, counter
-            )
-            assert countdict['rx'] >= 0
-            assert countdict['tx'] >= 0
+            ctl = pynvml.nvmlDeviceGetNvLinkUtilizationControl(handles[i], j, counter)
+            assert countdict["rx"] >= 0
+            assert countdict["tx"] >= 0
             assert ctl == control
-            assert pynvml.nvmlDeviceFreezeNvLinkUtilizationCounter(
-                handles[i], j, counter, 1
-            ) == pynvml.NVML_SUCCESS
-            assert pynvml.nvmlDeviceFreezeNvLinkUtilizationCounter(
-                handles[i], j, counter, 0
-            ) == pynvml.NVML_SUCCESS
+            assert (
+                pynvml.nvmlDeviceFreezeNvLinkUtilizationCounter(
+                    handles[i], j, counter, 1
+                )
+                == pynvml.NVML_SUCCESS
+            )
+            assert (
+                pynvml.nvmlDeviceFreezeNvLinkUtilizationCounter(
+                    handles[i], j, counter, 0
+                )
+                == pynvml.NVML_SUCCESS
+            )
+
 
 # Test pynvml.nvmlDeviceResetNvLinkErrorCounters
 # Test pynvml.nvmlDeviceGetNvLinkErrorCounter
-@pytest.mark.parametrize('error_type', [
-    pynvml.NVML_NVLINK_ERROR_DL_REPLAY,
-    pynvml.NVML_NVLINK_ERROR_DL_RECOVERY,
-    pynvml.NVML_NVLINK_ERROR_DL_CRC_FLIT,
-    pynvml.NVML_NVLINK_ERROR_DL_CRC_DATA])
-def test_nvml_nvlink_error_counters(ngpus, handles, error_type):
+@pytest.mark.parametrize(
+    "error_type",
+    [
+        pynvml.NVML_NVLINK_ERROR_DL_REPLAY,
+        pynvml.NVML_NVLINK_ERROR_DL_RECOVERY,
+        pynvml.NVML_NVLINK_ERROR_DL_CRC_FLIT,
+        pynvml.NVML_NVLINK_ERROR_DL_CRC_DATA,
+    ],
+)
+def test_nvml_nvlink_error_counters(ngpus, handles, error_type, driver):
+
+    if driver > 450.0:
+        pytest.xfail(XFAIL_LEGACY_NVLINK_MSG)
+
     for i in range(ngpus):
         for j in range(pynvml.NVML_NVLINK_MAX_LINKS):
-            assert pynvml.nvmlDeviceResetNvLinkErrorCounters(
-                handles[i], j
-            ) == pynvml.NVML_SUCCESS
+            assert (
+                pynvml.nvmlDeviceResetNvLinkErrorCounters(handles[i], j)
+                == pynvml.NVML_SUCCESS
+            )
             error_count = pynvml.nvmlDeviceGetNvLinkErrorCounter(
                 handles[i], j, error_type
             )
diff --git a/pynvml/tests/test_smi.py b/pynvml/tests/test_smi.py
index 8dd9df7..cbfa745 100644
--- a/pynvml/tests/test_smi.py
+++ b/pynvml/tests/test_smi.py
@@ -3,48 +3,58 @@
 import os
 
 # Fixture to initialize and finalize nvml
-@pytest.fixture(scope='module')
+@pytest.fixture(scope="module")
 def smi(request):
     return nvidia_smi.getInstance()
 
+
 @pytest.fixture
 def ngpus(smi):
-    result = smi.DeviceQuery('count')['count']
+    result = smi.DeviceQuery("count")["count"]
     assert result > 0
-    print('['+str(result)+' GPUs]', end =' ')
+    print("[" + str(result) + " GPUs]", end=" ")
     return result
 
+
 ## ---------------------------- ##
 ##       Start Query Tests      ##
 ## ---------------------------- ##
 
 # Test free-memory query
 def test_query_memory(ngpus, smi):
-    result = smi.DeviceQuery('memory.free')
+    result = smi.DeviceQuery("memory.free")
     for i in range(ngpus):
-        assert result['gpu'][i]['fb_memory_usage']['free'] >= 0
+        assert result["gpu"][i]["fb_memory_usage"]["free"] >= 0
+
 
 # Test gpu-utilization query
 def test_gpu_utilization(ngpus, smi):
     for i in range(ngpus):
-        result = smi.DeviceQuery('utilization.gpu')['gpu'][i]['utilization']['gpu_util']
+        result = smi.DeviceQuery("utilization.gpu")["gpu"][i]["utilization"]["gpu_util"]
         assert result >= 0
 
+
 # Test memory-utilization query
 def test_memory_utilization(ngpus, smi):
     for i in range(ngpus):
-        result = smi.DeviceQuery('utilization.memory')['gpu'][i]['utilization']['memory_util']
+        result = smi.DeviceQuery("utilization.memory")["gpu"][i]["utilization"][
+            "memory_util"
+        ]
         assert result >= 0
 
+
 # Test pstate query
 def test_pstate(ngpus, smi):
     for i in range(ngpus):
-        result = smi.DeviceQuery('pstate')['gpu'][i]['performance_state']
-        assert result[0] == 'P'
+        result = smi.DeviceQuery("pstate")["gpu"][i]["performance_state"]
+        assert result[0] == "P"
+
 
 # Test temperature query
 def test_temperature(ngpus, smi):
     for i in range(ngpus):
-        temp = smi.DeviceQuery('temperature.gpu')['gpu'][i]['temperature']['gpu_temp']
-        max_temp = smi.DeviceQuery('temperature.gpu')['gpu'][i]['temperature']['gpu_temp_max_threshold']
+        temp = smi.DeviceQuery("temperature.gpu")["gpu"][i]["temperature"]["gpu_temp"]
+        max_temp = smi.DeviceQuery("temperature.gpu")["gpu"][i]["temperature"][
+            "gpu_temp_max_threshold"
+        ]
         assert (temp > 0) and (temp < max_temp)