peter-toth commented on code in PR #58763:
URL: https://github.com/apache/spark/pull/58763#discussion_r3999303060
##########
core/src/main/scala/org/apache/spark/memory/MemoryManager.scala:
##########
@@ -128,6 +258,24 @@ private[spark] abstract class MemoryManager(
taskAttemptId: Long,
memoryMode: MemoryMode): Long
+ /**
+ * Reserve all `numBytes` for optional task work without waiting for
capacity or reclaiming
+ * memory.
+ *
+ * Managers must opt in to this policy; the default rejects the request
without changing their
+ * accounting. Implementations may still contend on bookkeeping locks. A
successful reservation
+ * must be task-attributed and released through the existing
execution-memory release methods.
+ *
+ * @return `numBytes` on success, or zero with no reservation on denial
+ */
+ private[memory] def tryAcquireExecutionMemory(
Review Comment:
**Finding 1.** An optional reservation is charged to the task in
`ExecutionMemoryPool.memoryForTask` but is owned by no `MemoryConsumer`, so
`TaskMemoryManager` cannot see it. Two consequences, both reachable by the
first consumer:
- `cleanUpAllAllocatedMemory()` returns the leftover bytes, so `Executor`
reports `Managed memory leak detected`
(`core/src/main/scala/org/apache/spark/executor/Executor.scala:917-925`) for
any reservation still held when the task body returns.
`spark.unsafe.exceptionOnMemoryLeak` turns that warning into a thrown
`SparkException`. It defaults to `false`, but every Spark test run sets it to
`true` (`project/SparkBuild.scala:2038`, `pom.xml:2903`). So a consumer whose
prefetch outlives the task body only warns in production and fails outright in
CI.
- `snapshotMemoryUsage()` computes `memoryNotAccountedFor` as the task's
pool charge minus the sum over consumers
(`core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:490`).
Optional bytes therefore land in the unattributed bucket of the
`UNABLE_TO_ACQUIRE_MEMORY` breakdown, which is the diagnostic someone reads
when a task OOMs next to a prefetcher.
The description defers task-completion integration to a follow-up. The
contract as written makes it a prerequisite instead. Nothing ties the
registration or the reservation to task completion, so a task that dies before
its `finally` leaves a live registration. Its counter then keeps every later
allocation on the preflight-and-drain path for the life of the executor,
invoking a callback into a dead reader.
Two ways out. Add the completion-listener integration here, releasing the
outstanding bytes and unregistering. Or model the owner as a `MemoryConsumer`,
as suggested at
[r3998074909](https://github.com/apache/spark/pull/58763#discussion_r3998074909),
which gets the lifecycle, the breakdown and `showMemoryUsage` for free. One
caveat on that second option: `ExecutionMemoryPool.acquireMemory` is
`lock.synchronized`, so hooking the reclaim into its loop would run callbacks
under the manager monitor. That is the thing this design spends the RW gate to
avoid.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -414,20 +438,47 @@ private[spark] class MemoryStore(
def freeMemoryEntry[T <: MemoryEntry[_]](entry: T): Unit = {
entry match {
case SerializedMemoryEntry(buffer, _, _) => buffer.dispose()
- case e: DeserializedMemoryEntry[_] => e.value.foreach {
- case o: AutoCloseable =>
- try {
- o.close()
- } catch {
- case NonFatal(e) =>
- logWarning("Fail to close a memory entry", e)
- }
- case _ =>
- }
+ case e: DeserializedMemoryEntry[_] => freeValues(e.value.iterator)
}
}
- def remove(blockId: BlockId): Boolean = memoryManager.synchronized {
+ private def freeValues(values: Iterator[_]): Unit = values.foreach {
+ case o: AutoCloseable =>
+ try {
+ o.close()
+ } catch {
+ case NonFatal(e) =>
+ logWarning("Fail to close a memory entry", e)
+ }
+ case _ =>
+ }
+
+ /** Dispose only values consumed by this unroll operation. */
+ private def freeUnrolledValues(valuesHolder: ValuesHolder[_]): Unit =
valuesHolder match {
+ case holder: DeserializedValuesHolder[_] =>
+ // Final sizing has already moved the vector into arrayValues; do not
build it again.
+ freeValues(if (holder.vector != null) holder.vector.iterator else
holder.arrayValues.iterator)
+ case holder: SerializedValuesHolder[_] =>
+ Utils.tryWithSafeFinally {
+ // As in PartiallySerializedBlock.discard, closing must not allocate
or flush more data.
+
holder.redirectableStream.setOutputStream(OutputStream.nullOutputStream())
+ holder.serializationStream.close()
+ } {
+ holder.bbos.dispose()
+ }
+ }
+
+ /**
+ * Remove a block and release its storage charge; optional admission skips
object close callbacks.
+ */
+ def remove(blockId: BlockId): Boolean = {
Review Comment:
**Finding 2.** The class comment at `MemoryStore.scala:91-92` still reads
"all changes to memory allocations, notably putting blocks, evicting blocks,
and acquiring or releasing unroll memory, must be synchronized on
`memoryManager`". After this PR that recipe hangs.
I ran it on this head. One reclaimer registered, thread A in
`mm.synchronized { mm.acquireStorageMemory(block, 100, ON_HEAP) }`, thread B in
`mm.tryAcquireExecutionMemory(100, 2, ON_HEAP)` once A holds the monitor.
Result: `threadA=WAITING threadB=BLOCKED`, neither returns in 15 seconds. That
is the deadlock described at
[r3998074901](https://github.com/apache/spark/pull/58763#discussion_r3998074901).
The new rule is the opposite of the comment: gate first, monitor second. It
is stated only in the scaladoc of `withMemoryReclamation`, which a
`MemoryStore` contributor is less likely to read than the comment at the top of
the class they are editing.
Please rewrite lines 91-92 to state the ordering and name the methods that
establish it. The runtime check asked for on that thread catches the accident.
This catches the person writing the next `MemoryStore` method.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -589,17 +659,23 @@ private[spark] class MemoryStore(
blockId: BlockId,
memory: Long,
memoryMode: MemoryMode): Boolean = {
- memoryManager.synchronized {
- val success = memoryManager.acquireUnrollMemory(blockId, memory,
memoryMode)
- if (success) {
- val taskAttemptId = currentTaskAttemptId()
- val unrollMemoryMap = memoryMode match {
- case MemoryMode.ON_HEAP => onHeapUnrollMemoryMap
- case MemoryMode.OFF_HEAP => offHeapUnrollMemoryMap
+ if (memoryManager.isStorageMemoryRequestTooLarge(memory, memoryMode)) {
Review Comment:
**Finding 4.** This early return drops the diagnostic that used to fire for
an impossible unroll request. Before, the call reached
`acquireStorageMemoryInternal` and logged "Will not store `<block>` as the
required space (N bytes) exceeds our memory limit (M bytes)" on its way to
`false` (`UnifiedMemoryManager.scala:401`). Same result now, one fewer line
explaining why nothing got cached. A `logInfo` here, or reusing that message,
keeps it.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -414,20 +438,47 @@ private[spark] class MemoryStore(
def freeMemoryEntry[T <: MemoryEntry[_]](entry: T): Unit = {
entry match {
case SerializedMemoryEntry(buffer, _, _) => buffer.dispose()
- case e: DeserializedMemoryEntry[_] => e.value.foreach {
- case o: AutoCloseable =>
- try {
- o.close()
- } catch {
- case NonFatal(e) =>
- logWarning("Fail to close a memory entry", e)
- }
- case _ =>
- }
+ case e: DeserializedMemoryEntry[_] => freeValues(e.value.iterator)
}
}
- def remove(blockId: BlockId): Boolean = memoryManager.synchronized {
+ private def freeValues(values: Iterator[_]): Unit = values.foreach {
+ case o: AutoCloseable =>
+ try {
+ o.close()
+ } catch {
+ case NonFatal(e) =>
+ logWarning("Fail to close a memory entry", e)
+ }
+ case _ =>
+ }
+
+ /** Dispose only values consumed by this unroll operation. */
+ private def freeUnrolledValues(valuesHolder: ValuesHolder[_]): Unit =
valuesHolder match {
Review Comment:
**Finding 5.** This match is not exhaustive and `ValuesHolder` is not sealed
(`MemoryStore.scala:760`). A third implementation would fail with a
`MatchError` at the worst moment: inside the cleanup of an already-failing
unroll, replacing the original error with a confusing one. The trait and both
implementations live in this file, so `private sealed trait ValuesHolder[T]`
turns that into a compile error. `MemoryEntry` next door is already declared
that way (`MemoryStore.scala:44`), and `freeMemoryEntry` relies on it.
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -134,7 +180,75 @@ private[spark] class UnifiedMemoryManager(
override private[memory] def acquireExecutionMemory(
numBytes: Long,
taskAttemptId: Long,
- memoryMode: MemoryMode): Long = synchronized {
+ memoryMode: MemoryMode): Long = {
+ val gate = optionalAdmissionGate.readLock()
+ gate.lock()
+ try {
+ // A first registration can race this check, but cannot admit optional
bytes while this
+ // read gate is held. Keep the no-owner path free of reentrant-monitor
checks/preflight.
+ if (!hasOptionalMemoryReclaimers(memoryMode)) {
+ return synchronized {
+ acquireExecutionMemoryInternal(numBytes, taskAttemptId, memoryMode)
+ }
+ }
+ val enteredWithMonitor = Thread.holdsLock(this)
Review Comment:
**Finding 3.** No in-tree caller reaches `acquireExecutionMemory` with this
monitor held. The only production caller is
`TaskMemoryManager.acquireExecutionMemory`, which holds its own monitor, not
this one, and a reclaimer callback is forbidden from allocating. The storage
twin at line 335 is covered by "outer storage marker drains before inherited
monitor and excludes new optional admission", but nothing covers this one.
The branch is load-bearing where it does fire: without it, a nested call
whose full request exceeds the caller's fair share falls through to
`reclaimOptionalMemory`, whose `require(!Thread.holdsLock(this))` then fails.
So the test needs a registered owner that the drain does not fully release,
plus a second task holding ordinary bytes so the task count stays put, and then
an execution request under the monitor that must come back as a partial grant
rather than an exception.
If there is no caller and none planned, dropping the branch and asserting
`!Thread.holdsLock(this)` here is smaller than keeping an untested one.
##########
core/src/main/scala/org/apache/spark/util/io/ChunkedByteBufferOutputStream.scala:
##########
@@ -53,6 +53,16 @@ private[spark] class ChunkedByteBufferOutputStream(
def size: Long = _size
+ /** Release untransferred chunks without allocating a compact final chunk. */
+ def dispose(): Unit = {
Review Comment:
**Finding 6.** `dispose()` reuses `toChunkedByteBufferWasCalled` to mean
"the chunks are gone", so a later `toChunkedByteBuffer` fails on
`require(!toChunkedByteBufferWasCalled, "toChunkedByteBuffer() can only be
called once")` with a message saying the opposite of what happened. The new
test asserts the `IllegalArgumentException` without checking its text, so the
wrong message is locked in. A separate `disposed` flag with its own
`require(!disposed, "cannot call toChunkedByteBuffer() after dispose()")` keeps
both messages honest for the price of one field.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
