andygrove commented on code in PR #5998: URL: https://github.com/apache/datafusion-comet/pull/5998#discussion_r4040978448
########## 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: Yes, and it does now. `getUsed` reads an `AtomicLong` and a volatile flag, and `spill` returns `0` without touching anything, so neither waits on the reservation monitor. That matches what `CometTaskMemoryManager.NativeMemoryConsumer` already does, and for the same reason. The lock order is now stated on the class: the reservation monitor may be held across `acquireExecutionMemory`, so nothing Spark can call while holding the `TaskMemoryManager` monitor is allowed to need it. The regression test is deterministic rather than a stress loop. It holds the listener's monitor on the test thread and calls `getUsed` and `spill` from another one with a timeout, which fails outright if either becomes synchronized again. There is a bounded two-thread Arrow/native reservation test alongside it, but the monitor probe is the one that actually pins the property. ########## 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: They can, but not by looking harder at the callback, since Arrow gives it nothing but a size. The fix is to stop asking the thread and make the allocator the owner: `CometTaskArrowAllocator.forCurrentTask()` now hands each task a child of the root carrying its own listener instance, so `onAllocation` and `onRelease` both land on the allocating task no matter which thread drops the last reference. That is the same shape Gluten uses. `BaseAllocator.newChildAllocator` passes the listener down, so the Python runner and `CometNativeArrowSource`, which cut their own children, are covered too. Falling out of that: the root no longer carries a listener at all, so FFI imports are unaccounted simply by using it, and `CometImportedArrowAllocator` is gone. Off a task and in on-heap mode the accessor returns the root, so those paths behave exactly as they do on main. The awkward part is lifetime. The process-wide allocator exists because buffers can outlive their task, and Arrow treats closing a non-empty allocator as a leak, so at task completion the reservation is dropped and the allocator is closed only if it is drained; otherwise it is parked and reaped by a later task. It cannot just be left open, because `BaseAllocator` keeps every child in a map until it closes. Tests cover a release with no task context, a release under a different task, completion driven through `markTaskCompleted` rather than `cleanUpAllAllocatedMemory`, and a buffer outliving its task and being reaped afterwards. ########## 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: Agreed, and this is worse than "should not throw": Arrow's `AllocationListener` javadoc says an exception cannot be thrown from `onAllocation` or `onRelease`, and `BaseAllocator.buffer` sets `success = true` before the call, so a throw skips the `finally { releaseBytes }` and the buffer is neither returned nor freed. Rather than move the acquisition to a different boundary, the listener now simply never lets anything out. Both callbacks wrap the memory-manager call and log once. `SparkOutOfMemoryError` is caught explicitly because it extends `OutOfMemoryError` and so slips past `NonFatal`, which is exactly the `IOException`-from-spill path you describe. Two tests, one for each way `trySpillAndAcquire` fails: a consumer holding the only block that throws `IOException`, and one that throws `InterruptedIOException`. Both assert the Arrow allocation succeeds, that the allocator reports the buffer's bytes while it is open and zero after closing it, and that nothing was reserved, which is what proves the throw really happened rather than the acquisition quietly succeeding and making the test vacuous. -- 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]
