forked from dask/distributed
-
Notifications
You must be signed in to change notification settings - Fork 0
/
nvml.py
28 lines (21 loc) · 731 Bytes
/
nvml.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import pynvml
handles = None
def _pynvml_handles():
global handles
if handles is None:
pynvml.nvmlInit()
count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
return handles
def real_time():
handles = _pynvml_handles()
return {
"utilization": [pynvml.nvmlDeviceGetUtilizationRates(h).gpu for h in handles],
"memory-used": [pynvml.nvmlDeviceGetMemoryInfo(h).used for h in handles],
}
def one_time():
handles = _pynvml_handles()
return {
"memory-total": [pynvml.nvmlDeviceGetMemoryInfo(h).total for h in handles],
"name": [pynvml.nvmlDeviceGetName(h).decode() for h in handles],
}