expanding non-default version tests ; adding utility function to verify pyspark ; disk size correction

Author: cjac

spark-rapids/test_spark_rapids.py

@@ -20,6 +20,10 @@ class SparkRapidsTestCase(DataprocTestCase):
   def verify_spark_instance(self, name):
     self.assert_instance_command(name, "nvidia-smi")
 
+  def verify_pyspark(self, name):
+    # Verify that pyspark works
+    self.assert_instance_command(name, "echo 'from pyspark.sql import SparkSession ; SparkSession.builder.getOrCreate()' | pyspark -c spark.executor.resource.gpu.amount=1 -c spark.task.resource.gpu.amount=0.01", 1)
+
   def verify_mig_instance(self, name):
     self.assert_instance_command(name,
         "/usr/bin/nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader | uniq | xargs -I % test % = 'Enabled'")
@@ -114,13 +118,22 @@ def test_spark_rapids_sql(self, configuration, machine_suffixes, accelerator):
     # Only need to do this once
     self.verify_spark_job_sql()
 
-  @parameterized.parameters(("STANDARD", ["w-0"], GPU_T4, "12.4.0", "550.54.14"))
+  @parameterized.parameters(
+    ("STANDARD", ["w-0"], GPU_T4, "11.8.0", "525.147.05"),
+    ("STANDARD", ["w-0"], GPU_T4, "12.4.0", "550.54.14"),
+    ("STANDARD", ["w-0"], GPU_T4, "12.6.2", "560.35.03")
+  )
   def test_non_default_cuda_versions(self, configuration, machine_suffixes,
                                      accelerator, cuda_version, driver_version):
 
     if self.getImageOs() == "rocky":
       self.skipTest("Not supported for Rocky OS")
 
+    if pkg_resources.parse_version(cuda_version) > pkg_resources.parse_version("12.0") \
+    and ( ( self.getImageOs() == 'ubuntu' and self.getImageVersion() <= pkg_resources.parse_version("2.0") ) or \
+          ( self.getImageOs() == 'debian' and self.getImageVersion() <= pkg_resources.parse_version("2.1") ) ):
+      self.skipTest("CUDA > 12.0 not supported on older debian/ubuntu releases")
+
     if self.getImageVersion() <= pkg_resources.parse_version("2.0"):
       self.skipTest("Not supported in 2.0 and earlier images")