From 4fc31f65690d440c391f1c61b0ac71c4aadfc6bc Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 9 Oct 2026 00:15:17 -0700 Subject: [PATCH 1/2] refactor: share the missing task entry retry between the shuffle allocator and the native pool #6403 added a retry for Spark's missing task entry (SPARK-59444) to CometUnifiedShuffleMemoryAllocator, and #6310 added its own copy to CometTaskMemoryManager.acquireMemory. Move both into MissingTaskEntryRetry under org.apache.comet. CometTaskMemoryManager passes its task id, which keeps #6310's exact "key not found: " match; the shuffle allocator keeps #6403's prefix match. After the retries run out the native pool still gets a zero grant and the shuffle allocator still throws SparkOutOfMemoryError, and the warning now carries the last exception for both. The fixtures and the waiting-request harness of both suites move to TaskMemoryTestUtils, and CometTaskMemoryManagerSuite keeps one race test next to the refusal and rethrow tests. memory_management.md describes the retry once for both call sites, says what a zero grant does to try_grow and grow, and drops the stale #6304 reference. Co-authored-by: Isaac --- .../contributor-guide/memory_management.md | 31 +-- .../apache/comet/MissingTaskEntryRetry.java | 132 +++++++++++ .../apache/spark/CometTaskMemoryManager.java | 61 +---- .../CometUnifiedShuffleMemoryAllocator.java | 62 +---- .../spark/CometTaskMemoryManagerSuite.scala | 218 +++-------------- .../apache/spark/TaskMemoryTestUtils.scala | 219 ++++++++++++++++++ ...etUnifiedShuffleMemoryAllocatorSuite.scala | 198 +++------------- 7 files changed, 450 insertions(+), 471 deletions(-) create mode 100644 spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java create mode 100644 spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index af84e4a2786..c23e1eb465e 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -137,15 +137,10 @@ allocator before reading anything into the split. returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the same task like any other allocation. A page or pointer array allocation that waits in Spark for -memory can wake up to find the task's entry gone from Spark's execution pool, because another -consumer of the task released its last bytes in the meantime, and Spark then throws a -`NoSuchElementException` ("key not found") instead of a grant -([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). The allocator retries that -case, since the failed call was granted nothing, and after three attempts throws -`SparkOutOfMemoryError`, which the shuffle writers handle like any other refused page. Only the -shuffle allocator's callers are guarded this way. Spark's operators in the same task, such as its -sorters and aggregates, can still hit the exception until Spark re-registers a waiting task in -`ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). +memory can wake up to find the task's entry gone from Spark's execution pool; see +[Constraints on a Comet memory consumer](#constraints-on-a-comet-memory-consumer). The allocator +retries it, and after three attempts throws `SparkOutOfMemoryError`, which the shuffle writers +handle like any other refused page. Which allocator each call site uses, and who ends up charged for the bytes: @@ -216,12 +211,18 @@ that grant, and the only figure available for doing so is the task-wide one abov **A parked acquire can wake up to a missing task entry.** Spark removes the task's `memoryForTask` entry when its balance reaches zero, and an acquire that was waiting in `lock.wait()` and wakes -afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant. -`CometTaskMemoryManager` retries that one case. A missing entry means the task held nothing from -Spark at that moment, so the failed call has no partial grant to reconcile, and the retry registers -the task again and waits for its share as the first call would have. After a few attempts it returns -a zero grant, which the native side treats as a refusal and spills. JVM consumers that call -`allocatePage` directly are not covered; that gap is tracked in #6304. +afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant +([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). Releases take no task monitor, +so another native thread or plan, or a JVM consumer of the same task, can empty the balance while an +acquire waits. Both of Comet's consumers that ask Spark directly, `CometTaskMemoryManager` and the +JVM shuffle allocator, retry that one case through `MissingTaskEntryRetry`. A missing entry means +the task held nothing from Spark at that moment, so the failed call has no partial grant to +reconcile, and the retry registers the task again and waits for its share as the first call would +have. After three attempts `CometTaskMemoryManager` returns a zero grant. `try_grow` treats that as +a refusal and the operator spills, while `grow` records the request as overcommit. The shuffle +allocator throws `SparkOutOfMemoryError` instead. Spark's own operators in the same task, such as +its sorters and aggregates, are not guarded and can still hit the exception until Spark re-registers +a waiting task in `ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). **A consumer whose `spill` returns zero takes budget it can never give back.** `NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim diff --git a/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java new file mode 100644 index 00000000000..b0faa2ed8ba --- /dev/null +++ b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java @@ -0,0 +1,132 @@ +/* + * 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.comet; + +import java.util.NoSuchElementException; +import java.util.OptionalLong; +import java.util.function.Function; +import java.util.function.Supplier; + +import org.slf4j.Logger; + +/** + * Retries a request for execution memory that Spark failed because the task's entry was gone from + * its execution pool. + * + *

Spark registers a task's entry in {@code ExecutionMemoryPool} once, before its wait loop, and + * removes it when the task's balance reaches zero. A request that waits for memory below the task's + * minimum share then fails when it wakes up with a NoSuchElementException ("key not found: " and + * the task id). Releases take no task monitor, so any other consumer of the same task, JVM or + * native, can empty the balance while a request waits. The failed call was granted nothing, because + * Spark only drops the entry at a zero balance, which cannot hold while the call has a partial + * grant. So trying again is safe: it registers the task again and waits for its share as Spark + * would have. Nothing here holds a lock across attempts, so a retry parks in Spark exactly as the + * first attempt did. The Spark issue is SPARK-59444. + */ +public final class MissingTaskEntryRetry { + + /** Attempts at a request whose task entry Spark keeps losing. */ + public static final int MAX_ATTEMPTS = 3; + + private MissingTaskEntryRetry() {} + + /** + * Runs {@code request}, a call that asks Spark for {@code required} bytes for task {@code + * taskAttemptId}, and retries it when it lost that task's entry. See {@link #retry(Logger, long, + * Supplier, Function)}. + */ + public static T retry( + Logger logger, + long taskAttemptId, + long required, + Supplier request, + Function refuse) { + return retry(logger, OptionalLong.of(taskAttemptId), required, request, refuse); + } + + /** + * Runs {@code request}, a call that asks Spark for {@code required} bytes, and retries it when it + * lost the task's entry. Without the task id, any task's missing entry counts. Other exceptions + * are rethrown. After {@link #MAX_ATTEMPTS} attempts that all lost the entry, returns what {@code + * refuse} makes of the last attempt's exception, or throws what it throws. Both the retries and + * the refusal are logged to {@code logger}. + */ + public static T retry( + Logger logger, + long required, + Supplier request, + Function refuse) { + return retry(logger, OptionalLong.empty(), required, request, refuse); + } + + private static T retry( + Logger logger, + OptionalLong taskAttemptId, + long required, + Supplier request, + Function refuse) { + for (int attempt = 1; ; attempt++) { + try { + return request.get(); + } catch (NoSuchElementException e) { + if (!isMissingTaskEntry(e, taskAttemptId)) { + throw e; + } + String task = taskAttemptId.isPresent() ? "Task " + taskAttemptId.getAsLong() : "The task"; + if (attempt >= MAX_ATTEMPTS) { + logger.warn( + "{} lost its execution memory entry in Spark on {} attempts to acquire {} bytes, " + + "refusing the request", + task, + attempt, + required, + e); + return refuse.apply(e); + } + logger.info( + "{} lost its execution memory entry in Spark while waiting to acquire {} bytes, " + + "trying again", + task, + required); + } + } + } + + /** + * Whether {@code e} is Spark's execution pool failing to find the task's entry: the message names + * the task, or any task when its id is unknown, and the exception comes from the pool. + */ + static boolean isMissingTaskEntry(NoSuchElementException e, OptionalLong taskAttemptId) { + String message = e.getMessage(); + boolean matches = + taskAttemptId.isPresent() + ? ("key not found: " + taskAttemptId.getAsLong()).equals(message) + : message != null && message.startsWith("key not found: "); + if (!matches) { + return false; + } + for (StackTraceElement frame : e.getStackTrace()) { + if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { + return true; + } + } + return false; + } +} diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index a468446bae7..a0eccb85ffb 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -20,7 +20,6 @@ package org.apache.spark; import java.io.IOException; -import java.util.NoSuchElementException; import java.util.concurrent.atomic.AtomicLong; import org.slf4j.Logger; @@ -30,6 +29,8 @@ import org.apache.spark.memory.MemoryMode; import org.apache.spark.memory.TaskMemoryManager; +import org.apache.comet.MissingTaskEntryRetry; + /** * A adapter class that is used by Comet native to acquire & release memory through Spark's unified * memory manager. This assumes Spark's off-heap memory mode is enabled. All the native plans in a @@ -39,9 +40,6 @@ public class CometTaskMemoryManager { private static final Logger logger = LoggerFactory.getLogger(CometTaskMemoryManager.class); - /** Attempts at an acquire whose task entry Spark keeps losing, see acquireFromSpark. */ - private static final int MAX_ACQUIRE_ATTEMPTS = 3; - /** The id of the native plan this memory manager was created for, the first in its task. */ private final long id; @@ -99,54 +97,17 @@ public long acquireMemory(long size) { } /** - * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry. Spark - * removes a task's entry from its execution pool when the task's balance reaches zero, and an - * acquire that was waiting for memory then fails when it wakes up with a NoSuchElementException - * ("key not found: " and the task id). Releases take no task monitor, so another native thread or - * a JVM consumer of the same task can empty the balance while an acquire waits. The failed call - * was granted nothing, because Spark only drops the entry at a zero balance, which cannot hold - * while this call has a partial grant. So trying again is safe: it registers the task again and - * waits for its share as Spark would have. After a few failed attempts the grant is refused. The - * helper holds no lock of its own across attempts, so a retry parks in Spark exactly as the first - * attempt did. + * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry; see + * {@link MissingTaskEntryRetry}. When the retries run out the grant is refused, which the native + * pools treat like any other refusal. */ private long acquireFromSpark(long size) { - for (int attempt = 1; ; attempt++) { - try { - return internal.acquireExecutionMemory(size, nativeMemoryConsumer); - } catch (NoSuchElementException e) { - if (!isMissingTaskEntry(e)) { - throw e; - } - if (attempt >= MAX_ACQUIRE_ATTEMPTS) { - logger.warn( - "Task {} lost its execution memory entry in Spark on {} attempts to acquire {} " - + "bytes, refusing the request", - taskAttemptId, - attempt, - size); - return 0; - } - logger.info( - "Task {} lost its execution memory entry in Spark while waiting to acquire {} bytes, " - + "trying again", - taskAttemptId, - size); - } - } - } - - /** Whether {@code e} is Spark's execution pool failing to find this task's entry. */ - private boolean isMissingTaskEntry(NoSuchElementException e) { - if (!("key not found: " + taskAttemptId).equals(e.getMessage())) { - return false; - } - for (StackTraceElement frame : e.getStackTrace()) { - if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { - return true; - } - } - return false; + return MissingTaskEntryRetry.retry( + logger, + taskAttemptId, + size, + () -> internal.acquireExecutionMemory(size, nativeMemoryConsumer), + e -> 0L); } // Called by Comet native through JNI diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java index e3a2269cacc..82c3ddd48a8 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java @@ -21,7 +21,6 @@ import java.io.IOException; import java.util.Map; -import java.util.NoSuchElementException; import java.util.function.Supplier; import org.slf4j.Logger; @@ -34,6 +33,8 @@ import org.apache.spark.unsafe.array.LongArray; import org.apache.spark.unsafe.memory.MemoryBlock; +import org.apache.comet.MissingTaskEntryRetry; + /** * A simple memory allocator used by `CometShuffleExternalSorter` to allocate memory blocks which * store serialized rows. This class is simply an implementation of `MemoryConsumer` that delegates @@ -49,9 +50,6 @@ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemory private static final Logger logger = LoggerFactory.getLogger(CometUnifiedShuffleMemoryAllocator.class); - /** Attempts at an allocation whose task entry Spark keeps losing, see retryMissingTaskEntry. */ - private static final int MAX_ALLOCATE_ATTEMPTS = 3; - CometUnifiedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) { super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) { @@ -83,31 +81,16 @@ public LongArray allocateArray(long size) { /** * Runs an allocation through Spark, retrying when an allocation waiting in Spark lost the task's - * entry. Spark removes a task's entry from its execution pool when the task's balance reaches - * zero, and an allocation that was waiting for memory then fails when it wakes up with a - * NoSuchElementException ("key not found: " and the task id). Releases take no task monitor, so a - * native consumer or another JVM consumer of the same task can empty the balance while an - * allocation waits. The failed call was granted nothing and took no page, because Spark only - * drops the entry at a zero balance, so trying again is safe: it registers the task again and - * waits for its share as Spark would have. After a few failed attempts the allocation fails with - * a SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused - * by the last attempt's exception. The Spark issue is SPARK-59444. + * entry; see {@link MissingTaskEntryRetry}. When the retries run out the allocation fails with a + * SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused by + * the last attempt's exception. */ private T retryMissingTaskEntry(long required, Supplier allocation) { - for (int attempt = 1; ; attempt++) { - try { - return allocation.get(); - } catch (NoSuchElementException e) { - if (!isMissingTaskEntry(e)) { - throw e; - } - if (attempt >= MAX_ALLOCATE_ATTEMPTS) { - logger.warn( - "Lost the task's execution memory entry in Spark on {} attempts to allocate {} " - + "bytes, refusing the allocation", - attempt, - required, - e); + return MissingTaskEntryRetry.retry( + logger, + required, + allocation, + e -> { SparkOutOfMemoryError error = new SparkOutOfMemoryError( "UNABLE_TO_ACQUIRE_MEMORY", @@ -116,30 +99,7 @@ private T retryMissingTaskEntry(long required, Supplier allocation) { "receivedBytes", String.valueOf(0))); error.initCause(e); throw error; - } - logger.info( - "Lost the task's execution memory entry in Spark while waiting to allocate {} bytes, " - + "trying again", - required); - } - } - } - - /** - * Whether {@code e} is Spark's execution pool failing to find the task's entry. The task id is - * not available here, so this matches the message prefix and the pool's stack frame. - */ - private static boolean isMissingTaskEntry(NoSuchElementException e) { - String message = e.getMessage(); - if (message == null || !message.startsWith("key not found: ")) { - return false; - } - for (StackTraceElement frame : e.getStackTrace()) { - if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { - return true; - } - } - return false; + }); } public synchronized long free(MemoryBlock block) { diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index efd899252ac..a81efd7eb05 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -19,16 +19,11 @@ package org.apache.spark -import java.util.Properties -import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference} +import org.apache.logging.log4j.Level +import org.apache.logging.log4j.core.LogEvent +import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager} -import org.apache.logging.log4j.{Level, LogManager} -import org.apache.logging.log4j.core.{LogEvent, LoggerContext} -import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager} - -class CometTaskMemoryManagerSuite extends SparkFunSuite { +class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils { test("native memory usage is visible to Spark's memory consumer") { withTaskMemoryManager { taskMemoryManager => @@ -136,62 +131,29 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } } - /** The events that Comet's and Spark's task memory managers log at `level` or above. */ - private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = { - val appender = new LogAppender("task memory manager") - appender.setThreshold(level) - val loggers = Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName) - val context = LogManager.getContext(false).asInstanceOf[LoggerContext] - val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) - try withLogAppender(appender, loggers, Some(level))(f) - finally { - // For a logger with no config of its own, withLogAppender adds one and never removes it. - // It copies the root config's additivity, which is off, so left in place it would keep the - // logger's events out of the test log for the rest of the run. - unconfigured.foreach(context.getConfiguration.removeLogger) - context.updateLoggers() - } - appender.loggingEvents.toSeq - } - test("an acquire waiting in Spark survives a release that empties the task's balance") { - val memoryManager = offHeapMemoryManager() - val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) - - withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + checkRequestWhileBalanceEmpties(loggers) { taskMemoryManager => val manager = new CometTaskMemoryManager(1L, 0L) - assert(otherTask.acquireMemory(90L) == 90L) - assert(manager.acquireMemory(10L) == 10L) - // Another native thread releases the task's last 10 bytes while the acquire waits. - acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { - manager.releaseMemory(10L) - } - } - } - - test("an acquire waiting in Spark survives a sibling consumer freeing the task's last bytes") { - val memoryManager = offHeapMemoryManager() - val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) - - withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => - val manager = new CometTaskMemoryManager(1L, 0L) - val sibling = new OffHeapConsumer(taskMemoryManager) - assert(otherTask.acquireMemory(90L) == 90L) - assert(sibling.acquireMemory(10L) == 10L) - - // A JVM consumer of the same task frees with freeMemory, which takes no task monitor. - acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { - sibling.freeMemory(10L) - } + WaitingRequest( + manager.acquireMemory(_), + manager.releaseMemory(_), + () => manager.acquireMemory(20L), + granted => { + assert(granted == 20L) + assert(manager.getUsed == 20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L) + manager.releaseMemory(20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + }) } } test("a NoSuchElementException other than Spark's missing task entry is rethrown") { - val wrongTask = missingEntryError("key not found: 7", fromExecutionMemoryPool = true) - val notFromPool = missingEntryError("key not found: 0", fromExecutionMemoryPool = false) + val wrongTask = noSuchElement("key not found: 7", fromExecutionMemoryPool = true) + val notFromPool = noSuchElement("key not found: 0", fromExecutionMemoryPool = false) for (error <- Seq(wrongTask, notFromPool)) { - val taskMemoryManager = new FailingTaskMemoryManager(error) + val taskMemoryManager = new FailingTaskMemoryManager(() => error, failures = Int.MaxValue) withTaskContext(taskMemoryManager) { _ => val manager = new CometTaskMemoryManager(1L, 0L) val thrown = intercept[NoSuchElementException](manager.acquireMemory(10L)) @@ -203,108 +165,30 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } test("an acquire that keeps losing its task entry gives up with a zero grant") { - val error = missingEntryError("key not found: 0", fromExecutionMemoryPool = true) - val taskMemoryManager = new FailingTaskMemoryManager(error) + // A new exception on every call, so the test can tell which one is logged. + val taskMemoryManager = new FailingTaskMemoryManager( + () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true), + failures = Int.MaxValue) withTaskContext(taskMemoryManager) { _ => val manager = new CometTaskMemoryManager(1L, 0L) - assert(manager.acquireMemory(10L) == 0L) - assert(taskMemoryManager.calls.get == maxAcquireAttempts) + val events = logEvents(Level.INFO) { + // A refusal rather than an exception, so native spills as for any refused reservation. + assert(manager.acquireMemory(10L) == 0L) + } + assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events)) + assert(events.last.getThrown eq taskMemoryManager.errors.last) + assert(taskMemoryManager.calls.get == MaxAcquireAttempts) assert(manager.getUsed == 0L) assert(nativeMemoryConsumer(manager).getUsed == 0L) assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) } } - private val TimeoutSeconds = 10L - - /** - * The task holds 10 bytes of a 100 byte pool and another task holds 90. An acquire of 20 bytes - * waits in Spark below its minimum share of 25 while `release` takes the task's balance to - * zero, which removes the task's entry from Spark's pool. Once the other task frees its memory - * the acquire must be granted all 20 bytes. - */ - private def acquireWhileBalanceEmpties( - manager: CometTaskMemoryManager, - taskMemoryManager: TaskMemoryManager, - otherTask: OffHeapConsumer)(release: => Unit): Unit = { - val granted = new AtomicLong(-1L) - val failure = new AtomicReference[Throwable]() - val acquire = new Thread(() => - try granted.set(manager.acquireMemory(20L)) - catch { case t: Throwable => failure.set(t) }) - acquire.setDaemon(true) - - try { - acquire.start() - val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) - while (!waitingInSpark(acquire) && System.nanoTime() < deadline) Thread.sleep(10) - assert(waitingInSpark(acquire), s"the acquire is ${acquire.getState}") - release - } finally { - // Free the other task's memory so that the acquire does not outlive a failed test. - otherTask.freeMemory(otherTask.getUsed) - acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) - } - - assert(!acquire.isAlive, s"the acquire is ${acquire.getState}") - assert(failure.get == null, s"the acquire failed: ${failure.get}") - assert(granted.get == 20L) - assert(manager.getUsed == 20L) - assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L) - manager.releaseMemory(20L) - assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) - } - - private def waitingInSpark(thread: Thread): Boolean = - thread.getState == Thread.State.WAITING && - thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") - - /** A 100 byte off-heap execution pool. */ - private def offHeapMemoryManager(): UnifiedMemoryManager = { - val conf = new SparkConf() - .set("spark.memory.offHeap.enabled", "true") - .set("spark.memory.offHeap.size", "100") - .set("spark.memory.storageFraction", "0") - new UnifiedMemoryManager(conf, 1000L, 500L, 1) - } - - /** An off-heap consumer that never spills. */ - private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) - extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { - override def spill(size: Long, trigger: MemoryConsumer): Long = 0L - } - - private val MaxFailingCalls = 10 - - /** A task memory manager whose every acquire throws `error`. */ - private class FailingTaskMemoryManager(error: Throwable) - extends TaskMemoryManager(new TestMemoryManager(new SparkConf()), 0L) { - val calls = new AtomicInteger() - - override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { - // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. - if (calls.incrementAndGet() > MaxFailingCalls) { - throw new IllegalStateException( - s"acquireExecutionMemory called more than $MaxFailingCalls times") - } - throw error - } - } + private val loggers = + Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName) - private def missingEntryError( - message: String, - fromExecutionMemoryPool: Boolean): NoSuchElementException = { - val error = new NoSuchElementException(message) - if (fromExecutionMemoryPool) { - val poolFrame = new StackTraceElement( - "org.apache.spark.memory.ExecutionMemoryPool", - "acquireMemory", - "ExecutionMemoryPool.scala", - 115) - error.setStackTrace(poolFrame +: error.getStackTrace) - } - error - } + /** The events that Comet's and Spark's task memory managers log at `level` or above. */ + private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = logEvents(level, loggers)(f) private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = { val memoryManager = new TestMemoryManager(new SparkConf()) @@ -312,40 +196,6 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { withTaskContext(new TaskMemoryManager(memoryManager, 0L))(f) } - private def withTaskContext(taskMemoryManager: TaskMemoryManager)( - f: TaskMemoryManager => Unit): Unit = { - val taskContext = new TaskContextImpl( - stageId = 0, - stageAttemptNumber = 0, - partitionId = 0, - numPartitions = 1, - taskAttemptId = 0L, - attemptNumber = 0, - taskMemoryManager = taskMemoryManager, - localProperties = new Properties, - metricsSystem = null, - taskMetrics = TaskMetrics.empty, - cpus = 1, - resources = Map.empty) - - TaskContext.setTaskContext(taskContext) - try { - f(taskMemoryManager) - } finally { - try { - taskMemoryManager.cleanUpAllAllocatedMemory() - } finally { - TaskContext.unset() - } - } - } - - private def maxAcquireAttempts: Int = { - val field = classOf[CometTaskMemoryManager].getDeclaredField("MAX_ACQUIRE_ATTEMPTS") - field.setAccessible(true) - field.getInt(null) - } - private def nativeMemoryConsumer(manager: CometTaskMemoryManager): MemoryConsumer = { val field = classOf[CometTaskMemoryManager].getDeclaredField("nativeMemoryConsumer") field.setAccessible(true) diff --git a/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala b/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala new file mode 100644 index 00000000000..a82e2377e82 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala @@ -0,0 +1,219 @@ +/* + * 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 + +import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.collection.mutable.ArrayBuffer + +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.{LogEvent, LoggerContext} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, UnifiedMemoryManager} + +import org.apache.comet.MissingTaskEntryRetry + +/** + * Helpers for suites that drive Spark's task memory manager from Comet's memory consumers, among + * them a reproduction of a request that waits in Spark's execution pool and loses the task's + * entry there (SPARK-59444). + */ +trait TaskMemoryTestUtils extends SparkFunSuite { + + protected val TimeoutSeconds = 10L + protected val MaxAcquireAttempts: Int = MissingTaskEntryRetry.MAX_ATTEMPTS + + /** A 100 byte off-heap execution pool. */ + protected def offHeapMemoryManager(): UnifiedMemoryManager = { + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "100") + .set("spark.memory.storageFraction", "0") + new UnifiedMemoryManager(conf, 1000L, 500L, 1) + } + + /** An off-heap consumer that never spills. */ + protected class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + private val MaxFailingCalls = 10 + + /** + * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the + * result of `newError`, recording each one in `errors`. + */ + protected class FailingTaskMemoryManager(newError: () => Throwable, failures: Int) + extends TaskMemoryManager(offHeapMemoryManager(), 0L) { + val calls = new AtomicInteger() + val errors = new ArrayBuffer[Throwable]() + + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + val call = calls.incrementAndGet() + // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. + if (call > MaxFailingCalls) { + throw new IllegalStateException( + s"acquireExecutionMemory called more than $MaxFailingCalls times") + } + if (call <= failures) { + val error = newError() + errors += error + throw error + } + super.acquireExecutionMemory(required, consumer) + } + } + + protected def noSuchElement( + message: String, + fromExecutionMemoryPool: Boolean): NoSuchElementException = { + val error = new NoSuchElementException(message) + if (fromExecutionMemoryPool) { + val poolFrame = new StackTraceElement( + "org.apache.spark.memory.ExecutionMemoryPool", + "acquireMemory", + "ExecutionMemoryPool.scala", + 115) + error.setStackTrace(poolFrame +: error.getStackTrace) + } + error + } + + /** Whether `thread` is parked in Spark's execution pool waiting for memory. */ + protected def waitingInSpark(thread: Thread): Boolean = + thread.getState == Thread.State.WAITING && + thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") + + /** Waits until `thread` is parked in Spark's execution pool, failing after the timeout. */ + protected def awaitWaitingInSpark(thread: Thread): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) + while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10) + assert(waitingInSpark(thread), s"the request is ${thread.getState}") + } + + /** + * A request for memory made while the task's balance empties; see + * [[checkRequestWhileBalanceEmpties]]. `hold` and `release` take and give back the task's bytes + * through a consumer other than the one `request` asks through, and `request` returns what it + * was granted. `check` runs on the test thread once the request has returned. + */ + protected case class WaitingRequest( + hold: Long => Long, + release: Long => Unit, + request: () => Long, + check: Long => Unit) + + /** + * The task holds 10 bytes of a 100 byte pool through `hold` and another task holds 90. The + * request is made on another thread and waits in Spark below the task's minimum share of 25 + * while `release` gives back the task's last 10 bytes, which removes the task's entry from + * Spark's pool. Once the other task frees its memory the request must not fail, and `check` + * sees what it was granted. The first attempt lost the task's entry, so `loggers` must have + * logged a single retry at INFO. + */ + protected def checkRequestWhileBalanceEmpties(loggers: Seq[String])( + setup: TaskMemoryManager => WaitingRequest): Unit = { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val waiting = setup(taskMemoryManager) + // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50. + assert(otherTask.acquireMemory(90L) == 90L) + assert(waiting.hold(10L) == 10L) + + val granted = new AtomicReference[java.lang.Long]() + val failure = new AtomicReference[Throwable]() + val thread = new Thread(() => + try granted.set(waiting.request()) + catch { case t: Throwable => failure.set(t) }) + thread.setDaemon(true) + + val events = logEvents(Level.INFO, loggers) { + try { + thread.start() + awaitWaitingInSpark(thread) + waiting.release(10L) + } finally { + // Free the other task's memory so that the request does not outlive a failed test. + otherTask.freeMemory(otherTask.getUsed) + thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + } + + assert(!thread.isAlive, s"the request is ${thread.getState}") + assert(failure.get == null, s"the request failed: ${failure.get}") + waiting.check(granted.get) + assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + } + } + + /** The events that `loggers` log at `level` or above while `f` runs. */ + protected def logEvents(level: Level, loggers: Seq[String])(f: => Unit): Seq[LogEvent] = { + val appender = new LogAppender(loggers.mkString(", ")) + appender.setThreshold(level) + val context = LogManager.getContext(false).asInstanceOf[LoggerContext] + val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) + try withLogAppender(appender, loggers, Some(level))(f) + finally { + // For a logger with no config of its own, withLogAppender adds one and never removes it. + // It copies the root config's additivity, which is off, so left in place it would keep the + // logger's events out of the test log for the rest of the run. + unconfigured.foreach(context.getConfiguration.removeLogger) + context.updateLoggers() + } + appender.loggingEvents.toSeq + } + + protected def messages(events: Seq[LogEvent]): String = + events.map(_.getMessage.getFormattedMessage).mkString("\n") + + /** Runs `f` as task 0 with `taskMemoryManager`, then frees whatever the task still holds. */ + protected def withTaskContext(taskMemoryManager: TaskMemoryManager)( + f: TaskMemoryManager => Unit): Unit = { + val taskContext = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = 0L, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + + TaskContext.setTaskContext(taskContext) + try { + f(taskMemoryManager) + } finally { + try { + taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + TaskContext.unset() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala index b978c272802..90aeae8dc1c 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala @@ -19,23 +19,14 @@ package org.apache.spark.shuffle.comet -import java.util.Properties -import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} +import org.apache.logging.log4j.Level +import org.apache.logging.log4j.core.LogEvent +import org.apache.spark.{CometTaskMemoryManager, TaskMemoryTestUtils} +import org.apache.spark.memory.SparkOutOfMemoryError -import scala.collection.mutable.ArrayBuffer - -import org.apache.logging.log4j.{Level, LogManager} -import org.apache.logging.log4j.core.{LogEvent, LoggerContext} -import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkFunSuite, TaskContext, TaskContextImpl} -import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, UnifiedMemoryManager} - -class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { +class CometUnifiedShuffleMemoryAllocatorSuite extends TaskMemoryTestUtils { private val PageSize = 20L - private val TimeoutSeconds = 10L - private val MaxAllocateAttempts = 3 test("a page waiting in Spark survives a release that empties the task's balance") { checkWaitingAllocationSurvives(20L)(allocatePage) @@ -79,8 +70,8 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { } assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events)) assert(events.last.getThrown eq taskMemoryManager.errors.last) - assert(taskMemoryManager.calls.get == MaxAllocateAttempts) - assert(taskMemoryManager.errors.distinct.size == MaxAllocateAttempts) + assert(taskMemoryManager.calls.get == MaxAcquireAttempts) + assert(taskMemoryManager.errors.distinct.size == MaxAcquireAttempts) assert(allocator.getUsed == 0L) assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) @@ -91,13 +82,13 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { val error = noSuchElement("key not found: 0", fromExecutionMemoryPool = true) for ((allocation, requested) <- allocations.zip(Seq(20L, 24L))) { val taskMemoryManager = - new FailingTaskMemoryManager(() => error, failures = MaxAllocateAttempts - 1) + new FailingTaskMemoryManager(() => error, failures = MaxAcquireAttempts - 1) val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) val events = logEvents(Level.INFO) { assert(allocation(allocator) == requested) } assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO), messages(events)) - assert(taskMemoryManager.calls.get == MaxAllocateAttempts) + assert(taskMemoryManager.calls.get == MaxAcquireAttempts) assert(allocator.getUsed == 0L) assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) @@ -124,162 +115,27 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { Seq(allocatePage _, allocatePointerArray _) /** - * The task holds 10 bytes of a 100 byte pool through a native consumer and another task holds - * 90. `allocate` asks for `requested` bytes on another thread and waits in Spark below the - * task's minimum share of 25 while the native consumer releases the task's last 10 bytes, which - * removes the task's entry from Spark's pool. Once the other task frees its memory the - * allocation must be granted in full. + * An allocation of `requested` bytes waits in Spark while a native consumer of the task + * releases the task's last bytes, and must then be granted in full. */ private def checkWaitingAllocationSurvives(requested: Long)( - allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = { - val memoryManager = offHeapMemoryManager() - val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) - - withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => - val native = new CometTaskMemoryManager(1L, 0L) - val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) - // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50. - assert(otherTask.acquireMemory(90L) == 90L) - assert(native.acquireMemory(10L) == 10L) - - val granted = new AtomicReference[java.lang.Long]() - val failure = new AtomicReference[Throwable]() - val thread = new Thread(() => - try granted.set(allocate(allocator)) - catch { case t: Throwable => failure.set(t) }) - thread.setDaemon(true) - - val events = logEvents(Level.INFO) { - try { - thread.start() - val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) - while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10) - assert(waitingInSpark(thread), s"the allocation is ${thread.getState}") - native.releaseMemory(10L) - } finally { - // Free the other task's memory so that the allocation does not outlive a failed test. - otherTask.freeMemory(otherTask.getUsed) - thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) - } - } - - assert(!thread.isAlive, s"the allocation is ${thread.getState}") - assert(failure.get == null, s"the allocation failed: ${failure.get}") - assert(granted.get == requested) - assert(allocator.getUsed == 0L) - assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) - // The first attempt lost the task's entry and the allocator tried again. - assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = + checkRequestWhileBalanceEmpties(Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName)) { + taskMemoryManager => + val native = new CometTaskMemoryManager(1L, 0L) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + WaitingRequest( + native.acquireMemory(_), + native.releaseMemory(_), + () => allocate(allocator), + granted => { + assert(granted == requested) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + }) } - } - - private def waitingInSpark(thread: Thread): Boolean = - thread.getState == Thread.State.WAITING && - thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") - - /** A 100 byte off-heap execution pool. */ - private def offHeapMemoryManager(): UnifiedMemoryManager = { - val conf = new SparkConf() - .set("spark.memory.offHeap.enabled", "true") - .set("spark.memory.offHeap.size", "100") - .set("spark.memory.storageFraction", "0") - new UnifiedMemoryManager(conf, 1000L, 500L, 1) - } - - /** An off-heap consumer that never spills. */ - private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) - extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { - override def spill(size: Long, trigger: MemoryConsumer): Long = 0L - } - - private val MaxFailingCalls = 10 - - /** - * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the - * result of `newError`, recording each one in `errors`. - */ - private class FailingTaskMemoryManager(newError: () => Throwable, failures: Int) - extends TaskMemoryManager(offHeapMemoryManager(), 0L) { - val calls = new AtomicInteger() - val errors = new ArrayBuffer[Throwable]() - - override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { - val call = calls.incrementAndGet() - // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. - if (call > MaxFailingCalls) { - throw new IllegalStateException( - s"acquireExecutionMemory called more than $MaxFailingCalls times") - } - if (call <= failures) { - val error = newError() - errors += error - throw error - } - super.acquireExecutionMemory(required, consumer) - } - } - - private def noSuchElement( - message: String, - fromExecutionMemoryPool: Boolean): NoSuchElementException = { - val error = new NoSuchElementException(message) - if (fromExecutionMemoryPool) { - val poolFrame = new StackTraceElement( - "org.apache.spark.memory.ExecutionMemoryPool", - "acquireMemory", - "ExecutionMemoryPool.scala", - 115) - error.setStackTrace(poolFrame +: error.getStackTrace) - } - error - } /** The events that the allocator logs at `level` or above. */ - private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = { - val appender = new LogAppender("unified shuffle memory allocator") - appender.setThreshold(level) - val loggers = Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName) - val context = LogManager.getContext(false).asInstanceOf[LoggerContext] - val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) - try withLogAppender(appender, loggers, Some(level))(f) - finally { - // For a logger with no config of its own, withLogAppender adds one and never removes it. - // It copies the root config's additivity, which is off, so left in place it would keep the - // logger's events out of the test log for the rest of the run. - unconfigured.foreach(context.getConfiguration.removeLogger) - context.updateLoggers() - } - appender.loggingEvents.toSeq - } - - private def messages(events: Seq[LogEvent]): String = - events.map(_.getMessage.getFormattedMessage).mkString("\n") - - private def withTaskContext(taskMemoryManager: TaskMemoryManager)( - f: TaskMemoryManager => Unit): Unit = { - val taskContext = new TaskContextImpl( - stageId = 0, - stageAttemptNumber = 0, - partitionId = 0, - numPartitions = 1, - taskAttemptId = 0L, - attemptNumber = 0, - taskMemoryManager = taskMemoryManager, - localProperties = new Properties, - metricsSystem = null, - taskMetrics = TaskMetrics.empty, - cpus = 1, - resources = Map.empty) - - TaskContext.setTaskContext(taskContext) - try { - f(taskMemoryManager) - } finally { - try { - taskMemoryManager.cleanUpAllAllocatedMemory() - } finally { - TaskContext.unset() - } - } - } + private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = + logEvents(level, Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName))(f) } From 55e62af29dcce9c18631b7c2b535658ba93bbca0 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 10 Oct 2026 12:12:56 -0700 Subject: [PATCH 2/2] fix: don't retry a missing task entry once the task has ended releasePlan drops a plan's stream when the task ends, while a Tokio worker may still wait in Spark's acquireMemory for the memory that stream holds. When the release empties the task's balance, the worker loses the task's entry. Before #6310 it failed at once; with the retry it registered the ended task again and, below its minimum share, waited for other tasks to release memory while holding the TaskMemoryManager monitor that the executor's cleanUpAllAllocatedMemory needs. MissingTaskEntryRetry now refuses without retrying when the task has completed or been killed: CometTaskMemoryManager returns a zero grant and the shuffle allocator throws SparkOutOfMemoryError. Both callers pass the TaskContext they were created in, since native threads have none of their own. Co-authored-by: Isaac --- .../contributor-guide/memory_management.md | 6 ++- .../apache/comet/MissingTaskEntryRetry.java | 42 ++++++++++++--- .../apache/spark/CometTaskMemoryManager.java | 11 ++-- .../CometUnifiedShuffleMemoryAllocator.java | 12 +++-- .../spark/CometTaskMemoryManagerSuite.scala | 54 +++++++++++++++++++ ...etUnifiedShuffleMemoryAllocatorSuite.scala | 22 +++++++- 6 files changed, 133 insertions(+), 14 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index c23e1eb465e..5456c8e8c33 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -220,7 +220,11 @@ the task held nothing from Spark at that moment, so the failed call has no parti reconcile, and the retry registers the task again and waits for its share as the first call would have. After three attempts `CometTaskMemoryManager` returns a zero grant. `try_grow` treats that as a refusal and the operator spills, while `grow` records the request as overcommit. The shuffle -allocator throws `SparkOutOfMemoryError` instead. Spark's own operators in the same task, such as +allocator throws `SparkOutOfMemoryError` instead. Both refuse without retrying once the task has +completed or been killed: releasing a native plan at the end of the task can empty the balance under +a Tokio worker still waiting in Spark, and a retry would register the ended task again and could wait +for other tasks while holding the `TaskMemoryManager` monitor that Spark's final cleanup needs. +Spark's own operators in the same task, such as its sorters and aggregates, are not guarded and can still hit the exception until Spark re-registers a waiting task in `ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). diff --git a/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java index b0faa2ed8ba..ae35a4be923 100644 --- a/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java +++ b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java @@ -26,6 +26,8 @@ import org.slf4j.Logger; +import org.apache.spark.TaskContext; + /** * Retries a request for execution memory that Spark failed because the task's entry was gone from * its execution pool. @@ -39,6 +41,12 @@ * grant. So trying again is safe: it registers the task again and waits for its share as Spark * would have. Nothing here holds a lock across attempts, so a retry parks in Spark exactly as the * first attempt did. The Spark issue is SPARK-59444. + * + *

A request of a task that has completed or been killed is refused instead of retried. Releasing + * a native plan when the task ends gives back the plan's memory while a Tokio worker may still wait + * in Spark for it, and when that empties the task's balance the worker loses the entry. Retried, it + * would register the ended task again and could wait for other tasks to release memory, holding the + * TaskMemoryManager monitor that Spark's final cleanup of the task needs. */ public final class MissingTaskEntryRetry { @@ -49,35 +57,40 @@ private MissingTaskEntryRetry() {} /** * Runs {@code request}, a call that asks Spark for {@code required} bytes for task {@code - * taskAttemptId}, and retries it when it lost that task's entry. See {@link #retry(Logger, long, - * Supplier, Function)}. + * taskAttemptId}, and retries it when it lost that task's entry. See {@link #retry(Logger, + * TaskContext, long, Supplier, Function)}. */ public static T retry( Logger logger, + TaskContext context, long taskAttemptId, long required, Supplier request, Function refuse) { - return retry(logger, OptionalLong.of(taskAttemptId), required, request, refuse); + return retry(logger, context, OptionalLong.of(taskAttemptId), required, request, refuse); } /** * Runs {@code request}, a call that asks Spark for {@code required} bytes, and retries it when it * lost the task's entry. Without the task id, any task's missing entry counts. Other exceptions * are rethrown. After {@link #MAX_ATTEMPTS} attempts that all lost the entry, returns what {@code - * refuse} makes of the last attempt's exception, or throws what it throws. Both the retries and - * the refusal are logged to {@code logger}. + * refuse} makes of the last attempt's exception, or throws what it throws. It does the same + * without retrying once {@code context}, the task's context or null when there is none, says the + * task has completed or been killed. Both the retries and the refusal are logged to {@code + * logger}. */ public static T retry( Logger logger, + TaskContext context, long required, Supplier request, Function refuse) { - return retry(logger, OptionalLong.empty(), required, request, refuse); + return retry(logger, context, OptionalLong.empty(), required, request, refuse); } private static T retry( Logger logger, + TaskContext context, OptionalLong taskAttemptId, long required, Supplier request, @@ -90,6 +103,14 @@ private static T retry( throw e; } String task = taskAttemptId.isPresent() ? "Task " + taskAttemptId.getAsLong() : "The task"; + if (hasEnded(context)) { + logger.info( + "{} lost its execution memory entry in Spark while waiting to acquire {} bytes after " + + "it ended, refusing the request", + task, + required); + return refuse.apply(e); + } if (attempt >= MAX_ATTEMPTS) { logger.warn( "{} lost its execution memory entry in Spark on {} attempts to acquire {} bytes, " @@ -109,6 +130,15 @@ private static T retry( } } + /** + * Whether the task has completed or been killed. TaskContextImpl.isCompleted takes the context's + * monitor, but Spark runs the task's completion listeners, which release native plans, without + * holding it. + */ + private static boolean hasEnded(TaskContext context) { + return context != null && (context.isInterrupted() || context.isCompleted()); + } + /** * Whether {@code e} is Spark's execution pool failing to find the task's entry: the message names * the task, or any task when its id is unknown, and the exception comes from the pool. diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index a0eccb85ffb..1c76627c5b5 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -45,6 +45,9 @@ public class CometTaskMemoryManager { private final long taskAttemptId; + /** The task's context, read from native threads, which have none of their own. */ + private final TaskContext taskContext; + public final TaskMemoryManager internal; private final NativeMemoryConsumer nativeMemoryConsumer; private final AtomicLong used = new AtomicLong(); @@ -52,7 +55,8 @@ public class CometTaskMemoryManager { public CometTaskMemoryManager(long id, long taskAttemptId) { this.id = id; this.taskAttemptId = taskAttemptId; - this.internal = TaskContext$.MODULE$.get().taskMemoryManager(); + this.taskContext = TaskContext$.MODULE$.get(); + this.internal = taskContext.taskMemoryManager(); this.nativeMemoryConsumer = new NativeMemoryConsumer(); } @@ -98,12 +102,13 @@ public long acquireMemory(long size) { /** * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry; see - * {@link MissingTaskEntryRetry}. When the retries run out the grant is refused, which the native - * pools treat like any other refusal. + * {@link MissingTaskEntryRetry}. When the retries run out, or the task has ended, the grant is + * refused, which the native pools treat like any other refusal. */ private long acquireFromSpark(long size) { return MissingTaskEntryRetry.retry( logger, + taskContext, taskAttemptId, size, () -> internal.acquireExecutionMemory(size, nativeMemoryConsumer), diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java index 82c3ddd48a8..f7bc13ea152 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java @@ -26,6 +26,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.apache.spark.TaskContext; import org.apache.spark.memory.MemoryConsumer; import org.apache.spark.memory.MemoryMode; import org.apache.spark.memory.SparkOutOfMemoryError; @@ -50,8 +51,12 @@ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemory private static final Logger logger = LoggerFactory.getLogger(CometUnifiedShuffleMemoryAllocator.class); + /** The task's context, or null outside a task; see retryMissingTaskEntry. */ + private final TaskContext taskContext; + CometUnifiedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) { super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); + this.taskContext = TaskContext.get(); if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) { throw new IllegalArgumentException( "CometUnifiedShuffleMemoryAllocator should be used with off-heap " @@ -81,13 +86,14 @@ public LongArray allocateArray(long size) { /** * Runs an allocation through Spark, retrying when an allocation waiting in Spark lost the task's - * entry; see {@link MissingTaskEntryRetry}. When the retries run out the allocation fails with a - * SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused by - * the last attempt's exception. + * entry; see {@link MissingTaskEntryRetry}. When the retries run out, or the task has ended, the + * allocation fails with a SparkOutOfMemoryError, which the shuffle writers handle like any other + * refused page, caused by the last attempt's exception. */ private T retryMissingTaskEntry(long required, Supplier allocation) { return MissingTaskEntryRetry.retry( logger, + taskContext, required, allocation, e -> { diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index a81efd7eb05..82c5731da7a 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -19,6 +19,9 @@ package org.apache.spark +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicReference + import org.apache.logging.log4j.Level import org.apache.logging.log4j.core.LogEvent import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager} @@ -149,6 +152,57 @@ class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils { } } + for ((state, end) <- Seq[(String, TaskContextImpl => Unit)]( + ("completed", _.markTaskCompleted(None)), + ("was killed", _.markInterrupted("killed")))) { + test(s"an acquire that loses its task entry after the task $state is not made again") { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val manager = new CometTaskMemoryManager(1L, 0L) + assert(otherTask.acquireMemory(90L) == 90L) + assert(manager.acquireMemory(10L) == 10L) + + val granted = new AtomicReference[java.lang.Long]() + val failure = new AtomicReference[Throwable]() + val acquire = new Thread(() => + try granted.set(manager.acquireMemory(20L)) + catch { case t: Throwable => failure.set(t) }) + acquire.setDaemon(true) + // Spark's executor frees what the task still holds once the task has ended, under the + // task memory manager's monitor, which a parked acquire holds. + val cleanUp = new Thread(() => taskMemoryManager.cleanUpAllAllocatedMemory()) + cleanUp.setDaemon(true) + + try { + acquire.start() + awaitWaitingInSpark(acquire) + // The task ends and releasing its plan gives back the task's last bytes while a native + // thread still waits in Spark, as releasePlan does when it drops a plan's stream. + end(TaskContext.get().asInstanceOf[TaskContextImpl]) + manager.releaseMemory(10L) + acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + // Parked again, the acquire would wait for the other task while holding the monitor. + assert(!acquire.isAlive, s"the acquire is ${acquire.getState} after the task $state") + cleanUp.start() + cleanUp.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + assert(!cleanUp.isAlive, s"the clean up is ${cleanUp.getState}") + } finally { + // Free the other task's memory so that neither thread outlives a failed test. + otherTask.freeMemory(otherTask.getUsed) + acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + cleanUp.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + + assert(failure.get == null, s"the acquire failed: ${failure.get}") + assert(granted.get == 0L) + assert(manager.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + } + test("a NoSuchElementException other than Spark's missing task entry is rethrown") { val wrongTask = noSuchElement("key not found: 7", fromExecutionMemoryPool = true) val notFromPool = noSuchElement("key not found: 0", fromExecutionMemoryPool = false) diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala index 90aeae8dc1c..ecf2e47be21 100644 --- a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala @@ -21,7 +21,7 @@ package org.apache.spark.shuffle.comet import org.apache.logging.log4j.Level import org.apache.logging.log4j.core.LogEvent -import org.apache.spark.{CometTaskMemoryManager, TaskMemoryTestUtils} +import org.apache.spark.{CometTaskMemoryManager, TaskContext, TaskContextImpl, TaskMemoryTestUtils} import org.apache.spark.memory.SparkOutOfMemoryError class CometUnifiedShuffleMemoryAllocatorSuite extends TaskMemoryTestUtils { @@ -78,6 +78,26 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends TaskMemoryTestUtils { } } + test("an allocation that loses its task entry after the task ended is not made again") { + for ((allocation, requested) <- allocations.zip(Seq("20", "24"))) { + val taskMemoryManager = new FailingTaskMemoryManager( + () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true), + failures = Int.MaxValue) + withTaskContext(taskMemoryManager) { _ => + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + TaskContext.get().asInstanceOf[TaskContextImpl].markTaskCompleted(None) + val events = logEvents(Level.INFO) { + val thrown = intercept[SparkOutOfMemoryError](allocation(allocator)) + assert(thrown.getMessageParameters.get("requestedBytes") == requested) + assert(thrown.getCause eq taskMemoryManager.errors.last) + } + assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + assert(taskMemoryManager.calls.get == 1) + assert(allocator.getUsed == 0L) + } + } + } + test("an allocation that loses its task entry twice is granted on the last attempt") { val error = noSuchElement("key not found: 0", fromExecutionMemoryPool = true) for ((allocation, requested) <- allocations.zip(Seq(20L, 24L))) {