Skip to content

Fix direct NVML compatibility and optional metadata handling #116

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Qubitium merged 3 commits into from
Aug 23, 2026
+521 −38
Merged

Fix direct NVML compatibility and optional metadata handling #116

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
459a44e
Fix direct NVML compatibility paths
Qubitium Aug 23, 2026
da1cba1
Correct legacy NVML PCI ABI
Qubitium Aug 23, 2026
d0187ec
Harden NVML compatibility regressions
Qubitium Aug 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 118 additions & 38 deletions device_smi/nvidia.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,13 @@

logger = logging.getLogger(__name__)

_NVML_ERROR_NOT_SUPPORTED = 3
_NVML_ERROR_FUNCTION_NOT_FOUND = 13
_NVML_COMPATIBILITY_ERRORS = {
_NVML_ERROR_NOT_SUPPORTED,
_NVML_ERROR_FUNCTION_NOT_FOUND,
}


class NvidiaGPUMetrics(BaseMetrics):
pass
Expand All @@ -34,7 +41,7 @@ class _nvmlUtilization_t(ctypes.Structure):
]


class _nvmlPciInfo_t(ctypes.Structure):
class _nvmlPciInfo_v3_t(ctypes.Structure):
_fields_ = [
("busIdLegacy", ctypes.c_char * 16),
("domain", ctypes.c_uint),
Expand All @@ -46,6 +53,19 @@ class _nvmlPciInfo_t(ctypes.Structure):
]


class _nvmlPciInfo_v2_t(ctypes.Structure):
"""Legacy PCI info ABI used by both the v2 and unversioned symbols."""

_fields_ = [
("busId", ctypes.c_char * 16),
("domain", ctypes.c_uint),
("bus", ctypes.c_uint),
("device", ctypes.c_uint),
("pciDeviceId", ctypes.c_uint),
("pciSubSystemId", ctypes.c_uint),
]


def _nvml_library_candidates():
system = platform.system()
if system == "Linux":
Expand Down Expand Up @@ -101,6 +121,7 @@ def __new__(cls):
cls._instance._initialized = False
cls._instance._available = False
cls._instance._lib = None
cls._instance._shutdown_registered = False
return cls._instance

def __init__(self):
Expand All @@ -123,7 +144,9 @@ def __init__(self):
if self._available:
import atexit

atexit.register(self._shutdown)
if not self._shutdown_registered:
atexit.register(self._shutdown)
self._shutdown_registered = True

@property
def available(self) -> bool:
Expand Down Expand Up @@ -163,7 +186,15 @@ def _bind_functions(self):
ctypes.c_int,
),
"nvmlDeviceGetPciInfo_v3": (
[ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_t)],
[ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_v3_t)],
ctypes.c_int,
),
"nvmlDeviceGetPciInfo_v2": (
[ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_v2_t)],
ctypes.c_int,
),
"nvmlDeviceGetPciInfo": (
[ctypes.c_void_p, ctypes.POINTER(_nvmlPciInfo_v2_t)],
ctypes.c_int,
),
"nvmlDeviceGetMaxPcieLinkGeneration": (
Expand Down Expand Up @@ -209,13 +240,14 @@ def _init(self):
return self._call(["nvmlInit_v2", "nvmlInit"])

def _shutdown(self):
if not self._available or self._lib is None:
return
try:
if self._nvmlShutdown is not None:
self._nvmlShutdown()
except (RuntimeError, AttributeError) as exc:
logger.debug("NVML shutdown failed: %s", exc)
with self._lock:
if not self._available or self._lib is None:
return
self._available = False
try:
self._check(self._call(["nvmlShutdown"]))
except (OSError, RuntimeError, AttributeError) as exc:
logger.debug("NVML shutdown failed: %s", exc)

def _check(self, code):
if code == 0:
Expand Down Expand Up @@ -248,6 +280,7 @@ def device_handle(self, gpu_id: str):
if not self._available:
raise RuntimeError("NVML is not available")
handle = ctypes.c_void_p()
gpu_id = gpu_id.strip()
if gpu_id.isdigit():
self._check(
self._call(
Expand All @@ -256,7 +289,7 @@ def device_handle(self, gpu_id: str):
ctypes.byref(handle),
)
)
elif gpu_id.startswith("GPU-"):
elif gpu_id.startswith(("GPU-", "MIG-")):
self._check(
self._call(
["nvmlDeviceGetHandleByUUID"],
Expand Down Expand Up @@ -287,60 +320,91 @@ def device_handle(self, gpu_id: str):

def device_name(self, handle):
buf = ctypes.create_string_buffer(96)
self._check(self._nvmlDeviceGetName(handle, buf, 96))
self._check(self._call(["nvmlDeviceGetName"], handle, buf, 96))
return buf.value.decode("utf-8").strip()

def pci_info(self, handle):
info = _nvmlPciInfo_t()
self._check(self._nvmlDeviceGetPciInfo_v3(handle, ctypes.byref(info)))
bus_id = info.busId.decode("utf-8").strip()
if not bus_id:
bus_id = info.busIdLegacy.decode("utf-8").strip()
if not re.match(r"^[0-9a-fA-F]{4,8}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$", bus_id):
raise RuntimeError(f"Invalid NVML PCI bus id: {bus_id!r}")
return bus_id
attempts = (
("nvmlDeviceGetPciInfo_v3", _nvmlPciInfo_v3_t),
("nvmlDeviceGetPciInfo_v2", _nvmlPciInfo_v2_t),
("nvmlDeviceGetPciInfo", _nvmlPciInfo_v2_t),
)
attempted = []
for name, info_type in attempts:
fn = getattr(self, f"_{name}", None)
if fn is None:
continue
attempted.append(name)
info = info_type()
code = self._call([name], handle, ctypes.byref(info))
if code in _NVML_COMPATIBILITY_ERRORS:
continue
self._check(code)
bus_id = info.busId.decode("utf-8").strip()
if not bus_id and isinstance(info, _nvmlPciInfo_v3_t):
bus_id = info.busIdLegacy.decode("utf-8").strip()
if not re.match(
r"^[0-9a-fA-F]{4,8}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$",
bus_id,
):
raise RuntimeError(f"Invalid NVML PCI bus id: {bus_id!r}")
return bus_id
raise RuntimeError(f"NVML PCI info functions unavailable: {attempted}")

def max_pcie_generation(self, handle):
gen = ctypes.c_uint()
self._check(self._nvmlDeviceGetMaxPcieLinkGeneration(handle, ctypes.byref(gen)))
self._check(
self._call(
["nvmlDeviceGetMaxPcieLinkGeneration"], handle, ctypes.byref(gen)
)
)
return gen.value

def current_pcie_generation(self, handle):
gen = ctypes.c_uint()
self._check(self._nvmlDeviceGetCurrPcieLinkGeneration(handle, ctypes.byref(gen)))
self._check(
self._call(
["nvmlDeviceGetCurrPcieLinkGeneration"], handle, ctypes.byref(gen)
)
)
return gen.value

def memory_total(self, handle):
mem = _nvmlMemory_t()
self._check(self._nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem)))
self._check(self._call(["nvmlDeviceGetMemoryInfo"], handle, ctypes.byref(mem)))
return mem.total

def memory_used(self, handle):
mem = _nvmlMemory_t()
self._check(self._nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem)))
self._check(self._call(["nvmlDeviceGetMemoryInfo"], handle, ctypes.byref(mem)))
return mem.used

def utilization_gpu(self, handle):
util = _nvmlUtilization_t()
self._check(self._nvmlDeviceGetUtilizationRates(handle, ctypes.byref(util)))
self._check(
self._call(["nvmlDeviceGetUtilizationRates"], handle, ctypes.byref(util))
)
return float(util.gpu)

def driver_version(self):
buf = ctypes.create_string_buffer(80)
self._check(self._nvmlSystemGetDriverVersion(buf, 80))
self._check(self._call(["nvmlSystemGetDriverVersion"], buf, 80))
return buf.value.decode("utf-8").strip()

def vbios_version(self, handle):
buf = ctypes.create_string_buffer(32)
self._check(self._nvmlDeviceGetVbiosVersion(handle, buf, 32))
self._check(self._call(["nvmlDeviceGetVbiosVersion"], handle, buf, 32))
return buf.value.decode("utf-8").strip()

def compute_capability(self, handle):
major = ctypes.c_int()
minor = ctypes.c_int()
self._check(
self._nvmlDeviceGetCudaComputeCapability(
handle, ctypes.byref(major), ctypes.byref(minor)
self._call(
["nvmlDeviceGetCudaComputeCapability"],
handle,
ctypes.byref(major),
ctypes.byref(minor),
)
)
return f"{major.value}.{minor.value}"
Expand Down Expand Up @@ -425,7 +489,9 @@ def _get_gpu_id(self):
gpu_count = 0

cudas = os.environ.get("CUDA_VISIBLE_DEVICES", "")
cuda_list = cudas.split(",") if cudas else []
cuda_list = [cuda.strip() for cuda in cudas.split(",")] if cudas.strip() else []
if any(not cuda for cuda in cuda_list):
raise ValueError("CUDA_VISIBLE_DEVICES contains an empty device entry")
if gpu_count > 0 and os.environ.get("CUDA_DEVICE_ORDER", "") != "PCI_BUS_ID":
warnings.warn(
"Detected different devices in the system. Please make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to avoid unexpected behavior.",
Expand All @@ -442,23 +508,37 @@ def _init_from_nvml(self, cls):

model = _nvml.device_name(handle)
total_memory = _nvml.memory_total(handle)
pci_bus_id = _nvml.pci_info(handle)
pcie_gen_max = _nvml.max_pcie_generation(handle)
pcie_gen_current = _nvml.current_pcie_generation(handle)
driver = _nvml.driver_version()
firmware = _nvml.vbios_version(handle)
compute_cap = _nvml.compute_capability(handle)
driver = self._optional_nvml(_nvml.driver_version, "")
firmware = self._optional_nvml(_nvml.vbios_version, "", handle)
compute_cap = self._optional_nvml(_nvml.compute_capability, None, handle)
pci_bus_id = self._optional_nvml(_nvml.pci_info, None, handle)
pcie_gen_max = self._optional_nvml(_nvml.max_pcie_generation, None, handle)
pcie_gen_current = self._optional_nvml(
_nvml.current_pcie_generation, None, handle
)

if model.lower().startswith("nvidia"):
model = model[len("nvidia"):]

cls.model = model.strip().lower()
cls.memory_total = int(total_memory)
cls.vendor = "nvidia"
cls.features = [compute_cap]
cls.pcie = Pcie(gen=int(pcie_gen_max), speed=int(pcie_gen_current), id=pci_bus_id)
cls.features = [compute_cap] if compute_cap else []
cls.pcie = None
if None not in (pci_bus_id, pcie_gen_max, pcie_gen_current):
cls.pcie = Pcie(
gen=int(pcie_gen_max), speed=int(pcie_gen_current), id=pci_bus_id
)
cls.gpu = GPU(driver=driver, firmware=firmware)

@staticmethod
def _optional_nvml(fetch, default, *args):
try:
return fetch(*args)
except (OSError, RuntimeError, ValueError, AttributeError) as exc:
logger.debug("Optional NVML metadata unavailable: %s", exc)
return default

def _init_from_nvidia_smi(self, cls):
args = [
"nvidia-smi",
Expand Down
Loading
Loading