GPU Switcher v2

nvmgr.py

#!/bin/python3
import os
import sys
import time
import libvirt
import argparse
import subprocess
from primesel import Switcher

NVIDIA_PCI_ID = '0000:01:00.0'
NVIDIA_PCI_VID = '10de:1bb5'
NVIDIA_MODULES = [
    'nvidia_uvm',
    'nvidia_drm',
    'nvidia_modeset'
]

STATE_FILE = '/root/.nvmgr_state'
VFIO_FILE = '/etc/modprobe.d/vfio.conf'
VM_ID = '17f8d8bc-2707-43de-8122-995613d1a7ba'

def detach(cleanup = False):
    try:
        pid = os.fork()
        if pid > 0:
            exit(0)
    except OSError as e:
        sys.stderr.write(f'Unable to detach from parent process: {e.strerror} {e.errno}')
    
    os.setsid()
    if cleanup:
        detach()

def read_state():
    if not os.path.isfile(STATE_FILE):
        return None
    with open(STATE_FILE, 'r') as file:
        return file.read().replace('\n', '')

def save_state(state):
    with open(STATE_FILE, 'w') as file:
        file.write(state)

def _device_cmd(device, cmd):
    with open(device, 'a') as file:
        file.write(cmd)

def _load_module(module, args = []):
    result = subprocess.run(['modprobe', module] + args)
    if result.returncode != 0:
        sys.stderr.write(f'Unable to load kernel module {module}! modprobe returned code {result.returncode}')
        exit(1)

def _unload_module(module, args = []):
    result = subprocess.run(['modprobe', '-r', module] + args)
    if result.returncode != 0:
        sys.stderr.write(f'Unable to unload kernel module {module}! modprobe returned code {result.returncode}')
        exit(1)

def load_gpu_modules(enabled = True):
    for module in NVIDIA_MODULES:
        if enabled:
            _load_module(module)
        else:
            _unload_module(module)

def set_libvirt_vmbind(bound = True):
    # try to open libvirt
    virt = libvirt.open('qemu:///system')
    vm = virt.lookupByName('win10')
    if not vm:
        sys.stderr.write(f'Unable to find virtual machine with UUID {VM_ID}!')
        return

    if bound and vm.state == 'paused':
        vm.suspend()
    elif not bound and vm.state == 'running':
        vm.resume()
    else:
        sys.stdout.write(f'Unsupported VM state {vm.state}, not doing anything')

def set_gpu_vmbind(bound = True):
    if bound:
        ids = f'ids={NVIDIA_PCI_VID}'
        _load_module('vfio_pci', [ids])

        # persist
        with open(VFIO_FILE, 'w') as file:
            file.write(f'options vfio-pci {ids}')

        _device_cmd(f'/sys/bus/pci/devices/{NVIDIA_PCI_ID}/driver/unbind', NVIDIA_PCI_ID)
        _device_cmd(f'/sys/bus/pci/drivers/vfio-pci/bind', NVIDIA_PCI_ID)
    else:
        _device_cmd(f'/sys/bus/pci/devices/{NVIDIA_PCI_ID}/remove', '1')
        _device_cmd(f'/sys/bus/pci/rescan', '1')
        _unload_module('vfio_pci')

        # persist
        os.remove(VFIO_FILE)

def switch_gpu(state_src, state_dst):
    detach_gpu = False
    update_initramfs = False

    if state_dst == 'xserver':
        detach_gpu = False
    elif state_dst == 'detached':
        detach_gpu = True
    elif state_dst == 'virtual':
        detach_gpu = True
    
    # process should now be a daemon and so we can shutdown the display manager
    result = subprocess.run(['service', 'gdm3', 'stop'])
    if result.returncode != 0:
        sys.stderr.write(f'Unable to stop the display manager!')
        exit(1)

    # switch the gpu profile
    switcher = Switcher()
    if detach_gpu:
        load_gpu_modules(False)
        switcher.enable_profile('intel')

        if state_dst == 'virtual':
            set_gpu_vmbind(True)
            set_libvirt_vmbind(True)
            update_initramfs = True
    else:
        if state_src == 'virtual':
            print('unbinding from virtual')
            set_libvirt_vmbind(False)
            set_gpu_vmbind(False)
            update_initramfs = True

        switcher.enable_profile('on-demand')
        load_gpu_modules(True)

    # restart gdm3
    result = subprocess.run(['service', 'gdm3', 'start'])
    if result.returncode != 0:
        sys.stderr.write(f'Unable to start the display manager! Perhaps there was a configuration problem?')
        exit(1)

    sys.stdout.write('GPU has been switched and X restarted.')
    save_state(state_dst)

    # need to do this if we modify a modprobe.d file
    # run in background as we shouldn't delay starting dm for this
    if update_initramfs:
        subprocess.run(['update-initramfs', '-u'])
    
    exit(0)

def print_help():
    print('Nvidia VFIO Manager 1.0')
    print('Usage:\n')
    print('switch [xserver|virtual|detached] - Switches the nVidia GPU to the specified destination')

def _switch_gpu_cmd(argv):
    if not os.geteuid() == 0:
        sys.stderr.write("This operation requires root privileges\n")
        exit(1)

    state_src = read_state()
    state_dest = argv[2]
    if not state_dest in ['xserver', 'virtual', 'detached']:
        print_help()
        exit(1)
    if state_src == state_dest:
        sys.stderr.write(f'GPU is already attached to {state_src}!')
        exit(1)

    # daemonize
    detach(True)
    switch_gpu(state_src, state_dest)

commands = {
    'switch': _switch_gpu_cmd
}

if len(sys.argv) < 2:
    print_help()
    exit(1)

fun = commands.get(sys.argv[1])
if not fun:
    print_help()
    exit(1)

fun(sys.argv)