diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index f45fb9beaf..c1f53b138c 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -132,7 +132,7 @@ JVM Heap: Native Memory: │ ColumnarBatch │ │ FFI_ArrowArray │ │ ┌──────────────┐ │ │ ┌──────────────┐ │ │ │ ArrowBuf │─┼──────────────>│ │ buffers[0] │ │ -│ │ (off-heap) │ │ │ │ (pointer) │ │ +│ │ (handle) │ │ │ │ (pointer) │ │ │ └──────────────┘ │ │ └──────────────┘ │ └──────────────────┘ └──────────────────┘ │ │ @@ -253,7 +253,7 @@ pub extern "system" fn Java_..._exportVector( ### Wrapper Object Lifecycle (Native → JVM) ``` -Time Native Memory JVM Heap Off-heap/Native +Time Native Memory JVM Heap Data location ──────────────────────────────────────────────────────────────────────── t0 RecordBatch produced - Data in native in DataFusion diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index c252afedc4..91199dec35 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -53,10 +53,10 @@ calls `System.exit(SparkExitCode.OOM)` (exit code 52). Real heap exhaustion ther the whole executor, not one task. An executor loss on its own does not tell you which budget was exceeded; the exit code does. -Comet's difficulty is that its allocations are made by Rust code, so they are invisible to the JVM -heap and to Spark's own off-heap accounting, yet they land squarely in container RSS. Comet -therefore maintains its own budget that is meant to shadow the physical one, and the accuracy of -that shadow is the central problem this page is about. +Comet's difficulty is that its allocations are made by Rust code, so no JVM allocator produces them +and no JVM metric measures them, yet they land squarely in container RSS. Comet therefore maintains +its own budget that is meant to shadow the physical one, and declares it to Spark so that the two +compete for a single number. The accuracy of that shadow is the central problem this page is about. ## Who allocates what @@ -68,11 +68,32 @@ page: | --------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- | | Spark execution + storage (on-heap) | JVM heap | `spark.executor.memory` and the unified memory manager | Yes | | Spark Tungsten (off-heap) | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | -| Comet native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | No | +| Comet native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | Reservations only | | Comet JVM Arrow (`CometArrowAllocator`) | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | | Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | -Two observations follow. +Several observations follow. + +**"Off-heap" and "native heap" are not the same thing.** Both sit outside the JVM heap and both +count toward container RSS, which is what makes them easy to conflate, but they differ in who +allocates the bytes and in who is able to see them. + +- **Off-heap** is allocated by JVM code: `sun.misc.Unsafe` for Spark's Tungsten pages, and Arrow's + Unsafe-backed allocator for Comet's Java Arrow buffers. A JVM-side allocator makes the call, so a + JVM-side consumer is in a position to report it, and `spark.memory.offHeap.size` together with + `TaskMemoryManager` arbitrates it. +- **Native heap** is allocated by Rust through its global allocator, in the same process. No JVM + code makes the call, so no JVM-side allocator or metric ever measures these bytes. The only layer + that sees any of them is Comet's own memory pool, and then only the portion that operators + explicitly reserve. That portion is still charged to Spark, as the next paragraph describes; what + is never reserved is measured by nothing and budgeted by nobody. + +**The budget is shared even though the memory is not.** `memory_limit` is derived from +`spark.memory.offHeap.size` (see [Where Comet's budget comes from](#where-comets-budget-comes-from)), +and in both off-heap pool types every native reservation is forwarded to Spark's off-heap execution +pool over JNI. A native allocation therefore consumes the same accounting budget as a Tungsten page +while occupying entirely different memory. Raising `spark.memory.offHeap.size` raises the ceiling +for both at once, and raises the pod's memory request by the same amount. **Comet's JVM-side Arrow allocator is unbounded and accounted by nobody.** `CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single process-wide @@ -87,6 +108,78 @@ returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing f `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the same task like any other allocation. +Which allocator each call site uses, and who ends up charged for the bytes: + +```mermaid +flowchart LR + subgraph SITES["JVM Arrow allocation sites"] + NU["NativeUtil
FFI structs, imports, exports"] + UDF["CometUdfBridge
JVM UDF inputs and result"] + CGO["CometBatchKernelCodegenOutput
codegen UDF output"] + SR["StreamReader
shuffle and IPC reads"] + NAS["CometNativeArrowSource
stream and readerBatchIter"] + CACHE["ArrowCachedBatchSerializer"] + PY["CometArrowPythonRunnerBase"] + BC["Utils broadcast-coalesce"] + end + + ROOT["CometArrowAllocator
RootAllocator, no limit
no allocation listener"] + SHUF["Comet JVM shuffle pages
CometUnifiedShuffleMemoryAllocator"] + NPOOL["Comet native memory pool
declared reservations only"] + TMM["Spark off-heap execution pool
TaskMemoryManager"] + NOBODY["accounted by nobody"] + + NU --> ROOT + UDF --> ROOT + CGO --> ROOT + SR --> ROOT + NAS --> ROOT + CACHE --> ROOT + PY --> ROOT + BC --> ROOT + ROOT --> NOBODY + SHUF --> TMM + NPOOL -->|"CometTaskMemoryManager over JNI"| TMM +``` + +### Constraints on a Comet memory consumer + +`CometTaskMemoryManager` is the one place where Comet code acts as a Spark `MemoryConsumer`, and the +rules below are why it is written the way it is. Each is easy to break by accident. + +**`getUsed` and `spill` must stay lock-free.** Spark calls both while already holding the +`TaskMemoryManager` monitor, which is why `NativeMemoryConsumer.getUsed` reads an `AtomicLong` and +`spill` returns without touching guarded state. A consumer that took a monitor of its own in either +method would risk a deadlock: a native reservation arriving over JNI on a Comet Tokio worker holds +Spark's monitor and would wait for the consumer's, while whatever held the consumer's waited for +Spark's. For the same reason, nothing reachable from a `spill` callback may allocate memory that +routes back through the same consumer. + +**`scala.util.control.NonFatal` does not contain an acquisition.** `acquireExecutionMemory` fails in +three ways and `NonFatal` catches only the first. It runs other consumers' `spill`, where +`TaskMemoryManager` turns an interrupted spill into a `RuntimeException` and an `IOException` into a +`SparkOutOfMemoryError`, which extends `OutOfMemoryError` and is therefore an `Error`. Separately, +`ExecutionMemoryPool.acquireMemory` parks in `lock.wait()` when a task is below its fair share, so +killing a task raises a plain `InterruptedException`, which `NonFatal` excludes by name. Comet +reaches this method from native over JNI, so whatever escapes crosses the JNI boundary. + +**Spark exposes no per-consumer usage figure.** `TaskMemoryManager.getMemoryConsumptionForThisTask` +is task-wide, and a consumer that reaches `acquireExecutionMemory` directly rather than through +`MemoryConsumer.acquireMemory` keeps an inherited `used` of zero. That is why `NativeMemoryConsumer` +overrides `getUsed` to report Comet's own tally instead: without it, Spark's spill-victim ordering +and `showMemoryUsage` would believe the consumer held nothing. + +**A partial grant can be stranded.** `acquireExecutionMemory` takes its first grant from the pool +before asking other consumers to spill, so when a spill throws, the task has been charged for bytes +the call never returns. Nothing releases them until Spark's final task cleanup, so they are headroom +nobody can use for the rest of the task. Any caller that swallows the exception has to reconcile +that grant, and the only figure available for doing so is the task-wide one above. + +**A consumer whose `spill` returns zero takes budget it can never give back.** +`NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim +nothing from it; it is only ever a spill trigger. The bytes it holds are real, so other consumers in +the same task see correspondingly less headroom and can spill earlier than they otherwise would. + ## Where Comet's budget comes from `CometExecIterator.getMemoryConfig` computes the budget once per executor and passes it across JNI @@ -253,7 +346,15 @@ diverge for several structural reasons: retained/dirty page cache all add resident bytes that no layer above the allocator can see. Freeing memory does not necessarily return pages to the OS. - **Non-Rust allocations.** Memory allocated by C dependencies through libc `malloc`, and anything - `mmap`ed, never passes through Rust's `GlobalAlloc`. + `mmap`ed, never passes through Rust's `GlobalAlloc`, so neither the memory pool nor the + `jemalloc_allocated` metric sees it. In a default build the C dependencies are libzstd + (`zstd-sys`, behind the Parquet `zstd` codec), libhdfs (`hdfs-sys`, pulled in by the default + `hdfs-opendal` feature), and the TLS stack used for cloud object stores (`aws-lc-sys`). Building + with the `jemalloc` or `mimalloc` feature adds the allocator itself (`tikv-jemalloc-sys`, + `libmimalloc-sys`). It is worth knowing which dependencies are _not_ C, because several names + suggest otherwise: the other Parquet codecs are pure Rust in this build, `snap` for Snappy, + `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its name, so those + allocations do pass through `GlobalAlloc` and are counted. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the @@ -261,7 +362,7 @@ diverge for several structural reasons: The practical consequence is that `reserved()` is a lower bound on Comet's real footprint, and the gap is workload-dependent. `spark.comet.exec.memoryPool.fraction` exists purely so operators can -hand-tune a haircut that covers the gap for their workload. +hand-tune a margin that covers the gap for their workload. To measure the gap on a real query, enable tracing with the `jemalloc` feature and compare `jemalloc_allocated` against the summed `thread_NNN_comet_memory_reserved` values; see @@ -289,11 +390,38 @@ hard ceiling on the sum of everything in the container. That cgroup counts, amon - Comet's JVM-side Arrow buffers (`CometArrowAllocator`), - page cache charged to the cgroup by the container's file I/O, including spill files. -Only the first and a portion of the third are visible to Spark's accounting. When the total crosses -`memory.max`, the kernel OOM killer kills the process. The failure mode is significantly worse than -a task-level OOM: every task running on that executor dies, every cached block it held is lost and -must be recomputed, and the shuffle files it produced become unavailable to downstream fetches. -Spark's driver sees only `ExecutorLostFailure` with exit code 137. +Everything the cgroup counts, and who accounts for each part: + +```mermaid +flowchart TB + subgraph CG["pod cgroup memory.max, kernel OOM kill above this"] + subgraph SEEN["visible to Spark's accounting"] + HEAP["JVM heap
execution and storage
spark.executor.memory"] + TUNG["Spark Tungsten off-heap
TaskMemoryManager"] + SHUFP["Comet JVM shuffle pages
CometUnifiedShuffleMemoryAllocator"] + NATRES["Comet native heap, reserved
operators that call try_grow
declared to Spark over JNI, never measured"] + end + subgraph NONE["accounted by nobody"] + NATUND["Comet native heap, undeclared
kernels, array builders, decompression
Parquet metadata, object_store, tokio"] + ARROWR["Comet JVM Arrow
CometArrowAllocator, unbounded"] + NONHEAP["JVM non-heap
metaspace, code cache, thread stacks
GC structures, Netty direct buffers"] + PAGEC["page cache charged to the cgroup
file I/O, including spill files"] + FRAG["allocator overhead
fragmentation, padding
jemalloc retained and dirty pages"] + end + end +``` + +Spark's accounting covers the first group, though not in the same sense throughout it. The JVM +heap, Tungsten pages and Comet's shuffle pages are allocated by JVM code that reports what it +allocated. A native reservation is a number an operator declared before allocating: `try_grow` +succeeds only once `CometTaskMemoryManager` has charged Spark's off-heap execution pool over JNI, so +the budget really is spent, but nothing measured the bytes and the reservation is only a lower bound +on them. The second group is outside every accounting layer. + +When the total crosses `memory.max`, the kernel OOM killer kills the process. The failure mode is +significantly worse than a task-level OOM: every task running on that executor dies, every cached +block it held is lost and must be recomputed, and the shuffle files it produced become unavailable +to downstream fetches. Spark's driver sees only `ExecutorLostFailure` with exit code 137. Two facts follow that are easy to get wrong: @@ -320,7 +448,7 @@ much they matter: with the `jemalloc` feature and compare `jemalloc_allocated` against summed reservations after the fact. There is no runtime value that an operator, a metric, or a policy could read. - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to - guess a per-workload haircut rather than measuring anything. + guess a per-workload margin rather than measuring anything. - **`CometArrowAllocator` is unbounded** and participates in no budget. - **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 diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 25b805d1cc..205721408d 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -83,12 +83,27 @@ than before. See the [Determining How Much Memory to Allocate] section for more ### Configuring Comet Memory Comet shares an off-heap memory pool with Spark. The size of the pool is -specified by `spark.memory.offHeap.size`. - -Comet's memory accounting isn't 100% accurate and this can result in Comet using more memory than it reserves, -leading to out-of-memory exceptions. To work around this issue, it is possible to -set `spark.comet.exec.memoryPool.fraction` to a value less than `1.0` to restrict the amount of memory that can be -reserved by Comet. +specified by `spark.memory.offHeap.size`. The pool is a shared _budget_ rather than a shared +allocator: Comet's native operators allocate from the Rust heap rather than from JVM off-heap +memory, but every reservation they make is charged against this same pool, so Comet and Spark's own +off-heap consumers draw down one number. + +Comet's memory pool only tracks memory that an operator explicitly reserves, which in practice means the batches +an operator deliberately accumulates: the sort buffer, the build side of a hash join, hash aggregation state, and the +shuffle writer's buffered partitions. Memory that is not reserved is invisible to the pool no matter how much of it +there is. That includes: + +- per-batch working memory in expression kernels and Arrow array builders, +- decompression buffers and Parquet reader structures, +- object store request buffers and the async runtime's own machinery, +- Arrow buffers allocated on the JVM side, which no budget covers at all, +- allocator overhead: buffer padding, size-class rounding, fragmentation, and pages the allocator retains after a + free rather than returning to the operating system. + +Reserved memory is therefore a lower bound on what Comet really uses, and how far below it sits depends on the +workload. This is why Comet can stay within the pool's limit and still push the executor past its container limit. +To leave room for the part that is not counted, set `spark.comet.exec.memoryPool.fraction` to a value less than +`1.0`, which restricts the amount of memory Comet is allowed to reserve. For more details about Spark off-heap memory mode, please refer to [Spark documentation].