You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/12 20:33:36 UTC
(datafusion-comet) 01/02: fix: share the missing task entry retry and stop retrying it after the task has ended (#6813)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit c937df09c2cbd302772310a85f14ca52f32455c7
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Mon Oct 12 19:25:01 2026 +0000
fix: share the missing task entry retry and stop retrying it after the task has ended (#6813)
* 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: <id>" 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 <no...@databricks.com>
* 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 <no...@databricks.com>
---------
Co-authored-by: Isaac <no...@databricks.com>
---
docs/source/contributor-guide/memory_management.md | 35 +--
.../org/apache/comet/MissingTaskEntryRetry.java | 162 +++++++++++++
.../org/apache/spark/CometTaskMemoryManager.java | 68 ++----
.../comet/CometUnifiedShuffleMemoryAllocator.java | 66 ++----
.../apache/spark/CometTaskMemoryManagerSuite.scala | 258 +++++++--------------
.../org/apache/spark/TaskMemoryTestUtils.scala | 219 +++++++++++++++++
.../CometUnifiedShuffleMemoryAllocatorSuite.scala | 218 ++++-------------
7 files changed, 562 insertions(+), 464 deletions(-)
diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 9c31d3fecb..184bf0de27 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,22 @@ 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. 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)).
**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 0000000000..ae35a4be92
--- /dev/null
+++ b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java
@@ -0,0 +1,162 @@
+/*
+ * 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;
+
+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.
+ *
+ * <p>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.
+ *
+ * <p>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 {
+
+ /** 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,
+ * TaskContext, long, Supplier, Function)}.
+ */
+ public static <T> T retry(
+ Logger logger,
+ TaskContext context,
+ long taskAttemptId,
+ long required,
+ Supplier<T> request,
+ Function<NoSuchElementException, T> 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. 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> T retry(
+ Logger logger,
+ TaskContext context,
+ long required,
+ Supplier<T> request,
+ Function<NoSuchElementException, T> refuse) {
+ return retry(logger, context, OptionalLong.empty(), required, request, refuse);
+ }
+
+ private static <T> T retry(
+ Logger logger,
+ TaskContext context,
+ OptionalLong taskAttemptId,
+ long required,
+ Supplier<T> request,
+ Function<NoSuchElementException, T> 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 (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, "
+ + "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 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.
+ */
+ 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 a468446bae..1c76627c5b 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.MemoryConsumer;
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,14 +40,14 @@ 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;
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();
@@ -54,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();
}
@@ -99,54 +101,18 @@ public class CometTaskMemoryManager {
}
/**
- * 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, or the task has ended, 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,
+ taskContext,
+ 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 e3a2269cac..f7bc13ea15 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,12 +21,12 @@ package org.apache.spark.shuffle.comet;
import java.io.IOException;
import java.util.Map;
-import java.util.NoSuchElementException;
import java.util.function.Supplier;
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;
@@ -34,6 +34,8 @@ import org.apache.spark.memory.TaskMemoryManager;
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,11 +51,12 @@ 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;
+ /** 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 "
@@ -83,31 +86,17 @@ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemory
/**
* 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, 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> T retryMissingTaskEntry(long required, Supplier<T> 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,
+ taskContext,
+ required,
+ allocation,
+ e -> {
SparkOutOfMemoryError error =
new SparkOutOfMemoryError(
"UNABLE_TO_ACQUIRE_MEMORY",
@@ -116,30 +105,7 @@ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemory
"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 efd899252a..82c5731da7 100644
--- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
@@ -19,16 +19,14 @@
package org.apache.spark
-import java.util.Properties
import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference}
+import java.util.concurrent.atomic.AtomicReference
-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}
+import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.core.LogEvent
+import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager}
-class CometTaskMemoryManagerSuite extends SparkFunSuite {
+class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils {
test("native memory usage is visible to Spark's memory consumer") {
withTaskMemoryManager { taskMemoryManager =>
@@ -136,62 +134,80 @@ 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)
- }
+ 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("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))
+ 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)
- 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)
+ 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 = 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 +219,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()
+ private val loggers =
+ Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName)
- 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 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 +250,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 0000000000..a82e2377e8
--- /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 b978c27280..ecf2e47be2 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, TaskContext, TaskContextImpl, 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,25 +70,45 @@ 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)
}
}
+ 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))) {
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 +135,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)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org