Move nvidia extension to the rootfs image

truenas · Nov 20, 2024 · 75cc110 · 75cc110
1 parent 85f8a44
commit 75cc110
Show file tree

Hide file tree

Showing 7 changed files with 108 additions and 81 deletions.
diff --git 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
@@ -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,77 +17,71 @@
 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) /")
 
 
 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
@@ -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/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
@@ -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/utils/kernel.py 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-*'))[1].split('/')[-1][len('vmlinuz-'):]
diff --git 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',