sunchao commented on code in PR #5998:
URL: https://github.com/apache/datafusion-comet/pull/5998#discussion_r4041278468
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala:
##########
@@ -87,21 +87,23 @@ private[codegen] object CometBatchKernelCodegenOutput
extends CometTypeShim {
* Closes the vector on any failure so a partially-initialized tree doesn't
leak buffers.
*/
def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int):
FieldVector = {
+ // JVM-owned codegen output, accounted to the task that is running the
kernel.
+ val allocator = CometTaskArrowAllocator.forCurrentTask()
Review Comment:
### Correctness
[P2] Give exported JVM output a single Spark reservation owner
This charges codegen output to the task allocator, but `CometUdfBridge`
exports that vector through C Data and closes only its original reference.
Export retains the same buffers and their JVM charge until native code releases
them. With a unified native memory pool, the existing DataFusion 55.1.0
hash-join build path then [counts and reserves the incoming batch's
buffers](https://github.com/apache/datafusion/blob/7d3835c71f30cbd3c3ae4041732267f1f453097a/datafusion/physical-plan/src/joins/hash_join/exec.rs#L2266-L2276)
through Comet's unified pool, charging the same Spark task again. Arrow's
foreign-buffer capacity is nonzero. [The custom allocation retains its
size](https://github.com/apache/arrow-rs/blob/f90e061326bd821a7af09281d9e92de6f3b603d9/arrow-buffer/src/bytes.rs#L100-L107).
I verified the export lifetime with real Arrow C Data export/import and the
current Java memory bridge: an unchanged data address retained a 2 MiB JVM
reservation while the corresponding 1,064,960-byt
e native-style reservation also succeeded. This is a component reproduction
plus a source trace, not an end-to-end join run. Under a bounded pool, this
reduces available headroom twice for one allocation and can reject the join
build. Please transfer or otherwise coordinate reservation ownership at export
and cover JVM-UDF output feeding a native hash join. This path already exists
without unmerged #5027.
##########
spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala:
##########
@@ -0,0 +1,217 @@
+/*
+ * 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.comet
+
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
+
+import scala.util.control.NonFatal
+
+import org.apache.arrow.memory.AllocationListener
+import org.apache.spark.internal.Logging
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode,
SparkOutOfMemoryError, TaskMemoryManager}
+
+import org.apache.comet.CometConf
+
+/**
+ * Accounts one task's JVM-side Arrow allocations against Spark's off-heap
execution pool.
+ *
+ * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, so
until now the
+ * off-heap bytes it hands out were counted by nobody: not Spark's
`TaskMemoryManager`, and not
+ * Comet's native memory pool. They are still resident in the container, which
makes them a blind
+ * spot when an executor is killed for exceeding its memory limit. This closes
the reporting half
+ * of that gap: the bytes appear in `TaskMemoryManager.showMemoryUsage` and
are arbitrated against
+ * Spark's other off-heap consumers.
+ *
+ * '''Ownership.''' One instance is created per task and attached to that
task's Arrow allocator
+ * by [[CometTaskArrowAllocator]]. Arrow reports an allocation and its
matching release to the
+ * listener of the allocator that '''owns''' the buffer, on whichever thread
happens to drop the
+ * last reference, and `AllocationListener` is handed nothing but a size.
Binding the listener to
+ * an allocator is therefore the only way to attribute a release, and reading
`TaskContext` inside
+ * the callbacks would get it wrong: the JVM UDF path exports a JVM-owned
vector to native, which
+ * drops it later from a Tokio worker with no task context installed. That
release would be lost,
+ * leaving the task charged for memory it had already freed, batch after batch.
+ *
+ * '''Reporting only.''' A short grant is logged and the allocation proceeds,
because Arrow
+ * allocation on these paths cannot fail today and making it fail is a
behavioural change that
+ * belongs in its own commit. Enforcement belongs in `onPreAllocation`, the
only callback
+ * permitted to throw, and in `onFailedAllocation`, not here. See
+ * [[https://github.com/apache/datafusion-comet/issues/5997]].
+ *
+ * '''Neither callback may throw.''' Arrow's `AllocationListener` documents
that, and
+ * `BaseAllocator.buffer` marks the allocation successful before calling
`onAllocation`, so
+ * throwing from here loses the buffer Arrow has already created and never
hands back. Spark's
+ * acquisition is fallible -- it runs other consumers' `spill`, which turns a
task interrupt into
+ * a `RuntimeException` and an I/O failure into a `SparkOutOfMemoryError` --
so every call into
+ * the memory manager is wrapped and reported rather than propagated.
+ *
+ * '''Lock order.''' [[getUsed]] and [[spill]] must stay lock-free, because
Spark calls both while
+ * holding the `TaskMemoryManager` monitor, and [[adjust]] holds this
listener's monitor across
+ * `acquireExecutionMemory`, which takes that monitor. Were the snapshot to
take this monitor
+ * instead, a native reservation arriving through `CometTaskMemoryManager` on
a Comet Tokio thread
+ * could hold Spark's monitor and wait for ours while an Arrow allocation on
the same task held
+ * ours and waited for Spark's.
+ */
+private[comet] class CometArrowAllocationListener(taskMemoryManager:
TaskMemoryManager)
+ extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP)
+ with AllocationListener {
+
+ import CometArrowAllocationListener._
+
+ /**
+ * Bytes Arrow currently holds on this task's behalf. An atomic rather than
a guarded field so
+ * that [[getUsed]] can read it without taking this listener's monitor; see
the lock order note
+ * above.
+ */
+ private val live = new AtomicLong(0L)
+
+ /** Bytes currently reserved with Spark. Guarded by this listener's monitor.
*/
+ private var reserved = 0L
+
+ /** Set once the owning task has finished. Volatile so [[getUsed]] can read
it lock-free. */
+ @volatile private var completed = false
+
+ override def onAllocation(size: Long): Unit = {
+ live.addAndGet(size)
+ adjustQuietly()
+ }
+
+ override def onRelease(size: Long): Unit = {
+ live.addAndGet(-size)
+ adjustQuietly()
+ }
+
+ /**
+ * Reports our own tally. Spark reads this for spill-victim ordering,
`showMemoryUsage` and
+ * end-of-task leak reporting. The inherited `used` counter stays at zero
because this consumer
+ * never calls `acquireMemory` or `allocatePage`; Arrow has already obtained
the memory and we
+ * are only accounting for it.
+ *
+ * Reports zero once the task has finished, so that buffers deliberately
allowed to outlive
+ * their task are not reported by `cleanUpAllAllocatedMemory` as a Spark
memory leak.
+ */
+ override def getUsed: Long = if (completed) 0L else math.max(0L, live.get())
+
+ /** Comet's native operators cannot be made to spill from here. See issue
#5997. */
+ override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+
+ /**
+ * Drops the whole reservation and stops accounting.
+ *
+ * Called from the owning task's completion listener. Anything still alive
afterwards is a
+ * buffer that outlives its task, which the process-wide allocator exists to
allow; those
+ * releases are ignored rather than charged to whichever task happens to be
running by then.
+ */
+ private[comet] def taskCompleted(): Unit = {
+ try {
+ synchronized {
+ completed = true
+ if (reserved > 0L) {
+ taskMemoryManager.releaseExecutionMemory(reserved, this)
+ reserved = 0L
+ }
+ }
+ } catch {
+ case NonFatal(e) => warnOnMemoryManagerFailure(e)
+ case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
+ }
+ }
+
+ /** Bytes Arrow currently holds on this task's behalf. Visible for testing.
*/
+ private[comet] def liveBytes: Long = live.get()
+
+ /** Bytes currently reserved with Spark on this task's behalf. Visible for
testing. */
+ private[comet] def reservedBytes: Long = synchronized(reserved)
+
+ private def adjustQuietly(): Unit = {
+ try {
+ adjust()
+ } catch {
+ // Both of these are reachable: `acquireExecutionMemory` runs other
consumers' `spill`, and
+ // `TaskMemoryManager` rethrows an interrupt as a RuntimeException and
an IOException as a
+ // SparkOutOfMemoryError, which is an Error and so slips past NonFatal.
+ case NonFatal(e) => warnOnMemoryManagerFailure(e)
+ case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e)
+ }
+ }
+
+ private def adjust(): Unit = synchronized {
+ if (!completed) {
+ val liveBytes = math.max(0L, live.get())
+ if (reserved < liveBytes) {
+ // Round up so `reserved` stays a block multiple and growth always
leaves headroom.
+ // Requesting the bare deficit would land exactly on `liveBytes` for
any buffer at or above
+ // the block size, sending the very next allocation straight back into
Spark's lock.
+ val request = roundUpToBlock(liveBytes - reserved)
+ val granted = taskMemoryManager.acquireExecutionMemory(request, this)
+ reserved += granted
Review Comment:
### Correctness
[P2] Account for a partial grant when Spark's later spill attempt throws
`TaskMemoryManager.acquireExecutionMemory` can acquire some bytes from the
execution pool and then throw while asking another consumer to spill. In that
case it never returns `granted`, so `reserved` stays unchanged, while
`adjustQuietly` swallows the exception and lets the Arrow allocation succeed.
With a 2 MiB pool, another consumer holding 1 MiB, and a 2 MiB Arrow request, I
reproduced this with both `IOException` and `InterruptedIOException`: after
closing the Arrow buffer and freeing the other consumer, Arrow and
`reservedBytes` are zero but Spark still charges the task 1 MiB.
`taskCompleted()` cannot return that unrecorded grant. It lasts until Spark's
final task cleanup, reducing headroom for the rest of the task. Could
acquisition failure reconcile or roll back partial grants? A regression with
partially available capacity would cover this. The current full-pool failure
tests only exercise a zero-byte initial grant.
--
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]
