build: split NVIDIA rootfs into base-nvidia image + gpu addon

Carve the monolithic NVIDIA GPU rootfs into a driver-agnostic
base-nvidia image (NVRC init + agent + base libs + in-tree modules)
and a driver-versioned gpu addon (GPU userspace, configs, firmware,
NVIDIA modules) laid out for /run/kata-addons/gpu. Both halves are
produced from the same chiseled tree via a partition-after step, so
the monolith build stays byte-identical and continues to ship.

The shared stage-one (driver install) is reused across the monolith,
base and addon variants; base-nvidia is cached without the driver/ctk
versions so one base image can back multiple driver addons.

tools/osbuilder/rootfs-builder/nvidia/nvidia_rootfs.sh

@@ -37,7 +37,29 @@ else
     die "Unsupported architecture: ${machine_arch}"
 fi
 
-readonly stage_one="${BUILD_DIR:?}/rootfs-${BUILD_VARIANT:?}-stage-one"
+# The base-nvidia and gpu-addon images are carved out of the very same chiseled
+# tree as the monolith, so they share the (expensive) driver stage-one with the
+# monolithic "nvidia-gpu" build instead of rebuilding it per layout.
+nvidia_stage_one_variant() {
+	case "${BUILD_VARIANT}" in
+		nvidia-gpu-base|nvidia-gpu-addon) echo "nvidia-gpu" ;;
+		*) echo "${BUILD_VARIANT}" ;;
+	esac
+}
+
+readonly stage_one="${BUILD_DIR:?}/rootfs-$(nvidia_stage_one_variant)-stage-one"
+
+# Image layout produced from the chiseled tree:
+#   monolith    - the full GPU image (default; unchanged behaviour)
+#   base        - driver-agnostic base-nvidia (NVRC init + agent + base libs)
+#   gpu-addon   - GPU userspace only, laid out for /run/kata-addons/gpu
+nvidia_image_layout() {
+	case "${BUILD_VARIANT}" in
+		nvidia-gpu-base) echo "base" ;;
+		nvidia-gpu-addon) echo "gpu-addon" ;;
+		*) echo "monolith" ;;
+	esac
+}
 
 setup_nvidia-nvrc() {
 	# NVRC is built from a pinned git ref by the Kata "nvrc" static-build
@@ -380,10 +402,122 @@ coco_guest_components() {
 	copy_cdh_runtime_deps
 }
 
+# GPU userspace owned by the gpu addon. Anything not listed here stays in the
+# driver-agnostic base-nvidia image. Paths are relative to the chiseled rootfs.
+readonly nvidia_gpu_addon_bins=(
+	bin/nvidia-smi
+	bin/nvidia-ctk
+	bin/nvidia-cdi-hook
+	bin/nvidia-persistenced
+	bin/nv-hostengine
+	bin/nv-fabricmanager
+	sbin/nvlsm
+)
+
+# GPU shared-library globs (inside the multiarch lib dir, plus libgrpc_mgr in
+# /lib) owned by the gpu addon.
+readonly nvidia_gpu_addon_lib_globs=(
+	'libnv*'
+	'libcuda.so*'
+	'libdcgm.*'
+	'libnvidia-nscq.so*'
+)
+
+# Lay the GPU userspace out for the addon mount (/run/kata-addons/gpu) and drop
+# everything else. Matches how NVRC consumes the addon: bins from <root>/bin and
+# <root>/sbin, libraries via LD_LIBRARY_PATH=<root>/lib:<root>/usr/lib, kernel
+# modules via `modprobe --dirname <root>`, configs from <root>/usr/share/nvidia,
+# and <root>/lib/firmware/nvidia bind-mounted onto /lib/firmware/nvidia. Runs
+# inside the (full, chiseled) ${stage_two}/${ROOTFS_DIR}.
+partition_gpu_addon() {
+	echo "nvidia: building gpu addon layout"
+
+	local addon
+	addon="$(mktemp -d "${BUILD_DIR}/.nvidia-gpu-addon.XXXX")"
+	mkdir -p "${addon}/bin" "${addon}/sbin" "${addon}/lib" \
+		 "${addon}/usr/share/nvidia" "${addon}/lib/firmware" "${addon}/lib/modules"
+
+	local f
+	for f in "${nvidia_gpu_addon_bins[@]}"; do
+		[[ -e "${f}" ]] && install -D -m0755 "${f}" "${addon}/${f}"
+	done
+
+	# Flatten the GPU shared libraries into <root>/lib so NVRC's LD_LIBRARY_PATH
+	# (<root>/lib:<root>/usr/lib) resolves them; libc/loader stay in the base.
+	local md="lib/${machine_arch}-linux-gnu"
+	local g
+	for g in "${nvidia_gpu_addon_lib_globs[@]}"; do
+		find "${md}" -maxdepth 1 -name "${g}" -exec cp -a {} "${addon}/lib/" \;
+	done
+	[[ -e lib/libgrpc_mgr.so ]] && cp -a lib/libgrpc_mgr.so "${addon}/lib/"
+
+	# GPU configs (fabricmanager.cfg, nvlsm.conf, ...).
+	[[ -d usr/share/nvidia ]] && cp -a usr/share/nvidia/. "${addon}/usr/share/nvidia/"
+
+	# GPU firmware (GSP, ...); NVRC binds this onto /lib/firmware/nvidia.
+	[[ -d lib/firmware/nvidia ]] && cp -a lib/firmware/nvidia "${addon}/lib/firmware/"
+
+	# Ship a self-contained module tree so `modprobe --dirname <root>` resolves
+	# the NVIDIA modules and their dependencies.
+	cp -a lib/modules/. "${addon}/lib/modules/"
+	if command -v depmod >/dev/null 2>&1; then
+		local kdir kver
+		for kdir in "${addon}"/lib/modules/*/; do
+			[[ -d "${kdir}" ]] || continue
+			kver="$(basename "${kdir}")"
+			depmod -b "${addon}" "${kver}" || true
+		done
+	fi
+
+	# Replace the rootfs with the addon-only content.
+	find . -mindepth 1 -maxdepth 1 -exec rm -rf {} +
+	cp -a "${addon}/." .
+	rm -rf "${addon}"
+}
+
+# Strip the GPU userspace from the chiseled tree, leaving a driver-agnostic
+# base-nvidia: NVRC init + kata-agent + busybox + loader/libc + in-tree kernel
+# modules. The empty /lib/firmware/nvidia directory is kept as the bind
+# mountpoint NVRC uses for the addon firmware. Runs inside ${ROOTFS_DIR}.
+partition_base() {
+	echo "nvidia: building driver-agnostic base layout"
+
+	local f
+	for f in "${nvidia_gpu_addon_bins[@]}"; do
+		rm -f "${f}"
+	done
+
+	local md="lib/${machine_arch}-linux-gnu"
+	local g
+	for g in "${nvidia_gpu_addon_lib_globs[@]}"; do
+		find "${md}" -maxdepth 1 -name "${g}" -delete
+	done
+	rm -f lib/libgrpc_mgr.so
+
+	# GPU configs live in the addon; keep usr/share/nvidia as an empty stub.
+	rm -rf usr/share/nvidia
+	mkdir -p usr/share/nvidia
+
+	# Keep /lib/firmware/nvidia as an empty mountpoint for NVRC's firmware bind.
+	rm -rf lib/firmware/nvidia
+	mkdir -p lib/firmware/nvidia
+
+	# Reset kernel modules to the in-tree set (no NVIDIA driver modules) so the
+	# base stays driver-agnostic; the gpu addon ships the driver modules.
+	rm -rf lib/modules
+	mkdir -p lib/modules
+	local appendix=""
+	echo "${NVIDIA_GPU_STACK}" | grep -q '\<dragonball\>' && appendix="-dragonball-experimental"
+	tar --zstd -xf "${BUILD_DIR}"/kata-static-kernel-nvidia-gpu"${appendix}"-modules.tar.zst \
+		-C ./lib/modules/
+}
+
 setup_nvidia_gpu_rootfs_stage_two() {
 	readonly stage_two="${ROOTFS_DIR:?}"
 	readonly stack="${NVIDIA_GPU_STACK:?}"
 	readonly type=${1:-""}
+	local layout
+	layout="$(nvidia_image_layout)"
 
 	# If devkit flag is set, skip chisseling, use stage_one
 	if echo "${stack}" | grep -q '\<devkit\>'; then
@@ -428,10 +562,19 @@ setup_nvidia_gpu_rootfs_stage_two() {
 
 		coco_guest_components
 		chisseled_nvat
+
+		# Carve the freshly chiseled (monolith) tree into the requested layout.
+		# The monolith path is left untouched.
+		case "${layout}" in
+			base) partition_base ;;
+			gpu-addon) partition_gpu_addon ;;
+		esac
 	fi
 
 	compress_rootfs
-	chroot . ldconfig
+	# The gpu addon has no loader/ldconfig of its own; its libraries are found
+	# via NVRC's LD_LIBRARY_PATH, so skip the ld.so cache rebuild there.
+	[[ "${layout}" != "gpu-addon" ]] && chroot . ldconfig
 
 	popd >> /dev/null
 }

tools/osbuilder/rootfs-builder/rootfs.sh

@@ -937,7 +937,13 @@ main()
 	init="${ROOTFS_DIR}/sbin/init"
 	setup_rootfs
 
-	if [[ "${BUILD_VARIANT}" = "nvidia-gpu" ]]; then
+	# The base-nvidia and gpu-addon layouts are carved from the same chiseled
+	# tree as the monolith (which they share the driver stage-one with), so they
+	# drive the non-confidential stage-one/two and only differ in the final
+	# partition step (see nvidia_image_layout in nvidia_rootfs.sh).
+	if [[ "${BUILD_VARIANT}" = "nvidia-gpu" || \
+	      "${BUILD_VARIANT}" = "nvidia-gpu-base" || \
+	      "${BUILD_VARIANT}" = "nvidia-gpu-addon" ]]; then
 		setup_nvidia_gpu_rootfs_stage_one
 		setup_nvidia_gpu_rootfs_stage_two
 		return $?

tools/packaging/guest-image/build_image.sh

@@ -97,7 +97,11 @@ build_image() {
 		PAUSE_IMAGE_TARBALL="${PAUSE_IMAGE_TARBALL:-}" \
 		GUEST_HOOKS_TARBALL="${GUEST_HOOKS_TARBALL}"
 
-	if [[ "${image_initrd_suffix}" == "nvidia-gpu"* ]]; then
+	# The driver-versioned NVIDIA artefacts (the monolith, the confidential
+	# monolith and the gpu addon) carry the driver version in their filename so
+	# multiple driver builds can coexist.  The base-nvidia image is
+	# driver-agnostic, so it must not get the driver-version suffix.
+	if [[ "${image_initrd_suffix}" == "nvidia-gpu"* && "${image_initrd_suffix}" != "nvidia-gpu-base" ]]; then
 		nvidia_driver_version=$(get_from_kata_deps .externals.nvidia.driver.version)
 		artifact_name=${artifact_name/.image/"-${nvidia_driver_version}".image}
 	fi

tools/packaging/kata-deploy/local-build/kata-deploy-binaries.sh

@@ -189,7 +189,7 @@ cleanup_and_fail_shim_v2_specifics() {
 	local extra_tarballs="${3:-}"
 	local tarball_dir="${repo_root_dir}/tools/packaging/kata-deploy/local-build/build"
 
-	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-addon; do
 		local root_hash_file="${tarball_dir}/${component}-root_hash_${variant}.txt"
 		[[ -f "${root_hash_file}" ]] && rm -f "${root_hash_file}"
 	done
@@ -219,7 +219,7 @@ install_cached_shim_v2_tarball_get_root_hash() {
 	local tarball_dir="${repo_root_dir}/tools/packaging/kata-deploy/local-build/build"
 	local root_hash_basedir="./opt/kata/share/kata-containers/"
 
-	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-addon; do
 		# The measured base image ships as kata-static-rootfs-image.tar.zst
 		# (no variant suffix), but carries its root hash under the "base" label.
 		local image_conf_tarball
@@ -247,7 +247,7 @@ install_cached_shim_v2_tarball_compare_root_hashes() {
 	local found_any=""
 	local tarball_dir="${repo_root_dir}/tools/packaging/kata-deploy/local-build/build"
 
-	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-addon; do
 		local image_root_hash="${tarball_dir}/root_hash_${variant}.txt"
 		local cached_root_hash="${component}-root_hash_${variant}.txt"
 
@@ -483,14 +483,26 @@ install_image() {
 		fi
 	fi
 
-	if [[ "${variant}" == "nvidia-gpu" ]]; then
-		# If we bump the kernel we need to rebuild the image
+	if [[ "${variant}" == "nvidia-gpu" || "${variant}" == "nvidia-gpu-addon" ]]; then
+		# If we bump the kernel we need to rebuild the image.  The gpu addon
+		# carries the driver userspace carved out of the same chiseled tree,
+		# so it is driver-versioned just like the monolith.
 		latest_artefact+="-$(get_latest_kernel_nvidia_artefact_and_builder_image_version)"
 		latest_artefact+="-$(get_latest_nvidia_driver_version)"
 		latest_artefact+="-$(get_latest_nvidia_ctk_version)"
 		latest_artefact+="-$(get_latest_nvidia_nvrc_version)"
 	fi
 
+	if [[ "${variant}" == "nvidia-gpu-base" ]]; then
+		# The base-nvidia image strips all driver userspace and resets the
+		# kernel modules to in-tree only, so it is driver-agnostic: it depends
+		# on the NVIDIA kernel build and NVRC (its init), but not on the driver
+		# or container-toolkit versions.  That lets a single base image back
+		# multiple driver-specific gpu addons.
+		latest_artefact+="-$(get_latest_kernel_nvidia_artefact_and_builder_image_version)"
+		latest_artefact+="-$(get_latest_nvidia_nvrc_version)"
+	fi
+
 	# The base guest image (empty variant) is built as a measured rootfs so
 	# that confidential configurations can dm-verity-protect it; non-confidential
 	# configurations simply boot its data partition unverified.  Reflect the
@@ -874,6 +886,45 @@ install_image_nvidia_gpu_confidential() {
 	install_image "nvidia-gpu-confidential"
 }
 
+# Install the driver-agnostic base-nvidia image.
+#
+# This is the NVRC-init half of the chiseled NVIDIA tree: NVRC, the kata-agent,
+# busybox, the base libc/loader and the in-tree kernel modules.  The driver
+# userspace is stripped out and shipped separately in the gpu addon, so this
+# image is reused across driver versions.  The driver still has to be installed
+# to build the shared stage-one (the GPU files are carved out afterwards), so we
+# keep the same NVIDIA_GPU_STACK as the monolith.
+install_image_nvidia_gpu_base() {
+	export AGENT_POLICY
+	export MEASURED_ROOTFS="yes"
+	export FS_TYPE="erofs"
+	export SKIP_DAX_HEADER="yes"
+	local version
+	version=$(get_latest_nvidia_driver_version)
+	EXTRA_PKGS="apt curl ${EXTRA_PKGS}"
+	NVIDIA_GPU_STACK=${NVIDIA_GPU_STACK:-"driver=${version},compute,dcgm,nvswitch"}
+	install_image "nvidia-gpu-base"
+}
+
+# Install the gpu addon image.
+#
+# This is the driver half of the chiseled NVIDIA tree, laid out for
+# /run/kata-addons/gpu: the GPU binaries, libraries, configs, firmware and the
+# NVIDIA kernel modules.  It is an erofs+verity image (MEASURED_ROOTFS) and is
+# driver-versioned, so multiple driver addons can coexist against a single
+# base-nvidia image.
+install_image_nvidia_gpu_addon() {
+	export AGENT_POLICY
+	export MEASURED_ROOTFS="yes"
+	export FS_TYPE="erofs"
+	export SKIP_DAX_HEADER="yes"
+	local version
+	version=$(get_latest_nvidia_driver_version)
+	EXTRA_PKGS="apt curl ${EXTRA_PKGS}"
+	NVIDIA_GPU_STACK=${NVIDIA_GPU_STACK:-"driver=${version},compute,dcgm,nvswitch"}
+	install_image "nvidia-gpu-addon"
+}
+
 install_se_image() {
 	# shellcheck disable=SC2154
 	info "Create IBM SE image configured with AA_KBC=${AA_KBC}"
@@ -1225,7 +1276,7 @@ install_nydus() {
 # Shared helper: extract measured-rootfs root hashes from confidential image tarballs.
 # These are needed by the Rust runtime (runtime-rs) at build time for dm-verity.
 _collect_root_hashes() {
-	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential; do
+	for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-addon; do
 		# The measured base image ships as kata-static-rootfs-image.tar.zst
 		# (no variant suffix), but carries its root hash under the "base" label.
 		local tarball_glob="kata-static-rootfs-image-${variant}.tar.zst"
@@ -1708,6 +1759,10 @@ handle_build() {
 
 	rootfs-image-nvidia-gpu-confidential) install_image_nvidia_gpu_confidential ;;
 
+	rootfs-image-nvidia-gpu-base) install_image_nvidia_gpu_base ;;
+
+	rootfs-image-nvidia-gpu-addon) install_image_nvidia_gpu_addon ;;
+
 	rootfs-cca-confidential-image) install_image_confidential ;;
 
 	rootfs-cca-confidential-initrd) install_initrd_confidential ;;
@@ -1767,7 +1822,7 @@ handle_build() {
 			;;
 		shim-v2-go|shim-v2-rust)
 			if [[ "${MEASURED_ROOTFS}" == "yes" ]]; then
-				for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential; do
+				for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-addon; do
 					[[ -f "${workdir}/root_hash_${variant}.txt" ]] && mv "${workdir}/root_hash_${variant}.txt" "${workdir}/${build_target}-root_hash_${variant}.txt"
 				done
 			fi
@@ -1836,7 +1891,7 @@ handle_build() {
 		shim-v2-go|shim-v2-rust)
 			if [[ "${MEASURED_ROOTFS}" == "yes" ]]; then
 				local found_any=""
-				for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential; do
+				for variant in base confidential coco-addon nvidia-gpu nvidia-gpu-confidential nvidia-gpu-base nvidia-gpu-addon; do
 					# The variants could be built independently we need to check if
 					# they exist and then push them to the registry
 					[[ -f "${workdir}/${build_target}-root_hash_${variant}.txt" ]] && files_to_push+=("${build_target}-root_hash_${variant}.txt") && found_any="yes"

versions.yaml

@@ -137,6 +137,12 @@ assets:
         nvidia-gpu-confidential:
           name: "ubuntu"
           version: "noble"  # 24.04 LTS
+        nvidia-gpu-base:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
+        nvidia-gpu-addon:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
       ppc64le:
         name: "ubuntu"
         version: "noble"  # 24.04 LTS
@@ -161,6 +167,12 @@ assets:
         nvidia-gpu-confidential:
           name: "ubuntu"
           version: "noble"  # 24.04 LTS
+        nvidia-gpu-base:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
+        nvidia-gpu-addon:
+          name: "ubuntu"
+          version: "noble"  # 24.04 LTS
 
   initrd:
     description: |