Merge pull request #289 from NVIDIA/precompile-kernel-mod-type

tariq1890 · web-flow · commit a2ca287f22fe · 2025-03-26T15:19:32.000-07:00
diff --git a/ubuntu22.04/precompiled/Dockerfile b/ubuntu22.04/precompiled/Dockerfile
@@ -2,6 +2,9 @@ FROM nvcr.io/nvidia/cuda:12.8.1-base-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
+ARG BASE_URL=https://us.download.nvidia.com/tesla
+ARG TARGETARCH
+ENV TARGETARCH=$TARGETARCH
 ARG DRIVER_BRANCH=535
 ENV DRIVER_BRANCH=$DRIVER_BRANCH
 ARG DRIVER_VERSION=535.230.02
@@ -60,6 +63,7 @@ ADD local-repo.sh /tmp
 RUN mkdir -p /usr/local/repos && \
     /tmp/local-repo.sh download_driver_package_deps && \
     /tmp/local-repo.sh build_local_apt_repo && \
+    /tmp/local-repo.sh fetch_nvidia_installer && \
     # Remove cuda repository to avoid GPG errors
     rm -f /etc/apt/sources.list.d/cuda*
 
diff --git a/ubuntu22.04/precompiled/local-repo.sh b/ubuntu22.04/precompiled/local-repo.sh
@@ -3,6 +3,8 @@
 set -eu
 
 LOCAL_REPO_DIR=/usr/local/repos
+DRIVER_ARCH=${TARGETARCH/amd64/x86_64} && DRIVER_ARCH=${DRIVER_ARCH/arm64/aarch64}
+DRIVER_RUN_FILE=NVIDIA-Linux-$DRIVER_ARCH-$DRIVER_VERSION
 
 download_apt_with_dep () {
   local package="$1"
@@ -40,10 +42,21 @@ build_local_apt_repo () {
   apt-get update
 }
 
+fetch_nvidia_installer () {
+  curl -fSsl -O $BASE_URL/$DRIVER_VERSION/$DRIVER_RUN_FILE.run
+  chmod +x $DRIVER_RUN_FILE.run
+  sh $DRIVER_RUN_FILE.run -x
+  mv $DRIVER_RUN_FILE/nvidia-installer /usr/bin/
+  rm -rf $DRIVER_RUN_FILE
+  rm $DRIVER_RUN_FILE.run
+}
+
 if [ "$1" = "download_driver_package_deps" ]; then
   download_driver_package_deps
 elif [ "$1" = "build_local_apt_repo" ]; then
   build_local_apt_repo
+elif [ "$1" = "fetch_nvidia_installer" ]; then
+  fetch_nvidia_installer
 else
   echo "Unknown function: $1"
   exit 1
diff --git a/ubuntu22.04/precompiled/nvidia-driver b/ubuntu22.04/precompiled/nvidia-driver
@@ -4,7 +4,6 @@
 set -eu
 
 KERNEL_VERSION=$(uname -r)
-OPEN_KERNEL_MODULES_ENABLED="${OPEN_KERNEL_MODULES_ENABLED:-false}"
 RUN_DIR=/run/nvidia
 PID_FILE=${RUN_DIR}/${0##*/}.pid
 DRIVER_BRANCH=${DRIVER_BRANCH:?"Missing driver version"}
@@ -14,6 +13,8 @@ NVIDIA_MODULE_PARAMS=()
 NVIDIA_UVM_MODULE_PARAMS=()
 NVIDIA_MODESET_MODULE_PARAMS=()
 NVIDIA_PEERMEM_MODULE_PARAMS=()
+TARGETARCH=${TARGETARCH:?"Missing TARGETARCH env"}
+KERNEL_MODULE_TYPE=${KERNEL_MODULE_TYPE:-auto}
 
 _update_package_cache() {
     if [ "${PACKAGE_TAG:-}" != "builtin" ]; then
@@ -232,6 +233,36 @@ _unload_driver() {
     return 0
 }
 
+_resolve_kernel_type_from_driver_branch() {
+  [[ "${DRIVER_BRANCH}" -lt 560 ]] && KERNEL_TYPE=kernel || KERNEL_TYPE=kernel-open
+}
+
+# _resolve_kernel_type determines which kernel module type, open or proprietary, to install.
+# This function assumes that the nvidia-installer binary is in the PATH, so this function
+# should only be invoked after the userspace driver components have been installed.
+#
+# KERNEL_MODULE_TYPE is the frontend interface that users can use to configure which module
+# to install. Valid values for KERNEL_MODULE_TYPE are 'auto' (default), 'open', and 'proprietary'.
+# When 'auto' is configured, we use the nvidia-installer to recommend the module type to install.
+_resolve_kernel_type() {
+  if [ "${KERNEL_MODULE_TYPE}" == "proprietary" ]; then
+    KERNEL_TYPE=kernel
+  elif [ "${KERNEL_MODULE_TYPE}" == "open" ]; then
+    KERNEL_TYPE=kernel-open
+  elif [ "${KERNEL_MODULE_TYPE}" == "auto" ]; then
+    kernel_module_type=$(nvidia-installer --print-recommended-kernel-module-type 2> /dev/null)
+    if [ $? -ne 0 ]; then
+      echo "failed to retrieve the recommended kernel module type from nvidia-installer, falling back to using the driver branch"
+      _resolve_kernel_type_from_driver_branch
+      return 0
+    fi
+    [[ "${kernel_module_type}" == "open" ]] && KERNEL_TYPE=kernel-open || KERNEL_TYPE=kernel
+  else
+    echo "invalid value for the KERNEL_MODULE_TYPE variable: ${KERNEL_MODULE_TYPE}"
+    return 1
+  fi
+}
+
 # Link and install the kernel modules from a precompiled packages
 _install_driver() {
     # Install necessary driver userspace packages
@@ -244,7 +275,7 @@ _install_driver() {
         libnvidia-fbc1-${DRIVER_BRANCH}-server
 
     # Now install the precompiled kernel module packages signed by Canonical
-    if [ "$OPEN_KERNEL_MODULES_ENABLED" = true ]; then
+    if [ "$KERNEL_TYPE" = "kernel-open" ]; then
         echo "Installing Open NVIDIA driver kernel modules..."
         apt-get install --no-install-recommends -y \
             linux-signatures-nvidia-${KERNEL_VERSION} \
@@ -276,6 +307,9 @@ _unmount_rootfs() {
 }
 
 init() {
+    # Determine the kernel module type
+    _resolve_kernel_type || exit 1
+
     echo -e "\n========== NVIDIA Software Installer ==========\n"
     echo -e "Starting installation of NVIDIA driver branch ${DRIVER_BRANCH} for Linux kernel version ${KERNEL_VERSION}\n"
 
diff --git a/ubuntu24.04/precompiled/Dockerfile b/ubuntu24.04/precompiled/Dockerfile
@@ -2,6 +2,9 @@ FROM nvcr.io/nvidia/cuda:12.8.1-base-ubuntu24.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
+ARG BASE_URL=https://us.download.nvidia.com/tesla
+ARG TARGETARCH
+ENV TARGETARCH=$TARGETARCH
 ARG DRIVER_BRANCH=550
 ENV DRIVER_BRANCH=$DRIVER_BRANCH
 ARG DRIVER_VERSION=550.90.12
@@ -55,6 +58,7 @@ ADD local-repo.sh /tmp
 RUN mkdir -p /usr/local/repos && \
     /tmp/local-repo.sh download_driver_package_deps && \
     /tmp/local-repo.sh build_local_apt_repo && \
+    /tmp/local-repo.sh fetch_nvidia_installer && \
     # Remove all other ubuntu apt sources to ensure we only pull from the local apt repo
     rm /etc/apt/sources.list.d/*
 
diff --git a/ubuntu24.04/precompiled/local-repo.sh b/ubuntu24.04/precompiled/local-repo.sh
@@ -3,6 +3,8 @@
 set -eu
 
 LOCAL_REPO_DIR=/usr/local/repos
+DRIVER_ARCH=${TARGETARCH/amd64/x86_64} && DRIVER_ARCH=${DRIVER_ARCH/arm64/aarch64}
+DRIVER_RUN_FILE=NVIDIA-Linux-$DRIVER_ARCH-$DRIVER_VERSION
 
 download_apt_with_dep () {
   local package="$1"
@@ -40,10 +42,21 @@ build_local_apt_repo () {
   apt-get update
 }
 
+fetch_nvidia_installer () {
+  curl -fSsl -O $BASE_URL/$DRIVER_VERSION/$DRIVER_RUN_FILE.run
+  chmod +x $DRIVER_RUN_FILE.run
+  sh $DRIVER_RUN_FILE.run -x
+  mv $DRIVER_RUN_FILE/nvidia-installer /usr/bin/
+  rm -rf $DRIVER_RUN_FILE
+  rm $DRIVER_RUN_FILE.run
+}
+
 if [ "$1" = "download_driver_package_deps" ]; then
   download_driver_package_deps
 elif [ "$1" = "build_local_apt_repo" ]; then
   build_local_apt_repo
+elif [ "$1" = "fetch_nvidia_installer" ]; then
+  fetch_nvidia_installer
 else
   echo "Unknown function: $1"
   exit 1
diff --git a/ubuntu24.04/precompiled/nvidia-driver b/ubuntu24.04/precompiled/nvidia-driver
@@ -4,7 +4,6 @@
 set -eu
 
 KERNEL_VERSION=$(uname -r)
-OPEN_KERNEL_MODULES_ENABLED="${OPEN_KERNEL_MODULES_ENABLED:-false}"
 RUN_DIR=/run/nvidia
 PID_FILE=${RUN_DIR}/${0##*/}.pid
 DRIVER_BRANCH=${DRIVER_BRANCH:?"Missing driver version"}
@@ -14,6 +13,8 @@ NVIDIA_MODULE_PARAMS=()
 NVIDIA_UVM_MODULE_PARAMS=()
 NVIDIA_MODESET_MODULE_PARAMS=()
 NVIDIA_PEERMEM_MODULE_PARAMS=()
+TARGETARCH=${TARGETARCH:?"Missing TARGETARCH env"}
+KERNEL_MODULE_TYPE=${KERNEL_MODULE_TYPE:-auto}
 
 _update_package_cache() {
     if [ "${PACKAGE_TAG:-}" != "builtin" ]; then
@@ -232,6 +233,36 @@ _unload_driver() {
     return 0
 }
 
+_resolve_kernel_type_from_driver_branch() {
+  [[ "${DRIVER_BRANCH}" -lt 560 ]] && KERNEL_TYPE=kernel || KERNEL_TYPE=kernel-open
+}
+
+# _resolve_kernel_type determines which kernel module type, open or proprietary, to install.
+# This function assumes that the nvidia-installer binary is in the PATH, so this function
+# should only be invoked after the userspace driver components have been installed.
+#
+# KERNEL_MODULE_TYPE is the frontend interface that users can use to configure which module
+# to install. Valid values for KERNEL_MODULE_TYPE are 'auto' (default), 'open', and 'proprietary'.
+# When 'auto' is configured, we use the nvidia-installer to recommend the module type to install.
+_resolve_kernel_type() {
+  if [ "${KERNEL_MODULE_TYPE}" == "proprietary" ]; then
+    KERNEL_TYPE=kernel
+  elif [ "${KERNEL_MODULE_TYPE}" == "open" ]; then
+    KERNEL_TYPE=kernel-open
+  elif [ "${KERNEL_MODULE_TYPE}" == "auto" ]; then
+    kernel_module_type=$(nvidia-installer --print-recommended-kernel-module-type 2> /dev/null)
+    if [ $? -ne 0 ]; then
+      echo "failed to retrieve the recommended kernel module type from nvidia-installer, falling back to using the driver branch"
+      _resolve_kernel_type_from_driver_branch
+      return 0
+    fi
+    [[ "${kernel_module_type}" == "open" ]] && KERNEL_TYPE=kernel-open || KERNEL_TYPE=kernel
+  else
+    echo "invalid value for the KERNEL_MODULE_TYPE variable: ${KERNEL_MODULE_TYPE}"
+    return 1
+  fi
+}
+
 # Link and install the kernel modules from a precompiled packages
 _install_driver() {
     # Install necessary driver userspace packages
@@ -244,7 +275,7 @@ _install_driver() {
         libnvidia-fbc1-${DRIVER_BRANCH}-server
 
     # Now install the precompiled kernel module packages signed by Canonical
-    if [ "$OPEN_KERNEL_MODULES_ENABLED" = true ]; then
+    if [ "$KERNEL_TYPE" = "kernel-open" ]; then
         echo "Installing Open NVIDIA driver kernel modules..."
         apt-get install --no-install-recommends -y \
             linux-signatures-nvidia-${KERNEL_VERSION} \
@@ -276,6 +307,9 @@ _unmount_rootfs() {
 }
 
 init() {
+    # Determine the kernel module type
+    _resolve_kernel_type || exit 1
+
     echo -e "\n========== NVIDIA Software Installer ==========\n"
     echo -e "Starting installation of NVIDIA driver branch ${DRIVER_BRANCH} for Linux kernel version ${KERNEL_VERSION}\n"
 
