# library of small functions
import logging
import subprocess
import numpy as np
from ibllib.exceptions import NvidiaDriverNotReady
_logger = logging.getLogger(__name__)
def _parametrized(dec):
def layer(*args, **kwargs):
def repl(f):
return dec(f, *args, **kwargs)
return repl
return layer
[docs]
def structarr(names, shape=None, formats=None):
if not formats:
formats = ['f8'] * len(names)
dtyp = np.dtype({'names': names, 'formats': formats})
return np.zeros(shape, dtype=dtyp)
[docs]
def check_nvidia_driver():
"""
Checks if the GPU driver reacts and otherwise raises a custom error.
Useful to check before long GPU-dependent processes.
"""
process = subprocess.Popen('nvidia-smi', shell=True, stdout=subprocess.PIPE,
stderr=subprocess.PIPE, executable="/bin/bash")
info, error = process.communicate()
if process.returncode != 0:
raise NvidiaDriverNotReady(f"{error.decode('utf-8')}")
_logger.info("nvidia-smi command successful")