Skip to content

Instantly share code, notes, and snippets.

@ekardon
Last active May 4, 2021 18:05
Show Gist options
  • Save ekardon/865554650b763dcdb7324ae32ff8712c to your computer and use it in GitHub Desktop.
Save ekardon/865554650b763dcdb7324ae32ff8712c to your computer and use it in GitHub Desktop.
This is a partial snippet, so there may be some missing/undefined parts
from ctypes import CFUNCTYPE, c_int, POINTER, byref, c_char_p, c_har, c_vooid_p, c_uint
# SOURCE: https://docs.nvidia.com/deploy/nvml-api/group__nvmlDeviceQueries.html
class nvmlPciInfo_t(Structure):
_fields_ = [
("bus", c_uint), # 0 to 0xff
("busId", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), # string + 0x00
("busIdLegacy", c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), # string + 0x00
("device", c_uint), # 0 to 31
("domain", c_uint), # 0 to 0xffffffff
("pciDeviceId", c_uint), # 16-bit device id + 16-bit vendor id
("pciSubSystemId", c_uint), # 32-bit id
]
nvmlDeviceInfo_t = c_void_p
def nvml_device_get_handle_by_pci_bus_id_v2(pci_bus_id):
proto = CFUNCTYPE(
c_int,
c_char_p,
nvmlDeviceInfo_t
)
params = (1, "pci_bus_id", 0), (1, "device_handler", 0),
func = proto(("nvmlDeviceGetHandleByPciBusId_v2", nvml_lib), params)
#pci_bus_id = c_char_p(pci_bus_id) # + b"\0"
device_handler = nvmlDeviceInfo_t()
ret_code = func(pci_bus_id, byref(device_handler))
if ret_code == NVML_SUCCESS:
return device_handler
else:
raise Exception(f"Err {ret_code}: {nvml_error_string(ret_code)}")
def nvml_device_get_pci_info_v3(device_handler):
proto = CFUNCTYPE(
c_int,
nvmlDeviceInfo_t,
POINTER(nvmlPciInfo_t)
)
params = (1, "device_handler", 0), (1, "pci_info_struct", 0),
func = proto(("nvmlDeviceGetPciInfo_v3", nvml_lib), params)
pci_info = nvmlPciInfo_t()
ret_code = func(device_handler, byref(pci_info))
if ret_code == NVML_SUCCESS:
return pci_info
else:
raise Exception(f"Err {ret_code}: {nvml_error_string(ret_code)}")
def nvml_device_get_handle_by_index_v2(device_index):
proto = CFUNCTYPE(
c_int,
c_uint,
POINTER(nvmlDeviceInfo_t)
)
params = (1, "device_index", 0), (1, "device_handler", 0),
func = proto(("nvmlDeviceGetHandleByIndex_v2", nvml_lib), params)
device_handler = nvmlDeviceInfo_t()
ret_code = func(c_uint(device_index), byref(device_handler))
if ret_code == NVML_SUCCESS:
return device_handler
else:
raise Exception(f"Err {ret_code}: {nvml_error_string(ret_code)}")
if __name__ == "__main__":
device_handler = nvml_device_get_handle_by_index_v2(0)
pci_info = nvml_device_get_pci_info_v3(device_handler)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment