decode pynvml bytes to str

Signed-off-by: Jonathan Nitisastro <jonathancn@anyscale.com>
ray-project · Nov 20, 2023 · 1e0f5ad · 1e0f5ad
1 parent f23603d
commit 1e0f5ad
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/dashboard/modules/reporter/reporter_agent.py b/dashboard/modules/reporter/reporter_agent.py
@@ -9,7 +9,7 @@
 
 import psutil
 
-from typing import List, Optional, Tuple, TypedDict
+from typing import List, Optional, Tuple, TypedDict, Union
 from collections import defaultdict
 
 import ray
@@ -400,6 +400,12 @@ def _get_gpu_usage():
         if not enable_gpu_usage_check:
             return []
         gpu_utilizations = []
+
+        def decode(b: Union[str, bytes]) -> str:
+            if isinstance(b, bytes):
+                return b.decode("utf-8")  # for python3, to unicode
+            return b
+
         try:
             pynvml.nvmlInit()
             num_gpus = pynvml.nvmlDeviceGetCount()
@@ -423,8 +429,8 @@ def _get_gpu_usage():
                     )
                     processes_pids = [
                         ProcessGPUInfo(
-                            pid=nv_process.pid,
-                            gpu_memory_usage=nv_process.usedGpuMemory // MB
+                            pid=int(nv_process.pid),
+                            gpu_memory_usage=int(nv_process.usedGpuMemory) // MB
                             if nv_process.usedGpuMemory
                             else 0,
                         )
@@ -435,8 +441,8 @@ def _get_gpu_usage():
 
                 info = GpuUtilizationInfo(
                     index=i,
-                    name=pynvml.nvmlDeviceGetName(gpu_handle),
-                    uuid=pynvml.nvmlDeviceGetUUID(gpu_handle),
+                    name=decode(pynvml.nvmlDeviceGetName(gpu_handle)),
+                    uuid=decode(pynvml.nvmlDeviceGetUUID(gpu_handle)),
                     utilization_gpu=utilization,
                     memory_used=int(pynvml.nvmlDeviceGetMemoryInfo(gpu_handle).used)
                     // MB,

diff --git a/python/ray/_private/accelerators/nvidia_gpu.py b/python/ray/_private/accelerators/nvidia_gpu.py
@@ -63,10 +63,11 @@ def get_current_node_accelerator_type() -> Optional[str]:
         cuda_device_type = None
         if device_count > 0:
             handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+            device_name = pynvml.nvmlDeviceGetName(handle)
+            if isinstance(device_name, bytes):
+                device_name = device_name.decode("utf-8")
             cuda_device_type = (
-                NvidiaGPUAcceleratorManager._gpu_name_to_accelerator_type(
-                    pynvml.nvmlDeviceGetName(handle)
-                )
+                NvidiaGPUAcceleratorManager._gpu_name_to_accelerator_type(device_name)
             )
         pynvml.nvmlShutdown()
         return cuda_device_type