def get_gpus()

in optimum_benchmark/system_utils.py [0:0]

• 36 lines of code
• 12 McCabe index (conditional complexity)

def get_gpus():
    if is_nvidia_system():
        if not is_pynvml_available():
            raise ValueError(
                "The library PyNVML is required to get available GPUs, but is not installed. "
                "Please install the official and NVIDIA maintained PyNVML library through `pip install nvidia-ml-py`."
            )

        gpus = []

        pynvml.nvmlInit()
        for i in range(pynvml.nvmlDeviceGetCount()):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            gpu = pynvml.nvmlDeviceGetName(handle)
            # Older pynvml versions may return bytes
            gpu = gpu.decode("utf-8") if isinstance(gpu, bytes) else gpu
            gpus.append(gpu)
        pynvml.nvmlShutdown()

    elif is_rocm_system():
        if not is_amdsmi_available() and not is_pyrsmi_available():
            raise ValueError(
                "Either the library AMD SMI or PyRSMI is required to get available GPUs, but neither is installed. "
                "Please install the official and AMD maintained AMD SMI library from https://github.com/ROCm/amdsmi "
                "or PyRSMI library from https://github.com/ROCm/pyrsmi."
            )

        gpus = []

        if is_amdsmi_available():
            amdsmi.amdsmi_init()
            for processor_handles in amdsmi.amdsmi_get_processor_handles():
                gpus.append(amdsmi.amdsmi_get_gpu_vendor_name(processor_handles))
            amdsmi.amdsmi_shut_down()

        elif is_pyrsmi_available():
            rocml.smi_initialize()
            for i in range(rocml.smi_get_device_count()):
                gpus.append(rocml.smi_get_device_name(i))
            rocml.smi_shutdown()

    else:
        raise ValueError("No NVIDIA or ROCm GPUs found.")

    return gpus