GPU Switcher

nvmgr.py

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

nvidia_modules = [
    'nvidia_uvm',
    'nvidia_drm',
    'nvidia_modeset'
]

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 load_gpu_modules(enabled = True):
    for module in nvidia_modules:
        if enabled:
            result = subprocess.run(['modprobe', module])
            if result.returncode != 0:
                sys.stderr.write(f'Unable to load kernel module {module}! modprobe returned code {result.returncode}')
                exit(1)
        else:
            result = subprocess.run(['rmmod', module])
            if result.returncode != 0:
                sys.stderr.write(f'Unable to unload kernel module {module}! rmmod returned code {result.returncode}')
                exit(1)

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

    # daemonize
    detach(True)

    vmid = None
    to_vm = False
    if argv[2] != 'host':
        vmid = argv[2]
        to_vm = 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 to_vm:
        load_gpu_modules(False)
        switcher.enable_profile('intel')
    else:
        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.')
    exit(0)

def print_help():
    print('Usage:\n')
    print('switch [host/(vmid)]')

commands = {
    'switch': switch_gpu
}

fun = commands.get(sys.argv[1])
if not fun:
    print(f'Invalid command: {sys.argv[1]}')
    print_help()
    exit(1)

fun(sys.argv)