2026-02-05 18:05:06 -05:00

172 lines
4.9 KiB
Python

"""NVIDIA driver installation and configuration."""
import shutil
from pathlib import Path
from .utils import info, success, error, run, emerge, prompt
def emerge_nvidia() -> None:
"""Install NVIDIA drivers."""
info("=== Installing NVIDIA Drivers ===")
emerge("x11-drivers/nvidia-drivers")
success("NVIDIA drivers installed.")
def blacklist_nouveau() -> None:
"""Blacklist nouveau driver."""
info("=== Blacklisting Nouveau ===")
modprobe_dir = Path("/etc/modprobe.d")
blacklist_conf = modprobe_dir / "blacklist-nouveau.conf"
# Check if already configured (idempotency)
if blacklist_conf.exists():
info("Nouveau already blacklisted")
return
modprobe_dir.mkdir(parents=True, exist_ok=True)
blacklist_conf.write_text(
"# Blacklist nouveau in favor of proprietary nvidia driver\n"
"blacklist nouveau\n"
"options nouveau modeset=0\n"
)
success(f"Created {blacklist_conf}")
def configure_nvidia_drm() -> None:
"""Configure NVIDIA DRM modeset."""
info("=== Configuring NVIDIA DRM ===")
modprobe_dir = Path("/etc/modprobe.d")
nvidia_conf = modprobe_dir / "nvidia.conf"
# Check if already configured (idempotency)
if nvidia_conf.exists():
info("NVIDIA DRM already configured")
return
modprobe_dir.mkdir(parents=True, exist_ok=True)
nvidia_conf.write_text(
"# Enable NVIDIA DRM kernel mode setting\n"
"options nvidia_drm modeset=1 fbdev=1\n"
)
success(f"Created {nvidia_conf}")
def copy_dracut_config(source_dir: Path) -> None:
"""Copy NVIDIA dracut configuration."""
info("=== Copying Dracut NVIDIA Config ===")
dracut_src = source_dir / "dracut.conf.d" / "nvidia.conf"
dracut_dst = Path("/etc/dracut.conf.d")
dracut_dst.mkdir(parents=True, exist_ok=True)
if dracut_src.exists():
shutil.copy2(dracut_src, dracut_dst / "nvidia.conf")
success("Copied nvidia.conf to /etc/dracut.conf.d/")
else:
# Create default if not in source
nvidia_conf = dracut_dst / "nvidia.conf"
nvidia_conf.write_text(
"# NVIDIA early KMS\n"
'add_drivers+=" nvidia nvidia_modeset nvidia_uvm nvidia_drm "\n'
'force_drivers+=" nvidia nvidia_modeset nvidia_uvm nvidia_drm "\n'
)
success("Created default /etc/dracut.conf.d/nvidia.conf")
def rebuild_initramfs() -> None:
"""Rebuild initramfs to include NVIDIA modules."""
info("=== Rebuilding Initramfs ===")
# Find installed kernel version
modules_dir = Path("/lib/modules")
if not modules_dir.exists():
error("No kernel modules found. Install kernel first.")
return
kernels = sorted(modules_dir.iterdir(), reverse=True)
if not kernels:
error("No kernel versions found in /lib/modules")
return
kernel_version = kernels[0].name
info(f"Rebuilding initramfs for kernel {kernel_version}")
run("dracut", "--force", f"/boot/initramfs-{kernel_version}.img", kernel_version)
success("Initramfs rebuilt.")
def verify_initramfs() -> None:
"""Verify NVIDIA modules are in initramfs."""
info("=== Verifying Initramfs ===")
# Find latest initramfs
boot = Path("/boot")
initramfs_files = sorted(boot.glob("initramfs-*.img"), reverse=True)
if not initramfs_files:
error("No initramfs found in /boot")
return
initramfs = initramfs_files[0]
info(f"Checking {initramfs.name}...")
# Check for NVIDIA modules
result = run("lsinitrd", str(initramfs), capture=True, check=False)
if result.ok and result.stdout:
nvidia_modules = [
line for line in result.stdout.split("\n")
if "nvidia" in line.lower()
]
if nvidia_modules:
success("NVIDIA modules found in initramfs:")
for module in nvidia_modules[:10]: # Show first 10
print(f" {module}")
else:
error("WARNING: No NVIDIA modules found in initramfs!")
print(" Initramfs may need to be rebuilt after kernel installation.")
else:
error("Could not inspect initramfs")
def setup_nvidia(source_dir: Path | None = None) -> None:
"""Full NVIDIA setup workflow."""
if source_dir is None:
source_dir = Path("/root/gentoo")
emerge_nvidia()
print()
blacklist_nouveau()
print()
configure_nvidia_drm()
print()
copy_dracut_config(source_dir)
print()
# Ask about rebuilding initramfs
print()
info("Initramfs should be rebuilt to include NVIDIA modules.")
response = prompt("Rebuild initramfs now? (y/n): ").strip().lower()
if response == "y":
rebuild_initramfs()
print()
verify_initramfs()
print()
success("=== NVIDIA Setup Complete ===")
print()
info("Next steps:")
print(" 1. Install and configure GRUB bootloader")
print(" 2. Verify initramfs has NVIDIA and crypt modules")
print(" 3. Reboot and test")