Detect: Add multiple GPU handling

This commit is contained in:
Ross Patterson 2019-08-09 20:48:34 -07:00
parent 7781c9c7b4
commit d7ae12f04b

View File

@ -59,33 +59,41 @@ GPU_PRODUCT_RE = re.compile(r'(?P<chip>[^[]+)(\[(?P<board>[^]]+)\]|)')
DATA = None DATA = None
def detect_gpu(): def detect_gpus():
""" """
Detect the GPU vendor, generation and model if possible. Detect the vendor, generation and model for each GPU if possible.
""" """
plat_sys = platform.system() plat_sys = platform.system()
gpu = None gpus = []
if plat_sys == 'Linux': if plat_sys == 'Linux':
# TODO: Android and other Linux'es that don't have `lshw` # TODO: Android and other Linux'es that don't have `lshw`
display_output = subprocess.check_output( display_output = subprocess.check_output(
['lshw', '-class', 'display', '-json']) ['lshw', '-class', 'display', '-json'])
display_data = json.loads(display_output.decode().strip().strip(',')) displays_data = json.loads(display_output.decode().strip().strip(','))
gpu = dict( if not isinstance(displays_data, list):
vendor=display_data['vendor'].replace(' Corporation', '')) # TODO: Confirm this is how `lshw` handles multiple GPUs
displays_data = [displays_data]
for display_data in displays_data:
gpu = dict(
vendor=display_data['vendor'].replace(' Corporation', ''))
# TODO get multiple GPUs from lshw
gpus.append(gpu)
product_match = GPU_PRODUCT_RE.search(display_data['product']) product_match = GPU_PRODUCT_RE.search(display_data['product'])
if product_match: if product_match:
gpu.update(**product_match.groupdict()) gpu.update(**product_match.groupdict())
if not gpu['board']: if not gpu['board']:
gpu['board'] = gpu.pop('chip') gpu['board'] = gpu.pop('chip')
else: else:
# TODO Other platforms # TODO Other platforms
raise NotImplementedError( raise NotImplementedError(
'GPU detection for {0!r} not supported yet'.format(plat_sys)) 'GPU detection for {0!r} not supported yet'.format(plat_sys))
return gpu if not gpus:
raise ValueError('No GPUs detected')
return gpus
def detect_hwaccels(hwaccels=None, cmd='ffmpeg'): def detect_hwaccels(hwaccels=None, cmd='ffmpeg'):
@ -106,8 +114,11 @@ def detect_hwaccels(hwaccels=None, cmd='ffmpeg'):
# Filter against which APIs are available on this OS+GPU # Filter against which APIs are available on this OS+GPU
data = _get_data() data = _get_data()
plat_sys = platform.system() plat_sys = platform.system()
gpu = detect_gpu() gpus = detect_gpus()
api_avail = data['hwaccels']['api_avail'][plat_sys][gpu['vendor']] api_avail = set()
for gpu in gpus:
api_avail.update(
data['hwaccels']['api_avail'][plat_sys][gpu['vendor']])
hwaccels = [ hwaccels = [
hwaccel for hwaccel in hwaccels if hwaccel['name'] in api_avail] hwaccel for hwaccel in hwaccels if hwaccel['name'] in api_avail]
@ -179,7 +190,7 @@ def detect_codecs(decoder, encoder, hwaccels=None, cmd='ffmpeg'):
__all__ = [ __all__ = [
'detect_gpu', 'detect_gpus',
'detect_hwaccels', 'detect_hwaccels',
'detect_codecs', 'detect_codecs',
] ]
@ -203,7 +214,7 @@ def main(args=None):
""" """
args = parser.parse_args(args) args = parser.parse_args(args)
data = dict( data = dict(
gpu=detect_gpu(), gpus=detect_gpus(),
hwaccels=detect_hwaccels(cmd=args.ffmpeg), hwaccels=detect_hwaccels(cmd=args.ffmpeg),
codecs=detect_codecs(cmd=args.ffmpeg)) codecs=detect_codecs(cmd=args.ffmpeg))
json.dump(data, sys.stdout, indent=2) json.dump(data, sys.stdout, indent=2)