Try to gather NVidia GPU device information via libcudart.
()
| 116 | |
| 117 | |
| 118 | def _gather_gpu_devices_cudart(): |
| 119 | """Try to gather NVidia GPU device information via libcudart.""" |
| 120 | dev_info = [] |
| 121 | |
| 122 | system = platform.system() |
| 123 | if system == "Linux": |
| 124 | libcudart = ct.cdll.LoadLibrary("libcudart.so") |
| 125 | elif system == "Darwin": |
| 126 | libcudart = ct.cdll.LoadLibrary("libcudart.dylib") |
| 127 | elif system == "Windows": |
| 128 | libcudart = ct.windll.LoadLibrary("libcudart.dll") |
| 129 | else: |
| 130 | raise NotImplementedError("Cannot identify system.") |
| 131 | |
| 132 | version = ct.c_int() |
| 133 | rc = libcudart.cudaRuntimeGetVersion(ct.byref(version)) |
| 134 | if rc != 0: |
| 135 | raise ValueError("Could not get version") |
| 136 | if version.value < 6050: |
| 137 | raise NotImplementedError("CUDA version must be between >= 6.5") |
| 138 | |
| 139 | device_count = ct.c_int() |
| 140 | libcudart.cudaGetDeviceCount(ct.byref(device_count)) |
| 141 | |
| 142 | for i in range(device_count.value): |
| 143 | properties = CUDADeviceProperties() |
| 144 | rc = libcudart.cudaGetDeviceProperties(ct.byref(properties), i) |
| 145 | if rc != 0: |
| 146 | raise ValueError("Could not get device properties") |
| 147 | pci_bus_id = " " * 13 |
| 148 | rc = libcudart.cudaDeviceGetPCIBusId(ct.c_char_p(pci_bus_id), 13, i) |
| 149 | if rc != 0: |
| 150 | raise ValueError("Could not get device PCI bus id") |
| 151 | |
| 152 | info = test_log_pb2.GPUInfo() # No UUID available |
| 153 | info.model = properties.name |
| 154 | info.bus_id = pci_bus_id |
| 155 | dev_info.append(info) |
| 156 | |
| 157 | del properties |
| 158 | |
| 159 | return dev_info |
| 160 | |
| 161 | |
| 162 | def gather_gpu_devices(): |
no test coverage detected