forked from dashingsoft/pyarmor
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcheck_gpu.py
More file actions
52 lines (46 loc) · 1.82 KB
/
check_gpu.py
File metadata and controls
52 lines (46 loc) · 1.82 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
50
51
52
import re
import subprocess
from pytransform import get_expired_days, get_user_data
def get_expiration_info():
try:
# license_info = get_license_info()
left_days = get_expired_days()
if left_days == -1:
print('This license for %s is never expired')
else:
print(f'This license for will be expired in {left_days} days')
except Exception as e:
print(e)
def check_gpu():
gpu_uuids = get_gpu_list()
assert len(gpu_uuids) == 1 # 1 GPU per instance
if gpu_uuids:
if len(gpu_uuids) > 1:
raise RuntimeError(f'This license is issued for one particular GPU, {len(gpu_uuids)} GPUs detected')
else:
assert len(gpu_uuids) == 1
gpu_uuid = gpu_uuids[0]
if gpu_uuid != get_user_data().decode('utf-8').lower():
print(f'User data {get_user_data()}')
raise RuntimeError('A GPU matching the license is not found')
else:
print('A GPU license check is passed')
get_expiration_info()
else:
raise RuntimeError('No GPUs detected, this license is issued to particular GPU')
def get_gpu_list():
try:
out = subprocess.run(["nvidia-smi", "-L"], check=True, capture_output=True)
out.check_returncode()
out = out.stdout.decode('utf-8').lower()
# output example
# gpu 0: geforce gtx 1080 ti (uuid: gpu-70ef1701-4072-9722-cc0b-7c7e75ff76db)
# gpu 1: geforce gtx 1080 ti (uuid: gpu-5b8df9cc-3b3c-d07a-8bd1-e2a51af4cfa9)
# ...
uuid_list = re.findall(r'\(uuid: \S+\)', out)
uuid_list = [uuid.replace('(uuid: ', '').replace(')', '')
for uuid in uuid_list]
return uuid_list
except:
# any error means no GPUs available
return None