Update spark rapids version to 24.02.0 (#1140)

* update rapids 24.02.0 jar Signed-off-by: Suraj Aralihalli <[email protected]> * add support for Dataproc 2.2 Signed-off-by: Suraj Aralihalli <[email protected]> --------- Signed-off-by: Suraj Aralihalli <[email protected]>
GoogleCloudDataproc · Mar 28, 2024 · 2d10d8e · 2d10d8e
1 parent 8c3698d
commit 2d10d8e
Showing 1 changed file with 23 additions and 2 deletions.
diff --git a/spark-rapids/spark-rapids.sh b/spark-rapids/spark-rapids.sh
@@ -41,7 +41,7 @@ else
 fi
 
 # Update SPARK RAPIDS config
-readonly DEFAULT_SPARK_RAPIDS_VERSION="23.12.1"
+readonly DEFAULT_SPARK_RAPIDS_VERSION="24.02.0"
 readonly SPARK_RAPIDS_VERSION=$(get_metadata_attribute 'spark-rapids-version' ${DEFAULT_SPARK_RAPIDS_VERSION})
 readonly XGBOOST_VERSION=$(get_metadata_attribute 'xgboost-version' ${DEFAULT_XGBOOST_VERSION})
 
@@ -55,6 +55,7 @@ CUDA_VERSION=$(get_metadata_attribute 'cuda-version' '12.2.2')  #12.2.2
 NVIDIA_DRIVER_VERSION=$(get_metadata_attribute 'driver-version' '535.104.05') #535.104.05
 CUDA_VERSION_MAJOR="${CUDA_VERSION%.*}"  #12.2
 
+# EXCEPTIONS
 # Change CUDA version for Ubuntu 18 (Cuda 12.1.1 - Driver v530.30.02 is the latest version supported by Ubuntu 18)
 if [[ "${OS_NAME}" == "ubuntu" ]]; then
     UBUNTU_VERSION=$(lsb_release -r | awk '{print $2}') # 20.04
@@ -65,6 +66,15 @@ if [[ "${OS_NAME}" == "ubuntu" ]]; then
       CUDA_VERSION_MAJOR="${CUDA_VERSION%.*}"  #12.1
     fi
 fi
+# Change CUDA version for Debian 12 (Cuda 12.3.2 - Driver v545.23.08 is the latest version supported by Debian 12)
+if [[ "${OS_NAME}" == "debian" ]]; then
+    DEBIAN_VERSION=$(lsb_release -r | awk '{print $2}') # 12
+    if [[ "${DEBIAN_VERSION}" == "12" ]]; then
+      CUDA_VERSION=$(get_metadata_attribute 'cuda-version' '12.3.2')  #12.3.2
+      NVIDIA_DRIVER_VERSION=$(get_metadata_attribute 'driver-version' '545.23.08') #545.23.08
+      CUDA_VERSION_MAJOR="${CUDA_VERSION%.*}"  #12.3
+    fi
+fi
 
 # Verify Secure boot
 SECURE_BOOT="disabled"
@@ -184,13 +194,24 @@ function install_nvidia_gpu_driver() {
 
     dpkg -i /tmp/local-installer.deb
     cp /var/cuda-repo-debian${DEBIAN_VERSION}-${CUDA_VERSION_MAJOR//./-}-local/cuda-*-keyring.gpg /usr/share/keyrings/
+
+    ## EXCEPTION
+    if [[ ${DEBIAN_VERSION} == 12 ]]; then
+      sed -i '0,/Components: main/s//& contrib/' /etc/apt/sources.list.d/debian.sources
+    fi
+
     add-apt-repository contrib
     execute_with_retries "apt-get update"
 
+    ## EXCEPTION
     if [[ ${DEBIAN_VERSION} == 10 ]]; then
       apt remove -y libglvnd0
       apt install -y ca-certificates-java
+    fi
 
+    ## EXCEPTION
+    if [[ ${DEBIAN_VERSION} == 12 ]]; then
+      execute_with_retries "apt-get install -y -q nvidia-kernel-open-dkms"
     fi
 
     execute_with_retries "apt-get install -y -q --no-install-recommends cuda-drivers-${NVIDIA_DRIVER_VERSION_PREFIX}"
@@ -544,7 +565,7 @@ function upgrade_kernel() {
 function check_os_and_secure_boot() {
   if [[ "${OS_NAME}" == "debian" ]]; then
     DEBIAN_VERSION=$(lsb_release -r | awk '{print $2}') # 10 or 11
-    if [[ "${DEBIAN_VERSION}" != "10" && "${DEBIAN_VERSION}" != "11" ]]; then
+    if [[ "${DEBIAN_VERSION}" != "10" && "${DEBIAN_VERSION}" != "11" && "${DEBIAN_VERSION}" != "12" ]]; then
       echo "Error: The Debian version (${DEBIAN_VERSION}) is not supported. Please use a compatible Debian version."
       exit 1
     fi