sunchao commented on code in PR #5998: URL: https://github.com/apache/datafusion-comet/pull/5998#discussion_r4039113437
########## spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala: ########## @@ -0,0 +1,225 @@ +/* + * 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.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.arrow.memory.AllocationListener +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager} + +import org.apache.comet.CometConf + +/** + * Reports JVM-side Arrow allocations to Spark's memory manager. + * + * `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 listener closes the reporting half of that gap. Every allocation is charged to a + * [[MemoryConsumer]] belonging to the task that made it, so the bytes appear in + * `TaskMemoryManager.showMemoryUsage` and are arbitrated against Spark's other off-heap + * consumers. + * + * It deliberately does not enforce. A short grant from Spark 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. Note that enforcement belongs in + * `onPreAllocation`, the only callback permitted to throw, and `onFailedAllocation`, not here. + * See [[https://github.com/apache/datafusion-comet/issues/5997]]. + * + * Buffers imported over the C Data Interface never reach this listener at all. They wrap memory + * the native side owns, so Comet imports them through `CometImportedArrowAllocator`, a child with + * no listener. Charging them here would double count bytes already reserved in Comet's native + * pool. Arrow notifies only the allocating allocator's own listener, which is what makes that + * separation work. + * + * Three cases are handled by doing nothing, each for a different reason: + * - No active task. Broadcast coalescing and the cached batch serializer can allocate from the + * driver or a non-task thread, where there is no task to charge. + * - On-heap mode. Comet's on-heap mode exists so the Spark SQL suite can run without off-heap + * memory configured; charging an off-heap consumer there would be wrong. + * - A buffer released after its allocating task has finished. The allocator is process-wide + * precisely because buffers can outlive the task that created them, so the task's reservation + * is dropped at task end and later releases are ignored rather than double-counted. + */ +class CometArrowAllocationListener extends AllocationListener { + + import CometArrowAllocationListener._ + + private val reservations = new ConcurrentHashMap[Long, TaskReservation]() + + override def onAllocation(size: Long): Unit = { + val reservation = reservationForCurrentTask() + if (reservation != null) { + reservation.allocated(size) + } + } + + override def onRelease(size: Long): Unit = { + val reservation = reservationForCurrentTask() Review Comment: ### Correctness [P2] Can releases be associated with the allocating task rather than the task currently on the releasing thread? The existing JVM UDF path installs a task context, allocates/exports JVM Arrow buffers, then restores or unsets the context before Rust retains the exported result through `from_ffi`. Its later release callback can run on a Tokio thread without a task context, so this lookup ignores the release while the allocating task is still running. Repeated batches leave reservations charged for memory already freed and can starve the native pool. In a component probe, allocating 1 MiB under task A and closing it without a context leaves A charged 1 MiB with zero Arrow bytes alive. Closing A's buffer under task B instead subtracts B's accounting. The size-only callback cannot recover ownership from the current task. Please bind the listener/accounting to the allocation owner and cover cross-thread release plus release after task completion. ########## spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala: ########## @@ -0,0 +1,225 @@ +/* + * 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.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.arrow.memory.AllocationListener +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager} + +import org.apache.comet.CometConf + +/** + * Reports JVM-side Arrow allocations to Spark's memory manager. + * + * `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 listener closes the reporting half of that gap. Every allocation is charged to a + * [[MemoryConsumer]] belonging to the task that made it, so the bytes appear in + * `TaskMemoryManager.showMemoryUsage` and are arbitrated against Spark's other off-heap + * consumers. + * + * It deliberately does not enforce. A short grant from Spark 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. Note that enforcement belongs in + * `onPreAllocation`, the only callback permitted to throw, and `onFailedAllocation`, not here. + * See [[https://github.com/apache/datafusion-comet/issues/5997]]. + * + * Buffers imported over the C Data Interface never reach this listener at all. They wrap memory + * the native side owns, so Comet imports them through `CometImportedArrowAllocator`, a child with + * no listener. Charging them here would double count bytes already reserved in Comet's native + * pool. Arrow notifies only the allocating allocator's own listener, which is what makes that + * separation work. + * + * Three cases are handled by doing nothing, each for a different reason: + * - No active task. Broadcast coalescing and the cached batch serializer can allocate from the + * driver or a non-task thread, where there is no task to charge. + * - On-heap mode. Comet's on-heap mode exists so the Spark SQL suite can run without off-heap + * memory configured; charging an off-heap consumer there would be wrong. + * - A buffer released after its allocating task has finished. The allocator is process-wide + * precisely because buffers can outlive the task that created them, so the task's reservation + * is dropped at task end and later releases are ignored rather than double-counted. + */ +class CometArrowAllocationListener extends AllocationListener { + + import CometArrowAllocationListener._ + + private val reservations = new ConcurrentHashMap[Long, TaskReservation]() + + override def onAllocation(size: Long): Unit = { + val reservation = reservationForCurrentTask() + if (reservation != null) { + reservation.allocated(size) + } + } + + override def onRelease(size: Long): Unit = { + val reservation = reservationForCurrentTask() + if (reservation != null) { + reservation.released(size) + } + } + + /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */ + private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = { + val reservation = reservations.get(taskAttemptId) + if (reservation == null) 0L else reservation.reservedBytes + } + + private[comet] def trackedTaskCount: Int = reservations.size() + + private def reservationForCurrentTask(): TaskReservation = { + // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast + // coalescing and the cached batch serializer all allocate with no task in scope. Reading the + // config before this would also mean re-reading `SparkEnv` on every allocation in a process + // that never has one. + val taskContext = TaskContext.get() + if (taskContext == null) return null + if (!accountingEnabled) return null + + val taskMemoryManager = taskContext.taskMemoryManager() + if (taskMemoryManager == null || + taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) { + return null + } + + val taskAttemptId = taskContext.taskAttemptId() + val existing = reservations.get(taskAttemptId) + if (existing != null) return existing + + // Deliberately not `computeIfAbsent`: `addTaskCompletionListener` runs the callback inline if + // the task has already completed, and that callback removes from this same map, which is a + // recursive update inside a mapping function. Registering outside the map operation avoids it. + val created = new TaskReservation(taskMemoryManager) + val previous = reservations.putIfAbsent(taskAttemptId, created) + if (previous != null) return previous + + taskContext.addTaskCompletionListener[Unit] { _ => + val finished = reservations.remove(taskAttemptId) + if (finished != null) { + finished.close() + } + } + created + } +} + +object CometArrowAllocationListener extends Logging { + + /** + * Batching granularity for reservations. Arrow allocates per buffer and + * `acquireExecutionMemory` takes an executor-wide lock, so the reservation is grown and shrunk + * in whole blocks and only block-crossing changes reach Spark. Deliberately not configurable: + * it trades lock chatter against reservation slack and has no plausible per-workload tuning. + */ + private val BLOCK_SIZE = 1024L * 1024L + + private val shortGrantLogged = new AtomicBoolean(false) + + /** + * Resolved once per JVM. The listener is attached to a `val` in a package object, so it is + * constructed on first touch of `CometArrowAllocator`, which can happen before any + * `SparkSession` exists and on executors where `SQLConf` does not carry Comet's settings. This + * is only read once a `TaskContext` exists, by which point an executor has a `SparkEnv`; the + * `Option` guard covers tests that install a task context without one. + */ + private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env => + env.conf.getBoolean( + CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key, + CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get) + } + + private def roundUpToBlock(bytes: Long): Long = + ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE + + private def warnOnShortGrant(requested: Long, granted: Long): Unit = { + if (shortGrantLogged.compareAndSet(false, true)) { + logWarning( + s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " + + "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " + + s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " + + "reporting these allocations to Spark.") + } + } + + /** One task's reservation against Spark's off-heap pool. */ + private class TaskReservation(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + + // Named `usedBytes` rather than `used` on purpose: `MemoryConsumer` already declares a + // `protected long used`, and a private field of that name narrows the inherited member, which + // the compiler rejects as weaker access privileges in overriding. + private var usedBytes: Long = 0L + private var reserved: Long = 0L + + /** Comet's native operators cannot be made to spill from here. See issue #5997. */ + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + + /** + * 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. + */ + override def getUsed: Long = synchronized(usedBytes) + + def reservedBytes: Long = synchronized(reserved) + + def allocated(size: Long): Unit = synchronized { + usedBytes += size + if (reserved < usedBytes) { + // Round up so `reserved` stays a block multiple and growth always leaves headroom. + // Requesting the bare deficit would land exactly on `usedBytes` for any buffer at or above + // the block size, sending the very next allocation straight back into Spark's lock. + val request = roundUpToBlock(usedBytes - reserved) + val granted = taskMemoryManager.acquireExecutionMemory(request, this) Review Comment: ### Correctness [P2] Please keep potentially throwing Spark acquisition out of Arrow's `onAllocation` callback. This call can invoke another consumer's spill method, and maintained Spark 3.5/4.0 propagate spill interruption or I/O failure as exceptions. Arrow 18.3.0 sets `success = true` and creates the buffer before invoking `onAllocation`, so an exception here escapes without returning or closing that buffer. With the unchanged listener, a 1 MiB pool occupied by a spill consumer that throws `InterruptedIOException` makes `root.buffer(1 MiB)` throw `RuntimeException` while the root still owns 1 MiB that the caller never received. Closing the allocator then reports that leak. Please move fallible work to a safe allocation boundary with rollback, and add a failing-spill/cancellation test that checks both Arrow live bytes and Spark reservations. ########## spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala: ########## @@ -0,0 +1,225 @@ +/* + * 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.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.arrow.memory.AllocationListener +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager} + +import org.apache.comet.CometConf + +/** + * Reports JVM-side Arrow allocations to Spark's memory manager. + * + * `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 listener closes the reporting half of that gap. Every allocation is charged to a + * [[MemoryConsumer]] belonging to the task that made it, so the bytes appear in + * `TaskMemoryManager.showMemoryUsage` and are arbitrated against Spark's other off-heap + * consumers. + * + * It deliberately does not enforce. A short grant from Spark 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. Note that enforcement belongs in + * `onPreAllocation`, the only callback permitted to throw, and `onFailedAllocation`, not here. + * See [[https://github.com/apache/datafusion-comet/issues/5997]]. + * + * Buffers imported over the C Data Interface never reach this listener at all. They wrap memory + * the native side owns, so Comet imports them through `CometImportedArrowAllocator`, a child with + * no listener. Charging them here would double count bytes already reserved in Comet's native + * pool. Arrow notifies only the allocating allocator's own listener, which is what makes that + * separation work. + * + * Three cases are handled by doing nothing, each for a different reason: + * - No active task. Broadcast coalescing and the cached batch serializer can allocate from the + * driver or a non-task thread, where there is no task to charge. + * - On-heap mode. Comet's on-heap mode exists so the Spark SQL suite can run without off-heap + * memory configured; charging an off-heap consumer there would be wrong. + * - A buffer released after its allocating task has finished. The allocator is process-wide + * precisely because buffers can outlive the task that created them, so the task's reservation + * is dropped at task end and later releases are ignored rather than double-counted. + */ +class CometArrowAllocationListener extends AllocationListener { + + import CometArrowAllocationListener._ + + private val reservations = new ConcurrentHashMap[Long, TaskReservation]() + + override def onAllocation(size: Long): Unit = { + val reservation = reservationForCurrentTask() + if (reservation != null) { + reservation.allocated(size) + } + } + + override def onRelease(size: Long): Unit = { + val reservation = reservationForCurrentTask() + if (reservation != null) { + reservation.released(size) + } + } + + /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */ + private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = { + val reservation = reservations.get(taskAttemptId) + if (reservation == null) 0L else reservation.reservedBytes + } + + private[comet] def trackedTaskCount: Int = reservations.size() + + private def reservationForCurrentTask(): TaskReservation = { + // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast + // coalescing and the cached batch serializer all allocate with no task in scope. Reading the + // config before this would also mean re-reading `SparkEnv` on every allocation in a process + // that never has one. + val taskContext = TaskContext.get() + if (taskContext == null) return null + if (!accountingEnabled) return null + + val taskMemoryManager = taskContext.taskMemoryManager() + if (taskMemoryManager == null || + taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) { + return null + } + + val taskAttemptId = taskContext.taskAttemptId() + val existing = reservations.get(taskAttemptId) + if (existing != null) return existing + + // Deliberately not `computeIfAbsent`: `addTaskCompletionListener` runs the callback inline if + // the task has already completed, and that callback removes from this same map, which is a + // recursive update inside a mapping function. Registering outside the map operation avoids it. + val created = new TaskReservation(taskMemoryManager) + val previous = reservations.putIfAbsent(taskAttemptId, created) + if (previous != null) return previous + + taskContext.addTaskCompletionListener[Unit] { _ => + val finished = reservations.remove(taskAttemptId) + if (finished != null) { + finished.close() + } + } + created + } +} + +object CometArrowAllocationListener extends Logging { + + /** + * Batching granularity for reservations. Arrow allocates per buffer and + * `acquireExecutionMemory` takes an executor-wide lock, so the reservation is grown and shrunk + * in whole blocks and only block-crossing changes reach Spark. Deliberately not configurable: + * it trades lock chatter against reservation slack and has no plausible per-workload tuning. + */ + private val BLOCK_SIZE = 1024L * 1024L + + private val shortGrantLogged = new AtomicBoolean(false) + + /** + * Resolved once per JVM. The listener is attached to a `val` in a package object, so it is + * constructed on first touch of `CometArrowAllocator`, which can happen before any + * `SparkSession` exists and on executors where `SQLConf` does not carry Comet's settings. This + * is only read once a `TaskContext` exists, by which point an executor has a `SparkEnv`; the + * `Option` guard covers tests that install a task context without one. + */ + private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env => + env.conf.getBoolean( + CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key, + CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get) + } + + private def roundUpToBlock(bytes: Long): Long = + ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE + + private def warnOnShortGrant(requested: Long, granted: Long): Unit = { + if (shortGrantLogged.compareAndSet(false, true)) { + logWarning( + s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " + + "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " + + s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " + + "reporting these allocations to Spark.") + } + } + + /** One task's reservation against Spark's off-heap pool. */ + private class TaskReservation(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + + // Named `usedBytes` rather than `used` on purpose: `MemoryConsumer` already declares a + // `protected long used`, and a private field of that name narrows the inherited member, which + // the compiler rejects as weaker access privileges in overriding. + private var usedBytes: Long = 0L + private var reserved: Long = 0L + + /** Comet's native operators cannot be made to spill from here. See issue #5997. */ + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + + /** + * 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. + */ + override def getUsed: Long = synchronized(usedBytes) Review Comment: ### Correctness [P1] Could `getUsed` expose a lock-free snapshot instead of taking the reservation monitor? `allocated` holds this monitor while entering `TaskMemoryManager.acquireExecutionMemory`, while Spark 3.5/4.0 hold the task-memory-manager monitor when calling consumers' `getUsed` during a short grant. A concurrent native reservation through `CometTaskMemoryManager` can therefore hold the Spark monitor and wait here while the Arrow allocation holds this monitor and waits for Spark. Comet's shared Tokio execution permits those requests to overlap. A two-thread component probe using the unchanged listener and the real Spark task memory manager reproduces the lock cycle. Please establish a consistent lock order and add a bounded concurrent-reservation regression test. -- 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]
