From 1e57ff583be26cc2ebec9925d898c05cbfa6ae46 Mon Sep 17 00:00:00 2001 From: themylogin Date: Thu, 21 Nov 2024 14:07:42 +0100 Subject: [PATCH] Move nvidia extension to the rootfs image (#762) --- Makefile | 2 - conf/build.manifest | 6 ++ scale_build/extensions.py | 107 +++++++++++++++------------------ scale_build/image/bootstrap.py | 2 +- scale_build/image/manifest.py | 6 +- scale_build/image/update.py | 50 +++++++++------ scale_build/image/utils.py | 3 +- scale_build/main.py | 5 -- scale_build/utils/kernel.py | 6 ++ scale_build/utils/manifest.py | 15 +++++ 10 files changed, 110 insertions(+), 92 deletions(-) create mode 100644 scale_build/utils/kernel.py diff --git a/Makefile b/Makefile index a1f14a33..022e89a9 100644 --- a/Makefile +++ b/Makefile @@ -31,8 +31,6 @@ check_upstream_package_updates: check . ./venv-${COMMIT_HASH}/bin/activate && scale_build check_upstream_package_updates iso: check . ./venv-${COMMIT_HASH}/bin/activate && scale_build iso -extensions: check - . ./venv-${COMMIT_HASH}/bin/activate && scale_build extensions packages: check ifeq ($(PACKAGES),"") . ./venv-${COMMIT_HASH}/bin/activate && scale_build packages diff --git a/conf/build.manifest b/conf/build.manifest index 9e5ffcd9..c4c42f73 100644 --- a/conf/build.manifest +++ b/conf/build.manifest @@ -637,3 +637,9 @@ sources: - "./pull.sh" deoptions: nocheck generate_version: false + +# Nvidia extensions versions +############################################################################ +extensions: + nvidia: + current: "550.127.05" diff --git a/scale_build/extensions.py b/scale_build/extensions.py index 45be627e..33fd61db 100644 --- a/scale_build/extensions.py +++ b/scale_build/extensions.py @@ -1,16 +1,14 @@ import errno -import json import logging import os import shutil import requests -from .exceptions import CallError -from .image.bootstrap import clean_mounts -from .image.manifest import update_file_path from .image.utils import run_in_chroot -from .utils.paths import CD_DIR, CHROOT_BASEDIR, CHROOT_OVERLAY, RELEASE_DIR +from .utils.kernel import get_kernel_version +from .utils.manifest import get_manifest +from .utils.paths import TMPFS, PKG_DIR from .utils.run import run logger = logging.getLogger(__name__) @@ -19,65 +17,61 @@ TEMPORARY_PACKAGES = ["gcc", "make", "pkg-config"] PERMANENT_PACKAGES = ["libvulkan1", "nvidia-container-toolkit", "vulkan-validationlayers"] HEADERS = {"User-Agent": "curl/7.88.1"} -EXTENSIONS_DIR = os.path.join(RELEASE_DIR, "extensions") +EXTENSIONS_CHROOT = os.path.join(TMPFS, "extensions_chroot") +EXTENSIONS_CHROOT_BASE = os.path.join(TMPFS, "extensions_chroot_base") -def build_extensions(): - clean_mounts() +def build_extensions(rootfs_image, dst_dir): + for path in [EXTENSIONS_CHROOT, EXTENSIONS_CHROOT_BASE]: + if os.path.exists(path): + shutil.rmtree(path) + os.makedirs(path) - if not os.path.exists(update_file_path()): - raise CallError("Missing rootfs image. Run `make update` first.") + run(["unsquashfs", "-dest", EXTENSIONS_CHROOT, rootfs_image]) + run(["unsquashfs", "-dest", EXTENSIONS_CHROOT_BASE, rootfs_image]) - if os.path.exists(CHROOT_BASEDIR): - shutil.rmtree(CHROOT_BASEDIR) - if os.path.exists(CHROOT_OVERLAY): - shutil.rmtree(CHROOT_OVERLAY) + kernel_version = get_kernel_version(EXTENSIONS_CHROOT) - run(["mount", "-o", "loop", update_file_path(), CD_DIR]) + os.makedirs(os.path.join(EXTENSIONS_CHROOT, "proc"), exist_ok=True) + run(["mount", "proc", os.path.join(EXTENSIONS_CHROOT, "proc"), "-t", "proc"]) + os.makedirs(os.path.join(EXTENSIONS_CHROOT, "sys"), exist_ok=True) + run(["mount", "sysfs", os.path.join(EXTENSIONS_CHROOT, "sys"), "-t", "sysfs"]) + os.makedirs(os.path.join(EXTENSIONS_CHROOT, "packages"), exist_ok=True) + run(["mount", "--bind", PKG_DIR, os.path.join(EXTENSIONS_CHROOT, "packages")]) try: - run(["unsquashfs", "-dest", CHROOT_BASEDIR, os.path.join(CD_DIR, "rootfs.squashfs")]) - run(["unsquashfs", "-dest", CHROOT_OVERLAY, os.path.join(CD_DIR, "rootfs.squashfs")]) - with open(f"{CD_DIR}/manifest.json") as f: - manifest = json.load(f) - finally: - run(["umount", CD_DIR]) - - run(["mount", "proc", os.path.join(CHROOT_BASEDIR, "proc"), "-t", "proc"]) - run(["mount", "sysfs", os.path.join(CHROOT_BASEDIR, "sys"), "-t", "sysfs"]) - try: - shutil.copyfile("/etc/resolv.conf", f"{CHROOT_BASEDIR}/etc/resolv.conf") + shutil.copyfile("/etc/resolv.conf", f"{EXTENSIONS_CHROOT}/etc/resolv.conf") for binary in BINARIES: - os.unlink(os.path.join(CHROOT_BASEDIR, f"usr/local/bin/{binary}")) - os.chmod(os.path.join(CHROOT_BASEDIR, f"usr/bin/{binary}"), 0o755) + os.unlink(os.path.join(EXTENSIONS_CHROOT, f"usr/local/bin/{binary}")) + os.chmod(os.path.join(EXTENSIONS_CHROOT, f"usr/bin/{binary}"), 0o755) add_nvidia_repository() - run_in_chroot(["apt", "update"]) - run_in_chroot(["apt", "-y", "install"] + TEMPORARY_PACKAGES + PERMANENT_PACKAGES) + run_in_chroot(["apt", "update"], chroot=EXTENSIONS_CHROOT) + run_in_chroot(["apt", "-y", "install"] + TEMPORARY_PACKAGES + PERMANENT_PACKAGES, chroot=EXTENSIONS_CHROOT) - install_nvidia_driver(manifest["kernel_version"]) + install_nvidia_driver(kernel_version) - run_in_chroot(["apt", "-y", "remove"] + TEMPORARY_PACKAGES) - run_in_chroot(["apt", "-y", "autoremove"]) + run_in_chroot(["apt", "-y", "remove"] + TEMPORARY_PACKAGES, chroot=EXTENSIONS_CHROOT) + run_in_chroot(["apt", "-y", "autoremove"], chroot=EXTENSIONS_CHROOT) finally: - run(["umount", os.path.join(CHROOT_BASEDIR, "sys")]) - run(["umount", os.path.join(CHROOT_BASEDIR, "proc")]) + run(["umount", os.path.join(EXTENSIONS_CHROOT, "packages")]) + run(["umount", os.path.join(EXTENSIONS_CHROOT, "sys")]) + run(["umount", os.path.join(EXTENSIONS_CHROOT, "proc")]) - if os.path.exists(EXTENSIONS_DIR): - shutil.rmtree(EXTENSIONS_DIR) - build_extension("nvidia") + build_extension("nvidia", f"{dst_dir}/nvidia.raw") def add_nvidia_repository(): r = requests.get("https://nvidia.github.io/libnvidia-container/gpgkey") r.raise_for_status() - with open(f"{CHROOT_BASEDIR}/key.gpg", "w") as f: + with open(f"{EXTENSIONS_CHROOT}/key.gpg", "w") as f: f.write(r.text) - run_in_chroot(["gpg", "-o", "/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg", "--dearmor", "/key.gpg"]) + run_in_chroot(["gpg", "-o", "/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg", "--dearmor", "/key.gpg"], + chroot=EXTENSIONS_CHROOT) - with open(f"{CHROOT_BASEDIR}/etc/apt/sources.list.d/nvidia-container-toolkit.list", "w") as f: + with open(f"{EXTENSIONS_CHROOT}/etc/apt/sources.list.d/nvidia-container-toolkit.list", "w") as f: f.write("deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] " "https://nvidia.github.io/libnvidia-container/stable/deb/$(ARCH) /") @@ -85,11 +79,9 @@ def add_nvidia_repository(): def download_nvidia_driver(): prefix = "https://download.nvidia.com/XFree86/Linux-x86_64" - r = requests.get(f"{prefix}/latest.txt", headers=HEADERS, timeout=10) - r.raise_for_status() - version = r.text.split()[0] + version = get_manifest()["extensions"]["nvidia"]["current"] filename = f"NVIDIA-Linux-x86_64-{version}-no-compat32.run" - result = f"{CHROOT_BASEDIR}/{filename}" + result = f"{EXTENSIONS_CHROOT}/{filename}" with requests.get(f"{prefix}/{version}/{filename}", headers=HEADERS, stream=True, timeout=10) as r: r.raise_for_status() @@ -104,29 +96,30 @@ def install_nvidia_driver(kernel_version): driver = download_nvidia_driver() run_in_chroot([f"/{os.path.basename(driver)}", "--skip-module-load", "--silent", f"--kernel-name={kernel_version}", - "--allow-installation-with-running-driver", "--no-rebuild-initramfs"]) + "--allow-installation-with-running-driver", "--no-rebuild-initramfs"], + chroot=EXTENSIONS_CHROOT) os.unlink(driver) -def build_extension(name): +def build_extension(name, dst_path): changed_files = [ - os.path.relpath(filename, CHROOT_BASEDIR) + os.path.relpath(filename, EXTENSIONS_CHROOT) for filename in map( lambda filename: os.path.join(os.getcwd(), filename), run( - ["rsync", "-avn", "--out-format=%f", f"{CHROOT_BASEDIR}/", f"{CHROOT_OVERLAY}/"], + ["rsync", "-avn", "--out-format=%f", f"{EXTENSIONS_CHROOT}/", f"{EXTENSIONS_CHROOT_BASE}/"], log=False, ).stdout.split("\n") ) - if os.path.abspath(filename).startswith(os.path.abspath(CHROOT_BASEDIR)) + if os.path.abspath(filename).startswith(os.path.abspath(EXTENSIONS_CHROOT)) ] sysext_files = [f for f in changed_files if f.startswith("usr/") and not (f.startswith("usr/src/"))] - for root, dirs, files in os.walk(CHROOT_BASEDIR, topdown=False): + for root, dirs, files in os.walk(EXTENSIONS_CHROOT, topdown=False): for f in files: - path = os.path.relpath(os.path.abspath(os.path.join(root, f)), CHROOT_BASEDIR) + path = os.path.relpath(os.path.abspath(os.path.join(root, f)), EXTENSIONS_CHROOT) if path not in sysext_files: os.unlink(os.path.join(root, f)) @@ -141,12 +134,8 @@ def build_extension(name): else: raise - os.makedirs(f"{CHROOT_BASEDIR}/usr/lib/extension-release.d", exist_ok=True) - with open(f"{CHROOT_BASEDIR}/usr/lib/extension-release.d/extension-release.{name}", "w") as f: + os.makedirs(f"{EXTENSIONS_CHROOT}/usr/lib/extension-release.d", exist_ok=True) + with open(f"{EXTENSIONS_CHROOT}/usr/lib/extension-release.d/extension-release.{name}", "w") as f: f.write("ID=_any\n") - os.makedirs(EXTENSIONS_DIR, exist_ok=True) - path = os.path.join(EXTENSIONS_DIR, f"{name}.raw") - run(["mksquashfs", CHROOT_BASEDIR, path, "-comp", "xz"]) - with open(f"{path}.sha256", "w") as f: - f.write(run(["sha256sum", path], log=False).stdout.split()[0]) + run(["mksquashfs", EXTENSIONS_CHROOT, dst_path, "-comp", "xz"]) diff --git a/scale_build/image/bootstrap.py b/scale_build/image/bootstrap.py index 359c7431..92a9a1c3 100644 --- a/scale_build/image/bootstrap.py +++ b/scale_build/image/bootstrap.py @@ -15,7 +15,7 @@ def setup_chroot_basedir(bootstrapdir_obj): if os.path.exists(CHROOT_BASEDIR): shutil.rmtree(CHROOT_BASEDIR) os.makedirs(TMPFS, exist_ok=True) - run(['mount', '-t', 'tmpfs', '-o', 'size=12G', 'tmpfs', TMPFS]) + run(['mount', '-t', 'tmpfs', '-o', 'size=16G', 'tmpfs', TMPFS]) bootstrapdir_obj.restore_cache(CHROOT_BASEDIR) run(['mount', 'proc', os.path.join(CHROOT_BASEDIR, 'proc'), '-t', 'proc']) run(['mount', 'sysfs', os.path.join(CHROOT_BASEDIR, 'sys'), '-t', 'sysfs']) diff --git a/scale_build/image/manifest.py b/scale_build/image/manifest.py index 4b90db31..525a915d 100644 --- a/scale_build/image/manifest.py +++ b/scale_build/image/manifest.py @@ -1,11 +1,11 @@ from datetime import datetime -import glob import json import os import shutil import subprocess from scale_build.exceptions import CallError +from scale_build.utils.kernel import get_kernel_version from scale_build.utils.paths import BUILDER_DIR, CHROOT_BASEDIR, RELEASE_DIR, UPDATE_DIR @@ -45,9 +45,7 @@ def build_manifest(): 'version': version, 'size': size, 'checksums': checksums, - 'kernel_version': glob.glob( - os.path.join(CHROOT_BASEDIR, 'boot/vmlinuz-*') - )[1].split('/')[-1][len('vmlinuz-'):], + 'kernel_version': get_kernel_version(CHROOT_BASEDIR), })) return version diff --git a/scale_build/image/update.py b/scale_build/image/update.py index f75183af..dcb2cc2c 100644 --- a/scale_build/image/update.py +++ b/scale_build/image/update.py @@ -8,6 +8,7 @@ import tempfile from scale_build.config import SIGNING_KEY, SIGNING_PASSWORD +from scale_build.extensions import build_extensions as do_build_extensions from scale_build.utils.manifest import get_manifest, get_apt_repos from scale_build.utils.run import run from scale_build.utils.paths import CHROOT_BASEDIR, RELEASE_DIR, UPDATE_DIR @@ -103,6 +104,8 @@ def install_rootfs_packages_impl(): # Do any pruning of rootfs clean_rootfs() + build_extensions() + with open(os.path.join(CHROOT_BASEDIR, 'etc/apt/sources.list'), 'w') as f: f.write('\n'.join(get_apt_sources())) @@ -183,26 +186,6 @@ def custom_rootfs_setup(): shutil.rmtree(local_cacerts, ignore_errors=True) os.symlink("/var/local/ca-certificates", local_cacerts) - # Build a systemd-sysext extension that, upon loading, will make `/usr/bin/dpkg` working. - # It is necessary for `update-initramfs` to function properly. - sysext_extensions_dir = os.path.join(CHROOT_BASEDIR, "usr/share/truenas/sysext-extensions") - os.makedirs(sysext_extensions_dir, exist_ok=True) - with tempfile.TemporaryDirectory() as td: - os.makedirs(f"{td}/usr/bin") - shutil.copy2(f"{CHROOT_BASEDIR}/usr/bin/dpkg", f"{td}/usr/bin/dpkg") - - os.makedirs(f"{td}/usr/local/bin") - with open(f"{td}/usr/local/bin/dpkg", "w") as f: - f.write("#!/bin/bash\n") - f.write("exec /usr/bin/dpkg \"$@\"") - os.chmod(f"{td}/usr/local/bin/dpkg", 0o755) - - os.makedirs(f"{td}/usr/lib/extension-release.d") - with open(f"{td}/usr/lib/extension-release.d/extension-release.functioning-dpkg", "w") as f: - f.write("ID=_any\n") - - run(["mksquashfs", td, f"{sysext_extensions_dir}/functioning-dpkg.raw"]) - def clean_rootfs(): to_remove = get_manifest()['base-prune'] @@ -227,3 +210,30 @@ def clean_rootfs(): ): shutil.rmtree(path) os.makedirs(path, exist_ok=True) + + +def build_extensions(): + # Build a systemd-sysext extension that, upon loading, will make `/usr/bin/dpkg` working. + # It is necessary for `update-initramfs` to function properly. + sysext_extensions_dir = os.path.join(CHROOT_BASEDIR, "usr/share/truenas/sysext-extensions") + os.makedirs(sysext_extensions_dir, exist_ok=True) + with tempfile.TemporaryDirectory() as td: + os.makedirs(f"{td}/usr/bin") + shutil.copy2(f"{CHROOT_BASEDIR}/usr/bin/dpkg", f"{td}/usr/bin/dpkg") + + os.makedirs(f"{td}/usr/local/bin") + with open(f"{td}/usr/local/bin/dpkg", "w") as f: + f.write("#!/bin/bash\n") + f.write("exec /usr/bin/dpkg \"$@\"") + os.chmod(f"{td}/usr/local/bin/dpkg", 0o755) + + os.makedirs(f"{td}/usr/lib/extension-release.d") + with open(f"{td}/usr/lib/extension-release.d/extension-release.functioning-dpkg", "w") as f: + f.write("ID=_any\n") + + run(["mksquashfs", td, f"{sysext_extensions_dir}/functioning-dpkg.raw"]) + + with tempfile.NamedTemporaryFile(suffix=".squashfs") as tf: + tf.close() + run(["mksquashfs", CHROOT_BASEDIR, tf.name, "-one-file-system"]) + do_build_extensions(tf.name, sysext_extensions_dir) diff --git a/scale_build/image/utils.py b/scale_build/image/utils.py index dc1d84a4..fd82b030 100644 --- a/scale_build/image/utils.py +++ b/scale_build/image/utils.py @@ -8,6 +8,7 @@ def run_in_chroot(command, exception_message=None, **kwargs): + chroot = kwargs.pop('chroot', CHROOT_BASEDIR) return run( - ['chroot', CHROOT_BASEDIR] + command, exception_msg=exception_message, env={**APT_ENV, **os.environ}, **kwargs + ['chroot', chroot] + command, exception_msg=exception_message, env={**APT_ENV, **os.environ}, **kwargs ) diff --git a/scale_build/main.py b/scale_build/main.py index 5828c61f..aad63fcf 100644 --- a/scale_build/main.py +++ b/scale_build/main.py @@ -10,7 +10,6 @@ from .upstream_package_updates import check_upstream_package_updates from .epoch import check_epoch from .exceptions import CallError -from .extensions import build_extensions from .iso import build_iso from .package import build_packages from .preflight import preflight_check @@ -63,7 +62,6 @@ def main(): ) subparsers.add_parser('update', help='Create TrueNAS Scale update image') subparsers.add_parser('iso', help='Create TrueNAS Scale iso installation file') - subparsers.add_parser('extensions', help='Create TrueNAS extensions files') branchout_parser = subparsers.add_parser('branchout', help='Checkout new branch for all packages') branchout_parser.add_argument( '--skip-push', '-sp', action='store_true', default=False, @@ -92,9 +90,6 @@ def main(): elif args.action == 'iso': validate() build_iso() - elif args.action == 'extensions': - validate() - build_extensions() elif args.action == 'clean': complete_cleanup() elif args.action == 'validate': diff --git a/scale_build/utils/kernel.py b/scale_build/utils/kernel.py new file mode 100644 index 00000000..4472e160 --- /dev/null +++ b/scale_build/utils/kernel.py @@ -0,0 +1,6 @@ +import glob +import os + + +def get_kernel_version(rootfs_path): + return glob.glob(os.path.join(rootfs_path, 'boot/vmlinuz-*-production+truenas'))[0].split('/')[-1][len('vmlinuz-'):] diff --git a/scale_build/utils/manifest.py b/scale_build/utils/manifest.py index cbeace08..892381ee 100644 --- a/scale_build/utils/manifest.py +++ b/scale_build/utils/manifest.py @@ -167,6 +167,21 @@ 'required': ['name', 'branch', 'repo'], } }, + 'extensions': { + 'type': 'object', + 'properties': { + 'nvidia': { + 'type': 'object', + 'properties': { + 'current': {'type': 'string'}, + }, + 'required': ['current'], + 'additionalProperties': False, + }, + }, + 'required': ['nvidia'], + 'additionalProperties': False, + }, }, 'required': [ 'code_name',