add gpu support

crates/openshell-driver-vm/runtime/kernel/openshell.kconfig

@@ -123,6 +123,28 @@ CONFIG_MEMCG=y
 CONFIG_POSIX_MQUEUE=y
 CONFIG_POSIX_MQUEUE_SYSCTL=y
 
+# ── PCI/PCIe (required for GPU passthrough via QEMU vfio-pci) ────────────
+# The libkrunfw base config disables CONFIG_PCI. GPU sandboxes using the
+# QEMU backend pass the GPU through as a PCIe device on a q35 machine.
+# Without PCI core support the guest kernel cannot see any PCI bus, so the
+# nvidia driver loads but finds zero devices.
+CONFIG_PCI=y
+CONFIG_PCI_MSI=y
+CONFIG_PCIEPORTBUS=y
+
+# ── Loadable kernel modules (required for GPU passthrough) ──────────────
+# The libkrunfw base config disables CONFIG_MODULES. GPU sandboxes need it
+# to load nvidia.ko, nvidia-uvm.ko, and nvidia-modeset.ko at boot via
+# modprobe. Without this, the guest kernel rejects all module loads.
+#
+# SECURITY NOTE: This enables module loading for ALL VMs (including
+# non-GPU), expanding the guest kernel attack surface. The sandbox
+# supervisor's seccomp profile must block init_module/finit_module
+# syscalls for the sandbox user to prevent arbitrary module loading.
+# Tracked: consider per-purpose kernel builds (GPU vs non-GPU).
+CONFIG_MODULES=y
+CONFIG_MODULE_UNLOAD=y
+
 # ── Security features required by the sandbox runtime ───────────────────
 CONFIG_SECURITY_LANDLOCK=y
 CONFIG_SECCOMP_FILTER=y

crates/openshell-driver-vm/scripts/openshell-vm-sandbox-init.sh

@@ -234,6 +234,17 @@ create_gpu_device_nodes_mknod() {
 setup_gpu() {
     ts "GPU_ENABLED=true — initializing GPU passthrough"
 
+    # Kernel modules are built for a specific guest kernel version.
+    # If the running kernel doesn't match, depmod/modprobe will silently fail.
+    local expected_kver="6.12.76"
+    local actual_kver
+    actual_kver="$(uname -r)"
+    if [ "${actual_kver}" != "${expected_kver}" ]; then
+        ts "WARNING: kernel version mismatch: expected ${expected_kver}, got ${actual_kver}"
+        ts "         GPU modules are installed under lib/modules/${expected_kver}/"
+        ts "         modprobe may fail to find them"
+    fi
+
     if ! command -v modprobe >/dev/null 2>&1; then
         ts "FATAL: modprobe not found; cannot load nvidia kernel modules"
         return 1
@@ -249,6 +260,11 @@ setup_gpu() {
         fi
     fi
 
+    ts "generating module dependency index"
+    if ! depmod -a "$(uname -r)" 2>/dev/null; then
+        ts "WARNING: depmod failed; modprobe may not find modules"
+    fi
+
     ts "loading nvidia kernel modules"
     modprobe nvidia || { ts "FATAL: modprobe nvidia failed"; return 1; }
     modprobe nvidia_uvm 2>/dev/null || true

crates/openshell-driver-vm/src/driver.rs

@@ -5,8 +5,8 @@ use crate::gpu::{
     GpuInventory, SubnetAllocator, allocate_vsock_cid, mac_from_sandbox_id, tap_device_name,
 };
 use crate::rootfs::{
-    create_rootfs_archive_from_dir, extract_rootfs_archive_to,
-    prepare_sandbox_rootfs_from_image_root, sandbox_guest_init_path,
+    create_rootfs_archive_from_dir, extract_rootfs_archive_to, inject_gpu_modules,
+    prepare_sandbox_rootfs_from_image_root, refresh_runtime_artifacts, sandbox_guest_init_path,
 };
 use bollard::Docker;
 use bollard::errors::Error as BollardError;
@@ -419,6 +419,28 @@ impl VmDriver {
                 return Err(err);
             }
         };
+        if is_gpu {
+            let rootfs_for_gpu = rootfs.clone();
+            let driver_state_dir = self.config.state_dir.clone();
+            if let Err(err) = tokio::task::spawn_blocking(move || {
+                inject_gpu_modules(&rootfs_for_gpu, &driver_state_dir)
+            })
+            .await
+            .map_err(|e| Status::internal(format!("GPU module injection panicked: {e}")))?
+            {
+                warn!(
+                    sandbox_id = %sandbox.id,
+                    error = %err,
+                    "vm driver: GPU module injection failed"
+                );
+                let _ = tokio::fs::remove_dir_all(&state_dir).await;
+                return Err(Status::failed_precondition(format!(
+                    "GPU module injection failed: {err}"
+                )));
+            }
+            info!(sandbox_id = %sandbox.id, "vm driver: GPU modules injected into rootfs");
+        }
+
         if let Some(tls_paths) = tls_paths.as_ref()
             && let Err(err) = prepare_guest_tls_materials(&rootfs, tls_paths).await
         {
@@ -740,10 +762,13 @@ impl VmDriver {
             .await?;
         let archive_path = image_cache_rootfs_archive(&self.config.state_dir, &image_identity);
         let rootfs_dest = rootfs.to_path_buf();
-        tokio::task::spawn_blocking(move || extract_rootfs_archive_to(&archive_path, &rootfs_dest))
-            .await
-            .map_err(|err| Status::internal(format!("sandbox rootfs extraction panicked: {err}")))?
-            .map_err(|err| Status::internal(format!("extract sandbox rootfs failed: {err}")))?;
+        tokio::task::spawn_blocking(move || {
+            extract_rootfs_archive_to(&archive_path, &rootfs_dest)?;
+            refresh_runtime_artifacts(&rootfs_dest)
+        })
+        .await
+        .map_err(|err| Status::internal(format!("sandbox rootfs extraction panicked: {err}")))?
+        .map_err(|err| Status::internal(format!("extract sandbox rootfs failed: {err}")))?;
 
         Ok(image_identity)
     }

crates/openshell-driver-vm/src/rootfs.rs

@@ -4,7 +4,7 @@
 use std::fs;
 use std::fs::File;
 use std::io::{BufWriter, Cursor};
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 const SUPERVISOR: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/openshell-sandbox.zst"));
 const ROOTFS_VARIANT_MARKER: &str = ".openshell-rootfs-variant";
@@ -29,6 +29,32 @@ pub fn prepare_sandbox_rootfs_from_image_root(
     Ok(())
 }
 
+/// Re-inject the init script and supervisor binary into an already-prepared
+/// rootfs. The image rootfs archive cache is keyed by image digest, so a
+/// driver rebuild does not invalidate it. Calling this after extraction
+/// ensures the guest always runs the init script and supervisor that match
+/// the running driver binary.
+pub fn refresh_runtime_artifacts(rootfs: &Path) -> Result<(), String> {
+    let init_path = rootfs.join("srv/openshell-vm-sandbox-init.sh");
+    if let Some(parent) = init_path.parent() {
+        fs::create_dir_all(parent).map_err(|e| format!("create {}: {e}", parent.display()))?;
+    }
+    fs::write(
+        &init_path,
+        include_str!("../scripts/openshell-vm-sandbox-init.sh"),
+    )
+    .map_err(|e| format!("write {}: {e}", init_path.display()))?;
+    #[cfg(unix)]
+    {
+        use std::os::unix::fs::PermissionsExt as _;
+        fs::set_permissions(&init_path, fs::Permissions::from_mode(0o755))
+            .map_err(|e| format!("chmod {}: {e}", init_path.display()))?;
+    }
+
+    ensure_supervisor_binary(rootfs)?;
+    Ok(())
+}
+
 pub fn extract_rootfs_archive_to(archive_path: &Path, dest: &Path) -> Result<(), String> {
     if dest.exists() {
         fs::remove_dir_all(dest)
@@ -182,6 +208,350 @@ pub fn validate_sandbox_rootfs(rootfs: &Path) -> Result<(), String> {
     Ok(())
 }
 
+/// Kernel version of the libkrunfw guest. Modules must be compiled against
+/// this exact version; a mismatch causes `modprobe` failures at boot.
+///
+/// Keep in sync with:
+///   - `tasks/scripts/vm/build-nvidia-modules.sh` (KERNEL_TREE path)
+///   - `openshell-vm-sandbox-init.sh` `setup_gpu()` expected version
+const GUEST_KERNEL_VERSION: &str = "6.12.76";
+
+/// Inject NVIDIA kernel modules, firmware, and `kmod` tooling into a prepared
+/// sandbox rootfs. Called by the driver when a sandbox requests GPU support.
+///
+/// Module source resolution order:
+///   1. `OPENSHELL_GPU_MODULES_DIR` environment variable
+///   2. `<state_dir>/gpu-modules/` (pre-provisioned by the operator)
+///
+/// Firmware source resolution (first match wins):
+///   0. Rootfs already contains `.bin` files (e.g. from the image's `.run`
+///      installer) — **skip injection entirely** to avoid version mismatch.
+///   1. `<modules_dir>/../nvidia-firmware/`
+///   2. Host `/lib/firmware/nvidia/`
+///
+/// Returns an error only if module injection is impossible (no source found
+/// or a write fails). Missing firmware emits a warning and continues.
+pub fn inject_gpu_modules(rootfs: &Path, state_dir: &Path) -> Result<(), String> {
+    let modules_dir = resolve_gpu_modules_dir(state_dir)?;
+
+    let ko_files: Vec<PathBuf> = fs::read_dir(&modules_dir)
+        .map_err(|e| format!("read GPU modules dir {}: {e}", modules_dir.display()))?
+        .filter_map(|entry| {
+            let entry = entry.ok()?;
+            let path = entry.path();
+            if path.extension().is_some_and(|ext| ext == "ko") {
+                Some(path)
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    if ko_files.is_empty() {
+        return Err(format!(
+            "GPU modules dir {} contains no .ko files",
+            modules_dir.display()
+        ));
+    }
+
+    let modules_dst = rootfs.join(format!(
+        "lib/modules/{GUEST_KERNEL_VERSION}/kernel/drivers/nvidia"
+    ));
+    fs::create_dir_all(&modules_dst)
+        .map_err(|e| format!("create {}: {e}", modules_dst.display()))?;
+
+    for ko in &ko_files {
+        let dest = modules_dst.join(ko.file_name().unwrap());
+        let bytes_copied = fs::copy(ko, &dest).map_err(|e| {
+            format!(
+                "copy {} -> {}: {e}",
+                ko.display(),
+                dest.display()
+            )
+        })?;
+        tracing::info!(
+            module = %ko.file_name().unwrap().to_string_lossy(),
+            size_bytes = bytes_copied,
+            src = %ko.display(),
+            "injected GPU kernel module"
+        );
+    }
+
+    inject_gpu_firmware(rootfs, &modules_dir);
+    ensure_kmod_symlinks(rootfs);
+    warn_missing_gpu_userspace(rootfs);
+
+    Ok(())
+}
+
+/// Check whether the rootfs contains essential GPU userspace binaries.
+/// Emits actionable warnings when the sandbox image lacks nvidia-smi
+/// or CUDA libraries — common when `--gpu` is used with a non-GPU base
+/// image like `ubuntu:latest` instead of the GPU sandbox Dockerfile.
+fn warn_missing_gpu_userspace(rootfs: &Path) {
+    let nvidia_smi_candidates = [
+        "usr/bin/nvidia-smi",
+        "usr/local/bin/nvidia-smi",
+        "bin/nvidia-smi",
+    ];
+    let has_nvidia_smi = nvidia_smi_candidates
+        .iter()
+        .any(|p| rootfs.join(p).exists());
+
+    if !has_nvidia_smi {
+        tracing::warn!(
+            "GPU sandbox image does not contain nvidia-smi. The sandbox will \
+             have GPU kernel modules but no NVIDIA userspace tools. Use a \
+             GPU-enabled image (e.g. --from ./sandboxes/nvidia-gpu/Dockerfile) \
+             or install the NVIDIA driver userspace in your image."
+        );
+    }
+}
+
+/// Locate the directory containing pre-built NVIDIA `.ko` files.
+///
+/// Resolution order:
+///   1. `OPENSHELL_GPU_MODULES_DIR` env var (explicit override)
+///   2. `<state_dir>/gpu-modules/` (operator pre-provisioned)
+///   3. `<project_root>/target/libkrun-build/nvidia-modules/` (build tree,
+///       discovered relative to the driver executable)
+///   4. Host `/lib/modules/<GUEST_KERNEL_VERSION>/kernel/drivers/nvidia/`
+fn resolve_gpu_modules_dir(state_dir: &Path) -> Result<PathBuf, String> {
+    if let Ok(dir) = std::env::var("OPENSHELL_GPU_MODULES_DIR") {
+        let p = PathBuf::from(&dir);
+        if p.is_dir() {
+            tracing::info!(path = %p.display(), "using GPU modules from OPENSHELL_GPU_MODULES_DIR");
+            return Ok(p);
+        }
+        return Err(format!(
+            "OPENSHELL_GPU_MODULES_DIR={dir} is not a directory"
+        ));
+    }
+
+    let provisioned = state_dir.join("gpu-modules");
+    if provisioned.is_dir() {
+        tracing::info!(path = %provisioned.display(), "using pre-provisioned GPU modules");
+        return Ok(provisioned);
+    }
+
+    // Auto-discover from the build tree. The driver binary lives at
+    // `target/{debug,release}/openshell-driver-vm`, so the project root
+    // is two levels up. The old GPU rootfs script places modules at
+    // `target/libkrun-build/nvidia-modules/`.
+    if let Some(build_tree_dir) = discover_build_tree_modules() {
+        return Ok(build_tree_dir);
+    }
+
+    // Check common host-installed module paths.
+    for candidate in [
+        format!("/lib/modules/{GUEST_KERNEL_VERSION}/kernel/drivers/nvidia"),
+        format!("/lib/modules/{GUEST_KERNEL_VERSION}/extra/nvidia"),
+    ] {
+        let p = PathBuf::from(&candidate);
+        if dir_has_ko_files(&p) {
+            tracing::info!(path = %p.display(), "using host-installed GPU modules");
+            return Ok(p);
+        }
+    }
+
+    Err(format!(
+        "No GPU kernel modules found. Searched: OPENSHELL_GPU_MODULES_DIR (unset), \
+         {}, build tree, host /lib/modules/{}. \
+         Build modules with `mise run vm:nvidia-modules` \
+         or set OPENSHELL_GPU_MODULES_DIR.",
+        provisioned.display(),
+        GUEST_KERNEL_VERSION,
+    ))
+}
+
+/// Walk up from the driver executable to find `target/libkrun-build/nvidia-modules/`.
+///
+/// This is a development convenience — production deployments should use
+/// `OPENSHELL_GPU_MODULES_DIR` or pre-provision `<state_dir>/gpu-modules/`.
+fn discover_build_tree_modules() -> Option<PathBuf> {
+    #[cfg(unix)]
+    if unsafe { libc::getuid() } == 0 {
+        tracing::debug!("build-tree GPU module discovery running as root; \
+                         prefer OPENSHELL_GPU_MODULES_DIR in production");
+    }
+    let exe = std::env::current_exe().ok()?;
+    // exe is typically target/{debug,release}/openshell-driver-vm
+    let target_dir = exe.parent()?.parent()?;
+    let modules_dir = target_dir.join("libkrun-build/nvidia-modules");
+    if dir_has_ko_files(&modules_dir) {
+        tracing::info!(
+            path = %modules_dir.display(),
+            "auto-discovered GPU modules in build tree"
+        );
+        return Some(modules_dir);
+    }
+
+    // Also try CWD-relative (for `cargo run` or `mise run` from project root).
+    let cwd_candidate = PathBuf::from("target/libkrun-build/nvidia-modules");
+    if dir_has_ko_files(&cwd_candidate) {
+        let abs = cwd_candidate.canonicalize().unwrap_or(cwd_candidate.clone());
+        tracing::info!(
+            path = %abs.display(),
+            "auto-discovered GPU modules relative to CWD"
+        );
+        return Some(abs);
+    }
+
+    None
+}
+
+fn dir_has_ko_files(dir: &Path) -> bool {
+    if !dir.is_dir() {
+        return false;
+    }
+    let Some(entries) = fs::read_dir(dir).ok() else {
+        return false;
+    };
+    let mut has_uncompressed = false;
+    let mut has_compressed = false;
+    for entry in entries.flatten() {
+        let path = entry.path();
+        match path.extension().and_then(|e| e.to_str()) {
+            Some("ko") => has_uncompressed = true,
+            Some("zst" | "xz") => {
+                if path.file_stem().and_then(|s| std::path::Path::new(s).extension()).is_some_and(|ext| ext == "ko") {
+                    has_compressed = true;
+                }
+            }
+            _ => {}
+        }
+    }
+    if !has_uncompressed && has_compressed {
+        tracing::warn!(
+            path = %dir.display(),
+            "directory contains compressed .ko.zst/.ko.xz modules but only uncompressed .ko files are supported"
+        );
+    }
+    has_uncompressed
+}
+
+/// Copy NVIDIA GSP firmware into the rootfs. Non-fatal on failure.
+///
+/// Skips injection if the rootfs already contains `.bin` firmware files
+/// (e.g. the sandbox Docker image installed them via the NVIDIA `.run`
+/// installer). Overwriting image-provided firmware with build-tree or
+/// host firmware causes version mismatches when the host driver differs
+/// from the image's driver version.
+fn inject_gpu_firmware(rootfs: &Path, modules_dir: &Path) {
+    let fw_dst = rootfs.join("lib/firmware/nvidia");
+
+    if rootfs_has_firmware_bins(&fw_dst) {
+        tracing::info!(
+            path = %fw_dst.display(),
+            "rootfs already contains GPU firmware; skipping injection"
+        );
+        return;
+    }
+
+    // Try version-matched firmware next to the modules directory.
+    let fw_parent = modules_dir
+        .parent()
+        .map(|p| p.join("nvidia-firmware"));
+
+    if let Some(ref fw_dir) = fw_parent {
+        if fw_dir.is_dir() {
+            if let Err(e) = copy_dir_contents(fw_dir, &fw_dst) {
+                tracing::warn!(error = %e, "failed to copy version-matched firmware");
+            } else {
+                tracing::info!(src = %fw_dir.display(), "injected GPU firmware (version-matched)");
+                return;
+            }
+        }
+    }
+
+    // Fallback: host firmware
+    for candidate in ["/lib/firmware/nvidia", "/usr/lib/firmware/nvidia"] {
+        let host_fw = Path::new(candidate);
+        if host_fw.is_dir() {
+            if let Err(e) = copy_dir_contents(host_fw, &fw_dst) {
+                tracing::warn!(error = %e, src = candidate, "failed to copy host firmware");
+            } else {
+                tracing::info!(src = candidate, "injected GPU firmware from host");
+                return;
+            }
+        }
+    }
+
+    tracing::warn!(
+        "no NVIDIA GSP firmware found; GPU guests may fail to initialize. \
+         Place firmware in {:?} or host /lib/firmware/nvidia/",
+        fw_parent.as_deref().unwrap_or(Path::new("(unknown)"))
+    );
+}
+
+/// Check whether a firmware directory (or any subdirectory) contains `.bin` files.
+fn rootfs_has_firmware_bins(fw_dir: &Path) -> bool {
+    if !fw_dir.is_dir() {
+        return false;
+    }
+    let Ok(entries) = fs::read_dir(fw_dir) else {
+        return false;
+    };
+    for entry in entries.flatten() {
+        let path = entry.path();
+        if path.extension().is_some_and(|ext| ext == "bin") {
+            return true;
+        }
+        if path.is_dir() && rootfs_has_firmware_bins(&path) {
+            return true;
+        }
+    }
+    false
+}
+
+/// Ensure `modprobe`, `insmod`, etc. symlinks exist. Many minimal container
+/// images install `kmod` but lack the convenience symlinks in `/usr/sbin`.
+fn ensure_kmod_symlinks(rootfs: &Path) {
+    let kmod_candidates = ["bin/kmod", "usr/bin/kmod", "sbin/kmod", "usr/sbin/kmod"];
+    let kmod_exists = kmod_candidates
+        .iter()
+        .any(|p| rootfs.join(p).exists());
+
+    if !kmod_exists {
+        tracing::warn!("kmod not found in rootfs; modprobe will fail. \
+                        Ensure the sandbox image installs the 'kmod' package.");
+        return;
+    }
+
+    let sbin = rootfs.join("usr/sbin");
+    let _ = fs::create_dir_all(&sbin);
+    for tool in ["modprobe", "insmod", "rmmod", "lsmod", "depmod"] {
+        let link = sbin.join(tool);
+        if !link.exists() {
+            #[cfg(unix)]
+            {
+                let _ = std::os::unix::fs::symlink("../../bin/kmod", &link)
+                    .or_else(|_| std::os::unix::fs::symlink("/usr/bin/kmod", &link));
+            }
+        }
+    }
+}
+
+/// Recursively copy all files from `src` to `dst`, preserving directory structure.
+fn copy_dir_contents(src: &Path, dst: &Path) -> Result<(), String> {
+    fs::create_dir_all(dst).map_err(|e| format!("create {}: {e}", dst.display()))?;
+
+    for entry in fs::read_dir(src).map_err(|e| format!("read {}: {e}", src.display()))? {
+        let entry = entry.map_err(|e| format!("read entry in {}: {e}", src.display()))?;
+        let src_path = entry.path();
+        let dst_path = dst.join(entry.file_name());
+
+        if src_path.is_dir() {
+            copy_dir_contents(&src_path, &dst_path)?;
+        } else {
+            fs::copy(&src_path, &dst_path).map_err(|e| {
+                format!("copy {} -> {}: {e}", src_path.display(), dst_path.display())
+            })?;
+        }
+    }
+    Ok(())
+}
+
 fn ensure_sandbox_guest_user(rootfs: &Path) -> Result<(), String> {
     const SANDBOX_UID: u32 = 10001;
     const SANDBOX_GID: u32 = 10001;
@@ -417,6 +787,195 @@ mod tests {
         let _ = fs::remove_dir_all(&dir);
     }
 
+    #[test]
+    fn refresh_runtime_artifacts_overwrites_stale_init_script() {
+        let dir = unique_temp_dir();
+        let rootfs = dir.join("rootfs");
+
+        fs::create_dir_all(rootfs.join("srv")).expect("create srv");
+        fs::create_dir_all(rootfs.join("opt/openshell/bin")).expect("create openshell bin");
+        fs::write(
+            rootfs.join("srv/openshell-vm-sandbox-init.sh"),
+            b"#!/bin/bash\n# stale placeholder",
+        )
+        .expect("write stale init");
+        fs::write(
+            rootfs.join("opt/openshell/bin/openshell-sandbox"),
+            b"old-supervisor",
+        )
+        .expect("write stale supervisor");
+
+        refresh_runtime_artifacts(&rootfs).expect("refresh runtime artifacts");
+
+        let init_content =
+            fs::read_to_string(rootfs.join("srv/openshell-vm-sandbox-init.sh")).expect("read init");
+        assert!(
+            init_content.contains("setup_gpu"),
+            "refreshed init script should contain GPU setup logic"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn inject_gpu_modules_copies_ko_files() {
+        let dir = unique_temp_dir();
+        let modules_dir = dir.join("modules");
+        let rootfs = dir.join("rootfs");
+
+        fs::create_dir_all(&modules_dir).expect("create modules dir");
+        fs::create_dir_all(&rootfs).expect("create rootfs dir");
+        fs::write(modules_dir.join("nvidia.ko"), b"\x7fELF-fake-module-1").expect("write nvidia.ko");
+        fs::write(modules_dir.join("nvidia-uvm.ko"), b"\x7fELF-fake-module-2")
+            .expect("write nvidia-uvm.ko");
+
+        unsafe { std::env::set_var("OPENSHELL_GPU_MODULES_DIR", &modules_dir) };
+        let result = inject_gpu_modules(&rootfs, Path::new("/dummy/state"));
+        unsafe { std::env::remove_var("OPENSHELL_GPU_MODULES_DIR") };
+
+        result.expect("inject_gpu_modules should succeed");
+
+        let dest = rootfs.join("lib/modules/6.12.76/kernel/drivers/nvidia");
+        assert!(dest.join("nvidia.ko").is_file());
+        assert!(dest.join("nvidia-uvm.ko").is_file());
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn inject_gpu_modules_fails_with_no_ko_files() {
+        let dir = unique_temp_dir();
+        let modules_dir = dir.join("modules");
+
+        fs::create_dir_all(&modules_dir).expect("create modules dir");
+        fs::write(modules_dir.join("readme.txt"), b"not a kernel module").expect("write txt");
+
+        unsafe { std::env::set_var("OPENSHELL_GPU_MODULES_DIR", &modules_dir) };
+        let result = inject_gpu_modules(Path::new("/dummy/rootfs"), Path::new("/dummy/state"));
+        unsafe { std::env::remove_var("OPENSHELL_GPU_MODULES_DIR") };
+
+        let err = result.expect_err("should fail with no .ko files");
+        assert!(
+            err.contains("no .ko files"),
+            "error should mention 'no .ko files', got: {err}"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn inject_gpu_modules_fails_with_missing_dir() {
+        let dir = unique_temp_dir();
+        let missing = dir.join("does-not-exist");
+
+        unsafe { std::env::set_var("OPENSHELL_GPU_MODULES_DIR", &missing) };
+        let result = inject_gpu_modules(Path::new("/dummy/rootfs"), Path::new("/dummy/state"));
+        unsafe { std::env::remove_var("OPENSHELL_GPU_MODULES_DIR") };
+
+        let err = result.expect_err("should fail with missing directory");
+        assert!(
+            err.contains("not a directory"),
+            "error should mention 'not a directory', got: {err}"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn inject_gpu_firmware_skips_when_rootfs_has_bins() {
+        let dir = unique_temp_dir();
+        let rootfs = dir.join("rootfs");
+        let modules_dir = dir.join("modules");
+        let fw_dir = rootfs.join("lib/firmware/nvidia");
+
+        fs::create_dir_all(&fw_dir).expect("create firmware dir");
+        fs::create_dir_all(&modules_dir).expect("create modules dir");
+        fs::write(fw_dir.join("gsp.bin"), b"original-firmware-content").expect("write gsp.bin");
+
+        inject_gpu_firmware(&rootfs, &modules_dir);
+
+        let content = fs::read(fw_dir.join("gsp.bin")).expect("read gsp.bin after injection");
+        assert_eq!(
+            content,
+            b"original-firmware-content",
+            "firmware should not be overwritten when rootfs already has .bin files"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn ensure_kmod_symlinks_creates_links() {
+        let dir = unique_temp_dir();
+        let rootfs = dir.join("rootfs");
+
+        fs::create_dir_all(rootfs.join("bin")).expect("create bin");
+        fs::write(rootfs.join("bin/kmod"), b"kmod-stub").expect("write kmod");
+
+        ensure_kmod_symlinks(&rootfs);
+
+        assert!(
+            rootfs.join("usr/sbin/modprobe").exists(),
+            "modprobe symlink should exist"
+        );
+        assert!(
+            rootfs.join("usr/sbin/insmod").exists(),
+            "insmod symlink should exist"
+        );
+        assert!(
+            rootfs.join("usr/sbin/depmod").exists(),
+            "depmod symlink should exist"
+        );
+        assert!(
+            fs::symlink_metadata(rootfs.join("usr/sbin/modprobe"))
+                .unwrap()
+                .file_type()
+                .is_symlink(),
+            "modprobe should be a symlink"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn ensure_kmod_symlinks_warns_without_kmod() {
+        let dir = unique_temp_dir();
+        let rootfs = dir.join("rootfs");
+
+        fs::create_dir_all(&rootfs).expect("create rootfs");
+
+        ensure_kmod_symlinks(&rootfs);
+
+        assert!(
+            !rootfs.join("usr/sbin/modprobe").exists(),
+            "modprobe should not exist when kmod is missing"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn rootfs_has_firmware_bins_detects_nested() {
+        let dir1 = unique_temp_dir();
+        fs::create_dir_all(dir1.join("subdir")).expect("create subdir");
+        fs::write(dir1.join("subdir/file.bin"), b"firmware").expect("write .bin");
+        assert!(
+            rootfs_has_firmware_bins(&dir1),
+            "should detect .bin in nested subdir"
+        );
+        let _ = fs::remove_dir_all(&dir1);
+
+        let dir2 = unique_temp_dir();
+        fs::create_dir_all(dir2.join("subdir")).expect("create subdir");
+        fs::write(dir2.join("subdir/file.txt"), b"not firmware").expect("write .txt");
+        assert!(
+            !rootfs_has_firmware_bins(&dir2),
+            "should not detect .txt as firmware"
+        );
+        let _ = fs::remove_dir_all(&dir2);
+    }
+
     fn unique_temp_dir() -> PathBuf {
         static COUNTER: AtomicU64 = AtomicU64::new(0);
         let nanos = SystemTime::now()

crates/openshell-driver-vm/src/runtime.rs

@@ -346,6 +346,8 @@ fn build_kernel_cmdline(config: &VmLaunchConfig) -> String {
 
     if config.gpu_bdf.is_some() {
         parts.push("firmware_class.path=/lib/firmware".to_string());
+        parts.push("modprobe.blacklist=nouveau".to_string());
+        parts.push("nouveau.modeset=0".to_string());
     }
 
     parts.join(" ")

sandboxes/nvidia-gpu/Dockerfile

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# GPU-enabled sandbox image for OpenShell VM driver.
+#
+# Provides userspace GPU tooling (nvidia-smi, NVML, CUDA driver libs, kmod)
+# on top of a minimal Ubuntu base with the full NVIDIA driver userspace
+# installed via the official .run installer (no kernel modules -- those are
+# injected at rootfs preparation time by the VM driver).
+#
+# Usage:
+#   openshell sandbox create --gpu --from ./sandboxes/nvidia-gpu/Dockerfile
+#   openshell sandbox create --gpu --from nvidia-gpu   # once published
+#
+# Build-time args:
+#   CUDA_VERSION          - CUDA toolkit version (default: 12.8.1)
+#   UBUNTU_VERSION        - Ubuntu release (default: 22.04)
+#   NVIDIA_DRIVER_VERSION - Must match the kernel modules built by
+#                           `mise run vm:nvidia-modules` (default: 580.159.03)
+
+ARG CUDA_VERSION=12.8.1
+ARG UBUNTU_VERSION=22.04
+
+FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION}
+
+# Must match NVIDIA_DRIVER_VERSION in sandboxes/nvidia-gpu/versions.env
+# and NVIDIA_OPEN_VERSION in tasks/scripts/vm/build-nvidia-modules.sh
+ARG NVIDIA_DRIVER_VERSION=580.159.03
+
+# ── System packages required by the sandbox init script ──────────────
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        bash \
+        busybox-static \
+        ca-certificates \
+        curl \
+        iproute2 \
+        iptables \
+        kmod \
+        pciutils \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN mkdir -p /usr/share/udhcpc && ln -sf /bin/busybox /sbin/udhcpc
+
+# ── NVIDIA driver userspace ──────────────────────────────────────────
+# The nvidia/cuda base image does NOT include the driver (nvidia-smi,
+# libcuda.so, libnvidia-ml.so). It relies on the NVIDIA Container
+# Runtime to mount them from the host. In a VM there is no container
+# runtime, so we install the driver userspace via the .run installer
+# with --no-kernel-module (kernel modules are injected separately).
+# TODO(gpu): Pin SHA-256 checksum for reproducible builds. Compute with:
+#   curl -fsSL <url> | sha256sum
+RUN curl -fsSL \
+        "https://us.download.nvidia.com/XFree86/Linux-x86_64/${NVIDIA_DRIVER_VERSION}/NVIDIA-Linux-x86_64-${NVIDIA_DRIVER_VERSION}.run" \
+        -o /tmp/nvidia.run \
+    && chmod +x /tmp/nvidia.run \
+    && /tmp/nvidia.run \
+        --silent \
+        --no-kernel-module \
+        --no-drm \
+        --no-x-check \
+        --no-systemd \
+        --no-nvidia-modprobe \
+        --no-distro-scripts \
+    && rm -f /tmp/nvidia.run
+
+# Ensure library paths are indexed for dlopen.
+RUN mkdir -p /etc/ld.so.conf.d \
+    && echo "/usr/local/cuda/lib64" > /etc/ld.so.conf.d/cuda.conf \
+    && echo "/usr/lib/x86_64-linux-gnu" >> /etc/ld.so.conf.d/cuda.conf \
+    && ldconfig 2>/dev/null || true
+
+# ── Kernel modules ───────────────────────────────────────────────────
+# NVIDIA kernel modules (.ko) must match the guest VM kernel (libkrunfw).
+# They are NOT in this image -- the VM driver injects them at rootfs
+# preparation time via `inject_gpu_modules`.
+#
+# GSP firmware (.bin) IS provided by the .run installer above. The VM
+# driver detects its presence and skips firmware injection, avoiding
+# version mismatches when the host driver differs from this image's.
+RUN mkdir -p /lib/modules
+
+LABEL org.opencontainers.image.title="OpenShell GPU Sandbox" \
+      org.opencontainers.image.description="GPU-enabled sandbox for OpenShell VM driver with CUDA support" \
+      io.openshell.sandbox.gpu="true"

sandboxes/nvidia-gpu/README.md

@@ -0,0 +1,98 @@
+<!-- SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!-- SPDX-License-Identifier: Apache-2.0 -->
+
+# GPU Sandbox Image
+
+GPU-enabled sandbox image for the OpenShell VM driver. Provides NVIDIA
+userspace tooling (nvidia-smi, NVML, CUDA driver libraries) on top of a
+minimal Ubuntu base. Kernel modules are injected separately by the VM
+driver at sandbox creation time.
+
+## Architecture
+
+The GPU sandbox splits responsibility between the container image and the
+VM driver:
+
+| Layer | Source | Contents |
+|-------|--------|----------|
+| **Userspace** | This Dockerfile | nvidia-smi, libcuda.so, libnvidia-ml.so, kmod, iproute2 |
+| **Kernel modules** | VM driver injection | nvidia.ko, nvidia_uvm.ko, nvidia_modeset.ko (built for guest kernel 6.12.76) |
+| **GSP firmware** | `.run` installer in image OR host fallback | gsp_ga10x.bin, gsp_tu10x.bin |
+
+The kernel modules must be compiled against the exact guest kernel version
+used by libkrunfw. The VM driver injects them into each sandbox's rootfs
+at creation time via `inject_gpu_modules()`.
+
+## Prerequisites
+
+- Linux x86_64 host with an NVIDIA GPU
+- IOMMU enabled (for VFIO GPU passthrough)
+- Docker (for building the sandbox image)
+- Guest kernel built with `CONFIG_MODULES=y` (`mise run vm:setup`)
+
+## Quick Start
+
+```shell
+# 1. One-time: build the VM runtime (includes guest kernel with module support)
+mise run vm:setup
+
+# 2. Build NVIDIA kernel modules for the guest kernel
+mise run vm:nvidia-modules
+
+# 3. Build the GPU sandbox image
+docker build -t nvidia-gpu:latest ./sandboxes/nvidia-gpu/
+
+# 4. Start the gateway with GPU support
+sudo mise run gateway:vm -- --gpu
+
+# 5. Create a GPU sandbox
+openshell sandbox create --gpu --from nvidia-gpu:latest
+```
+
+## Version Coupling
+
+The NVIDIA driver version must match across three components:
+
+| Component | Variable | Default |
+|-----------|----------|---------|
+| Dockerfile (userspace) | `NVIDIA_DRIVER_VERSION` | `580.159.03` |
+| Module build script | `NVIDIA_OPEN_VERSION` | `580.159.03` |
+| Shared reference | `sandboxes/nvidia-gpu/versions.env` | `580.159.03` |
+
+A mismatch causes `modprobe` "version magic" errors or nvidia-smi ABI
+failures at sandbox boot time.
+
+## Customization
+
+### Changing the CUDA version
+
+```shell
+docker build \
+  --build-arg CUDA_VERSION=12.6.0 \
+  --build-arg UBUNTU_VERSION=22.04 \
+  -t my-gpu-sandbox:latest \
+  ./sandboxes/nvidia-gpu/
+```
+
+### Changing the NVIDIA driver version
+
+Update all three locations:
+1. `sandboxes/nvidia-gpu/versions.env`
+2. `sandboxes/nvidia-gpu/Dockerfile` ARG `NVIDIA_DRIVER_VERSION`
+3. Rebuild kernel modules: `NVIDIA_OPEN_VERSION=<version> mise run vm:nvidia-modules`
+
+### Adding packages
+
+Add packages to the `apt-get install` line in the Dockerfile. The image
+must retain `bash`, `kmod`, `iproute2`, and `busybox-static` — the VM
+driver validates these at rootfs preparation time.
+
+## Troubleshooting
+
+| Symptom | Cause | Fix |
+|---------|-------|-----|
+| "No GPU kernel modules found" | Modules not built | `mise run vm:nvidia-modules` |
+| "kmod not found in rootfs" | Image missing kmod package | Add `kmod` to Dockerfile `apt-get install` |
+| `modprobe nvidia` fails | Kernel version mismatch | Rebuild modules after `mise run vm:setup` |
+| nvidia-smi "driver/library mismatch" | Userspace/module version mismatch | Ensure Dockerfile and module versions match |
+| "kernel version mismatch: expected X, got Y" | Guest kernel was rebuilt | Rebuild modules: `mise run vm:nvidia-modules` |

sandboxes/nvidia-gpu/versions.env

@@ -0,0 +1,6 @@
+# Shared NVIDIA driver/module version for GPU sandbox images.
+# Referenced by:
+#   - sandboxes/nvidia-gpu/Dockerfile (ARG NVIDIA_DRIVER_VERSION)
+#   - tasks/scripts/vm/build-nvidia-modules.sh (NVIDIA_OPEN_VERSION)
+# These MUST match for kernel modules and userspace to be compatible.
+NVIDIA_DRIVER_VERSION=580.159.03

tasks/scripts/gateway-vm.sh

@@ -32,6 +32,8 @@
 
 set -euo pipefail
 
+MISE="${__MISE_EXE:-$(command -v mise 2>/dev/null || echo mise)}"
+
 ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
 PORT="${OPENSHELL_SERVER_PORT:-18081}"
 GATEWAY_NAME="${OPENSHELL_VM_GATEWAY_NAME:-vm-dev}"
@@ -274,13 +276,13 @@ if [ ! -d "${COMPRESSED_DIR}" ] \
     || ! find "${COMPRESSED_DIR}" -maxdepth 1 -name 'libkrun*.zst' | grep -q . \
     || [ ! -f "${COMPRESSED_DIR}/gvproxy.zst" ]; then
   echo "==> Preparing embedded VM runtime (mise run vm:setup)"
-  mise run vm:setup
+  "$MISE" run vm:setup
 fi
 
 if [ ! -f "${COMPRESSED_DIR}/openshell-sandbox.zst" ]; then
   check_supervisor_cross_toolchain
   echo "==> Building bundled VM supervisor (mise run vm:supervisor)"
-  mise run vm:supervisor
+  "$MISE" run vm:supervisor
 fi
 
 export OPENSHELL_VM_RUNTIME_COMPRESSED_DIR="${COMPRESSED_DIR}"
@@ -291,8 +293,17 @@ if [[ -n "${CARGO_BUILD_JOBS:-}" ]]; then
 fi
 
 echo "==> Building openshell-gateway and openshell-driver-vm"
-cargo build ${CARGO_BUILD_JOBS_ARG[@]+"${CARGO_BUILD_JOBS_ARG[@]}"} \
-  -p openshell-server -p openshell-driver-vm
+CARGO_CMD=(cargo build ${CARGO_BUILD_JOBS_ARG[@]+"${CARGO_BUILD_JOBS_ARG[@]}"} \
+  -p openshell-server -p openshell-driver-vm)
+if [ "$(id -u)" = "0" ] && [ -n "${SUDO_USER:-}" ]; then
+  chown -R "${SUDO_USER}" "${ROOT}/target" 2>/dev/null || true
+  sudo -u "${SUDO_USER}" env \
+    "PATH=${PATH}" \
+    "OPENSHELL_VM_RUNTIME_COMPRESSED_DIR=${OPENSHELL_VM_RUNTIME_COMPRESSED_DIR}" \
+    "${CARGO_CMD[@]}"
+else
+  "${CARGO_CMD[@]}"
+fi
 
 if [ "$(uname -s)" = "Darwin" ]; then
   echo "==> Codesigning openshell-driver-vm (Hypervisor entitlement)"

tasks/scripts/vm/build-nvidia-modules.sh

@@ -0,0 +1,176 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Build NVIDIA kernel modules against the libkrunfw guest kernel.
+#
+# Uses the NVIDIA DKMS source already installed on the host (from the
+# nvidia-dkms-* or nvidia-kernel-source-* package) and compiles it
+# against the guest kernel tree produced by build-libkrun.sh.
+#
+# Prerequisites:
+#   - NVIDIA kernel source in /usr/src/nvidia-*/
+#   - Guest kernel built with CONFIG_MODULES=y (mise run vm:setup)
+#
+# Output:
+#   target/libkrun-build/nvidia-modules/*.ko
+#   target/libkrun-build/nvidia-firmware/<version>/*.bin (if available)
+#
+# Usage:
+#   ./build-nvidia-modules.sh
+
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+BUILD_DIR="${ROOT}/target/libkrun-build"
+OUTPUT_DIR="${BUILD_DIR}/nvidia-modules"
+
+# Guest kernel version — keep in sync with GUEST_KERNEL_VERSION in
+# crates/openshell-driver-vm/src/rootfs.rs and the init script.
+KERNEL_TREE="${BUILD_DIR}/libkrunfw/linux-6.12.76"
+if [ ! -d "${KERNEL_TREE}" ]; then
+    echo "ERROR: Guest kernel tree not found at ${KERNEL_TREE}" >&2
+    echo "       Run: mise run vm:setup" >&2
+    exit 1
+fi
+
+if ! grep -q 'CONFIG_MODULES=y' "${KERNEL_TREE}/.config" 2>/dev/null; then
+    echo "ERROR: Guest kernel was built without CONFIG_MODULES=y" >&2
+    echo "       Ensure openshell.kconfig includes CONFIG_MODULES=y and rebuild:" >&2
+    echo "       mise run vm:setup" >&2
+    exit 1
+fi
+
+if [ ! -f "${KERNEL_TREE}/Module.symvers" ]; then
+    echo "ERROR: Module.symvers not found — the kernel needs to be rebuilt" >&2
+    echo "       with CONFIG_MODULES=y so the build produces Module.symvers." >&2
+    echo "       Run: mise run vm:setup" >&2
+    exit 1
+fi
+
+# Detect the host NVIDIA driver version to pick a compatible module source.
+HOST_DRIVER_VERSION="${NVIDIA_DRIVER_VERSION:-}"
+if [ -z "${HOST_DRIVER_VERSION}" ]; then
+    HOST_DRIVER_VERSION="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null || true)"
+fi
+if [ -z "${HOST_DRIVER_VERSION}" ]; then
+    HOST_DRIVER_VERSION="$(modinfo -F version /lib/modules/$(uname -r)/updates/dkms/nvidia.ko 2>/dev/null || true)"
+fi
+
+# Use the open-gpu-kernel-modules release matching the host driver major
+# version. The open modules support newer kernels better than the
+# proprietary DKMS source shipped in /usr/src/.
+# Must match NVIDIA_DRIVER_VERSION in sandboxes/nvidia-gpu/versions.env
+# and sandboxes/nvidia-gpu/Dockerfile ARG NVIDIA_DRIVER_VERSION
+NVIDIA_OPEN_VERSION="${NVIDIA_OPEN_VERSION:-580.159.03}"
+NVIDIA_SRC_DIR="${BUILD_DIR}/open-gpu-kernel-modules-${NVIDIA_OPEN_VERSION}"
+
+if [ ! -d "${NVIDIA_SRC_DIR}/kernel-open" ]; then
+    echo "==> Downloading NVIDIA open kernel modules ${NVIDIA_OPEN_VERSION}"
+    TARBALL="${BUILD_DIR}/nvidia-open-${NVIDIA_OPEN_VERSION}.tar.gz"
+    if [ ! -f "${TARBALL}" ]; then
+        curl -fSL \
+            "https://github.com/NVIDIA/open-gpu-kernel-modules/archive/refs/tags/${NVIDIA_OPEN_VERSION}.tar.gz" \
+            -o "${TARBALL}"
+        # TODO(gpu): Add SHA-256 verification for supply chain integrity.
+        # echo "<expected-hash>  ${TARBALL}" | sha256sum -c -
+    fi
+    echo "    Extracting..."
+    tar -xzf "${TARBALL}" -C "${BUILD_DIR}"
+    echo "    Source: ${NVIDIA_SRC_DIR}"
+fi
+
+NVIDIA_SRC="${NVIDIA_SRC_DIR}"
+
+# Patch API incompatibilities with newer kernels.
+# __flush_tlb() was removed in kernel 6.12; use __flush_tlb_all() instead.
+NV_PAT="${NVIDIA_SRC}/kernel-open/nvidia/nv-pat.c"
+if [ -f "${NV_PAT}" ] && grep -q '__flush_tlb()' "${NV_PAT}"; then
+    echo "==> Patching nv-pat.c (__flush_tlb -> __flush_tlb_all)"
+    sed -i 's/__flush_tlb()/__flush_tlb_all()/g' "${NV_PAT}"
+fi
+
+echo "==> Building NVIDIA ${NVIDIA_OPEN_VERSION} open kernel modules for guest kernel 6.12.76"
+echo "    NVIDIA source: ${NVIDIA_SRC}"
+echo "    Kernel tree:   ${KERNEL_TREE}"
+echo "    Output:        ${OUTPUT_DIR}"
+if [ -n "${HOST_DRIVER_VERSION}" ]; then
+    echo "    Host driver:   ${HOST_DRIVER_VERSION}"
+fi
+echo ""
+
+mkdir -p "${OUTPUT_DIR}"
+
+NPROC="$(nproc 2>/dev/null || echo 4)"
+IGNORE_CC_MISMATCH=1 make -C "${NVIDIA_SRC}" \
+    SYSSRC="${KERNEL_TREE}" \
+    SYSOUT="${KERNEL_TREE}" \
+    -j"${NPROC}" \
+    modules 2>&1 | tail -30
+
+echo ""
+echo "==> Collecting .ko files"
+
+# Open modules build into kernel-open/<subdir>/
+for subdir in nvidia nvidia-uvm nvidia-modeset nvidia-drm nvidia-peermem; do
+    for search in "${NVIDIA_SRC}/kernel-open/${subdir}" "${NVIDIA_SRC}/${subdir}"; do
+        ko_file="${search}/${subdir//-/_}.ko"
+        if [ -f "${ko_file}" ]; then
+            cp "${ko_file}" "${OUTPUT_DIR}/"
+            echo "    $(basename "${ko_file}") ($(du -h "${ko_file}" | cut -f1))"
+            break
+        fi
+    done
+done
+
+# Also check for flat layouts.
+for ko in "${NVIDIA_SRC}"/*.ko "${NVIDIA_SRC}"/kernel-open/*.ko; do
+    [ -f "${ko}" ] || continue
+    base="$(basename "${ko}")"
+    [ -f "${OUTPUT_DIR}/${base}" ] && continue
+    cp "${ko}" "${OUTPUT_DIR}/"
+    echo "    ${base} ($(du -h "${ko}" | cut -f1))"
+done
+
+KO_COUNT=$(find "${OUTPUT_DIR}" -name '*.ko' | wc -l)
+if [ "${KO_COUNT}" -eq 0 ]; then
+    echo "ERROR: No .ko files produced. Check build output above." >&2
+    exit 1
+fi
+
+echo ""
+echo "==> Collecting firmware"
+
+# GSP firmware is included in the open-gpu-kernel-modules source tree.
+FW_SRC="${NVIDIA_SRC}/src/nvidia/firmware"
+FW_OUTPUT="${BUILD_DIR}/nvidia-firmware/${NVIDIA_OPEN_VERSION}"
+if [ -d "${FW_SRC}" ] && ls "${FW_SRC}"/*.bin >/dev/null 2>&1; then
+    mkdir -p "${FW_OUTPUT}"
+    cp "${FW_SRC}"/*.bin "${FW_OUTPUT}/"
+    FW_COUNT=$(find "${FW_OUTPUT}" -name '*.bin' 2>/dev/null | wc -l)
+    echo "    Copied ${FW_COUNT} firmware files from source tree"
+else
+    # Fall back to host firmware.
+    HOST_FW=""
+    for candidate in "/lib/firmware/nvidia/${HOST_DRIVER_VERSION}" /lib/firmware/nvidia; do
+        if [ -d "${candidate}" ] && ls "${candidate}"/*.bin >/dev/null 2>&1; then
+            HOST_FW="${candidate}"
+            break
+        fi
+    done
+    if [ -n "${HOST_FW}" ]; then
+        mkdir -p "${FW_OUTPUT}"
+        cp -r "${HOST_FW}"/* "${FW_OUTPUT}/" 2>/dev/null || true
+        FW_COUNT=$(find "${FW_OUTPUT}" -name '*.bin' 2>/dev/null | wc -l)
+        echo "    Copied ${FW_COUNT} firmware files from host ${HOST_FW}"
+    else
+        echo "    WARNING: No firmware found. GPU guests may fail without GSP firmware."
+    fi
+fi
+
+echo ""
+echo "==> Done! ${KO_COUNT} kernel modules built for guest kernel 6.12.76."
+echo "    The VM driver will auto-discover them at:"
+echo "    ${OUTPUT_DIR}"
+echo ""
+echo "    Next: mise run gateway:vm -- --gpu"

tasks/scripts/vm/build-supervisor-bundle.sh

@@ -133,14 +133,31 @@ run_supervisor_build() {
         cargo_prefix=(env -u RUSTC_WRAPPER)
     fi
 
-    if command -v cargo-zigbuild >/dev/null 2>&1; then
-        "${cargo_prefix[@]}" cargo zigbuild --release -p openshell-sandbox --target "${RUST_TARGET}" \
-            --manifest-path "${ROOT}/Cargo.toml"
-    else
-        echo "    cargo-zigbuild not found, falling back to cargo build..."
-        "${cargo_prefix[@]}" cargo build --release -p openshell-sandbox --target "${RUST_TARGET}" \
-            --manifest-path "${ROOT}/Cargo.toml"
+    # When running under sudo, de-escalate the build to the original user.
+    # The target/ dir is owned by that user and root may lack write access
+    # (e.g. NFS root_squash). Only the final gateway execution needs root.
+    # Pass PATH explicitly so cargo/rustc/sccache remain reachable.
+    # Also reclaim any root-owned artifacts left by prior sudo builds.
+    if [ "$(id -u)" = "0" ] && [ -n "${SUDO_USER:-}" ]; then
+        if [ -d "${ROOT}/target" ]; then
+            chown -R "${SUDO_USER}" "${ROOT}/target" 2>/dev/null || true
+        fi
+        cargo_prefix=(sudo -u "${SUDO_USER}" env "PATH=${PATH}" "${cargo_prefix[@]}")
     fi
+
+    local host_arch
+    host_arch="$(uname -m)"
+    local cargo_build_cmd="build"
+    local cargo_bin="cargo"
+
+    if [ "${host_arch}" != "${GUEST_ARCH}" ] && command -v cargo-zigbuild >/dev/null 2>&1; then
+        cargo_build_cmd="zigbuild"
+    elif [ "${host_arch}" != "${GUEST_ARCH}" ]; then
+        echo "    cargo-zigbuild not found, falling back to cargo build..."
+    fi
+
+    "${cargo_prefix[@]}" ${cargo_bin} ${cargo_build_cmd} --release -p openshell-sandbox \
+        --target "${RUST_TARGET}" --manifest-path "${ROOT}/Cargo.toml"
 }
 
 print_build_failure() {

tasks/vm.toml

@@ -24,6 +24,18 @@ run = "tasks/scripts/vm/vm-setup.sh"
 description = "Build the bundled openshell-sandbox supervisor for openshell-driver-vm"
 run = "tasks/scripts/vm/build-supervisor-bundle.sh"
 
+["vm:rootfs"]
+description = "Build the VM rootfs tarball (use -- --base for lightweight)"
+run = "tasks/scripts/vm/build-rootfs-tarball.sh"
+
+["vm:nvidia-modules"]
+description = "Build NVIDIA kernel modules for the guest VM kernel"
+run = "tasks/scripts/vm/build-nvidia-modules.sh"
+
+["vm:clean"]
+description = "Remove all VM cached artifacts (runtime, rootfs, builds)"
+run = "tasks/scripts/vm/vm-clean.sh"
+
 ["vm:smoke:orphan-cleanup"]
 description = "Smoke test: start gateway+driver, create a sandbox, signal the gateway, assert no orphaned processes survive"
 run = "tasks/scripts/vm/smoke-orphan-cleanup.sh"