From 500da14c227ae9a888fdbb6075af496c49dfce27 Mon Sep 17 00:00:00 2001
From: Hao Wu <tohaowu@google.com>
Date: Mon, 11 Sep 2023 17:23:53 -0700
Subject: [PATCH] Add CUDA 12.2 and fix version compatibility issues between
 Pandas and Dask.

PiperOrigin-RevId: 564547852
---
 .../linux_benchmarks/xgboost_benchmark.py              |  4 ++++
 perfkitbenchmarker/linux_packages/cuda_toolkit.py      | 10 +++++++++-
 perfkitbenchmarker/linux_packages/nvidia_driver.py     |  2 +-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/perfkitbenchmarker/linux_benchmarks/xgboost_benchmark.py b/perfkitbenchmarker/linux_benchmarks/xgboost_benchmark.py
index eb23d1fc2b..d9fb5ee55c 100644
--- a/perfkitbenchmarker/linux_benchmarks/xgboost_benchmark.py
+++ b/perfkitbenchmarker/linux_benchmarks/xgboost_benchmark.py
@@ -82,6 +82,10 @@ def Prepare(bm_spec: benchmark_spec.BenchmarkSpec) -> None:
       ' xgboost_ray/requirements/test-requirements.txt'
   )
   vm.RemoteCommand(f'{FLAGS.xgboost_env} python3 -m pip install xgboost_ray')
+  vm.RemoteCommand(
+      f'{FLAGS.xgboost_env} python3 -m pip install --upgrade pandas'
+      ' "dask[complete]"'
+  )
   vm.RemoteCommand(
       f'{FLAGS.xgboost_env} python3'
       ' xgboost_ray/xgboost_ray/tests/release/create_test_data.py'
diff --git a/perfkitbenchmarker/linux_packages/cuda_toolkit.py b/perfkitbenchmarker/linux_packages/cuda_toolkit.py
index efe8e91a64..b86725906e 100644
--- a/perfkitbenchmarker/linux_packages/cuda_toolkit.py
+++ b/perfkitbenchmarker/linux_packages/cuda_toolkit.py
@@ -52,6 +52,7 @@
         '11.8',
         '12.0',
         '12.1',
+        '12.2',
         'None',
         '',
     ],
@@ -70,6 +71,7 @@
 
 CUDA_PIN = 'https://developer.download.nvidia.com/compute/cuda/repos/{os}/{cpu_arch}/cuda-{os}.pin'
 
+CUDA_12_2_TOOLKIT = 'https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda-repo-{os}-12-2-local_12.2.2-535.104.05-1_{cpu_arch}.deb'
 CUDA_12_1_TOOLKIT = 'https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda-repo-{os}-12-1-local_12.1.1-530.30.02-1_{cpu_arch}.deb'
 CUDA_12_0_TOOLKIT = 'https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installers/cuda-repo-{os}-12-0-local_12.0.1-525.85.12-1_{cpu_arch}.deb'
 CUDA_11_8_TOOLKIT = 'https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-{os}-11-8-local_11.8.0-520.61.05-1_{cpu_arch}.deb'
@@ -174,7 +176,7 @@ def GetCudaToolkitVersion(vm):
 
 
 def EnrollSigningKey(vm):
-  if FLAGS.cuda_toolkit_version in ('11.7', '11.8', '12.0', '12.1'):
+  if FLAGS.cuda_toolkit_version in ('11.7', '11.8', '12.0', '12.1', '12.2'):
     version = FLAGS.cuda_toolkit_version.replace('.', '-')
     vm.RemoteCommand(
         'sudo cp'
@@ -359,6 +361,10 @@ def _InstallCuda12Point1(vm):
   _InstallCuda12Generic(vm, CUDA_12_1_TOOLKIT, '12-1')
 
 
+def _InstallCuda12Point2(vm):
+  _InstallCuda12Generic(vm, CUDA_12_2_TOOLKIT, '12-2')
+
+
 def _InstallCuda11Point0(vm):
   _InstallCuda11Generic(vm, CUDA_11_0_TOOLKIT, '11-0')
 
@@ -444,6 +450,8 @@ def AptInstall(vm):
     _InstallCuda12Point0(vm)
   elif version_to_install == '12.1':
     _InstallCuda12Point1(vm)
+  elif version_to_install == '12.2':
+    _InstallCuda12Point2(vm)
   else:
     raise UnsupportedCudaVersionError()
   DoPostInstallActions(vm)
diff --git a/perfkitbenchmarker/linux_packages/nvidia_driver.py b/perfkitbenchmarker/linux_packages/nvidia_driver.py
index 34cfbd2444..e11d543964 100644
--- a/perfkitbenchmarker/linux_packages/nvidia_driver.py
+++ b/perfkitbenchmarker/linux_packages/nvidia_driver.py
@@ -49,7 +49,7 @@
 
 flags.DEFINE_string(
     'nvidia_driver_version',
-    '525.105.17',
+    '535.104.05',
     (
         'The version of nvidia driver to install. '
         'For example, "418.67" or "418.87.01."'