From 295f2d69fb9feac049c64f22c4d8ad5558c7c938 Mon Sep 17 00:00:00 2001
From: r2k1 <yokree@gmail.com>
Date: Mon, 24 Feb 2025 10:36:53 +1300
Subject: [PATCH] clean installation scripts

---
 .../linux/cloud-init/artifacts/cse_config.sh  | 28 ++++++++++++++++++-
 parts/linux/cloud-init/artifacts/cse_main.sh  |  7 +----
 vhdbuilder/packer/install-dependencies.sh     | 16 +++++++++--
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/parts/linux/cloud-init/artifacts/cse_config.sh b/parts/linux/cloud-init/artifacts/cse_config.sh
index 55aebe30b0a..a47890a98e6 100755
--- a/parts/linux/cloud-init/artifacts/cse_config.sh
+++ b/parts/linux/cloud-init/artifacts/cse_config.sh
@@ -844,8 +844,17 @@ ensureGPUDrivers() {
     fi
 }
 
-# TODO: this is a temporary ubuntu-only HACK until we get a driver
 ensureAMDGPUDrivers() {
+    if [[ $OS == $UBUNTU_OS_NAME ]]; then
+        ensureAMDGPUDriversUbuntu
+    else
+        echo "os $OS not supported at this time. skipping ensureAMDGPUDrivers"
+        return
+    fi
+}
+
+# TODO: this is a temporary ubuntu-only HACK until we get a driver
+ensureAMDGPUDriversUbuntu() {
     echo "Installing AMD GPU drivers"
 
      # delete amdgpu module from blacklist
@@ -860,6 +869,23 @@ ensureAMDGPUDrivers() {
     echo "AMD GPU drivers installed"
 }
 
+cleanAMDGPUDriver() {
+    if [[ $OS == $UBUNTU_OS_NAME ]]; then
+      ensureAMDGPUDriversUbuntu
+    else
+      return
+    fi
+}
+
+cleanAMDGPUDriverUbuntu() {
+   # delete amd from a list of recognized vendors
+   sudo rm /etc/apt/keyrings/rocm.gpg
+   sudo rm /etc/apt/sources.list.d/amdgpu.list
+   # delete cached amd gpu packages to save disk space
+   sudo rm -rf /var/cache/amdgpu-apt/*
+}
+
+
 disableSSH() {
     systemctlDisableAndStop ssh || exit $ERR_DISABLE_SSH
 }
diff --git a/parts/linux/cloud-init/artifacts/cse_main.sh b/parts/linux/cloud-init/artifacts/cse_main.sh
index e896b12a3ca..ec6d288c0b0 100755
--- a/parts/linux/cloud-init/artifacts/cse_main.sh
+++ b/parts/linux/cloud-init/artifacts/cse_main.sh
@@ -208,13 +208,8 @@ fi
 
 if [[ "${AMD_GPU_NODE}" = true ]] && [[ "${skip_gpu_driver_install}" != "true" ]]; then
     logs_to_events "AKS.CSE.ensureAMDGPUDrivers" ensureAMDGPUDrivers
-else
-    # delete cached amd gpu packages to save disk space
-    sudo rm /etc/apt/keyrings/rocm.gpg
-    sudo rm /etc/apt/sources.list.d/amdgpu.list
-    sudo rm -rf /var/cache/amdgpu-apt/*
 fi
-
+cleanAMDGPUDrivers
 
 echo $(date),$(hostname), "End configuring GPU drivers"
 
diff --git a/vhdbuilder/packer/install-dependencies.sh b/vhdbuilder/packer/install-dependencies.sh
index 1a73af51bc6..7dd413b198a 100644
--- a/vhdbuilder/packer/install-dependencies.sh
+++ b/vhdbuilder/packer/install-dependencies.sh
@@ -621,7 +621,7 @@ capture_benchmark "${SCRIPT_NAME}_overall" true
 process_benchmarks
 
 
-download_amdgpu_drivers() {
+downloadAMDGPUDriversUbuntu() {
   if [[ $OS != $UBUNTU_OS_NAME ]]; then
     echo "Skipping AMD GPU driver setup: Unsupported OS (${OS})"
     return
@@ -641,7 +641,8 @@ download_amdgpu_drivers() {
   wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
       gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
   sudo chmod 0644 /etc/apt/keyrings/rocm.gpg
-  echo "deb [arch=amd64,i386 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/6.3.2/ubuntu ${DISTRO} main" \
+
+  echo "deb [arch=amd64,i386 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/6.3.3/ubuntu ${DISTRO} main" \
       | sudo tee /etc/apt/sources.list.d/amdgpu.list
   sudo apt-get update
 
@@ -655,4 +656,13 @@ download_amdgpu_drivers() {
   sudo apt-get install -o Dir::Cache::Archives="/var/cache/amdgpu-apt" --download-only --reinstall -y m4 amdgpu-dkms autoconf automake autotools-dev amdgpu-dkms-firmware
 }
 
-download_amdgpu_drivers
\ No newline at end of file
+downloadAMDGPUDrivers() {
+  if [[ $OS == $UBUNTU_OS_NAME ]]; then
+    downloadAMDGPUDriversUbuntu
+  else
+    echo "os $OS not supported at this time. skipping ensureAMDGPUDrivers"
+    return
+  fi
+}
+
+downloadAMDGPUDrivers
\ No newline at end of file