Refactor and update comments

Signed-off-by: Partho Sarthi <[email protected]>
NVIDIA · Jan 3, 2025 · 2a84eda · 2a84eda
1 parent 782e82a
commit 2a84eda
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 16 deletions.
diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala
@@ -771,10 +771,10 @@ class AutoTuner(
           case Some(smVersion) =>
             Right(autoTunerConfigsProvider.buildShuffleManagerClassName(smVersion))
           case None =>
-            Left(autoTunerConfigsProvider.shuffleManagerCommentForUnsupportedVersion(sparkVersion))
+            Left(autoTunerConfigsProvider.shuffleManagerComments("unsupported")(sparkVersion))
         }
       case None =>
-        Left(autoTunerConfigsProvider.shuffleManagerCommentForMissingVersion)
+        Left(autoTunerConfigsProvider.shuffleManagerComments("missing")(""))
     }
   }
 
@@ -1346,6 +1346,16 @@ trait AutoTunerConfigsProvider extends Logging {
   private val shuffleManagerDocUrl = "https://docs.nvidia.com/spark-rapids/user-guide/latest/" +
     "additional-functionality/rapids-shuffle.html#rapids-shuffle-manager"
 
+  val shuffleManagerComments: Map[String, String => String] = Map(
+    "unsupported" -> ((sparkVersion: String) =>
+      s"Cannot recommend RAPIDS Shuffle Manager for unsupported '$sparkVersion' version.\n" +
+        "  To enable RAPIDS Shuffle Manager, use a supported Spark version and set \n" +
+        "  'spark.shuffle.manager' to a valid RAPIDS Shuffle Manager version. \n" +
+        s"  See supported versions: $shuffleManagerDocUrl."),
+    "missing" -> (_ =>
+      "Could not recommend RapidsShuffleManager as Spark version cannot be determined.")
+  )
+
   /**
    * Abstract method to create an instance of the AutoTuner.
    */
@@ -1466,17 +1476,6 @@ trait AutoTunerConfigsProvider extends Logging {
   def buildShuffleManagerClassName(smVersion: String): String = {
     s"com.nvidia.spark.rapids.spark$smVersion.RapidsShuffleManager"
   }
-
-  def shuffleManagerCommentForUnsupportedVersion(sparkVersion: String): String = {
-    s"Cannot recommend RAPIDS Shuffle Manager for unsupported '$sparkVersion' version.\n" +
-     "  To enable RAPIDS Shuffle Manager, set 'spark.shuffle.manager' to a value\n" +
-     "  from the supported versions. \n" +
-    s"  See supported versions: $shuffleManagerDocUrl."
-  }
-
-  def shuffleManagerCommentForMissingVersion: String = {
-    "Could not recommend RapidsShuffleManager as Spark version cannot be determined."
-  }
 }
 
 /**

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
@@ -2254,8 +2254,8 @@ We recommend using nodes/workers with more memory. Need at least 7796MB memory."
       case Right(smClassName) =>
         fail(s"Expected error comment but got valid RapidsShuffleManager: $smClassName")
       case Left(comment) =>
-        assert(comment == ProfilingAutoTunerConfigsProvider
-          .shuffleManagerCommentForUnsupportedVersion(sparkVersion))
+        assert(comment ==
+          ProfilingAutoTunerConfigsProvider.shuffleManagerComments("unsupported")(sparkVersion))
     }
   }
 
@@ -2314,7 +2314,8 @@ We recommend using nodes/workers with more memory. Need at least 7796MB memory."
       case Right(smClassName) =>
         fail(s"Expected error comment but got valid RapidsShuffleManager: $smClassName")
       case Left(comment) =>
-        assert(comment == ProfilingAutoTunerConfigsProvider.shuffleManagerCommentForMissingVersion)
+        assert(comment ==
+          ProfilingAutoTunerConfigsProvider.shuffleManagerComments("missing")(""))
     }
   }