/
gpu.py
73 lines (62 loc) · 2.28 KB
/
gpu.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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import pyudev
import re
import subprocess
from middlewared.service_exception import CallError
RE_PCI_ADDR = re.compile(r'(?P<domain>.*):(?P<bus>.*):(?P<slot>.*)\.')
SENSITIVE_PCI_DEVICE_TYPES = (
'PCI Bridge',
'ISA Bridge',
'RAM memory',
'SMBus',
)
def get_gpus():
cp = subprocess.Popen(['lspci', '-D'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = cp.communicate()
if cp.returncode:
raise CallError(f'Unable to list available gpus: {stderr.decode()}')
gpus = []
gpu_slots = []
for line in stdout.decode().splitlines():
for k in (
'VGA compatible controller',
'Display controller',
'3D controller',
):
if k in line:
gpu_slots.append((line.strip(), k))
break
for gpu_line, key in gpu_slots:
addr = gpu_line.split()[0]
addr_re = RE_PCI_ADDR.match(addr)
gpu_dev = pyudev.Devices.from_name(pyudev.Context(), 'pci', addr)
# Let's normalise vendor for consistency
vendor = None
vendor_id_from_db = gpu_dev.get('ID_VENDOR_FROM_DATABASE', '').lower()
if 'nvidia' in vendor_id_from_db:
vendor = 'NVIDIA'
elif 'intel' in vendor_id_from_db:
vendor = 'INTEL'
elif 'amd' in vendor_id_from_db:
vendor = 'AMD'
devices = []
critical = False
for child in filter(lambda c: all(k in c for k in ('PCI_SLOT_NAME', 'PCI_ID')), gpu_dev.parent.children):
devices.append({
'pci_id': child['PCI_ID'],
'pci_slot': child['PCI_SLOT_NAME'],
'vm_pci_slot': f'pci_{child["PCI_SLOT_NAME"].replace(".", "_").replace(":", "_")}',
})
critical |= any(
k.lower() in child.get('ID_PCI_SUBCLASS_FROM_DATABASE', '').lower() for k in SENSITIVE_PCI_DEVICE_TYPES
)
gpus.append({
'addr': {
'pci_slot': addr,
**{k: addr_re.group(k) for k in ('domain', 'bus', 'slot')},
},
'description': gpu_line.split(f'{key}:')[-1].split('(rev')[0].strip(),
'devices': devices,
'vendor': vendor,
'uses_system_critical_devices': critical,
})
return gpus