path: root/scripts/stack_manager.sh

#!/bin/bash
# stack_manager.sh - Complete vGPU stack management (load, unload, test, debug)
# Usage: ./stack_manager.sh <command> [options]

set -euo pipefail

BASE_DIR="${OPEN_IOV_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
NVIDIA_KO="$BASE_DIR/nvidia-open-gpu-kernel-modules/kernel-open/nvidia.ko"
VFIO_KO="$BASE_DIR/kernel/open-iov-vfio.ko"
FIRMWARE_PATH="$BASE_DIR/firmware/57X-GSP/570.144/gsp-570.144.fwimage.bin"
FIRMWARE_DEST="/lib/firmware/gsp-570.144.fwimage.bin"

MDEV_SCRIPT="$BASE_DIR/scripts/create_mdev.sh"
# Legacy direct-QEMU script retained for reference only - VM launch MUST go
# through the harness (vgpu_vm.py). See cmd_launch_qemu.
QEMU_SCRIPT="$BASE_DIR/scripts/launch_qemu_vgpu.sh"
TEST_SCRIPT="$BASE_DIR/scripts/test_vgpu_stack.sh"

color_red='\033[0;31m'
color_green='\033[0;32m'
color_yellow='\033[1;33m'
color_blue='\033[0;34m'
color_nc='\033[0m'

log_info() { echo -e "${color_blue}[INFO]${color_nc} $*"; }
log_ok() { echo -e "${color_green}[OK]${color_nc} $*"; }
log_warn() { echo -e "${color_yellow}[WARN]${color_nc} $*"; }
log_err() { echo -e "${color_red}[ERROR]${color_nc} $*"; }

# ============================================================
# BUILD COMMANDS (correct order matters!)
# ============================================================

cmd_build_kernel() {
    log_info "Building open-iov-vfio.ko (kernel module)..."
    cd "$BASE_DIR"
    make kernel
    log_ok "open-iov-vfio.ko built"
}

cmd_build_nvidia() {
    log_info "Building NVIDIA OpenRM with vGPU symbols (KBUILD_EXTRA_SYMBOLS)..."
    cd "$BASE_DIR"
    
    # Step 1: Build nvidia kernel module (nv-kernel.o)
    log_info "Building nvidia kernel module (nv-kernel.o)..."
    cd "$BASE_DIR/nvidia-open-gpu-kernel-modules/src/nvidia"
    make -j$(nproc)
    log_ok "nvidia kernel module built"
    
    # Step 2: Build nvidia-modeset kernel module (nv-modeset-kernel.o)
    log_info "Building nvidia-modeset kernel module (nv-modeset-kernel.o)..."
    cd "$BASE_DIR/nvidia-open-gpu-kernel-modules/src/nvidia-modeset"
    make -j$(nproc)
    log_ok "nvidia-modeset kernel module built"
    
    # Step 3: Build kernel-open modules with KBUILD_EXTRA_SYMBOLS
    log_info "Building kernel-open modules with KBUILD_EXTRA_SYMBOLS..."
    cd "$BASE_DIR/nvidia-open-gpu-kernel-modules/kernel-open"
    KBUILD_EXTRA_SYMBOLS="$BASE_DIR/kernel/Module.symvers" make -j$(nproc) modules
    log_ok "NVIDIA OpenRM built"
}

cmd_build_userspace() {
    log_info "Building userspace (libiov.so + examples)..."
    cd "$BASE_DIR"
    make userspace
    make examples
    log_ok "Userspace built"
}

cmd_build_all() {
    log_info "=== BUILD ALL (correct order) ==="
    cmd_build_kernel
    cmd_build_userspace
    log_ok "=== ALL BUILDS COMPLETE ==="
    log_info "NOTE: NVIDIA submodule (610-based) NOT built by default."
    log_info "      Only needed for NVIDIA's native vGPU management (not QEMU)."
    log_info "      Build separately with: stack_manager.sh build-nvidia"
}

# ============================================================
# LOAD COMMANDS (correct order matters!)
# ============================================================

cmd_load_prereqs() {
    log_info "Loading kernel prerequisites (vfio, mdev)..."
    # vfio_mdev does not exist on kernels >= 6.4 (mdev core is standalone mdev.ko);
    # modprobe each individually so a missing one never blocks the rest.
    local m
    for m in vfio vfio_iommu_type1 vfio_pci mdev vfio_mdev; do
        if modinfo "$m" &>/dev/null; then
            sudo modprobe "$m" 2>/dev/null || true
        fi
    done
    lsmod | grep -E "vfio|mdev"
}

cmd_load_open_iov() {
    if [[ ! -f "$VFIO_KO" ]]; then
        log_err "Missing $VFIO_KO - run 'build' first"
        return 1
    fi
    if ! lsmod | grep -q "^open_iov_vfio "; then
        log_info "Loading open-iov-vfio.ko..."
        sudo insmod "$VFIO_KO" || { log_err "Failed to load open-iov-vfio.ko"; return 1; }
        log_ok "open-iov-vfio.ko loaded"
    else
        log_ok "open-iov-vfio.ko already loaded"
    fi
}

cmd_load_nvidia() {
    if [[ ! -f "$NVIDIA_KO" ]]; then
        log_warn "NVIDIA submodule not built - skipping (apt nvidia works for host GPU)"
        return 0
    fi
    if lsmod | grep -q "^nvidia "; then
        # A module named 'nvidia' is already loaded. For QEMU/vGPU with
        # open-iov-vfio.ko (emulated VFIO device), we do NOT need our
        # submodule nvidia.ko on the host - the apt nvidia driver manages
        # the physical GPU and open-iov-vfio handles the emulated vGPU.
        # Loading our submodule would conflict with the apt driver and
        # potentially crash the GSP (WPR2 stuck). Skip it.
        local loaded_ver
        loaded_ver=$(cat /sys/module/nvidia/version 2>/dev/null)
        log_ok "apt nvidia ($loaded_ver) already loaded - submodule not needed for QEMU"
        return 0
    fi
    log_info "Loading nvidia.ko (OpenRM)..."
    sudo insmod "$NVIDIA_KO" || { log_err "Failed to load nvidia.ko"; return 1; }
    log_ok "nvidia.ko loaded"
}

cmd_install_firmware() {
    if [[ ! -f "$FIRMWARE_PATH" ]]; then
        log_err "Missing firmware at $FIRMWARE_PATH"
        return 1
    fi
    if [[ ! -f "$FIRMWARE_DEST" ]] || ! cmp -s "$FIRMWARE_PATH" "$FIRMWARE_DEST"; then
        log_info "Installing firmware to $FIRMWARE_DEST..."
        sudo cp "$FIRMWARE_PATH" "$FIRMWARE_DEST"
        log_ok "Firmware installed"
    else
        log_ok "Firmware already installed"
    fi
}

cmd_load_stack() {
    log_info "=== LOAD STACK (correct order) ==="
    cmd_load_prereqs
    cmd_load_open_iov
    cmd_install_firmware
    # NOTE: NVIDIA submodule NOT loaded for QEMU/vGPU work.
    # apt nvidia (570.211.01) manages the physical GPU;
    # open-iov-vfio handles the emulated vGPU.
    log_ok "=== DRIVER STACK LOADED ==="
    lsmod | grep -E "nvidia|open_iov|vfio|mdev"
}

# Full cycle: build all + load stack + create mdev
cmd_up() {
    log_info "=== FULL STARTUP: build + load + mdev ==="
    cmd_build_all
    cmd_load_stack
    cmd_create_mdev
    log_ok "=== SYSTEM READY FOR QEMU ==="
    echo "Run: ./stack_manager.sh qemu"
}

# ============================================================
# UNLOAD COMMANDS
# ============================================================

cmd_unload_stack() {
    log_info "=== UNLOAD STACK ==="

    if [[ -e /sys/bus/mdev/devices/* ]]; then
        log_warn "MDEV devices exist - removing them first..."
        for mdev in /sys/bus/mdev/devices/*; do
            uuid=$(basename "$mdev")
            log_info "Removing mdev $uuid..."
            echo 1 | sudo tee "$mdev/remove" >/dev/null
        done
        sleep 2
    fi

    # Only unload OUR submodule nvidia.ko if it was loaded (not apt nvidia).
    # The apt nvidia driver (e.g. 570.211.01) manages the physical GPU and
    # must stay loaded for nvidia-smi to keep working.
    if [[ -f "$NVIDIA_KO" ]] && lsmod | grep -q "^nvidia "; then
        local loaded_ver want_ver
        loaded_ver=$(cat /sys/module/nvidia/version 2>/dev/null)
        want_ver=$(modinfo -F version "$NVIDIA_KO" 2>/dev/null)
        if [[ -n "$loaded_ver" && "$loaded_ver" == "$want_ver" ]]; then
            for mod in nvidia_uvm nvidia_modeset nvidia_drm nvidia_peermem nvidia; do
                if lsmod | grep -qE "^$mod\>"; then
                    log_info "Unloading submodule $mod..."
                    sudo rmmod "$mod" || { log_err "Failed to unload $mod"; return 1; }
                    log_ok "$mod unloaded"
                fi
            done
        else
            log_ok "apt nvidia ($loaded_ver) kept loaded for host GPU"
        fi
    fi

    if lsmod | grep -qE "^open_iov_vfio\>"; then
        log_info "Unloading open-iov-vfio.ko..."
        sudo rmmod open_iov_vfio || { log_err "Failed to unload open-iov-vfio.ko"; return 1; }
        log_ok "open-iov-vfio.ko unloaded"
    fi

    sudo modprobe -r vfio_mdev vfio_pci vfio_iommu_type1 vfio mdev 2>/dev/null || true
    log_ok "=== DRIVER STACK UNLOADED ==="
}

# ============================================================
# MDEV COMMANDS
# ============================================================

cmd_create_mdev() {
    log_info "Creating mdev device..."
    if [[ ! -f "$MDEV_SCRIPT" ]]; then
        log_err "create_mdev.sh not found"
        return 1
    fi
    "$MDEV_SCRIPT" "$@"
}

cmd_list_mdev() {
    log_info "MDEV devices:"
    ls -la /sys/bus/mdev/devices/ 2>/dev/null || log_warn "No mdev devices"
}

# ============================================================
# QEMU COMMANDS
# ============================================================

cmd_launch_qemu() {
    # WORKFLOW RULE: always use the QEMU harness (vgpu_vm.py), never raw
    # QEMU directly - the harness owns the mdev lifecycle, QMP,
    # guest-agent and serial logging.
    local harness="$BASE_DIR/scripts/vgpu_vm.py"
    if [[ ! -f "$harness" ]]; then
        log_err "vgpu_vm.py harness not found at $harness"
        return 1
    fi
    log_info "Launching QEMU with vGPU via harness (vgpu_vm.py)..."
    "$harness" start "${1:-vgpu-linux}"
}

# ============================================================
# TEST COMMANDS
# ============================================================

cmd_test_stack() {
    log_info "Running comprehensive test suite..."
    if [[ ! -f "$TEST_SCRIPT" ]]; then
        log_err "test_vgpu_stack.sh not found"
        return 1
    fi
    "$TEST_SCRIPT" "$@"
}

# ============================================================
# STATUS / DEBUG
# ============================================================

cmd_status() {
    echo "=== Kernel Modules ==="
    lsmod | grep -E "nvidia|open_iov|vfio|mdev" || echo "  (none)"
    echo ""
    echo "=== MDEV Devices ==="
    ls -la /sys/bus/mdev/devices/ 2>/dev/null || echo "  (none)"
    echo ""
    echo "=== VFIO Devices ==="
    ls -la /dev/vfio/ 2>/dev/null || echo "  (none)"
    echo ""
    echo "=== Firmware ==="
    ls -la "$FIRMWARE_DEST" 2>/dev/null || echo "  Not installed"
    echo ""
    echo "=== dmesg (last 20) ==="
    sudo dmesg -T | grep -i -e nvidia -e gsp -e vgpu -e open-iov -e mdev | tail -20
}

cmd_dmesg() {
    sudo dmesg -T -w | grep -i -e nvidia -e gsp -e vgpu -e open-iov -e mdev
}

cmd_debug_qemu() {
    # WORKFLOW RULE: harness only. The vgpu-linux.yaml already sets
    # serialLog=/tmp/vgpu-guest-serial.log and gdbPort=1234, so a normal
    # harness start IS the debug session: attach GDB with
    #   gdb -ex 'target remote localhost:1234'
    # and follow the guest serial log plus host dmesg.
    local harness="$BASE_DIR/scripts/vgpu_vm.py"
    if [[ ! -f "$harness" ]]; then
        log_err "vgpu_vm.py harness not found at $harness"
        return 1
    fi
    log_info "Starting debug session via harness (GDB stub tcp::1234, serial /tmp/vgpu-guest-serial.log)..."
    "$harness" start "${1:-vgpu-linux}"
    log_info "Attach GDB: gdb -ex 'target remote localhost:1234'"
    log_info "Serial log: tail -f /tmp/vgpu-guest-serial.log"
}

cmd_cleanup_crash() {
    log_info "Cleaning up after crash..."
    cmd_unload_stack
    log_info "Clearing any stuck mdev devices..."
    for mdev in /sys/bus/mdev/devices/*; do
        [[ -e "$mdev" ]] || continue
        uuid=$(basename "$mdev")
        echo 1 | sudo tee "$mdev/remove" >/dev/null 2>&1 || true
    done
    log_ok "Cleanup complete"
}

# ============================================================
# HELP
# ============================================================

cmd_help() {
    cat <<EOF
stack_manager.sh - Complete vGPU Stack Management

USAGE:
    ./stack_manager.sh <command> [options]

BUILD COMMANDS (run in order):
    build-kernel      Build open-iov-vfio.ko (generates Module.symvers)
    build-nvidia      Build NVIDIA OpenRM with KBUILD_EXTRA_SYMBOLS
    build-userspace   Build libiov.so + examples
    build-all         Run all three build steps in correct order

LOAD COMMANDS (run in order):
    load-prereqs      Load vfio, vfio_iommu_type1, vfio_pci, mdev, vfio_mdev
    load-open-iov     Load open-iov-vfio.ko
    load-nvidia       Load nvidia.ko (OpenRM)
    install-firmware  Copy GSP firmware to /lib/firmware/
    load-stack        Run all four load steps in correct order

FULL CYCLE:
    up                Build all + load stack + create mdev (one command!)

MDEV COMMANDS:
    mdev-create       Create mdev device
    mdev-list         List mdev devices

QEMU COMMANDS:
    qemu              Launch QEMU with vGPU
    qemu-debug        Launch QEMU with GDB + guest monitor

TEST COMMANDS:
    test              Run comprehensive test suite

STATUS / DEBUG:
    status            Show current stack status
    dmesg             Follow filtered dmesg output
    cleanup-crash     Emergency cleanup after crash/hang
    help              Show this help

EXAMPLES:
    # Full startup from clean state:
    ./stack_manager.sh up

    # Step by step:
    ./stack_manager.sh build-all
    ./stack_manager.sh load-stack
    ./stack_manager.sh mdev-create
    ./stack_manager.sh qemu

    # Or just load (if already built):
    ./stack_manager.sh load-stack
    ./stack_manager.sh mdev-create
    ./stack_manager.sh qemu

    # Testing:
    ./stack_manager.sh test

    # Cleanup:
    ./stack_manager.sh unload
    # Or force cleanup after crash:
    ./stack_manager.sh cleanup-crash

NOTE:
    - Requires root for module operations
    - BUILD ORDER: open-iov-vfio.ko FIRST (generates Module.symvers),
                   THEN nvidia.ko with KBUILD_EXTRA_SYMBOLS
    - LOAD ORDER: vfio/mdev prerequisites -> open-iov-vfio.ko -> nvidia.ko
    - UNLOAD ORDER: remove mdevs -> open-iov-vfio.ko -> nvidia.ko
    - NEVER REBOOT without explicit approval
EOF
}

main() {
    case "${1:-help}" in
        # Build
        build-kernel) cmd_build_kernel ;;
        build-nvidia) cmd_build_nvidia ;;
        build-userspace) cmd_build_userspace ;;
        build-all) cmd_build_all ;;

        # Load
        load-prereqs) cmd_load_prereqs ;;
        load-open-iov) cmd_load_open_iov ;;
        load-nvidia) cmd_load_nvidia ;;
        install-firmware) cmd_install_firmware ;;
        load-stack) cmd_load_stack ;;

        # Full cycle
        up) cmd_up ;;

        # Unload
        unload) cmd_unload_stack ;;

        # MDEV
        mdev-create) shift; cmd_create_mdev "$@" ;;
        mdev-list) cmd_list_mdev ;;

        # QEMU
        qemu) shift; cmd_launch_qemu "$@" ;;
        qemu-debug) shift; cmd_debug_qemu "$@" ;;

        # Test
        test) shift; cmd_test_stack "$@" ;;

        # Status/Debug
        status) cmd_status ;;
        dmesg) cmd_dmesg ;;
        cleanup-crash) cmd_cleanup_crash ;;
        help|*) cmd_help ;;
    esac
}

main "$@"