From f092605aeda8c1f81f4a047fb08e2adf614713ca Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 11 Sep 2026 10:06:18 -0700 Subject: [PATCH] [SPARK-59444][CORE] Retain execution-memory task registration while allocations wait --- .../spark/memory/ExecutionMemoryPool.scala | 111 +++++--- .../apache/spark/memory/MemoryManager.scala | 3 +- .../memory/ExecutionMemoryPoolSuite.scala | 242 ++++++++++++++++++ 3 files changed, 319 insertions(+), 37 deletions(-) create mode 100644 core/src/test/scala/org/apache/spark/memory/ExecutionMemoryPoolSuite.scala diff --git a/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala b/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala index 8b31d02cf97e3..c305aedd43e65 100644 --- a/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala +++ b/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala @@ -56,6 +56,15 @@ private[memory] class ExecutionMemoryPool( @GuardedBy("lock") private val memoryForTask = new mutable.HashMap[Long, Long]() + // Only acquisitions that actually wait need a retained zero-byte task entry. + // Lazily allocated under lock; the non-waiting admission path does no map work. + @GuardedBy("lock") + private var waitingAcquisitions: mutable.LongMap[Int] = null + + private def hasWaitingAcquisition(taskAttemptId: Long): Boolean = { + waitingAcquisitions != null && waitingAcquisitions.contains(taskAttemptId) + } + override def memoryUsed: Long = lock.synchronized { memoryForTask.values.sum } @@ -110,42 +119,71 @@ private[memory] class ExecutionMemoryPool( // task would have more than 1 / numActiveTasks of the memory) or we have enough free // memory to give it (we always let each task get at least 1 / (2 * numActiveTasks)). // TODO: simplify this to limit each task to its own slot - while (true) { - val numActiveTasks = memoryForTask.keys.size - val curMem = memoryForTask(taskAttemptId) - - // In every iteration of this loop, we should first try to reclaim any borrowed execution - // space from storage. This is necessary because of the potential race condition where new - // storage blocks may steal the free execution memory that this task was waiting for. - maybeGrowPool(numBytes - memoryFree) - - // Maximum size the pool would have after potentially growing the pool. - // This is used to compute the upper bound of how much memory each task can occupy. This - // must take into account potential free memory as well as the amount this pool currently - // occupies. Otherwise, we may run into SPARK-12155 where, in unified memory management, - // we did not take into account space that could have been freed by evicting cached blocks. - val maxPoolSize = computeMaxPoolSize() - val maxMemoryPerTask = maxPoolSize / numActiveTasks - val minMemoryPerTask = poolSize / (2 * numActiveTasks) - - // How much we can grant this task; keep its share within 0 <= X <= 1 / numActiveTasks - val maxToGrant = math.min(numBytes, math.max(0, maxMemoryPerTask - curMem)) - // Only give it as much memory as is free, which might be none if it reached 1 / numTasks - val toGrant = math.min(maxToGrant, memoryFree) - - // We want to let each task get at least 1 / (2 * numActiveTasks) before blocking; - // if we can't give it this much now, wait for other tasks to free up memory - // (this happens if older tasks allocated lots of memory before N grew) - if (toGrant < numBytes && curMem + toGrant < minMemoryPerTask) { - logInfo(log"TID ${MDC(TASK_ATTEMPT_ID, taskAttemptId)} waiting for at least 1/2N of" + - log" ${MDC(POOL_NAME, poolName)} pool to be free") - lock.wait() - } else { - memoryForTask(taskAttemptId) += toGrant - return toGrant + var registeredWaiter = false + try { + while (true) { + val numActiveTasks = memoryForTask.keys.size + val curMem = memoryForTask(taskAttemptId) + + // In every iteration of this loop, we should first try to reclaim any borrowed execution + // space from storage. This is necessary because of the potential race condition where new + // storage blocks may steal the free execution memory that this task was waiting for. + maybeGrowPool(numBytes - memoryFree) + + // Maximum size the pool would have after potentially growing the pool. + // This is used to compute the upper bound of how much memory each task can occupy. This + // must take into account potential free memory as well as the amount this pool currently + // occupies. Otherwise, we may run into SPARK-12155 where, in unified memory management, + // we did not take into account space that could have been freed by evicting cached blocks. + val maxPoolSize = computeMaxPoolSize() + val maxMemoryPerTask = maxPoolSize / numActiveTasks + val minMemoryPerTask = poolSize / (2 * numActiveTasks) + + // How much we can grant this task; keep its share within 0 <= X <= 1 / numActiveTasks + val maxToGrant = math.min(numBytes, math.max(0, maxMemoryPerTask - curMem)) + // Only give it as much memory as is free, which might be none if it reached 1 / numTasks + val toGrant = math.min(maxToGrant, memoryFree) + + // We want to let each task get at least 1 / (2 * numActiveTasks) before blocking; + // if we can't give it this much now, wait for other tasks to free up memory + // (this happens if older tasks allocated lots of memory before N grew) + if (toGrant < numBytes && curMem + toGrant < minMemoryPerTask) { + logInfo(log"TID ${MDC(TASK_ATTEMPT_ID, taskAttemptId)} waiting for at least 1/2N of" + + log" ${MDC(POOL_NAME, poolName)} pool to be free") + if (!registeredWaiter) { + if (waitingAcquisitions == null) { + waitingAcquisitions = mutable.LongMap.empty[Int] + } + waitingAcquisitions(taskAttemptId) = + waitingAcquisitions.getOrElse(taskAttemptId, 0) + 1 + registeredWaiter = true + } + lock.wait() + } else { + memoryForTask(taskAttemptId) += toGrant + return toGrant + } + } + 0L // Never reached + } finally { + if (registeredWaiter) { + val remaining = waitingAcquisitions(taskAttemptId) - 1 + if (remaining == 0) { + waitingAcquisitions.remove(taskAttemptId) + if (waitingAcquisitions.isEmpty) { + waitingAcquisitions = null + } + } else { + waitingAcquisitions(taskAttemptId) = remaining + } + // The last exiting waiter must not leave a zero-byte fairness participant. + if (!hasWaitingAcquisition(taskAttemptId) && + memoryForTask.get(taskAttemptId).contains(0L)) { + memoryForTask.remove(taskAttemptId) + lock.notifyAll() + } } } - 0L // Never reached } /** @@ -164,7 +202,7 @@ private[memory] class ExecutionMemoryPool( } if (memoryForTask.contains(taskAttemptId)) { memoryForTask(taskAttemptId) -= memoryToFree - if (memoryForTask(taskAttemptId) <= 0) { + if (memoryForTask(taskAttemptId) <= 0 && !hasWaitingAcquisition(taskAttemptId)) { memoryForTask.remove(taskAttemptId) } } @@ -172,7 +210,8 @@ private[memory] class ExecutionMemoryPool( } /** - * Release all memory for the given task and mark it as inactive (e.g. when a task ends). + * Release all memory for the given task. A task with a waiting acquisition remains active + * until that acquisition completes or is interrupted. * @return the number of bytes freed. */ def releaseAllMemoryForTask(taskAttemptId: Long): Long = lock.synchronized { diff --git a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala index dac82cc763577..1218b0754a147 100644 --- a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala +++ b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala @@ -143,7 +143,8 @@ private[spark] abstract class MemoryManager( } /** - * Release all memory for the given task and mark it as inactive (e.g. when a task ends). + * Release all memory for the given task. A task with a waiting acquisition remains active + * until that acquisition completes or is interrupted. * * @return the number of bytes freed. */ diff --git a/core/src/test/scala/org/apache/spark/memory/ExecutionMemoryPoolSuite.scala b/core/src/test/scala/org/apache/spark/memory/ExecutionMemoryPoolSuite.scala new file mode 100644 index 0000000000000..a4c1825dee001 --- /dev/null +++ b/core/src/test/scala/org/apache/spark/memory/ExecutionMemoryPoolSuite.scala @@ -0,0 +1,242 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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. + */ + +package org.apache.spark.memory + +import java.util.concurrent.{CompletableFuture, ExecutionException, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean + +import scala.collection.mutable.ArrayBuffer +import scala.util.control.NonFatal + +import org.scalatest.concurrent.Eventually +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.{SparkConf, SparkFunSuite} + +class ExecutionMemoryPoolSuite extends SparkFunSuite with Eventually { + private class WaitingAcquire(acquire: () => Long) { + val result = new CompletableFuture[Long]() + val thread = new Thread("execution-memory-waiter") { + override def run(): Unit = { + try { + result.complete(acquire()) + } catch { + case e: InterruptedException => result.completeExceptionally(e) + case NonFatal(e) => result.completeExceptionally(e) + } + } + } + thread.setDaemon(true) + + def awaitWaiting(): Unit = eventually(timeout(10.seconds)) { + assert(!result.isDone, "acquisition completed instead of waiting") + assert(thread.getState == Thread.State.WAITING) + assert(thread.getStackTrace.exists { frame => + frame.getClassName == classOf[ExecutionMemoryPool].getName && + frame.getMethodName == "acquireMemory" + }) + } + + def acquired(): Long = result.get(10, TimeUnit.SECONDS) + + def failure(): Throwable = { + val error = intercept[ExecutionException] { + result.get(10, TimeUnit.SECONDS) + } + error.getCause + } + + def interrupt(): Unit = { + thread.interrupt() + assert(failure().isInstanceOf[InterruptedException]) + } + } + + private val waiters = new ArrayBuffer[WaitingAcquire]() + + override protected def afterEach(): Unit = { + try { + waiters.foreach(_.thread.interrupt()) + waiters.foreach(_.thread.join(10000)) + assert(waiters.forall(!_.thread.isAlive), "memory acquisition thread did not terminate") + } finally { + waiters.clear() + super.afterEach() + } + } + + private def acquireAsync( + pool: ExecutionMemoryPool, + bytes: Long, + maybeGrowPool: Long => Unit = _ => ()): WaitingAcquire = { + acquireAsync(pool.acquireMemory(bytes, 1L, maybeGrowPool)) + } + + private def acquireAsync(acquire: => Long): WaitingAcquire = { + val waiter = new WaitingAcquire(() => acquire) + waiters += waiter + waiter.thread.start() + waiter.awaitWaiting() + waiter + } + + private def newPool(mode: MemoryMode): ExecutionMemoryPool = { + val pool = new ExecutionMemoryPool(new Object, mode) + pool.incrementPoolSize(1000L) + assert(pool.acquireMemory(900L, 2L) == 900L) + pool + } + + for (mode <- Seq(MemoryMode.ON_HEAP, MemoryMode.OFF_HEAP)) { + test(s"retain task registration across two consumers of one TaskMemoryManager ($mode)") { + val conf = new SparkConf(false) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "1000") + val memory = new UnifiedMemoryManager(conf, 1000L, 500L, 1) + val task = new TaskMemoryManager(memory, 1L) + val peerTask = new TaskMemoryManager(memory, 2L) + val owner = new TestMemoryConsumer(task, mode) + val requester = new TestMemoryConsumer(task, mode) + val peer = new TestMemoryConsumer(peerTask, mode) + assert(peer.acquireMemory(900L) == 900L) + assert(owner.acquireMemory(100L) == 100L) + val waiter = acquireAsync(requester.acquireMemory(300L)) + + // Release through another consumer of the waiting task, not through the pool directly. + owner.freeMemory(100L) + peer.freeMemory(300L) + assert(waiter.acquired() == 300L) + assert(owner.getUsed() == 0L) + assert(requester.getUsed() == 300L) + assert(task.getMemoryConsumptionForThisTask() == 300L) + assert(peerTask.getMemoryConsumptionForThisTask() == 600L) + assert(memory.executionMemoryUsed == 900L) + + requester.freeMemory(300L) + peer.freeMemory(600L) + assert(task.cleanUpAllAllocatedMemory() == 0L) + assert(peerTask.cleanUpAllAllocatedMemory() == 0L) + assert(memory.executionMemoryUsed == 0L) + } + + for (releaseAll <- Seq(false, true)) { + test(s"retain a waiting task after its last release ($mode, releaseAll=$releaseAll)") { + val pool = newPool(mode) + assert(pool.acquireMemory(100L, 1L) == 100L) + val waiter = acquireAsync(pool, 300L) + + if (releaseAll) { + assert(pool.releaseAllMemoryForTask(1L) == 100L) + } else { + pool.releaseMemory(100L, 1L) + } + pool.releaseMemory(300L, 2L) + + assert(waiter.acquired() == 300L) + assert(pool.getMemoryUsageForTask(1L) == 300L) + assert(pool.memoryUsed == 900L) + } + } + + test(s"retain a task until all its waiting acquisitions complete ($mode)") { + val pool = newPool(mode) + assert(pool.acquireMemory(100L, 1L) == 100L) + val first = acquireAsync(pool, 200L) + val second = acquireAsync(pool, 200L) + pool.releaseMemory(100L, 1L) + pool.releaseMemory(100L, 2L) + + eventually(timeout(10.seconds)) { + assert(first.result.isDone || second.result.isDone) + } + val (completed, remaining) = if (first.result.isDone) (first, second) else (second, first) + assert(completed.acquired() == 200L) + remaining.awaitWaiting() + pool.releaseMemory(200L, 1L) + + assert(remaining.acquired() == 200L) + assert(pool.getMemoryUsageForTask(1L) == 200L) + assert(pool.memoryUsed == 1000L) + } + + test(s"preserve a waiting task's remaining allocation after a partial release ($mode)") { + val pool = newPool(mode) + assert(pool.acquireMemory(100L, 1L) == 100L) + val waiter = acquireAsync(pool, 300L) + pool.releaseMemory(40L, 1L) + pool.releaseMemory(300L, 2L) + + assert(waiter.acquired() == 300L) + assert(pool.getMemoryUsageForTask(1L) == 360L) + assert(pool.memoryUsed == 960L) + } + + for (previousAllocation <- Seq(false, true)) { + test(s"remove an interrupted zero-byte task ($mode, previous=$previousAllocation)") { + val pool = newPool(mode) + if (previousAllocation) { + assert(pool.acquireMemory(100L, 1L) == 100L) + } + val waiter = acquireAsync(pool, 300L) + if (previousAllocation) { + pool.releaseMemory(100L, 1L) + } + waiter.interrupt() + + // The interrupted task must no longer reduce the remaining task's fair share. + assert(pool.acquireMemory(100L, 2L) == 100L) + assert(pool.getMemoryUsageForTask(1L) == 0L) + assert(pool.releaseAllMemoryForTask(2L) == 1000L) + assert(pool.memoryUsed == 0L) + } + } + + test(s"remove a zero-byte task if the pool-growth callback fails after waiting ($mode)") { + val pool = newPool(mode) + val failAfterWaiting = new AtomicBoolean(false) + val error = new IllegalStateException("pool-growth failure") + val waiter = acquireAsync(pool, 300L, _ => { + if (failAfterWaiting.get()) { + throw error + } + }) + failAfterWaiting.set(true) + pool.releaseMemory(0L, 2L) // Wake the waiter to run the callback again. + + assert(waiter.failure() eq error) + assert(pool.acquireMemory(100L, 2L) == 100L) + assert(pool.releaseAllMemoryForTask(2L) == 1000L) + assert(pool.memoryUsed == 0L) + } + + test(s"interrupting one acquisition must retain the same task's other waiter ($mode)") { + val pool = newPool(mode) + assert(pool.acquireMemory(100L, 1L) == 100L) + val interrupted = acquireAsync(pool, 300L) + val remaining = acquireAsync(pool, 300L) + pool.releaseMemory(100L, 1L) + interrupted.interrupt() + remaining.awaitWaiting() + pool.releaseMemory(300L, 2L) + + assert(remaining.acquired() == 300L) + assert(pool.getMemoryUsageForTask(1L) == 300L) + assert(pool.memoryUsed == 900L) + } + } +}