From e5927f205b99e864578d63e59b13866ad8b52c3d Mon Sep 17 00:00:00 2001
From: shenlan <manbuzhe2009@qq.com>
Date: Fri, 27 Jun 2025 16:25:09 +0800
Subject: [PATCH] fix offline gpu package script

---
 docs/gpu-offline-package.md               |  2 +-
 scripts/create-gpu-k8s-offline-package.sh | 11 +++++++--
 scripts/gpu-k8s.sh                        | 29 ++++++++++++++++++-----
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/docs/gpu-offline-package.md b/docs/gpu-offline-package.md
index ab9797d..0c902dc 100644
--- a/docs/gpu-offline-package.md
+++ b/docs/gpu-offline-package.md
@@ -17,7 +17,7 @@ bash scripts/create-gpu-k8s-offline-package.sh
 
 - Kubernetes 二进制镜像
 - Cilium、Helm 等依赖镜像
-- NVIDIA 驱动（nvidia-driver-535）及 nvidia-container-toolkit deb 包
+- NVIDIA 驱动（nvidia-driver-535）及 nvidia-container-toolkit 离线包（deb/rpm）
 - nerdctl CLI（v${NERDCTL_VERSION:-2.1.2}）
 - 必要的容器镜像，包括 `registry.k8s.io/pause:3.8`
 
diff --git a/scripts/create-gpu-k8s-offline-package.sh b/scripts/create-gpu-k8s-offline-package.sh
index e2b5683..5afebe0 100755
--- a/scripts/create-gpu-k8s-offline-package.sh
+++ b/scripts/create-gpu-k8s-offline-package.sh
@@ -40,12 +40,19 @@ curl -s -L "https://nvidia.github.io/nvidia-docker/${distribution}/nvidia-docker
   sudo tee /etc/apt/sources.list.d/nvidia-docker.list
 
 sudo apt-get update -y
-export NVIDIA_CONTAINER_TOOLKIT_VERSION=1.17.8-1
-sudo apt-get install --download-only -y \
+
+if [ -n "${NVIDIA_CONTAINER_TOOLKIT_VERSION:-}" ]; then
+  sudo apt-get install --download-only -y \
       nvidia-container-toolkit=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
       nvidia-container-toolkit-base=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
       libnvidia-container-tools=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
       libnvidia-container1=${NVIDIA_CONTAINER_TOOLKIT_VERSION}
+else
+  sudo apt-get install --download-only -y \
+      nvidia-container-toolkit nvidia-container-toolkit-base \
+      libnvidia-container-tools libnvidia-container1
+fi
+
 sudo apt-get install --download-only -y "${APT_PACKAGES[@]}"
 cp /var/cache/apt/archives/*.deb "$WORKDIR/packages/"
 sudo apt-get clean
diff --git a/scripts/gpu-k8s.sh b/scripts/gpu-k8s.sh
index 2144d6c..23e2717 100644
--- a/scripts/gpu-k8s.sh
+++ b/scripts/gpu-k8s.sh
@@ -79,12 +79,26 @@ load_offline_images() {
 }
 
 install_all_offline_packages() {
-  if [ -d "${OFFLINE_DIR}/packages" ]; then
-    echo "📦 Using offline deb packages"
-    sudo dpkg -i ${OFFLINE_DIR}/packages/*.deb 2>/dev/null || sudo apt-get -f install -y
-    return 0
+  if [ ! -d "${OFFLINE_DIR}/packages" ]; then
+    return 1
   fi
-  return 1
+
+  echo "📦 Using offline packages"
+
+  shopt -s nullglob
+  debs=("${OFFLINE_DIR}"/packages/*.deb)
+  rpms=("${OFFLINE_DIR}"/packages/*.rpm)
+
+  if [ ${#debs[@]} -gt 0 ]; then
+    sudo dpkg -i "${debs[@]}" 2>/dev/null || sudo apt-get -f install -y
+  fi
+
+  if [ ${#rpms[@]} -gt 0 ] && command -v dnf >/dev/null; then
+    sudo dnf install -y "${rpms[@]}"
+  fi
+
+  shopt -u nullglob
+  return 0
 }
 
 install_base() {
@@ -120,6 +134,7 @@ install_nerdctl() {
 install_nvidia() {
   echo "[3/8] 安装 NVIDIA 驱动和容器工具"
   distribution="ubuntu22.04"
+
   if [ -f "${OFFLINE_DIR}/nvidia-gpgkey" ]; then
     sudo install -m 0644 "${OFFLINE_DIR}/nvidia-gpgkey" /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
   else
@@ -127,7 +142,9 @@ install_nvidia() {
       sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
   fi
 
-  if ! install_all_offline_packages; then
+  if install_all_offline_packages; then
+    echo "✅ 离线安装 NVIDIA 相关包完成"
+  else
     proxy_curl -sL https://nvidia.github.io/nvidia-docker/${distribution}/nvidia-docker.list | \
       sed 's|^deb |deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] |' | \
       sudo tee /etc/apt/sources.list.d/nvidia-docker.list