Eagerly initialize RapidsShuffleManager for SPARK-45762 [databricks]

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuShuffleEnv.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.rapids
 import java.util.Locale
 
 import com.nvidia.spark.rapids._
+import com.nvidia.spark.rapids.shims.ShuffleManagerShimUtils
 
 import org.apache.spark.SparkEnv
 import org.apache.spark.internal.Logging
@@ -105,21 +106,25 @@ object GpuShuffleEnv extends Logging {
   //
   def initShuffleManager(): Unit = {
     val shuffleManager = SparkEnv.get.shuffleManager
-    shuffleManager match {
-      case rapidsShuffleManager: RapidsShuffleManagerLike =>
-        rapidsShuffleManager.initialize
-      case _ =>
-        val rsmLoaderViaShuffleManager = shuffleManager.getClass.getSuperclass.getInterfaces
-          .collectFirst {
-            case c if c.getName == classOf[RapidsShuffleManagerLike].getName => c.getClassLoader
-          }
-        val rsmLoaderDirect = classOf[RapidsShuffleManagerLike].getClassLoader
-
-        throw new IllegalStateException(s"Cannot initialize the RAPIDS Shuffle Manager " +
-          s"${shuffleManager}! Expected: an instance of RapidsShuffleManagerLike loaded by " +
-          s"${rsmLoaderDirect}. Actual: ${shuffleManager} tagged with RapidsShuffleManagerLike " +
-          s"loaded by: ${rsmLoaderViaShuffleManager}"
-        )
+    if (ShuffleManagerShimUtils.eagerlyInitialized) {
+      // skip deferred init
+    } else {
+      shuffleManager match {
+        case rapidsShuffleManager: RapidsShuffleManagerLike =>
+          rapidsShuffleManager.initialize
+        case _ =>
+          val rsmLoaderViaShuffleManager = shuffleManager.getClass.getSuperclass.getInterfaces
+            .collectFirst {
+              case c if c.getName == classOf[RapidsShuffleManagerLike].getName => c.getClassLoader
+            }
+          val rsmLoaderDirect = classOf[RapidsShuffleManagerLike].getClassLoader
+
+          throw new IllegalStateException(s"Cannot initialize the RAPIDS Shuffle Manager " +
+            s"${shuffleManager}! Expected: an instance of RapidsShuffleManagerLike loaded by " +
+            s"${rsmLoaderDirect}. Actual: ${shuffleManager} tagged with RapidsShuffleManagerLike " +
+            s"loaded by: ${rsmLoaderViaShuffleManager}"
+          )
+      }
     }
   }
 

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/RapidsShuffleManager.scala

@@ -37,11 +37,9 @@
 {"spark": "343"}
 {"spark": "344"}
 {"spark": "350"}
-{"spark": "350db143"}
 {"spark": "351"}
 {"spark": "352"}
 {"spark": "353"}
-{"spark": "400"}
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.$_spark.version.classifier_
 

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/ShuffleManagerShimUtils.scala

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+/*** spark-rapids-shim-json-lines
+{"spark": "320"}
+{"spark": "321"}
+{"spark": "321cdh"}
+{"spark": "322"}
+{"spark": "323"}
+{"spark": "324"}
+{"spark": "330"}
+{"spark": "330cdh"}
+{"spark": "330db"}
+{"spark": "331"}
+{"spark": "332"}
+{"spark": "332cdh"}
+{"spark": "332db"}
+{"spark": "333"}
+{"spark": "334"}
+{"spark": "340"}
+{"spark": "341"}
+{"spark": "341db"}
+{"spark": "342"}
+{"spark": "343"}
+{"spark": "344"}
+{"spark": "350"}
+{"spark": "351"}
+{"spark": "352"}
+{"spark": "353"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+object ShuffleManagerShimUtils {
+  def eagerlyInitialized = false  
+}

sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/RapidsShuffleManager.scala

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*** spark-rapids-shim-json-lines
+{"spark": "350db143"}
+{"spark": "400"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.$_spark.version.classifier_
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.rapids.ProxyRapidsShuffleInternalManagerBase
+
+/** A shuffle manager optimized for the RAPIDS Plugin for Apache Spark. */
+sealed class RapidsShuffleManager(
+    conf: SparkConf,
+    isDriver: Boolean
+) extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) {
+    initialize
+}

sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/ShuffleManagerShimUtils.scala

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*** spark-rapids-shim-json-lines
+{"spark": "350db143"}
+{"spark": "400"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+object ShuffleManagerShimUtils {
+  def eagerlyInitialized = true  
+}