This is an automated email from the ASF dual-hosted git repository. sunchao pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push: new 0ea318f3e23f [SPARK-46947][CORE] Delay memory manager initialization until Driver plugin is loaded 0ea318f3e23f is described below commit 0ea318f3e23fb2daf953c9129047e8716e4ebf9b Author: Chao Sun <sunc...@apache.org> AuthorDate: Mon Feb 26 10:17:38 2024 -0800 [SPARK-46947][CORE] Delay memory manager initialization until Driver plugin is loaded ### What changes were proposed in this pull request? This changes the initialization of `SparkEnv.memoryManager` to after the `DriverPlugin` is loaded, to allow the plugin to customize memory related configurations. A minor fix has been made to `Task` to make sure that it uses the same `BlockManager` through out the task execution. Previous a different `BlockManager` could be used in some corner cases. Also added a test for the fix. ### Why are the changes needed? Today, there is no way for a custom `DriverPlugin` to override memory configurations such as `spark.executor.memory`, `spark.executor.memoryOverhead`, `spark.memory.offheap.size` etc This is because the memory manager is initialized before `DriverPlugin` is loaded. A similar change has been made to `shuffleManager` in #43627. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. Also added new tests. ### Was this patch authored or co-authored using generative AI tooling? No Closes #45052 from sunchao/SPARK-46947. Authored-by: Chao Sun <sunc...@apache.org> Signed-off-by: Chao Sun <sunc...@apache.org> --- .../main/scala/org/apache/spark/SparkContext.scala | 1 + .../src/main/scala/org/apache/spark/SparkEnv.scala | 20 ++++-- .../scala/org/apache/spark/scheduler/Task.scala | 15 +++-- .../org/apache/spark/storage/BlockManager.scala | 20 ++++-- .../internal/plugin/PluginContainerSuite.scala | 53 ++++++++++++++++ .../apache/spark/scheduler/TaskContextSuite.scala | 71 +++++++++++++++++++++- 6 files changed, 163 insertions(+), 17 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala index 801b6dd85a2b..d519617c4095 100644 --- a/core/src/main/scala/org/apache/spark/SparkContext.scala +++ b/core/src/main/scala/org/apache/spark/SparkContext.scala @@ -577,6 +577,7 @@ class SparkContext(config: SparkConf) extends Logging { // Initialize any plugins before the task scheduler is initialized. _plugins = PluginContainer(this, _resources.asJava) _env.initializeShuffleManager() + _env.initializeMemoryManager(SparkContext.numDriverCores(master, conf)) // 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 index ca07c276fbff..84c0fa5840b7 100644 --- a/core/src/main/scala/org/apache/spark/SparkEnv.scala +++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala @@ -67,7 +67,6 @@ class SparkEnv ( val blockManager: BlockManager, val securityManager: SecurityManager, val metricsSystem: MetricsSystem, - val memoryManager: MemoryManager, val outputCommitCoordinator: OutputCommitCoordinator, val conf: SparkConf) extends Logging { @@ -77,6 +76,12 @@ class SparkEnv ( def shuffleManager: ShuffleManager = _shuffleManager + // We initialize the MemoryManager later in SparkContext after DriverPlugin is loaded + // to allow the plugin to overwrite executor memory configurations + private var _memoryManager: MemoryManager = _ + + def memoryManager: MemoryManager = _memoryManager + @volatile private[spark] var isStopped = false /** @@ -199,6 +204,12 @@ class SparkEnv ( "Shuffle manager already initialized to %s", _shuffleManager) _shuffleManager = ShuffleManager.create(conf, executorId == SparkContext.DRIVER_IDENTIFIER) } + + private[spark] def initializeMemoryManager(numUsableCores: Int): Unit = { + Preconditions.checkState(null == memoryManager, + "Memory manager already initialized to %s", _memoryManager) + _memoryManager = UnifiedMemoryManager(conf, numUsableCores) + } } object SparkEnv extends Logging { @@ -276,6 +287,8 @@ object SparkEnv extends Logging { numCores, ioEncryptionKey ) + // Set the memory manager since it needs to be initialized explicitly + env.initializeMemoryManager(numCores) SparkEnv.set(env) env } @@ -358,8 +371,6 @@ object SparkEnv extends Logging { new MapOutputTrackerMasterEndpoint( rpcEnv, mapOutputTracker.asInstanceOf[MapOutputTrackerMaster], conf)) - val memoryManager: MemoryManager = UnifiedMemoryManager(conf, numUsableCores) - val blockManagerPort = if (isDriver) { conf.get(DRIVER_BLOCK_MANAGER_PORT) } else { @@ -418,7 +429,7 @@ object SparkEnv extends Logging { blockManagerMaster, serializerManager, conf, - memoryManager, + _memoryManager = null, mapOutputTracker, _shuffleManager = null, blockTransferService, @@ -463,7 +474,6 @@ object SparkEnv extends Logging { blockManager, securityManager, metricsSystem, - memoryManager, outputCommitCoordinator, conf) diff --git a/core/src/main/scala/org/apache/spark/scheduler/Task.scala b/core/src/main/scala/org/apache/spark/scheduler/Task.scala index 1ecd185de557..6e449e4dc111 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/Task.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/Task.scala @@ -93,7 +93,12 @@ private[spark] abstract class Task[T]( require(cpus > 0, "CPUs per task should be > 0") - SparkEnv.get.blockManager.registerTask(taskAttemptId) + // Use the blockManager at start of the task through out the task - particularly in + // case of local mode, a SparkEnv can be initialized when spark context is restarted + // and we want to ensure the right env and block manager is used (given lazy initialization of + // block manager) + val blockManager = SparkEnv.get.blockManager + blockManager.registerTask(taskAttemptId) // TODO SPARK-24874 Allow create BarrierTaskContext based on partitions, instead of whether // the stage is barrier. val taskContext = new TaskContextImpl( @@ -143,15 +148,15 @@ private[spark] abstract class Task[T]( try { Utils.tryLogNonFatalError { // Release memory used by this thread for unrolling blocks - SparkEnv.get.blockManager.memoryStore.releaseUnrollMemoryForThisTask(MemoryMode.ON_HEAP) - SparkEnv.get.blockManager.memoryStore.releaseUnrollMemoryForThisTask( - MemoryMode.OFF_HEAP) + blockManager.memoryStore.releaseUnrollMemoryForThisTask(MemoryMode.ON_HEAP) + blockManager.memoryStore.releaseUnrollMemoryForThisTask(MemoryMode.OFF_HEAP) // Notify any tasks waiting for execution memory to be freed to wake up and try to // acquire memory again. This makes impossible the scenario where a task sleeps forever // because there are no other tasks left to notify it. Since this is safe to do but may // not be strictly necessary, we should revisit whether we can remove this in the // future. - val memoryManager = SparkEnv.get.memoryManager + + val memoryManager = blockManager.memoryManager memoryManager.synchronized { memoryManager.notifyAll() } } } finally { diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala index 42bbd025177b..228ec5752e1b 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala @@ -185,7 +185,7 @@ private[spark] class BlockManager( val master: BlockManagerMaster, val serializerManager: SerializerManager, val conf: SparkConf, - memoryManager: MemoryManager, + private val _memoryManager: MemoryManager, mapOutputTracker: MapOutputTracker, private val _shuffleManager: ShuffleManager, val blockTransferService: BlockTransferService, @@ -198,6 +198,12 @@ private[spark] class BlockManager( // (except for tests) and we ask for the instance from the SparkEnv. private lazy val shuffleManager = Option(_shuffleManager).getOrElse(SparkEnv.get.shuffleManager) + // Similarly, we also initialize MemoryManager later after DriverPlugin is loaded, to + // allow the plugin to overwrite certain memory configurations. The `_memoryManager` will be + // null here and we ask for the instance from SparkEnv + private[spark] lazy val memoryManager = + Option(_memoryManager).getOrElse(SparkEnv.get.memoryManager) + // same as `conf.get(config.SHUFFLE_SERVICE_ENABLED)` private[spark] val externalShuffleServiceEnabled: Boolean = externalBlockStoreClient.isDefined private val isDriver = executorId == SparkContext.DRIVER_IDENTIFIER @@ -224,17 +230,19 @@ private[spark] class BlockManager( ThreadUtils.newDaemonCachedThreadPool("block-manager-future", 128)) // Actual storage of where blocks are kept - private[spark] val memoryStore = - new MemoryStore(conf, blockInfoManager, serializerManager, memoryManager, this) + private[spark] lazy val memoryStore = { + val store = new MemoryStore(conf, blockInfoManager, serializerManager, memoryManager, this) + memoryManager.setMemoryStore(store) + store + } private[spark] val diskStore = new DiskStore(conf, diskBlockManager, securityManager) - memoryManager.setMemoryStore(memoryStore) // Note: depending on the memory manager, `maxMemory` may actually vary over time. // However, since we use this only for reporting and logging, what we actually want here is // the absolute maximum value that `maxMemory` can ever possibly reach. We may need // to revisit whether reporting this value as the "max" is intuitive to the user. - private val maxOnHeapMemory = memoryManager.maxOnHeapStorageMemory - private val maxOffHeapMemory = memoryManager.maxOffHeapStorageMemory + private lazy val maxOnHeapMemory = memoryManager.maxOnHeapStorageMemory + private lazy val maxOffHeapMemory = memoryManager.maxOffHeapStorageMemory private[spark] val externalShuffleServicePort = StorageUtils.externalShuffleServicePort(conf) diff --git a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala index 197c2f13d807..cdbe5553bc95 100644 --- a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala +++ b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala @@ -36,6 +36,7 @@ import org.apache.spark.TestUtils._ import org.apache.spark.api.plugin._ import org.apache.spark.internal.config._ import org.apache.spark.launcher.SparkLauncher +import org.apache.spark.memory.MemoryMode import org.apache.spark.resource.ResourceInformation import org.apache.spark.resource.ResourceUtils.GPU import org.apache.spark.resource.TestResourceIDs.{DRIVER_GPU_ID, EXECUTOR_GPU_ID, WORKER_GPU_ID} @@ -228,6 +229,58 @@ class PluginContainerSuite extends SparkFunSuite with LocalSparkContext { assert(driverResources.get(GPU).name === GPU) } } + + test("memory override in plugin") { + val conf = new SparkConf() + .setAppName(getClass().getName()) + .set(SparkLauncher.SPARK_MASTER, "local-cluster[2,1,1024]") + .set(PLUGINS, Seq(classOf[MemoryOverridePlugin].getName())) + + var sc: SparkContext = null + try { + sc = new SparkContext(conf) + val memoryManager = sc.env.memoryManager + + assert(memoryManager.tungstenMemoryMode == MemoryMode.OFF_HEAP) + assert(memoryManager.maxOffHeapStorageMemory == MemoryOverridePlugin.offHeapMemory) + + // Ensure all executors has started + TestUtils.waitUntilExecutorsUp(sc, 1, 60000) + + // Check executor memory is also updated + val execInfo = sc.statusTracker.getExecutorInfos.head + assert(execInfo.totalOffHeapStorageMemory() == MemoryOverridePlugin.offHeapMemory) + } finally { + if (sc != null) { + sc.stop() + } + } + } +} + +class MemoryOverridePlugin extends SparkPlugin { + override def driverPlugin(): DriverPlugin = { + new DriverPlugin { + override def init(sc: SparkContext, pluginContext: PluginContext): JMap[String, String] = { + // Take the original executor memory, and set `spark.memory.offHeap.size` to be the + // same value. Also set `spark.memory.offHeap.enabled` to true. + val originalExecutorMemBytes = + sc.conf.getSizeAsMb(EXECUTOR_MEMORY.key, EXECUTOR_MEMORY.defaultValueString) + sc.conf.set(MEMORY_OFFHEAP_ENABLED.key, "true") + sc.conf.set(MEMORY_OFFHEAP_SIZE.key, s"${originalExecutorMemBytes}M") + MemoryOverridePlugin.offHeapMemory = sc.conf.getSizeAsBytes(MEMORY_OFFHEAP_SIZE.key) + Map.empty[String, String].asJava + } + } + } + + override def executorPlugin(): ExecutorPlugin = { + new ExecutorPlugin {} + } +} + +object MemoryOverridePlugin { + var offHeapMemory: Long = _ } class NonLocalModeSparkPlugin extends SparkPlugin { diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala index 9aba41cea215..d08e75733abf 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.scheduler import java.util.Properties +import java.util.concurrent.Semaphore import java.util.concurrent.atomic.AtomicInteger import scala.collection.mutable.ArrayBuffer @@ -27,7 +28,9 @@ import org.mockito.Mockito._ import org.scalatest.BeforeAndAfter import org.apache.spark._ +import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, SparkPlugin} import org.apache.spark.executor.{Executor, TaskMetrics, TaskMetricsSuite} +import org.apache.spark.internal.Logging import org.apache.spark.internal.config.METRICS_CONF import org.apache.spark.memory.TaskMemoryManager import org.apache.spark.metrics.source.JvmSource @@ -680,9 +683,49 @@ class TaskContextSuite extends SparkFunSuite with BeforeAndAfter with LocalSpark context.markTaskCompleted(None) assert(isFailed) } + + test("SPARK-46947: ensure the correct block manager is used to unroll memory for task") { + import BlockManagerValidationPlugin._ + BlockManagerValidationPlugin.resetState() + + // run a task which ignores thread interruption when spark context is shutdown + sc = new SparkContext("local", "test") + + val rdd = new RDD[String](sc, List()) { + override def getPartitions = Array[Partition](StubPartition(0)) + + override def compute(split: Partition, context: TaskContext): Iterator[String] = { + context.addTaskCompletionListener(new TaskCompletionListener { + override def onTaskCompletion(context: TaskContext): Unit = { + try { + releaseTaskSem.acquire() + } catch { + case _: InterruptedException => + // ignore thread interruption + } + } + }) + taskStartedSem.release() + Iterator.empty + } + } + // submit the job, but don't block this thread + rdd.collectAsync() + // wait for task to start + taskStartedSem.acquire() + + sc.stop() + assert(sc.isStopped) + + // create a new SparkContext which will be blocked for certain amount of time + // during initializing the driver plugin below + val conf = new SparkConf() + conf.set("spark.plugins", classOf[BlockManagerValidationPlugin].getName) + sc = new SparkContext("local", "test", conf) + } } -private object TaskContextSuite { +private object TaskContextSuite extends Logging { @volatile var completed = false @volatile var lastError: Throwable = _ @@ -690,4 +733,30 @@ private object TaskContextSuite { class FakeTaskFailureException extends Exception("Fake task failure") } +class BlockManagerValidationPlugin extends SparkPlugin { + override def driverPlugin(): DriverPlugin = { + new DriverPlugin() { + // does nothing but block the current thread for certain time for the task thread + // to progress and reproduce the issue. + BlockManagerValidationPlugin.releaseTaskSem.release() + Thread.sleep(2500) + } + } + override def executorPlugin(): ExecutorPlugin = { + new ExecutorPlugin() { + // do nothing + } + } +} + +object BlockManagerValidationPlugin { + val releaseTaskSem = new Semaphore(0) + val taskStartedSem = new Semaphore(0) + + def resetState(): Unit = { + releaseTaskSem.drainPermits() + taskStartedSem.drainPermits() + } +} + private case class StubPartition(index: Int) extends Partition --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org