From 508358c06e731c85a6e8305c25632d7248106b6c Mon Sep 17 00:00:00 2001 From: Thomas Wynne Date: Wed, 9 Sep 2026 00:06:24 -0500 Subject: [PATCH 1/2] Add Databricks 14.3 and 17.3 to the AutoTuner shuffle manager map On a Databricks 14.3 or 17.3 event log the profiling AutoTuner commented "Cannot recommend RAPIDS Shuffle Manager for unsupported Databricks runtime" and advised moving to runtime 13.3, whose shim the plugin has dropped. The Databricks entries of `supportedShuffleManagerVersionMap` stopped at 13.3 while the plugin ships `spark350db143.RapidsShuffleManager` and `spark400db173.RapidsShuffleManager`, and Databricks sets `spark.shuffle.manager=SORT` explicitly, so the tuner's recommendation is the only place a user is told the class name. The map gains `14.3 -> 350db143` and `17.3 -> 400db173`. The 12.2 and 13.3 entries stay so old event logs keep their recommendation, and the "e.g." runtime in the unsupported-version comment becomes 17.3 through the existing `maxBy`. Tests: the per-entry Databricks test in `ProfilingAutoTunerSuite` covers the two entries; a 17.3 case with the `-scala2.13` runtime suffix is added. Fixes #2144 Signed-off-by: Thomas Wynne --- .../nvidia/spark/rapids/tool/Platform.scala | 4 ++- .../tool/tuning/ProfilingAutoTunerSuite.scala | 30 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) 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..cb4adc74f 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 @@ -2397,6 +2397,36 @@ class ProfilingAutoTunerSuite extends ProfilingAutoTunerSuiteBase { verifyUnsupportedSparkVersionForShuffleManager(autoTuner, databricksVersion) } + test("test shuffle manager version for supported databricks version - 17.3 on scala 2.13") { + 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 version without a plugin shim - 15.4") { + 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") + } + test("test shuffle manager version for unsupported spark version") { val sparkVersion = "3.1.2" val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), From 68edf95c01a4bf1aa38288b9ed023d17223218ec Mon Sep 17 00:00:00 2001 From: Thomas Wynne Date: Fri, 11 Sep 2026 09:43:20 -0500 Subject: [PATCH 2/2] Pin the Databricks shuffle manager entries in ProfilingAutoTunerSuiteV2 The two cases added for 14.3 and 17.3 sat in ProfilingAutoTunerSuite, whose header marks it deprecated and sends new AutoTuner cases to ProfilingAutoTunerSuiteV2. They move there, and a third case pins 14.3 to 350db143 with the value written out: the map-driven tests take both the runtime and the expected class from supportedShuffleManagerVersionMap, so a wrong or missing entry stays green in them, and 17.3 already had its own pin. verifyRecommendedShuffleManagerVersion and verifyUnsupportedSparkVersionForShuffleManager move from the deprecated suite into ProfilingAutoTunerSuiteBase, unchanged, so both suites share them and the map-driven tests keep running. With the two map entries reverted the three V2 cases fail; with them ProfilingAutoTunerSuite and ProfilingAutoTunerSuiteV2 pass under JDK 17, 198 succeeded and 0 failed across the two, scalastyle clean. Signed-off-by: Thomas Wynne --- .../tool/tuning/ProfilingAutoTunerSuite.scala | 94 +++++++------------ .../tuning/ProfilingAutoTunerSuiteV2.scala | 47 ++++++++++ 2 files changed, 79 insertions(+), 62 deletions(-) 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 cb4adc74f..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), @@ -2397,36 +2397,6 @@ class ProfilingAutoTunerSuite extends ProfilingAutoTunerSuiteBase { verifyUnsupportedSparkVersionForShuffleManager(autoTuner, databricksVersion) } - test("test shuffle manager version for supported databricks version - 17.3 on scala 2.13") { - 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 version without a plugin shim - 15.4") { - 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") - } - test("test shuffle manager version for unsupported spark version") { val sparkVersion = "3.1.2" 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") + } + }