xoscar 0.9.0__cp312-cp312-macosx_10_13_x86_64.whl

xoscar/nvutils.py

@@ -0,0 +1,717 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import os
+import subprocess
+import sys
+import uuid
+from collections import namedtuple
+from ctypes import (
+    CDLL,
+    POINTER,
+    Structure,
+    byref,
+    c_char,
+    c_char_p,
+    c_int,
+    c_uint,
+    c_ulonglong,
+    create_string_buffer,
+)
+from typing import List, Tuple, Union
+
+from .utils import parse_readable_size
+
+logger = logging.getLogger(__name__)
+
+# Some constants taken from cuda.h
+CUDA_SUCCESS = 0
+CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT = 16
+CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR = 39
+CU_DEVICE_ATTRIBUTE_CLOCK_RATE = 13
+CU_DEVICE_ATTRIBUTE_PCI_BUS_ID = 33
+CU_DEVICE_ATTRIBUTE_PCI_DEVICE_ID = 34
+CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE = 36
+
+CU_NO_CUDA_CAPABLE_DEVICE_DETECTED = 100
+
+# nvml constants
+NVML_SUCCESS = 0
+NVML_ERROR_UNINITIALIZED = 1
+NVML_ERROR_INVALID_ARGUMENT = 2
+NVML_ERROR_NOT_SUPPORTED = 3
+NVML_ERROR_NO_PERMISSION = 4
+NVML_ERROR_ALREADY_INITIALIZED = 5
+NVML_ERROR_NOT_FOUND = 6
+NVML_ERROR_INSUFFICIENT_SIZE = 7
+NVML_ERROR_INSUFFICIENT_POWER = 8
+NVML_ERROR_DRIVER_NOT_LOADED = 9
+NVML_ERROR_TIMEOUT = 10
+NVML_ERROR_IRQ_ISSUE = 11
+NVML_ERROR_LIBRARY_NOT_FOUND = 12
+NVML_ERROR_FUNCTION_NOT_FOUND = 13
+NVML_ERROR_CORRUPTED_INFOROM = 14
+NVML_ERROR_GPU_IS_LOST = 15
+NVML_ERROR_RESET_REQUIRED = 16
+NVML_ERROR_OPERATING_SYSTEM = 17
+NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18
+NVML_ERROR_IN_USE = 19
+NVML_ERROR_MEMORY = 20
+NVML_ERROR_NO_DATA = 21
+NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22
+NVML_ERROR_INSUFFICIENT_RESOURCES = 23
+NVML_ERROR_FREQ_NOT_SUPPORTED = 24
+NVML_ERROR_UNKNOWN = 999
+NVML_TEMPERATURE_GPU = 0
+NVML_DRIVER_NOT_LOADED = 9
+NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96
+NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1)
+NVML_DEVICE_MIG_DISABLE = 0x0
+NVML_DEVICE_MIG_ENABLE = 0x1
+
+
+class _CUuuid_t(Structure):
+    _fields_ = [("bytes", c_char * 16)]
+
+
+class _nvmlUtilization_t(Structure):
+    _fields_ = [
+        ("gpu", c_uint),
+        ("memory", c_uint),
+    ]
+
+
+class _struct_nvmlDevice_t(Structure):
+    pass  # opaque handle
+
+
+_nvmlDevice_t = POINTER(_struct_nvmlDevice_t)
+
+
+class _nvmlBAR1Memory_t(Structure):
+    _fields_ = [
+        ("total", c_ulonglong),
+        ("free", c_ulonglong),
+        ("used", c_ulonglong),
+    ]
+
+
+class _nvmlProcessInfo_t(Structure):
+    _fields_ = [
+        ("pid", c_uint),
+        ("usedGpuMemory", c_ulonglong),
+        ("gpuInstanceId", c_uint),
+        ("computeInstanceId", c_uint),
+    ]
+
+
+## Alternative object
+# Allows the object to be printed
+# Allows mismatched types to be assigned
+#  - like None when the Structure variant requires c_uint
+class nvmlFriendlyObject:
+    def __init__(self, dictionary):
+        for x in dictionary:
+            setattr(self, x, dictionary[x])
+
+    def __str__(self):
+        return self.__dict__.__str__()
+
+
+def nvmlStructToFriendlyObject(struct):
+    d = {}
+    for x in struct._fields_:
+        key = x[0]
+        value = getattr(struct, key)
+        # only need to convert from bytes if bytes, no need to check python version.
+        d[key] = value.decode() if isinstance(value, bytes) else value
+    obj = nvmlFriendlyObject(d)
+    return obj
+
+
+@dataclasses.dataclass
+class CudaDeviceInfo:
+    uuid: bytes | None = None
+    device_index: int | None = None
+    mig_index: int | None = None
+
+
+@dataclasses.dataclass
+class CudaContext:
+    has_context: bool
+    device_info: CudaDeviceInfo | None = None
+
+
+_is_windows: bool = sys.platform.startswith("win")
+_is_wsl: bool = "WSL_DISTRO_NAME" in os.environ
+
+
+def _load_nv_library(*libnames):
+    for lib in libnames:
+        try:
+            return CDLL(lib)
+        except OSError:
+            continue
+
+
+_cuda_lib = _nvml_lib = None
+
+_cu_device_info = namedtuple(
+    "_cu_device_info", "index uuid name multiprocessors cuda_cores threads"
+)
+_nvml_driver_info = namedtuple("_nvml_driver_info", "driver_version cuda_version")
+_nvml_device_status = namedtuple(
+    "_nvml_device_status",
+    "gpu_util mem_util temperature fb_total_mem fb_used_mem fb_free_mem",
+)
+
+_init_pid = None
+_gpu_count = None
+_driver_info = None
+_device_infos: dict[int, _cu_device_info] = dict()
+
+_no_device_warned = False
+
+
+class NVError(Exception):
+    def __init__(self, msg, *args, errno=None):
+        self._errno = errno
+        super().__init__(msg or "Unknown error", *args)
+
+    def __str__(self):
+        return f"({self._errno}) {super().__str__()}"
+
+    @property
+    def errno(self):
+        return self._errno
+
+    @property
+    def message(self):
+        return super().__str__()
+
+
+class NVDeviceAPIError(NVError):
+    pass
+
+
+class NVMLAPIError(NVError):
+    pass
+
+
+def _cu_check_error(result):
+    if result != CUDA_SUCCESS:
+        _error_str = c_char_p()
+        _cuda_lib.cuGetErrorString(result, byref(_error_str))
+        err_value = _error_str.value.decode() if _error_str.value is not None else None
+        raise NVDeviceAPIError(err_value, errno=result)
+
+
+_nvmlErrorString = None
+
+
+def _nvml_check_error(result):
+    global _nvmlErrorString
+    if _nvmlErrorString is None:
+        _nvmlErrorString = _nvml_lib.nvmlErrorString
+        _nvmlErrorString.restype = c_char_p
+
+    if result != NVML_SUCCESS:
+        _error_str = _nvmlErrorString(result)
+        raise NVMLAPIError(_error_str.decode(), errno=result)
+
+
+_cu_process_var_to_cores = {
+    (1, 0): 8,
+    (1, 1): 8,
+    (1, 2): 8,
+    (1, 3): 8,
+    (2, 0): 32,
+    (2, 1): 48,
+}
+
+
+def _cu_get_processor_cores(major, minor):
+    return _cu_process_var_to_cores.get((major, minor), 192)
+
+
+def _init_cp():
+    global _cuda_lib, _no_device_warned
+    if _init_pid == os.getpid():
+        return
+
+    libcuda_paths = ["libcuda.so", "libcuda.dylib", "cuda.dll", "nvcuda.dll"]
+    if _is_wsl:
+        libcuda_paths = ["/usr/lib/wsl/lib/libcuda.so"] + libcuda_paths
+    _cuda_lib = _load_nv_library(*libcuda_paths)
+
+    if _cuda_lib is None:
+        return
+    try:
+        _cu_check_error(_cuda_lib.cuInit(0))
+    except NVDeviceAPIError as ex:
+        if ex.errno == CU_NO_CUDA_CAPABLE_DEVICE_DETECTED:
+            _cuda_lib = None
+            if not _no_device_warned:
+                logger.warning("No CUDA device detected")
+                _no_device_warned = True
+        else:
+            logger.exception("Failed to initialize libcuda.")
+        return
+
+
+def _init_nvml():
+    global _nvml_lib, _no_device_warned
+    if _init_pid == os.getpid():
+        return
+
+    nvml_paths = [
+        "libnvidia-ml.so",
+        "libnvidia-ml.so.1",
+        "libnvidia-ml.dylib",
+        "nvml.dll",
+    ]
+    if _is_windows:
+        nvml_paths.append(
+            os.path.join(
+                os.getenv("ProgramFiles", "C:/Program Files"),
+                "NVIDIA Corporation/NVSMI/nvml.dll",
+            )
+        )
+    if _is_wsl:
+        nvml_paths = ["/usr/lib/wsl/lib/libnvidia-ml.so.1"] + nvml_paths
+    _nvml_lib = _load_nv_library(*nvml_paths)
+
+    if _nvml_lib is None:
+        return
+    try:
+        _nvml_check_error(_nvml_lib.nvmlInit_v2())
+    except NVMLAPIError as ex:
+        if ex.errno == NVML_DRIVER_NOT_LOADED:
+            _nvml_lib = None
+            if not _no_device_warned:
+                logger.warning(
+                    "Failed to load libnvidia-ml: %s, no CUDA device will be enabled",
+                    ex.message,
+                )
+                _no_device_warned = True
+        else:
+            logger.exception("Failed to initialize libnvidia-ml.")
+        return
+
+
+def _init():
+    global _init_pid
+
+    _init_cp()
+    _init_nvml()
+
+    if _nvml_lib is not None and _cuda_lib is not None:
+        _init_pid = os.getpid()
+
+
+def get_device_count() -> int | None:
+    global _gpu_count
+
+    if _gpu_count is not None:
+        return _gpu_count
+
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    if "CUDA_VISIBLE_DEVICES" in os.environ:
+        devices = os.environ["CUDA_VISIBLE_DEVICES"].strip()
+        if not devices or devices == "-1":
+            _gpu_count = 0
+        else:
+            _gpu_count = len(devices.split(","))
+    else:
+        n_gpus = c_uint()
+        _cu_check_error(_nvml_lib.nvmlDeviceGetCount(byref(n_gpus)))
+        _gpu_count = n_gpus.value
+    return _gpu_count
+
+
+def _get_all_device_count() -> int | None:
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    n_gpus = c_uint()
+    _cu_check_error(_nvml_lib.nvmlDeviceGetCount(byref(n_gpus)))
+    return n_gpus.value
+
+
+def get_driver_info() -> _nvml_driver_info | None:
+    global _driver_info
+
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+    if _driver_info is not None:
+        return _driver_info
+
+    version_buf = create_string_buffer(100)
+    cuda_version = c_uint()
+
+    _nvml_check_error(
+        _nvml_lib.nvmlSystemGetDriverVersion(version_buf, len(version_buf))
+    )
+    _nvml_check_error(_nvml_lib.nvmlSystemGetCudaDriverVersion(byref(cuda_version)))
+
+    _driver_info = _nvml_driver_info(
+        driver_version=version_buf.value.decode(),
+        cuda_version=".".join(str(v) for v in divmod(cuda_version.value, 1000)),
+    )
+    return _driver_info
+
+
+def get_device_info(dev_index: int) -> _cu_device_info | None:
+    try:
+        return _device_infos[dev_index]
+    except KeyError:
+        pass
+
+    _init()
+    if _init_pid is None:
+        return None
+
+    device = c_int()
+    name_buf = create_string_buffer(100)
+    uuid_t = _CUuuid_t()
+    cc_major = c_int()
+    cc_minor = c_int()
+    cores = c_int()
+    threads_per_core = c_int()
+
+    _cu_check_error(_cuda_lib.cuDeviceGet(byref(device), c_int(dev_index)))  # type: ignore
+    _cu_check_error(_cuda_lib.cuDeviceGetName(name_buf, len(name_buf), device))  # type: ignore
+    _cu_check_error(_cuda_lib.cuDeviceGetUuid(byref(uuid_t), device))  # type: ignore
+    _cu_check_error(
+        _cuda_lib.cuDeviceComputeCapability(byref(cc_major), byref(cc_minor), device)  # type: ignore
+    )
+    _cu_check_error(
+        _cuda_lib.cuDeviceGetAttribute(  # type: ignore
+            byref(cores), CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device
+        )
+    )
+    _cu_check_error(
+        _cuda_lib.cuDeviceGetAttribute(  # type: ignore
+            byref(threads_per_core),
+            CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR,
+            device,
+        )
+    )
+
+    if "CUDA_VISIBLE_DEVICES" in os.environ:
+        real_dev_index = [
+            int(s) for s in os.environ["CUDA_VISIBLE_DEVICES"].split(",")
+        ][dev_index]
+    else:
+        real_dev_index = dev_index
+
+    info = _device_infos[dev_index] = _cu_device_info(
+        index=real_dev_index,
+        uuid=uuid.UUID(bytes=uuid_t.bytes),
+        name=name_buf.value.decode(),
+        multiprocessors=cores.value,
+        cuda_cores=cores.value
+        * _cu_get_processor_cores(cc_major.value, cc_minor.value),
+        threads=cores.value * threads_per_core.value,
+    )
+    return info
+
+
+def get_device_status(dev_index: int) -> _nvml_device_status | None:
+    _init()
+    if _init_pid is None:
+        return None
+
+    c_device = _nvmlDevice_t()
+    c_utils = _nvmlUtilization_t()
+    c_temperature = c_uint()
+    c_memory_info = _nvmlBAR1Memory_t()
+
+    dev_uuid = get_device_info(dev_index).uuid  # type: ignore
+
+    uuid_str = ("GPU-" + str(dev_uuid)).encode()
+
+    if not _is_wsl:
+        _nvml_check_error(
+            _nvml_lib.nvmlDeviceGetHandleByUUID(uuid_str, byref(c_device))  # type: ignore
+        )
+
+        _nvml_check_error(
+            _nvml_lib.nvmlDeviceGetUtilizationRates(c_device, byref(c_utils))  # type: ignore
+        )
+        gpu_util = c_utils.gpu
+        mem_util = c_utils.memory
+
+        _nvml_check_error(
+            _nvml_lib.nvmlDeviceGetTemperature(  # type: ignore
+                c_device, NVML_TEMPERATURE_GPU, byref(c_temperature)
+            )
+        )
+        temperature = c_temperature.value
+
+        _nvml_check_error(
+            _nvml_lib.nvmlDeviceGetMemoryInfo(c_device, byref(c_memory_info))  # type: ignore
+        )
+        fb_total_mem = c_memory_info.total
+        fb_free_mem = c_memory_info.free
+        fb_used_mem = c_memory_info.used
+    else:
+        import defusedxml
+
+        proc = subprocess.Popen(
+            ["/usr/lib/wsl/lib/nvidia-smi", "-q", f"--id={dev_index}", "-x"],
+            stdout=subprocess.PIPE,
+        )
+        proc.wait()
+        xml_result = defusedxml.ElementTree.fromstring(proc.stdout.read())  # type: ignore
+        gpu_node = xml_result.find("gpu")
+
+        fb_node = gpu_node.find("fb_memory_usage")
+        fb_total_mem = int(parse_readable_size(fb_node.find("total").text)[0])
+        fb_free_mem = int(parse_readable_size(fb_node.find("free").text)[0])
+        fb_used_mem = int(parse_readable_size(fb_node.find("used").text)[0])
+
+        util_node = gpu_node.find("utilization")
+        if util_node.find("gpu_util").text == "N/A":
+            gpu_util = 0
+        else:
+            gpu_util = int(util_node.find("gpu_util"))
+        if util_node.find("memory_util").text == "N/A":
+            mem_util = 0
+        else:
+            mem_util = int(util_node.find("memory_util"))
+
+        temperature = int(gpu_node.find("temperature").find("gpu_temp").text[:-1])
+
+    return _nvml_device_status(
+        gpu_util=gpu_util,
+        mem_util=mem_util,
+        temperature=temperature,
+        fb_total_mem=fb_total_mem,
+        fb_free_mem=fb_free_mem,
+        fb_used_mem=fb_used_mem,
+    )
+
+
+def get_handle_by_index(index: int) -> _nvmlDevice_t:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_index = c_int(index)
+    device = _nvmlDevice_t()
+    _nvml_check_error(_nvml_lib.nvmlDeviceGetHandleByIndex_v2(c_index, byref(device)))
+    return device
+
+
+def get_handle_by_uuid(uuid: bytes) -> _nvmlDevice_t:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_uuid = c_char_p(uuid)
+    device = _nvmlDevice_t()
+    _nvml_check_error(_nvml_lib.nvmlDeviceGetHandleByUUID(c_uuid, byref(device)))
+    return device
+
+
+def get_mig_mode(device: _nvmlDevice_t) -> Tuple[int, int] | None:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_current_mode, c_pending_mode = c_uint(), c_uint()
+    _nvml_check_error(
+        _nvml_lib.nvmlDeviceGetMigMode(
+            device, byref(c_current_mode), byref(c_pending_mode)
+        )
+    )
+    return c_current_mode.value, c_pending_mode.value
+
+
+def get_max_mig_device_count(device: _nvmlDevice_t) -> int | None:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_count = c_uint()
+    _nvml_check_error(_nvml_lib.nvmlDeviceGetMaxMigDeviceCount(device, byref(c_count)))
+    return c_count.value
+
+
+def get_mig_device_handle_by_index(device: _nvmlDevice_t, index: int) -> _nvmlDevice_t:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_index = c_uint(index)
+    mig_device = _nvmlDevice_t()
+    _nvml_check_error(
+        _nvml_lib.nvmlDeviceGetMigDeviceHandleByIndex(
+            device, c_index, byref(mig_device)
+        )
+    )
+    return mig_device
+
+
+def get_index(handle: _nvmlDevice_t) -> int | None:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_index = c_uint()
+    _nvml_check_error(_nvml_lib.nvmlDeviceGetIndex(handle, byref(c_index)))
+    return c_index.value
+
+
+def get_uuid(handle: _nvmlDevice_t) -> bytes | None:  # type: ignore
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE)
+    _nvml_check_error(
+        _nvml_lib.nvmlDeviceGetUUID(
+            handle, c_uuid, c_uint(NVML_DEVICE_UUID_V2_BUFFER_SIZE)
+        )
+    )
+    return c_uuid.value
+
+
+def get_index_and_uuid(device: Union[int, bytes, str]) -> CudaDeviceInfo | None:
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    try:
+        device_index = int(device)
+        device_handle = get_handle_by_index(device_index)
+        uuid = get_uuid(device_handle)
+    except ValueError:
+        uuid = device if isinstance(device, bytes) else device.encode()  # type: ignore
+        uuid_handle = get_handle_by_uuid(uuid)
+        device_index = get_index(uuid_handle)  # type: ignore
+        uuid = get_uuid(uuid_handle)
+
+    return CudaDeviceInfo(uuid=uuid, device_index=device_index)
+
+
+def get_compute_running_processes(  # type: ignore
+    handle: _nvmlDevice_t,  # type: ignore
+) -> List[nvmlFriendlyObject] | None:
+    _init_nvml()
+    if _nvml_lib is None:
+        return None
+
+    c_count = c_uint(0)
+    func = getattr(_nvml_lib, "nvmlDeviceGetComputeRunningProcesses_v3", None)
+    if func is None:
+        func = getattr(_nvml_lib, "nvmlDeviceGetComputeRunningProcesses_v2")
+    ret = func(handle, byref(c_count), None)
+
+    if ret == NVML_SUCCESS:
+        # special case, no running processes
+        return []
+    elif ret == NVML_ERROR_INSUFFICIENT_SIZE:
+        # typical case
+        # oversize the array in case more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = _nvmlProcessInfo_t * c_count.value
+        c_procs = proc_array()
+
+        _nvml_check_error(func(handle, byref(c_count), c_procs))
+
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value:
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+
+        return procs
+    else:
+        # error case
+        _nvml_check_error(ret)
+
+
+def _running_process_matches(handle: _nvmlDevice_t) -> bool:  # type: ignore
+    """Check whether the current process is same as that of handle
+    Parameters
+    ----------
+    handle : _nvmlDevice_t
+        NVML handle to CUDA device
+    Returns
+    -------
+    out : bool
+        Whether the device handle has a CUDA context on the running process.
+    """
+    return any(os.getpid() == o.pid for o in get_compute_running_processes(handle))  # type: ignore
+
+
+def get_cuda_context() -> CudaContext:
+    """Check whether the current process already has a CUDA context created."""
+
+    _init()
+    if _init_pid is None:
+        return CudaContext(has_context=False)
+
+    for index in range(_get_all_device_count()):  # type: ignore
+        handle = get_handle_by_index(index)
+        try:
+            mig_current_mode, mig_pending_mode = get_mig_mode(handle)  # type: ignore
+        except NVMLAPIError as e:
+            if e.errno == NVML_ERROR_NOT_SUPPORTED:
+                mig_current_mode = NVML_DEVICE_MIG_DISABLE
+            else:
+                raise
+        if mig_current_mode == NVML_DEVICE_MIG_ENABLE:
+            for mig_index in range(get_max_mig_device_count(handle)):  # type: ignore
+                try:
+                    mig_handle = get_mig_device_handle_by_index(handle, mig_index)
+                except NVMLAPIError as e:
+                    if e.errno == NVML_ERROR_NOT_FOUND:
+                        # No MIG device with that index
+                        continue
+                    else:
+                        raise
+                if _running_process_matches(mig_handle):
+                    return CudaContext(
+                        has_context=True,
+                        device_info=CudaDeviceInfo(
+                            uuid=get_uuid(handle),
+                            device_index=index,
+                            mig_index=mig_index,
+                        ),
+                    )
+        else:
+            if _running_process_matches(handle):
+                return CudaContext(
+                    has_context=True,
+                    device_info=CudaDeviceInfo(
+                        uuid=get_uuid(handle), device_index=index
+                    ),
+                )
+
+    return CudaContext(has_context=False)