integrate gpu-driver-util into the driver images

Signed-off-by: Tariq Ibrahim <[email protected]>

Author: tariq1890

ubuntu20.04/Dockerfile

@@ -28,9 +28,10 @@ ENV PATH /usr/local/go/bin:$PATH
 WORKDIR /work
 
 RUN git clone https://github.com/NVIDIA/gpu-driver-container driver && \
-    cd driver/vgpu/src && \
-    go build -o vgpu-util && \
-    mv vgpu-util /work
+    go build -C driver/vgpu/src -o vgpu-util && \
+    mv driver/vgpu/src/vgpu-util /work && \
+    go build -C driver/gpu-driver-util -o gpu-driver-util && \
+    mv driver/gpu-driver-util/gpu-driver-util /work
 
 FROM nvcr.io/nvidia/cuda:12.6.3-base-ubuntu20.04
 
@@ -69,9 +70,15 @@ RUN /tmp/install.sh reposetup && /tmp/install.sh depinstall && \
     curl -fsSL -o /usr/local/bin/donkey https://github.com/3XX0/donkey/releases/download/v1.1.0/donkey && \
     chmod +x /usr/local/bin/donkey
 
+# Download the nvidia-driver-assistant to get the latest supported-gpus.json file
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-driver-assistant && \
+    rm -rf /var/lib/apt/lists/*
+
 COPY nvidia-driver /usr/local/bin
 
 COPY --from=build /work/vgpu-util /usr/local/bin
+COPY --from=build /work/gpu-driver-util /usr/local/bin
 
 ADD drivers drivers/
 

ubuntu22.04/Dockerfile

@@ -28,9 +28,10 @@ ENV PATH /usr/local/go/bin:$PATH
 WORKDIR /work
 
 RUN git clone https://github.com/NVIDIA/gpu-driver-container driver && \
-    cd driver/vgpu/src && \
-    go build -o vgpu-util && \
-    mv vgpu-util /work
+    go build -C driver/vgpu/src -o vgpu-util && \
+    mv driver/vgpu/src/vgpu-util /work && \
+    go build -C driver/gpu-driver-util -o gpu-driver-util && \
+    mv driver/gpu-driver-util/gpu-driver-util /work
 
 FROM nvcr.io/nvidia/cuda:12.6.3-base-ubuntu22.04
 
@@ -69,9 +70,15 @@ RUN /tmp/install.sh reposetup && /tmp/install.sh depinstall && \
     curl -fsSL -o /usr/local/bin/donkey https://github.com/3XX0/donkey/releases/download/v1.1.0/donkey && \
     chmod +x /usr/local/bin/donkey
 
+# Download the nvidia-driver-assistant to get the latest supported-gpus.json file
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-driver-assistant && \
+    rm -rf /var/lib/apt/lists/*
+
 COPY nvidia-driver /usr/local/bin
 
 COPY --from=build /work/vgpu-util /usr/local/bin
+COPY --from=build /work/gpu-driver-util /usr/local/bin
 
 ADD drivers drivers/
 

ubuntu22.04/nvidia-driver

@@ -15,9 +15,7 @@ NVIDIA_UVM_MODULE_PARAMS=()
 NVIDIA_MODESET_MODULE_PARAMS=()
 NVIDIA_PEERMEM_MODULE_PARAMS=()
 TARGETARCH=${TARGETARCH:?"Missing TARGETARCH env"}
-
-OPEN_KERNEL_MODULES_ENABLED=${OPEN_KERNEL_MODULES_ENABLED:-false}
-[[ "${OPEN_KERNEL_MODULES_ENABLED}" == "true" ]] && KERNEL_TYPE=kernel-open || KERNEL_TYPE=kernel
+KERNEL_MODULE_TYPE=${KERNEL_MODULE_TYPE:-auto}
 
 export DEBIAN_FRONTEND=noninteractive
 
@@ -481,6 +479,27 @@ _shutdown() {
     return 1
 }
 
+_resolve_kernel_type() {
+
+  if [ "${KERNEL_MODULE_TYPE}" == "proprietary" ]; then
+    KERNEL_TYPE=kernel
+    return 0
+  elif [ "${KERNEL_MODULE_TYPE}" == "open" ]; then
+    KERNEL_TYPE=kernel-open
+    return 0
+  elif [ "${KERNEL_MODULE_TYPE}" == "auto" ]; then
+      KERNEL_TYPE=$(gpu-driver-util get-kernel-module-type -b "${DRIVER_BRANCH}")
+      if [ $? -ne 0 ]; then
+          echo "cannot autodetect the kernel module type, please check /var/log/gpu-driver-util.log for more details..."
+          return 1
+      fi
+      return 0
+  else
+    echo "invalid value for the KERNEL_MODULE_TYPE variable: ${KERNEL_MODULE_TYPE}"
+    return 1
+  fi
+}
+
 _find_vgpu_driver_version() {
     local count=""
     local version=""
@@ -524,6 +543,8 @@ init() {
         _find_vgpu_driver_version || exit 1
     fi
 
+    _resolve_kernel_type || exit 1
+
     # Install the userspace components and copy the kernel module sources.
     sh NVIDIA-Linux-$DRIVER_ARCH-$DRIVER_VERSION.run -x && \
         cd NVIDIA-Linux-$DRIVER_ARCH-$DRIVER_VERSION && \
@@ -596,6 +617,8 @@ update() {
     fi
     exec 3>&-
 
+    _resolve_kernel_type || exit 1
+
     # vgpu driver version is choosen dynamically during runtime, so pre-compile modules for
     # only non-vgpu driver types
     if [ "${DRIVER_TYPE}" != "vgpu" ]; then

ubuntu24.04/Dockerfile

@@ -18,8 +18,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         git && \
     rm -rf /var/lib/apt/lists/*
 
-
-
 # download appropriate binary based on the target architecture for multi-arch builds
 RUN OS_ARCH=${TARGETARCH/x86_64/amd64} && OS_ARCH=${OS_ARCH/aarch64/arm64} && \
     curl https://storage.googleapis.com/golang/go${GOLANG_VERSION}.linux-${OS_ARCH}.tar.gz \
@@ -30,9 +28,10 @@ ENV PATH /usr/local/go/bin:$PATH
 WORKDIR /work
 
 RUN git clone https://github.com/NVIDIA/gpu-driver-container driver && \
-    cd driver/vgpu/src && \
-    go build -o vgpu-util && \
-    mv vgpu-util /work
+    go build -C driver/vgpu/src -o vgpu-util && \
+    mv driver/vgpu/src/vgpu-util /work && \
+    go build -C driver/gpu-driver-util -o gpu-driver-util && \
+    mv driver/gpu-driver-util/gpu-driver-util /work
 
 FROM nvcr.io/nvidia/cuda:12.6.3-base-ubuntu24.04
 
@@ -66,9 +65,15 @@ ADD install.sh /tmp
 RUN usermod -o -u 0 -g 0 _apt && \
     /tmp/install.sh depinstall && /tmp/install.sh setup_cuda_repo
 
+# Download the nvidia-driver-assistant to get the latest supported-gpus.json file
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends nvidia-driver-assistant && \
+    rm -rf /var/lib/apt/lists/*
+
 COPY nvidia-driver /usr/local/bin
 
 COPY --from=build /work/vgpu-util /usr/local/bin
+COPY --from=build /work/gpu-driver-util /usr/local/bin
 
 ADD drivers drivers/