Fix test failures

mridulm · Nov 8, 2023 · 585b493 · 585b493
1 parent d5ba0cc
commit 585b493
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -576,6 +576,7 @@ class SparkContext(config: SparkConf) extends Logging {
 
     // Initialize any plugins before the task scheduler is initialized.
     _plugins = PluginContainer(this, _resources.asJava)
+    _env.initiailzeShuffleManager()
 
     // Create and start the scheduler
     val (sched, ts) = SparkContext.createTaskScheduler(this, master)

diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -18,7 +18,6 @@
 package org.apache.spark
 
 import java.io.File
-import java.util.Locale
 
 import scala.collection.concurrent
 import scala.collection.mutable
@@ -195,11 +194,9 @@ class SparkEnv (
   }
 
   private[spark] def initiailzeShuffleManager(): Unit = {
-    Preconditions.checkState(null == _shuffleManager, "Shuffle manager already initialized")
-    // Must not be driver
-    Preconditions.checkState(executorId != SparkContext.DRIVER_IDENTIFIER,
-      "Should not be called in driver")
-    _shuffleManager = ShuffleManager.create(conf, isDriver = false)
+    Preconditions.checkState(null == _shuffleManager,
+      "Shuffle manager already initialized to %s", _shuffleManager)
+    _shuffleManager = ShuffleManager.create(conf, executorId == SparkContext.DRIVER_IDENTIFIER)
   }
 }
 
@@ -371,7 +368,8 @@ object SparkEnv extends Logging {
       new MapOutputTrackerMasterEndpoint(
         rpcEnv, mapOutputTracker.asInstanceOf[MapOutputTrackerMaster], conf))
 
-    val shuffleManager: ShuffleManager = if (isDriver) ShuffleManager.create(conf, true) else null
+    // val shuffleManager: ShuffleManager = if (isDriver) ShuffleManager.create(conf, true) else null
+    val shuffleManager: ShuffleManager = null
     val memoryManager: MemoryManager = UnifiedMemoryManager(conf, numUsableCores)
 
     val blockManagerPort = if (isDriver) {

diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -50,7 +50,7 @@ import org.apache.spark.resource.ResourceInformation
 import org.apache.spark.rpc.RpcTimeout
 import org.apache.spark.scheduler._
 import org.apache.spark.serializer.SerializerHelper
-import org.apache.spark.shuffle.{FetchFailedException, ShuffleBlockPusher, ShuffleManager}
+import org.apache.spark.shuffle.{FetchFailedException, ShuffleBlockPusher}
 import org.apache.spark.status.api.v1.ThreadStackTrace
 import org.apache.spark.storage.{StorageLevel, TaskResultBlockId}
 import org.apache.spark.util._
@@ -337,8 +337,11 @@ private[spark] class Executor(
       PluginContainer(env, resources.asJava)
     }
 
-  Utils.withContextClassLoader(defaultSessionState.replClassLoader) {
-    env.initiailzeShuffleManager()
+  // Already initialized in local mode
+  if (!isLocal) {
+    Utils.withContextClassLoader(defaultSessionState.replClassLoader) {
+      env.initiailzeShuffleManager()
+    }
   }
 
   metricsPoller.start()