From 1df72b2f5e1d7a39aa99b504789e8d61e7ec45bf Mon Sep 17 00:00:00 2001
From: Lujie Duan <lujieduan@google.com>
Date: Fri, 6 Oct 2023 13:49:10 +0000
Subject: [PATCH] Install drivers for K80 without causing vm to reboot

---
 .../applications/dcgm/debian_ubuntu/install          |  2 +-
 .../applications/nvml/debian_ubuntu/install          | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/integration_test/third_party_apps_data/applications/dcgm/debian_ubuntu/install b/integration_test/third_party_apps_data/applications/dcgm/debian_ubuntu/install
index 3009a02dc9..82babc83dd 100755
--- a/integration_test/third_party_apps_data/applications/dcgm/debian_ubuntu/install
+++ b/integration_test/third_party_apps_data/applications/dcgm/debian_ubuntu/install
@@ -9,7 +9,7 @@ DEVICE_CODE=$(lspci -n | grep -Po '10de:[\w\d]{4}')
 case $DEVICE_CODE in
     10de:102d)
         # Install a specific version for NVIDIA Tesla K80
-        DRIVER_VERSION=460.106.00
+        DRIVER_VERSION=470.82.01
         ;;
     *)
         DRIVER_VERSION=525.60.13
diff --git a/integration_test/third_party_apps_data/applications/nvml/debian_ubuntu/install b/integration_test/third_party_apps_data/applications/nvml/debian_ubuntu/install
index 3f01b91f6c..05b5279e36 100755
--- a/integration_test/third_party_apps_data/applications/nvml/debian_ubuntu/install
+++ b/integration_test/third_party_apps_data/applications/nvml/debian_ubuntu/install
@@ -10,12 +10,12 @@ DEVICE_CODE=$(lspci -n | grep -Po '10de:[\w\d]{4}')
 case $DEVICE_CODE in
     10de:102d)
         # Install a specific version for NVIDIA Tesla K80, R470 is the last supported version
-        echo "Installing NVIDIA CUDA 11.2.1 with driver 460.32.03"
-        curl https://raw.githubusercontent.com/GoogleCloudPlatform/compute-gpu-installation/main/linux/install_gpu_driver.py --output install_gpu_driver.py
-        sudo python3 install_gpu_driver.py
-        wget --no-verbose https://developer.download.nvidia.com/compute/cuda/11.2.1/local_installers/cuda_11.2.1_460.32.03_linux.run
-        sudo sh cuda_11.2.1_460.32.03_linux.run --toolkit --silent
-        sudo apt install -y libcublas10
+        DRIVER_VERSION=470.82.01
+        echo "Installing NVIDIA CUDA 11.2.1 with driver $DRIVER_VERSION"
+        curl -fSsl -O https://us.download.nvidia.com/tesla/$DRIVER_VERSION/NVIDIA-Linux-x86_64-$DRIVER_VERSION.run
+        sudo bash ./NVIDIA-Linux-x86_64-$DRIVER_VERSION.run --silent
+        wget --no-verbose https://developer.download.nvidia.com/compute/cuda/11.4.4/local_installers/cuda_11.4.4_${DRIVER_VERSION}_linux.run
+        sudo sh cuda_11.4.4_${DRIVER_VERSION}_linux.run --toolkit --silent
         ;;
     *)
         echo "Installing latest version of NVIDIA CUDA and driver"