Use nvidia-smi to detect nvidia GPU

This commit is contained in:
Jip-Hop 2023-03-04 11:36:43 +01:00
parent 3411021bcc
commit 1c62587403
1 changed files with 17 additions and 18 deletions

View File

@ -35,7 +35,7 @@ IT COMES WITHOUT WARRANTY AND IS NOT SUPPORTED BY IXSYSTEMS.{NORMAL}"""
DESCRIPTION = "Create persistent Linux 'jails' on TrueNAS SCALE, with full access to all files \ DESCRIPTION = "Create persistent Linux 'jails' on TrueNAS SCALE, with full access to all files \
via bind mounts, without modifying the host OS thanks to systemd-nspawn!" via bind mounts, without modifying the host OS thanks to systemd-nspawn!"
VERSION = '0.0.0' VERSION = '0.0.1'
JAILS_DIR_PATH = 'jails' JAILS_DIR_PATH = 'jails'
JAIL_CONFIG_NAME = 'config' JAIL_CONFIG_NAME = 'config'
@ -87,25 +87,18 @@ def passthrough_nvidia(gpu_passthrough_nvidia, systemd_nspawn_additional_args, j
os.remove(ld_so_conf_path) os.remove(ld_so_conf_path)
return return
# Detect nvidia GPU
if not len(glob.glob('/dev/nvidia*')):
eprint(dedent("""
No nvidia GPU seems to be present...
Skip passthrough of nvidia GPU."""))
return
if not os.path.exists('/dev/nvidia-uvm'): if not os.path.exists('/dev/nvidia-uvm'):
# Create the nvidia device files: /dev/nvidia* # Create the nvidia device files: /dev/nvidia*
for command in ( try:
['modprobe', 'nvidia-current-uvm'], # TODO: does nvidia-smi work on post-init?
['nvidia-modprobe', '-c0', '-u'], # Or do we need to explicitly call modprobe?
): subprocess.run(['nvidia-smi', '-f', '/dev/null'], check=True)
try: except:
subprocess.run(command, check=True) eprint("Skip passthrough of nvidia GPU.")
except: return
eprint("Skip passthrough of nvidia GPU.")
return
# TODO: do we really need all /dev/nvidia* devices?
# Perhaps starting with an empty set would be enough
nvidia_files = set(glob.glob('/dev/nvidia*')) nvidia_files = set(glob.glob('/dev/nvidia*'))
try: try:
nvidia_files.update([x for x in subprocess.check_output( nvidia_files.update([x for x in subprocess.check_output(
@ -538,7 +531,13 @@ def create_jail(jail_name):
gpu_passthrough_nvidia = 0 gpu_passthrough_nvidia = 0
if len(glob.glob('/dev/nvidia*')): try:
subprocess.run(['nvidia-smi'], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
nvidia_detected = True
except:
nvidia_detected = False
if nvidia_detected:
print("Detected the presence of an nvidia GPU.") print("Detected the presence of an nvidia GPU.")
if agree('Passthrough the nvidia GPU?', 'n'): if agree('Passthrough the nvidia GPU?', 'n'):
gpu_passthrough_nvidia = 1 gpu_passthrough_nvidia = 1