Loader for libnvidia-ml.so
| 50 | |
| 51 | |
| 52 | class NVML(object): |
| 53 | """ |
| 54 | Loader for libnvidia-ml.so |
| 55 | """ |
| 56 | |
| 57 | _nvmlLib = None |
| 58 | _lib_lock = threading.Lock() |
| 59 | |
| 60 | def load(self): |
| 61 | with self._lib_lock: |
| 62 | if self._nvmlLib is None: |
| 63 | self._nvmlLib = CDLL("libnvidia-ml.so.1") |
| 64 | |
| 65 | function_pointers = ["nvmlDeviceGetName", "nvmlDeviceGetUUID", "nvmlDeviceGetMemoryInfo", |
| 66 | "nvmlDeviceGetUtilizationRates", "nvmlInit_v2", "nvmlShutdown", |
| 67 | "nvmlDeviceGetCount_v2", "nvmlDeviceGetHandleByIndex_v2"] |
| 68 | |
| 69 | self.func_ptr = {n: self._function_pointer(n) for n in function_pointers} |
| 70 | |
| 71 | def _function_pointer(self, name): |
| 72 | try: |
| 73 | return getattr(self._nvmlLib, name) |
| 74 | except AttributeError: |
| 75 | raise NvmlException(NVML_ERROR_FUNCTION_NOT_FOUND) |
| 76 | |
| 77 | def get_function(self, name): |
| 78 | if name in self.func_ptr.keys(): |
| 79 | return self.func_ptr[name] |
| 80 | |
| 81 | |
| 82 | _NVML = NVML() |