diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/Platform.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/Platform.scala index 077230c99..e845d5575 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/Platform.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/Platform.scala @@ -844,7 +844,9 @@ abstract class DatabricksPlatform(gpuDevice: Option[GpuDevice], // TODO: Issue to automate this https://github.com/NVIDIA/cudf-spark-tools/issues/1676 override val supportedShuffleManagerVersionMap: Array[(String, String)] = Array( "12.2" -> "332db", - "13.3" -> "341db" + "13.3" -> "341db", + "14.3" -> "350db143", + "17.3" -> "400db173" ) override def createClusterInfo(coresPerExecutor: Int, diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala index 6c067fc44..dd09c9c56 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuite.scala @@ -56,6 +56,38 @@ abstract class ProfilingAutoTunerSuiteBase extends BaseAutoTunerSuite { } } + /** + * Helper method to verify that the recommended shuffle manager version matches the + * expected version. + */ + protected def verifyRecommendedShuffleManagerVersion( + autoTuner: AutoTuner, + expectedSmVersion: String): Unit = { + autoTuner.getShuffleManagerClassName match { + case Right(smClassName) => + assert(smClassName == ProfilingAutoTunerHelper + .buildShuffleManagerClassName(expectedSmVersion)) + case Left(comment) => + fail(s"Expected valid RapidsShuffleManager but got comment: $comment") + } + } + + /** + * Helper method to verify that the shuffle manager version is not recommended + * for the unsupported Spark version. + */ + protected def verifyUnsupportedSparkVersionForShuffleManager( + autoTuner: AutoTuner, + sparkVersion: String): Unit = { + autoTuner.getShuffleManagerClassName match { + case Right(smClassName) => + fail(s"Expected error comment but got valid RapidsShuffleManager: $smClassName") + case Left(comment) => + assert(comment == shuffleManagerCommentForUnsupportedVersion(sparkVersion, + autoTuner.platform)) + } + } + /** * Helper method to extract the AutoTuner results from the profile log content * TODO: We should store the AutoTuner results in a separate file. @@ -2310,22 +2342,6 @@ class ProfilingAutoTunerSuite extends ProfilingAutoTunerSuiteBase { compareOutput(expectedResults, autoTunerOutput) } - /** - * Helper method to verify that the recommended shuffle manager version matches the - * expected version. - */ - private def verifyRecommendedShuffleManagerVersion( - autoTuner: AutoTuner, - expectedSmVersion: String): Unit = { - autoTuner.getShuffleManagerClassName match { - case Right(smClassName) => - assert(smClassName == ProfilingAutoTunerHelper - .buildShuffleManagerClassName(expectedSmVersion)) - case Left(comment) => - fail(s"Expected valid RapidsShuffleManager but got comment: $comment") - } - } - val dbPlatform: Platform = PlatformFactory.createInstance(PlatformNames.DATABRICKS_AWS) dbPlatform.supportedShuffleManagerVersionMap.foreach { case (dbVersion, smVersion) => test(s"test shuffle manager version for supported databricks version - $dbVersion") { @@ -2367,22 +2383,6 @@ class ProfilingAutoTunerSuite extends ProfilingAutoTunerSuiteBase { verifyRecommendedShuffleManagerVersion(autoTuner, expectedSmVersion = "330") } - /** - * Helper method to verify that the shuffle manager version is not recommended - * for the unsupported Spark version. - */ - private def verifyUnsupportedSparkVersionForShuffleManager( - autoTuner: AutoTuner, - sparkVersion: String): Unit = { - autoTuner.getShuffleManagerClassName match { - case Right(smClassName) => - fail(s"Expected error comment but got valid RapidsShuffleManager: $smClassName") - case Left(comment) => - assert(comment == shuffleManagerCommentForUnsupportedVersion(sparkVersion, - autoTuner.platform)) - } - } - test("test shuffle manager version for unsupported databricks version") { val databricksVersion = "9.1.x-gpu-ml-scala2.12" val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala index b66fa4255..cbd1b36fb 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala @@ -19,6 +19,7 @@ package com.nvidia.spark.rapids.tool.tuning import scala.collection.mutable import com.nvidia.spark.rapids.tool.{DynamicAllocationInfo, GpuTypes, NodeInstanceMapKey, PlatformFactory, PlatformInstanceTypes, PlatformNames, ToolTestUtils} +import com.nvidia.spark.rapids.tool.planparser.db.DBVersionExtractor import com.nvidia.spark.rapids.tool.profiling.{Profiler, PySparkMemoryEvidence, RecommendedCommentResult, ShuffleStageInputAnalysis} import com.nvidia.spark.rapids.tool.tuning.config.{ConfTypeEnum, TuningConfigEntry, @@ -3763,4 +3764,50 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(!values.contains("spark.yarn.isPython")) } + // The Databricks entries of supportedShuffleManagerVersionMap. The map-driven tests in + // ProfilingAutoTunerSuite take both the runtime and the expected class from the map, so a + // wrong or missing entry stays green there; these pin the expected values independently. + test("test shuffle manager version for databricks 14.3 is 350db143") { + val databricksVersion = "14.3.x-gpu-ml-scala2.12" + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), + mutable.Map("spark.rapids.sql.enabled" -> "true", + "spark.plugins" -> "com.nvidia.spark.AnotherPlugin, com.nvidia.spark.SQLPlugin", + DBVersionExtractor.DB_SPARK_VERSION_KEY -> databricksVersion), + Some(databricksVersion), Seq()) + val autoTuner = buildAutoTunerForTests( + infoProvider, + PlatformFactory.createInstance(PlatformNames.DATABRICKS_AWS)) + verifyRecommendedShuffleManagerVersion(autoTuner, expectedSmVersion = "350db143") + } + + test("test shuffle manager version for databricks 17.3 on scala 2.13 is 400db173") { + val databricksVersion = "17.3.x-gpu-ml-scala2.13" + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), + mutable.Map("spark.rapids.sql.enabled" -> "true", + "spark.plugins" -> "com.nvidia.spark.AnotherPlugin, com.nvidia.spark.SQLPlugin", + DBVersionExtractor.DB_SPARK_VERSION_KEY -> databricksVersion), + Some(databricksVersion), Seq()) + val autoTuner = buildAutoTunerForTests( + infoProvider, + PlatformFactory.createInstance(PlatformNames.DATABRICKS_AZURE)) + verifyRecommendedShuffleManagerVersion(autoTuner, expectedSmVersion = "400db173") + } + + test("test shuffle manager version for databricks 15.4 without a plugin shim") { + val databricksVersion = "15.4.x-gpu-ml-scala2.12" + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), + mutable.Map("spark.rapids.sql.enabled" -> "true", + "spark.plugins" -> "com.nvidia.spark.AnotherPlugin, com.nvidia.spark.SQLPlugin", + DBVersionExtractor.DB_SPARK_VERSION_KEY -> databricksVersion), + Some(databricksVersion), Seq()) + val autoTuner = buildAutoTunerForTests( + infoProvider, + PlatformFactory.createInstance(PlatformNames.DATABRICKS_AWS)) + verifyUnsupportedSparkVersionForShuffleManager(autoTuner, databricksVersion) + // the comment points the user at the newest supported runtime, which must be 17.3 + val (latestVersion, latestSmVersion) = autoTuner.platform.latestSupportedShuffleManagerInfo + assert(latestVersion == "17.3") + assert(latestSmVersion == "400db173") + } + }