blob: 97c790986578dbe71a48b3a919f0a0a798da1cce (
plain)
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
|
# Copyright 2018 Lenovo, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
"""
Utils for GPU driver.
"""
import re
import subprocess
GPU_FLAGS = ["VGA compatible controller", "3D controller"]
GPU_INFO_PATTERN = re.compile("(?P<devices>[0-9]{2}:[0-9]{2}\.[0-9]) (?P"
"<name>.*) \[.* [\[](?P<vendor_id>[0-9a-fA-F]{4})"
":(?P<product_id>[0-9a-fA-F]{4})] .*")
def discover_vendors():
cmd = "sudo lspci -nnn | grep -E '%s'"
cmd = cmd % "|".join(GPU_FLAGS)
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
p.wait()
gpus = p.stdout.readlines()
vendors = set()
for gpu in gpus:
m = GPU_INFO_PATTERN.match(gpu)
if m:
vendor_id = m.groupdict().get("vendor_id")
vendors.add(vendor_id)
return vendors
def discover_gpus(vender_id=None):
cmd = "sudo lspci -nnn | grep -E '%s'"
cmd = cmd % "|".join(GPU_FLAGS)
if vender_id:
cmd = cmd + "| grep " + vender_id
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
p.wait()
gpus = p.stdout.readlines()
gpu_list = []
for gpu in gpus:
m = GPU_INFO_PATTERN.match(gpu)
if m:
gpu_dict = m.groupdict()
gpu_dict["function"] = "GPU"
gpu_dict["devices"] = _match_nova_addr(gpu_dict["devices"])
gpu_dict["assignable"] = True
gpu_list.append(gpu_dict)
return gpu_list
def _match_nova_addr(devices):
addr = '0000:'+devices.replace(".", ":")
return addr
|
