Get Nvidia GPU Name

get_gpu_name.py

from platform import machine
from re import sub as re_sub

from logging import getLogger


logger = getLogger(__name__)


__all__ = ["get_gpu_name"]


def _get_gpu_name_alt():
    """Returns the name of the GPU using nvidia-smi if available, otherwise None"""
    import subprocess

    gpu_name = None
    try:
        output_bytes = subprocess.check_output("nvidia-smi --query-gpu=gpu_name --format=csv,noheader", shell=True)
        output = output_bytes.decode("ascii").strip()
        gpu_name = re_sub(r"\s+|-", "_", output)  # replace spaces and dashes with underscores
    except Exception as ex:
        logger.exception(ex)

    logger.info("GPU name from nvidia-smi: {gpu_name}", gpu_name=gpu_name)
    return gpu_name


def _get_gpu_name_amd64() -> str | None:
    """Returns the name of the GPU if available, otherwise None"""
    from pynvml import (
        nvmlDeviceGetCount,
        nvmlDeviceGetHandleByIndex,
        nvmlDeviceGetName,
        NVMLError,
        nvmlInit,
        nvmlShutdown,
    )

    gpu_name = None
    nvml_initiated = False
    try:
        nvmlInit()
        logger.info("initialized NVML")
        nvml_initiated = True
    except NVMLError as err:
        logger.exception("Failed to initialize NVML: {err}", err=err)
        #  try alternate method
        gpu_name = _get_gpu_name_alt()
    else:
        try:
            deviceCount = nvmlDeviceGetCount()
            logger.info("GPU Count: {count}", count=deviceCount)
            if deviceCount > 0:
                logger.info("Getting GPU name")
                gpu_name = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(0))
                gpu_name = re_sub(r"\s+|-", "_", gpu_name)  # replace spaces and dashes with underscores
            else:
                gpu_name = None
        except NVMLError as err:
            logger.exception("Failed to get GPU name: {err}", err=err)
    finally:
        if nvml_initiated:
            nvmlShutdown()
            logger.info("shutdown NVML")

    logger.info("GPU name: {gpu_name}", gpu_name=gpu_name)

    return gpu_name


def _get_gpu_name_arm64() -> str | None:
    """Returns the name of the GPU if available, otherwise None"""
    # NOTE: using pynvml fails on current arm64 image due to libnividia-ml.so issues.
    #       Since Jetson Nano is the only supported arm device for now, the GPU name is hardcoded
    # TODO: Find a way to get the GPU name for ARM64 devices
    return "Xavier"


def get_gpu_name() -> str | None:
    """Returns the name of the GPU if available, otherwise None"""
    gpu_name = None
    arch = machine()
    if arch == "aarch64":
        gpu_name = _get_gpu_name_arm64()
    elif arch == "x86_64":
        gpu_name = _get_gpu_name_amd64()
    else:
        logger.error("Unsupported platform: {arch}", arch=arch)

    return gpu_name