From 8c8ce915d92013e80196092e39bbc6c449d75fc8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 26 Sep 2026 14:18:01 -0600 Subject: [PATCH] feat: report JVM Arrow memory in the executor memory usage log The executor's periodic memory usage log reported only what Comet's native allocator had handed out and what its memory pools had reserved. Arrow memory that Comet holds on the JVM side appeared in neither, so the tuning guide's recipe for sizing the memory overhead had to leave it to the margin. Each line now also reports the total charged to CometArrowAllocator and the part of it charged to CometArrowImportAllocator, the figures tracing already reports as jvm_arrow_allocated and jvm_arrow_imported. The container warning counts the Arrow memory the JVM allocated itself. It is added to the native allocation before the pools' reservations are subtracted, because a native operator that holds on to a JVM batch reserves it: those bytes are in both the JVM figure and the reservations, but not in the native allocation. The sizing recipe uses the same definition. --- .../contributor-guide/memory_management.md | 36 ++++-- docs/source/contributor-guide/tracing.md | 1 + .../source/user-guide/latest/tuning/memory.md | 74 +++++++------ .../scala/org/apache/comet/CometConf.scala | 7 +- .../org/apache/comet/CometExecIterator.scala | 104 ++++++++++++------ .../CometExecIteratorLifecycleSuite.scala | 87 ++++++++++++--- 6 files changed, 216 insertions(+), 93 deletions(-) diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 9768698b0d8..f768baf366e 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -113,7 +113,17 @@ for both at once, and raises the pod's memory request by the same amount. (`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. These are real off-heap bytes in container RSS that neither Spark's `TaskMemoryManager` nor Comet's native memory pool sees. In practice the volume is modest, a batch at a time per stream, but there is no -ceiling and no backpressure. +ceiling and no backpressure. The executor's memory usage log reports the allocator's total, and +the part of it charged to the import allocator described below, so that the overhead can be sized +for it. + +Charging these buffers to Spark's off-heap pool instead was tried in +[#5998](https://github.com/apache/datafusion-comet/pull/5998) and dropped. Spark cannot make +Comet's native consumer release anything, and native operators fill a task's share before they +spill, so the JVM allocation that decodes their next input batch is refused first. Refusing it +fails the task where native would have spilled; not refusing it bounds nothing when the pool is +full. It needs native reclaim first; see +[#3873](https://github.com/apache/datafusion-comet/issues/3873). One further child, `CometArrowImportAllocator` (`comet-ffi-imports`), is what the Arrow C Data Interface import path allocates from, so that tracing can report those charges apart from the rest @@ -430,7 +440,8 @@ gap is workload-dependent. The margin that covers it has to come from one; see [Where Comet's budget comes from](#where-comets-budget-comes-from). To measure the gap on a real workload, read the executor's periodic memory usage log, which -reports the bytes Rust's allocator has handed out next to the pools' reservations; see +reports the bytes Rust's allocator has handed out next to the pools' reservations, and the Arrow +memory Comet holds on the JVM side; see [Sizing the Overhead from the Memory Usage Log][memory-usage-log]. For a view per event rather than per interval, enable tracing and compare `native_allocated` against `comet_memory_reserved_total`; see [Tracing](tracing.md#analyzing-memory-usage). @@ -519,7 +530,8 @@ much they matter: its container is still stopped only by the kill. - **The memory overhead is sized by hand.** The gap has to fit in `spark.executor.memoryOverhead`, and the memory usage log measures it, but nothing sizes the overhead from it. -- **`CometArrowAllocator` is unbounded** and participates in no budget. +- **`CometArrowAllocator` is unbounded** and participates in no budget. The memory usage log + reports it, but nothing bounds it. - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be resident on either side with no reservation covering it, because reservations are made and withdrawn by individual operators while the bytes outlive them. @@ -531,14 +543,14 @@ much they matter: ## Debugging memory issues -| Tool | What it gives you | -| ------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- | -| `spark.comet.debug.memory=true` | `LoggingMemoryPool` logs every register/grow/shrink with the consumer name | -| `spark.comet.explain.native.enabled=true` | Native plan with per-operator metrics, including spill counts | -| [Memory usage log](../user-guide/latest/tuning/memory.md#sizing-the-overhead-from-the-memory-usage-log) | Executor-wide native allocation vs pool reservations, logged every 10 seconds by default | -| [Tracing](tracing.md#analyzing-memory-usage) | `native_allocated` vs `comet_memory_reserved_total` per event; the accounting gap over time | -| `TrackConsumersPool` | Names the top 10 consumers in `ResourcesExhausted` messages (always on) | -| [`thresher`](https://github.com/cetra3/thresher) | Third-party crate that dumps a jemalloc heap profile at a threshold | +| Tool | What it gives you | +| ------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------- | +| `spark.comet.debug.memory=true` | `LoggingMemoryPool` logs every register/grow/shrink with the consumer name | +| `spark.comet.explain.native.enabled=true` | Native plan with per-operator metrics, including spill counts | +| [Memory usage log](../user-guide/latest/tuning/memory.md#sizing-the-overhead-from-the-memory-usage-log) | Executor-wide native allocation vs pool reservations, and JVM Arrow memory, logged every 10 seconds by default | +| [Tracing](tracing.md#analyzing-memory-usage) | `native_allocated` vs `comet_memory_reserved_total` per event; the accounting gap over time | +| `TrackConsumersPool` | Names the top 10 consumers in `ResourcesExhausted` messages (always on) | +| [`thresher`](https://github.com/cetra3/thresher) | Third-party crate that dumps a jemalloc heap profile at a threshold | A checklist for triaging an executor OOM kill: @@ -550,7 +562,7 @@ A checklist for triaging an executor OOM kill: 2. Compare `allocated` against `reserved` in the executor's `Comet native memory usage` log lines leading up to the kill, or `native_allocated` against `comet_memory_reserved_total` in a trace. A large excess points at undeclared native allocations; a small excess points at the budget - simply being too small, or at the JVM side. + simply being too small, or at the JVM side, which the same lines report as `JVM Arrow allocated`. 3. Check `spark.comet.batchSize` against the schema width. Peak memory scales with `batch_size * columns`, and wide or deeply nested schemas amplify it. 4. Check whether the operators involved can spill at all. `ShuffledHashJoin` cannot, so diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 1b3a6bcb9e3..8acdbe841ae 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -134,6 +134,7 @@ Arrow memory on the JVM side is reported separately, because it is off-heap and `jvm_heap_used`. Comet imports batches from native over the Arrow C Data Interface, and Arrow charges a buffer to whichever allocator owns it, so those imports are taken against a dedicated child allocator and reported as `jvm_arrow_imported`, within the `jvm_arrow_allocated` total. +The executor's periodic memory usage log reports the same two figures. Both are allocator charges. They report what each allocator is accountable for, not where the bytes were allocated, and their difference is not a bound on the Arrow memory the JVM allocated itself. diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 0fb4288abac..95bd9178093 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -149,48 +149,58 @@ While Comet native plans are running, each executor logs its native memory usage one line every 10 seconds for the whole executor: ``` -Comet native memory usage: allocated 5412.3 MiB, reserved 3890.0 MiB (16 native plans, 8 memory pools) +Comet native memory usage: allocated 5412.3 MiB, reserved 3890.0 MiB (16 native plans, 8 memory pools); JVM Arrow allocated 310.4 MiB, 96.2 MiB of it imported from native ``` - `allocated` is the memory that Comet's native code has allocated and not yet freed, whether or not a pool tracks it. - `reserved` is the part that Comet's memory pools track. It is charged against `spark.memory.offHeap.size`, so the container already has room for it. - -The difference between the two, `allocated - reserved`, is Comet's untracked native memory. It is -the part of Comet's footprint that has to fit in `spark.executor.memoryOverhead`, alongside the -JVM's own non-heap memory. To size the overhead from it: - -1. Run a representative workload and find the line with the largest difference in each executor's - log. Take both figures from the same line: they are sampled together, and figures from different - lines describe different moments. Setting `spark.comet.memory.logInterval=1s` for this run makes a - short-lived peak less likely to fall between samples. +- `JVM Arrow allocated` is the Arrow memory Comet holds on the JVM side, such as batches read from + Comet's in-memory cache, broadcast data, and batches exchanged with native code or Python workers. + The part imported from native was allocated by Comet's native code, so `allocated` already counts + it, and the rest the JVM allocated itself. A pool reserves the JVM's part only while a native + operator holds on to one of its batches, such as a sort buffering its input. The split follows + which allocator a buffer is charged to rather than where it was allocated, so treat the JVM's + part as an estimate. + +Comet's untracked memory is what Comet holds outside the JVM heap that no pool reserves: `allocated`, +plus the JVM Arrow figure less the part imported from native, minus `reserved`. It is the part of +Comet's footprint that has to fit in `spark.executor.memoryOverhead`, alongside the JVM's own +non-heap memory. To size the overhead from it: + +1. Run a representative workload and find the line with the most untracked memory in each + executor's log. Take all the figures from the same line: they are sampled together, and figures + from different lines describe different moments. Setting `spark.comet.memory.logInterval=1s` for + this run makes a short-lived peak less likely to fall between samples. 2. Start from the overhead the executors had before Comet was enabled, which covers the JVM's own - non-heap memory, and add the largest difference seen on any executor. -3. Add a margin on top. The log can miss the true peak between samples, and neither figure includes - the allocator's fragmentation and retained pages, memory allocated by native C libraries such as - zstd, or Comet's Arrow buffers on the JVM side. + non-heap memory, and add the most untracked memory seen on any executor. +3. Add a margin on top. The log can miss the true peak between samples, and none of the figures + includes the allocator's fragmentation and retained pages, or memory allocated by native C + libraries such as zstd. -For example, a 16 GiB executor derives an overhead of 1638 MiB. If the largest difference in its -log is the 1522.3 MiB in the line above, the overhead needs to be at least 1638 + 1523 = 3161 MiB -before any margin, so `spark.executor.memoryOverhead=4g` would be a reasonable setting. +For example, a 16 GiB executor derives an overhead of 1638 MiB. If the line above has the most +untracked memory in its log, that is 5412.3 + (310.4 - 96.2) - 3890.0 = 1736.5 MiB. The overhead +then needs to be at least 1638 + 1737 = 3375 MiB before any margin, so +`spark.executor.memoryOverhead=4g` would be a reasonable setting. The executor also logs a warning when its native memory looks larger than its container allows: -when the difference, plus everything in use in Spark's off-heap memory pool (which includes Comet's -reservations), exceeds `spark.memory.offHeap.size` plus the memory overhead. This counts the part of -the off-heap pool that nothing has acquired at that moment, which untracked memory can occupy until -Spark hands it out, so a quiet log is not a sign that the overhead is large enough: size it from the -largest difference as described above. The overhead also has to hold the JVM's own non-heap memory, -so by the time the warning appears the executor has likely outgrown its container. It warns the first time this -happens, and again each time it happens after dropping back below. The overhead it uses is -`spark.executor.memoryOverhead` if set, otherwise `spark.executor.memoryOverheadFactor` of -`spark.executor.memory` with a minimum of `spark.executor.minMemoryOverhead`, as Spark sizes the -default container. There is no warning in local mode. - -Look more closely before raising the overhead if the difference keeps growing through a run rather -than levelling off: native memory that is not being released will exhaust any overhead eventually. -The executor logs one more line after its last native plan finishes, and an `allocated` figure there -that grows from one query to the next points the same way. +when the untracked memory, plus everything in use in Spark's off-heap memory pool (which includes +Comet's reservations), exceeds `spark.memory.offHeap.size` plus the memory overhead. This counts the +part of the off-heap pool that nothing has acquired at that moment, which untracked memory can +occupy until Spark hands it out, so a quiet log is not a sign that the overhead is large enough: +size it from the most untracked memory as described above. The overhead also has to hold the JVM's +own non-heap memory, so by the time the warning appears the executor has likely outgrown its +container. It warns the first time this happens, and again each time it happens after dropping back +below. The overhead it uses is `spark.executor.memoryOverhead` if set, otherwise +`spark.executor.memoryOverheadFactor` of `spark.executor.memory` with a minimum of +`spark.executor.minMemoryOverhead`, as Spark sizes the default container. There is no warning in +local mode. + +Look more closely before raising the overhead if the untracked memory keeps growing through a run +rather than levelling off: memory that is not being released will exhaust any overhead eventually. +The executor logs one more line after its last native plan finishes, and an `allocated` or +`JVM Arrow allocated` figure there that grows from one query to the next points the same way. `spark.comet.memory.logInterval` is read when an executor starts its first Comet native plan, so set it when the application is submitted. Set it to `0` to turn the log off. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index eeb5cf8045d..12a84af8cf5 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -362,9 +362,10 @@ object CometConf extends ShimCometConf { .category(CATEGORY_TUNING) .doc( "How often each executor logs its native memory usage at INFO level while Comet native " + - "plans are running: the bytes the native allocator has handed out, and the bytes " + - "reserved in Comet's memory pools. The difference is native memory that the pools are " + - "not accounting for. The executor logs one line per interval however many tasks are " + + "plans are running: the bytes the native allocator has handed out, the bytes reserved " + + "in Comet's memory pools, and the Arrow memory Comet holds on the JVM side. The " + + "difference between the first two is native memory that the pools are not accounting " + + "for. The executor logs one line per interval however many tasks are " + "running, and one more after the last plan finishes. It logs a warning when the " + "native memory looks larger than the executor's container allows. This is an executor " + "setting, read when an executor starts its first Comet native plan, so it must be set " + diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 5b70fbaaf24..a56408cb278 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -27,6 +27,7 @@ import java.util.concurrent.atomic.AtomicBoolean import scala.util.control.NonFatal import org.apache.arrow.c.ArrowArrayStream +import org.apache.arrow.memory.BufferAllocator import org.apache.hadoop.conf.Configuration import org.apache.spark._ import org.apache.spark.broadcast.Broadcast @@ -384,12 +385,13 @@ object CometExecIterator extends Logging { * Starts the executor's native memory usage log when the first native plan is created, unless * `spark.comet.memory.logInterval` is 0. * - * Both figures the log reports, the bytes the native allocator has handed out and the bytes - * reserved in Comet's memory pools, are executor-wide, so one daemon thread logs one line per - * interval for the whole executor, however many tasks are running. It runs on a timer rather - * than between batches, because a plan can spend its whole run inside one `executePlan` call: a - * plan rooted at a native shuffle writer consumes all of its input before it returns, and a - * plan fed directly by native scans parks the task thread until its next batch is ready. + * Every figure the log reports is executor-wide: the bytes the native allocator has handed out, + * the bytes reserved in Comet's memory pools, and the Arrow memory Comet holds on the JVM side. + * So one daemon thread logs one line per interval for the whole executor, however many tasks + * are running. It runs on a timer rather than between batches, because a plan can spend its + * whole run inside one `executePlan` call: a plan rooted at a native shuffle writer consumes + * all of its input before it returns, and a plan fed directly by native scans parks the task + * thread until its next batch is ready. */ private def startMemoryUsageLog(): Unit = { if (memoryUsageLogStarted.compareAndSet(false, true)) { @@ -505,10 +507,11 @@ object CometExecIterator extends Logging { private def logMemoryUsage(nativeLib: Native, limitBytes: Option[Long]): Unit = { try { val usage = nativeLib.getMemoryUsage() - memoryUsageMessage(usage, plansAtLastMemoryUsageLog).foreach(logInfo(_)) + val jvmArrow = JvmArrowMemory.current() + memoryUsageMessage(usage, jvmArrow, plansAtLastMemoryUsageLog).foreach(logInfo(_)) plansAtLastMemoryUsageLog = usage(3) val warning = limitBytes.flatMap( - nativeMemoryLimitWarning(usage, CometTaskMemoryManager.sparkOffHeapUsed(), _)) + nativeMemoryLimitWarning(usage, jvmArrow, CometTaskMemoryManager.sparkOffHeapUsed(), _)) // Warn when the footprint first exceeds the limit, not at every interval while it stays // there: the INFO line above keeps reporting it. if (!limitExceededAtLastLog) { @@ -525,19 +528,55 @@ object CometExecIterator extends Logging { } /** - * The memory usage log line for `usage`, as returned by [[Native.getMemoryUsage]], or None to - * stay quiet. The log reports while native plans are running, and once more after the last of - * them finishes, so that allocation that outlives them is visible, then waits for plans to run - * again. + * The Arrow memory Comet holds on the JVM side, read from its allocators' own accounting: + * `allocated` for the whole `CometArrowAllocator` tree, and `imported` for the part of it + * charged to `CometArrowImportAllocator`, which the Arrow C Data Interface import path + * allocates from. These are the figures tracing reports as `jvm_arrow_allocated` and + * `jvm_arrow_imported`. */ - def memoryUsageMessage(usage: Array[Long], plansAtLastLog: Long): Option[String] = { + case class JvmArrowMemory(allocated: Long, imported: Long) { + + /** + * An estimate of the Arrow memory the JVM allocated itself, which none of the native figures + * include. Imported buffers are left out because native code allocated them, so the native + * allocator already counts them. Both figures are allocator charges rather than measures of + * where the bytes came from, and ownership transfers move charges between allocators, so the + * split is not exact; see `CometArrowImportAllocator`. + */ + def allocatedByJvm: Long = math.max(allocated - imported, 0L) + } + + object JvmArrowMemory { + + def current(): JvmArrowMemory = of(CometArrowAllocator, CometArrowImportAllocator) + + def of(root: BufferAllocator, imports: BufferAllocator): JvmArrowMemory = { + // The import allocator is read first, so an import that lands between the two reads is + // counted in the total but not in the imported part. That overstates the JVM's own + // allocation rather than understating it. + val imported = imports.getAllocatedMemory + JvmArrowMemory(root.getAllocatedMemory, imported) + } + } + + /** + * The memory usage log line for `usage`, as returned by [[Native.getMemoryUsage]], and + * `jvmArrow`, or None to stay quiet. The log reports while native plans are running, and once + * more after the last of them finishes, so that allocation that outlives them is visible, then + * waits for plans to run again. + */ + def memoryUsageMessage( + usage: Array[Long], + jvmArrow: JvmArrowMemory, + plansAtLastLog: Long): Option[String] = { val (allocated, reserved, pools, plans) = (usage(0), usage(1), usage(2), usage(3)) if (plans == 0 && plansAtLastLog == 0) { None } else { Some( s"Comet native memory usage: allocated ${toMiB(allocated)}, reserved " + - s"${toMiB(reserved)} ($plans native plans, $pools memory pools)") + s"${toMiB(reserved)} ($plans native plans, $pools memory pools); JVM Arrow allocated " + + s"${toMiB(jvmArrow.allocated)}, ${toMiB(jvmArrow.imported)} of it imported from native") } } @@ -545,30 +584,33 @@ object CometExecIterator extends Logging { * A warning if the executor's native footprint exceeds `limitBytes`, the container's memory * outside the JVM heap; see [[nativeMemoryLimit]]. * - * The footprint is the native memory Comet's pools do not track, `allocated - reserved`, plus - * `sparkOffHeapUsed`, everything in use in Spark's off-heap pool, which includes Comet's - * reservations as well as Spark's own off-heap execution and storage memory. Comparing the sum - * rather than the untracked part against the overhead alone counts the part of - * `spark.memory.offHeap.size` that nothing has acquired at that moment, which untracked memory - * can occupy until Spark hands it out. The limit also has to hold the JVM's own non-heap - * memory, so by the time the footprint exceeds it the executor has likely outgrown its - * container. + * The footprint is the memory Comet holds outside the JVM heap that its pools do not track, + * `allocated + jvmArrow.allocatedByJvm - reserved`, plus `sparkOffHeapUsed`, everything in use + * in Spark's off-heap pool, which includes Comet's reservations as well as Spark's own off-heap + * execution and storage memory. The JVM's Arrow memory is added before the reservations are + * subtracted, because a native operator that holds on to a batch the JVM allocated, such as a + * sort buffering its input, reserves it: those bytes are in `reserved` and in the JVM figure + * but not in `allocated`, and are counted once. Comparing the sum rather than the untracked + * part against the overhead alone counts the part of `spark.memory.offHeap.size` that nothing + * has acquired at that moment, which untracked memory can occupy until Spark hands it out. The + * limit also has to hold the JVM's own non-heap memory, so by the time the footprint exceeds it + * the executor has likely outgrown its container. */ def nativeMemoryLimitWarning( usage: Array[Long], + jvmArrow: JvmArrowMemory, sparkOffHeapUsed: Long, limitBytes: Long): Option[String] = { - val untracked = math.max(usage(0) - usage(1), 0L) + val untracked = math.max(usage(0) + jvmArrow.allocatedByJvm - usage(1), 0L) val footprint = untracked + sparkOffHeapUsed if (footprint > limitBytes) { - Some( - s"Comet native memory not tracked by any memory pool (${toMiB(untracked)}) plus " + - s"Spark's off-heap memory in use (${toMiB(sparkOffHeapUsed)}, including Comet's " + - s"reservations) is ${toMiB(footprint)}, more than the ${toMiB(limitBytes)} the " + - "executor's container has outside the JVM heap (spark.memory.offHeap.size plus the " + - "memory overhead), which also has to hold the JVM's own non-heap memory. The cluster " + - "manager may kill this executor for exceeding its container limit. Raise " + - s"spark.executor.memoryOverhead. ${CometConf.TUNING_GUIDE}.") + Some(s"Comet memory not tracked by any memory pool (${toMiB(untracked)}, native and JVM " + + s"Arrow) plus Spark's off-heap memory in use (${toMiB(sparkOffHeapUsed)}, including " + + s"Comet's reservations) is ${toMiB(footprint)}, more than the ${toMiB(limitBytes)} the " + + "executor's container has outside the JVM heap (spark.memory.offHeap.size plus the " + + "memory overhead), which also has to hold the JVM's own non-heap memory. The cluster " + + "manager may kill this executor for exceeding its container limit. Raise " + + s"spark.executor.memoryOverhead. ${CometConf.TUNING_GUIDE}.") } else { None } diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index d16e7fe504f..65d26bbede3 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -24,6 +24,7 @@ import java.lang.ref.WeakReference import java.util.Properties import java.util.concurrent.atomic.AtomicBoolean +import org.apache.arrow.memory.RootAllocator import org.apache.spark.executor.TaskMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.sql.CometTestBase @@ -251,41 +252,97 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { } test("the memory usage log reports while plans run and once after the last one finishes") { + import CometExecIterator.JvmArrowMemory val mib = 1024L * 1024 val busy = Array(300 * mib, 100 * mib, 2L, 3L) + val jvmArrow = JvmArrowMemory(allocated = 40 * mib, imported = 10 * mib) assert( CometExecIterator - .memoryUsageMessage(busy, plansAtLastLog = 0) - .contains("Comet native memory usage: allocated 300.0 MiB, reserved 100.0 MiB " + - "(3 native plans, 2 memory pools)")) + .memoryUsageMessage(busy, jvmArrow, plansAtLastLog = 0) + .contains( + "Comet native memory usage: allocated 300.0 MiB, reserved 100.0 MiB " + + "(3 native plans, 2 memory pools); JVM Arrow allocated 40.0 MiB, 10.0 MiB of it " + + "imported from native")) // The line after the last plan finishes shows the allocation the plans left behind. val idle = Array(20 * mib, 0L, 0L, 0L) + val noArrow = JvmArrowMemory(0L, 0L) assert( CometExecIterator - .memoryUsageMessage(idle, plansAtLastLog = 3) + .memoryUsageMessage(idle, noArrow, plansAtLastLog = 3) .exists(_.contains("allocated 20.0 MiB, reserved 0.0 MiB (0 native plans"))) - assert(CometExecIterator.memoryUsageMessage(idle, plansAtLastLog = 0).isEmpty) + assert(CometExecIterator.memoryUsageMessage(idle, noArrow, plansAtLastLog = 0).isEmpty) + } + + test("the memory usage log reads JVM Arrow memory from the allocators, imports apart") { + import CometExecIterator.JvmArrowMemory + val root = new RootAllocator(Long.MaxValue) + try { + val imports = root.newChildAllocator("imports", 0, Long.MaxValue) + val others = root.newChildAllocator("others", 0, Long.MaxValue) + val owned = others.buffer(1024 * 1024) + val imported = imports.buffer(256 * 1024) + try { + val memory = JvmArrowMemory.of(root, imports) + assert(memory == JvmArrowMemory(allocated = 1280 * 1024, imported = 256 * 1024)) + assert(memory.allocatedByJvm == 1024 * 1024) + } finally { + imported.close() + owned.close() + imports.close() + others.close() + } + } finally { + root.close() + } + // The two figures are read one after the other, so an import can land in between. + assert(JvmArrowMemory(allocated = 10L, imported = 20L).allocatedByJvm == 0L) } test("the memory usage log warns when the native footprint exceeds the container") { - import CometExecIterator.nativeMemoryLimitWarning + import CometExecIterator.{nativeMemoryLimitWarning, JvmArrowMemory} val mib = 1024L * 1024 // A 4 GiB off-heap pool with a 1 GiB overhead, and a pool running at 0.8 of the off-heap size. val limit = 5120 * mib - val reserved = 3000 * mib - // 1500 MiB untracked is more than the overhead, but fits in what the pool left free, since - // Spark's off-heap pool holds only the reservation. - assert(nativeMemoryLimitWarning(Array(4500 * mib, reserved, 1L, 1L), reserved, limit).isEmpty) + val noArrow = JvmArrowMemory(0L, 0L) + // Spark's off-heap pool holds only Comet's reservation unless told otherwise. + def warning( + allocated: Long, + reserved: Long = 3000 * mib, + jvmArrow: JvmArrowMemory = noArrow, + sparkOffHeapUsed: Option[Long] = None): Option[String] = + nativeMemoryLimitWarning( + Array(allocated, reserved, 1L, 1L), + jvmArrow, + sparkOffHeapUsed.getOrElse(reserved), + limit) + + // 1500 MiB untracked is more than the overhead, but fits in what the pool left free. + assert(warning(4500 * mib).isEmpty) // 2500 MiB untracked does not fit: 2500 + 3000 = 5500 MiB. - val warning = nativeMemoryLimitWarning(Array(5500 * mib, reserved, 1L, 1L), reserved, limit) - assert(warning.exists(_.contains("(2500.0 MiB) plus Spark's off-heap memory in use (3000.0"))) - assert(warning.exists(_.contains("is 5500.0 MiB, more than the 5120.0 MiB"))) + val native = warning(5500 * mib) + assert(native.exists(_.contains("(2500.0 MiB, native and JVM Arrow) plus Spark's off-heap"))) + assert( + native.exists(_.contains("memory in use (3000.0 MiB, including Comet's reservations)"))) + assert(native.exists(_.contains("is 5500.0 MiB, more than the 5120.0 MiB"))) // Spark's own off-heap use counts against the same limit. + assert(warning(4500 * mib, sparkOffHeapUsed = Some(4000 * mib)).isDefined) + // So does Arrow memory the JVM allocated itself: 1500 + 700 + 3000 = 5200 MiB. + val jvm = warning(4500 * mib, jvmArrow = JvmArrowMemory(900 * mib, imported = 200 * mib)) + assert(jvm.exists(_.contains("(2200.0 MiB, native and JVM Arrow)"))) + assert(jvm.exists(_.contains("is 5200.0 MiB, more than the 5120.0 MiB"))) + // Imported buffers were allocated by native code, so the allocation already counts them. + assert( + warning(4500 * mib, jvmArrow = JvmArrowMemory(900 * mib, imported = 900 * mib)).isEmpty) + // A batch the JVM allocated and a native operator holds on to is reserved as well, so it + // counts once: 1000 + 2500 - 3500 leaves nothing untracked, and 3500 MiB in all fits. assert( - nativeMemoryLimitWarning(Array(4500 * mib, reserved, 1L, 1L), 4000 * mib, limit).isDefined) + warning( + 1000 * mib, + reserved = 3500 * mib, + jvmArrow = JvmArrowMemory(2500 * mib, 0L)).isEmpty) // Reservations can exceed the allocation, since operators reserve before they allocate. - assert(nativeMemoryLimitWarning(Array(100 * mib, reserved, 1L, 1L), reserved, limit).isEmpty) + assert(warning(100 * mib).isEmpty) } test("the memory pool limit reads a bare off-heap size as bytes, as Spark does") {