()
| 39 | |
| 40 | |
| 41 | def _init_nvml() -> bool: |
| 42 | global _nvml_inited, _nvml_handle |
| 43 | if _nvml_inited: |
| 44 | return _nvml_handle is not None |
| 45 | _nvml_inited = True |
| 46 | try: |
| 47 | from pynvml import ( # noqa: F401 |
| 48 | nvmlInit, |
| 49 | nvmlDeviceGetHandleByIndex, |
| 50 | nvmlDeviceGetName, |
| 51 | nvmlDeviceGetMemoryInfo, |
| 52 | ) |
| 53 | nvmlInit() |
| 54 | idx = int(os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]) |
| 55 | _nvml_handle = nvmlDeviceGetHandleByIndex(idx) |
| 56 | return True |
| 57 | except Exception: |
| 58 | return False |
| 59 | |
| 60 | |
| 61 | def gpu_device_used_mb() -> float: |
no test coverage detected