forked from snap-stanford/GreaseLM
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathnv_usage.py
More file actions
49 lines (41 loc) · 1.61 KB
/
Copy pathnv_usage.py
File metadata and controls
49 lines (41 loc) · 1.61 KB
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import pynvml
# 输入是需要并行计算的gpu数量，输出是符合条件的gpu标号list
def get_gpu_index(gpu_number):
# Initialize NVML
pynvml.nvmlInit()
gpu_index_list = []
deviceCount = pynvml.nvmlDeviceGetCount()
for i in range(deviceCount):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
# Get the GPU utilization
utilization = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
memrate = meminfo.used/meminfo.total
if utilization == 0 and memrate < 0.05:
gpu_index_list.append(i)
if len(gpu_index_list) == gpu_number:
return gpu_index_list
# Shutdown NVML
pynvml.nvmlShutdown()
raise ValueError(f"Only {len(gpu_index_list)} gpu are empty. No adequate gpu found !")
# 输入是需要并行计算的gpu数量，输出是符合条件的gpu标号list
def get_avail_gpu_index():
# Initialize NVML
pynvml.nvmlInit()
gpu_index_list = []
deviceCount = pynvml.nvmlDeviceGetCount()
for i in range(deviceCount):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
# Get the GPU utilization
utilization = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
memrate = meminfo.used/meminfo.total
if utilization == 0 and memrate < 0.05:
gpu_index_list.append(i)
# Shutdown NVML
pynvml.nvmlShutdown()
return gpu_index_list
if __name__ == "__main__":
print(get_gpu_index(1))
# print(get_gpu_index(2))
# print(get_gpu_index(8))