diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 701617cde3..21483ac54e 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -48,6 +48,7 @@ Comet Plugin Overview Arrow FFI JVM Shuffle Native Shuffle +Memory Management ANSI Error Propagation S3 Credential Provider Design ``` diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md new file mode 100644 index 0000000000..c252afedc4 --- /dev/null +++ b/docs/source/contributor-guide/memory_management.md @@ -0,0 +1,358 @@ + + +# Memory Management + +This page describes how memory is budgeted, accounted, and enforced across the JVM/native +boundary. It is aimed at contributors working on memory pools, operators that reserve memory, or +anyone debugging an out-of-memory report. For user-facing tuning advice, see the +[Tuning Guide](../user-guide/latest/tuning.md). + +This page covers off-heap mode (`spark.memory.offHeap.enabled=true`) only. Comet also has an +on-heap mode, but it exists so that the Spark SQL test suite can run against Comet without changing +Spark's memory configuration. It must not be used in production, and it is not described here. The +pool types that only on-heap mode exposes belong to the `CATEGORY_TESTING` config group for the +same reason. + +## Overview + +A Comet executor has to satisfy three separate memory budgets at once, and they are enforced by +three different parties: + +| Budget | Enforced by | What happens when it is exceeded | +| -------------------- | ----------------- | ----------------------------------------------------------------------------- | +| JVM heap | The JVM | `java.lang.OutOfMemoryError`; Spark treats it as fatal and exits the executor | +| Spark's memory pools | Spark bookkeeping | A consumer is asked to spill, or the task fails with `SparkOutOfMemoryError` | +| Container RSS | The OS / cgroup | `SIGKILL` of the whole executor process (exit 137, `OOMKilled`) | + +The first two are _accounting_: a running total of bytes that consumers have voluntarily declared. +The third is _physical_: the kernel measures resident pages and does not care what any accounting +layer believes. + +The first row is easy to get wrong. Spark's `Executor.isFatalError` exempts `SparkOutOfMemoryError`, +which is the error Spark's own memory manager throws when a consumer cannot get the bytes it asked +for, so that failure is confined to the task and the executor keeps running. An ordinary JVM +`OutOfMemoryError` is not exempt: the executor hands it to `SparkUncaughtExceptionHandler`, which +calls `System.exit(SparkExitCode.OOM)` (exit code 52). Real heap exhaustion therefore usually costs +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. + +## Who allocates what + +Enabling Comet does not add one new memory consumer, it adds several, and they are not all +accounted by the same party. This inventory is worth internalizing before reading the rest of the +page: + +| Allocator | Lives in | Bounded by | Visible to Spark? | +| --------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- | +| 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 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. + +**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 +`new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream export +(`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. + +**The JVM shuffle allocator is an ordinary Spark consumer.** `CometShuffleMemoryAllocator.getInstance` +returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from +`spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the +same task like any other allocation. + +## Where Comet's budget comes from + +`CometExecIterator.getMemoryConfig` computes the budget once per executor and passes it across JNI +to `Java_org_apache_comet_Native_createPlan` as `memory_limit`. Comet shares Spark's off-heap pool +rather than asking for a separate allocation: + +```text +memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction +``` + +`spark.comet.exec.memoryPool.fraction` defaults to `1.0`. Lowering it is the current workaround for +Comet's under-accounting (see [The accounting gap](#the-accounting-gap)). It holds back a slice of +the off-heap pool that Comet is not allowed to reserve, on the assumption that Comet's real usage +overshoots its reservations by roughly that slice. + +A second value, `memory_limit_per_task`, is computed and passed alongside it, but only the on-heap +pool types read it. + +### Resolving the pool type + +`parse_memory_pool_config` (`native/core/src/execution/memory_pools/config.rs`) turns the pool-type +string and the limit into a `MemoryPoolConfig`. Two pool types are valid in off-heap mode: + +| Pool type | Sized from | Notes | +| ------------------------ | ------------------- | ----------------------------------------------------- | +| `fair_unified` (default) | `memory_limit` | Delegates to Spark's `TaskMemoryManager`; task-shared | +| `greedy_unified` | n/a (pool size `0`) | Spark owns the limit entirely; task-shared | + +Any other pool type is rejected with a configuration error. + +## The pool stack + +`create_memory_pool` builds a base pool and `createPlan` then wraps it in decorators. Reading from +the inside out, a Comet plan in the default configuration sees: + +```text +[LoggingMemoryPool] <- only when spark.comet.debug.memory=true + [TaskSharedMemoryPool] <- RAII handle for the per-task registry + [TrackConsumersPool] <- DataFusion; names the top 10 consumers in error messages + [CometFairMemoryPool] <- delegates acquire/release to Spark over JNI +``` + +Each decorator forwards every `MemoryPool` method to its inner pool, so `reserved()` at any level +reports the base pool's number. + +### The unified pools + +`CometUnifiedMemoryPool` and `CometFairMemoryPool` (`unified_pool.rs`, `fair_pool.rs`) are the +bridge to Spark. Their `try_grow` calls `CometTaskMemoryManager.acquireMemory` over JNI, which goes +through Spark's ordinary `TaskMemoryManager`. That means: + +- Comet competes with Spark's own off-heap consumers (Tungsten sorters, `BytesToBytesMap`, and so + on) for the same `spark.memory.offHeap.size`, and Spark's unified memory manager arbitrates. +- Spark can force _Spark's_ consumers in the same task to spill to satisfy Comet's request. +- The reverse does not hold. Comet registers a `NativeMemoryConsumer` with the `TaskMemoryManager` + so that Spark has something to charge, but its `spill()` always returns `0` + (`CometTaskMemoryManager.java`). A Spark allocation can therefore never make a native sorter or + aggregate release its reservations. Native operators spill only when their _own_ `try_grow` + fails, so a JVM consumer that is blocked behind native reservations has no way to reclaim them. +- A partial grant (`acquired < additional`) is released immediately and reported as + `ResourcesExhausted`, which is the signal DataFusion uses to spill. + +`CometFairMemoryPool` additionally applies a local check before it asks Spark. It divides +`pool_size` by the number of consumers currently registered with the pool and rejects the request if +the pool's _total_ reserved bytes plus the request would exceed that quotient. Two details matter +for tuning: + +- The comparison is against the shared total (`state.used`), not against the requesting consumer's + own reservation. With an 8 GiB pool and two registered consumers, once one of them holds 3 GiB a + 2 GiB request from the other is rejected, even though neither would exceed 4 GiB. In effect the + usable pool shrinks to `pool_size / num_consumers` in aggregate as soon as more than one consumer + is registered. +- The pool is task-shared (see below), so `num_consumers` counts every registered consumer across + every native plan in the task, not just the plan making the request. + +This is why `fair_unified` spills earlier than `greedy_unified`. It is not a per-consumer quota, and +reading it as one overstates the memory a multi-operator task can use. + +### Task-shared pools and their lifetime + +A single Spark task can run more than one native plan concurrently: a shuffle runs the pre-shuffle +operators and the shuffle writer as separate native execution contexts. If each got its own pool, +the per-task limit would be enforced once per plan rather than once per task. + +`acquire_task_shared_pool` (`task_shared.rs`) keeps a process-wide +`HashMap>`. Plans in the same task upgrade the existing +`Weak` and share one pool; the returned `Arc` is the only lifetime handle, so the registry entry +disappears when the last plan (and its last reservation) drops. There is no explicit release call to +forget, and a `createPlan` that fails partway through cleans up on unwind. + +`TaskSharedMemoryPool::drop` has to handle one race: an `acquire` can observe an expired `Weak` and +insert a replacement before the dying pool reaches the registry lock. The drop therefore compares +pointers and only removes an entry that is still its own. + +## How DataFusion consumes the pool + +Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservation` API: + +- `try_grow(n)` may fail. Spillable operators (`ExternalSorter`, the grouped hash aggregate, + sort-merge join) respond to a `ResourcesExhausted` error by spilling to disk and retrying. This is + the only mechanism that turns memory pressure into progress rather than failure. +- `grow(n)` is infallible and panics if the pool refuses. It is used where the caller cannot spill. +- `shrink(n)` returns bytes to the pool. + +An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses. + +## Crossing the FFI boundary + +Batches move between the JVM and native over the Arrow C Data and C Stream interfaces, which are +zero-copy. Nothing is copied, so the _allocator_ that produced a batch and the _runtime_ that +decides when it dies can be on opposite sides of the boundary. See [Arrow FFI](ffi.md) for the +mechanics; what matters here is who is charged and who controls the lifetime. + +Two things are easy to conflate here. The _allocator_ of a buffer determines which process-level +accounting (if any) saw the bytes appear. Whether the buffer is _reserved_ in Comet's pool is a +separate decision, made by whichever operator holds it, and that operator neither knows nor cares +which side of the boundary the bytes came from. + +**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of +`CometArrowAllocator` and exports the whole per-partition iterator once as an `ArrowArrayStream`. +`ScanExec` imports each batch through `AlignedArrowStreamReader` with `CopyMode::UnpackOrClone`: +dictionary columns are unpacked into new native arrays, everything else is an `Arc` clone of the +imported buffers. Those bytes stay where Java Arrow put them and are pinned for as long as any native +reference survives. They are invisible to Spark's `TaskMemoryManager`, and `CometArrowAllocator` is +unbounded, so nobody charged for them at allocation time. Whether they are charged _later_ depends +on who holds them. DataFusion's `ExternalSorter` reserves `get_reserved_bytes_for_record_batch` for +every batch it retains, and the hash join build side reserves `get_record_batch_memory_size` for +each incoming batch; both read buffer sizes off the `ArrayData` and apply equally to imported +buffers. So an imported batch that a sort or join has buffered _is_ reserved in Comet's pool, and +through a unified pool is charged to Spark. The same batch held by an operator that does not reserve +(a projection or filter, or `ScanExec` itself between polls) is charged to nothing. + +**Native → JVM (`CometExecIterator`).** DataFusion produces the batch in Rust, and whether it is +still reserved when it reaches the boundary depends on the producing operator. The sort's output is +wrapped in a `ReservationStream` that shrinks the reservation by the batch's size as each batch is +emitted, so by the time `prepare_output` sees a sorted batch its bytes are no longer reserved. +`prepare_output` then calls `move_to_spark`, which writes an `FFI_ArrowArray` whose release callback +drops the Rust `ArrayData`; no reservation travels with it. The JVM wraps the pointers in +`ArrowBuf`s and the memory is freed when the JVM calls `close()`. An exported batch is therefore +usually _not_ pool-charged while the JVM holds it, but it is resident until the JVM releases it. A +slow or backed-up JVM consumer keeps native memory alive that no accounting layer is counting. + +The point is not that one direction is charged and the other is not. **Buffer lifetime and +reservation lifetime are independent.** A reservation is a number an operator chose to declare and +later withdraw; a buffer lives until its last reference drops, on whichever side of the boundary +that happens. The two are kept in step only inside operators written to do so, and only while that +operator holds the batch. Once the batch leaves the operator, in either direction, the reservation +stops and the bytes keep going. That is the shape of the gap to hold in mind when reading the next +section, and it is what any fix at the allocator level has to close. + +## The accounting gap + +The pool tracks _declared reservations_. Container RSS counts _pages the process touched_. The two +diverge for several structural reasons: + +- **Undeclared allocations.** Arrow array builders, expression kernels producing intermediate + arrays, decompression buffers, Parquet metadata structures, `object_store` request buffers, and + tokio's own machinery all allocate without reserving. Only operators that were explicitly written + to reserve show up in the pool. +- **Rounding and padding.** Arrow buffers are padded to 64-byte boundaries and builders grow by + doubling, so a reservation of exactly `n` bytes routinely corresponds to more than `n` bytes of + heap. +- **Allocator behavior.** `malloc`-level fragmentation, size-class rounding, and jemalloc's + 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`. +- **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 + JVM closes them (see [Crossing the FFI boundary](#crossing-the-ffi-boundary)). + +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. + +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 +[Tracing](tracing.md#analyzing-memory-usage). + +## What the container sees + +On Kubernetes, Spark sizes the executor pod from `ResourceProfile`: + +```text +pod memory request = pod memory limit + = spark.executor.memory + + spark.executor.memoryOverhead (default max(0.1 * executor.memory, 384 MiB)) + + spark.memory.offHeap.size + + pyspark memory (Python applications only) +``` + +Both the request and the limit are set to this same value, so the pod's cgroup `memory.max` is a +hard ceiling on the sum of everything in the container. That cgroup counts, among other things: + +- the JVM heap (`spark.executor.memory`), +- JVM non-heap: metaspace, code cache, thread stacks, GC structures, Netty direct buffers, +- Spark's own off-heap allocations, +- **all of Comet's native allocations**, +- 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. + +Two facts follow that are easy to get wrong: + +1. **`spark.memory.offHeap.size` is part of the pod limit, not extra headroom on top of it.** + Raising Comet's off-heap budget on Kubernetes raises the pod's memory request by the same + amount, so the scheduler will place fewer executors per node rather than silently giving Comet + more room. +2. **`spark.executor.memoryOverhead` is the only slack in the container**, and the JVM's own + non-heap usage already consumes a large part of it. Comet's overshoot beyond its declared + reservations eats into the same allowance. + +YARN behaves analogously. The container size is the same sum, and the NodeManager kills containers +that exceed it, but the kill is done by the NodeManager's monitor rather than the kernel, so it is +somewhat less abrupt. + +## Open problems + +Nothing described above prevents an executor from being OOM-killed. Every enforcement point Comet +has bounds _declared reservations_, and the sections above describe several structural reasons why +declared reservations are a lower bound on physical usage. The known gaps, roughly in order of how +much they matter: + +- **No signal for real native usage.** The only way to observe the gap today is to enable tracing + 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. +- **`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 + withdrawn by individual operators while the bytes outlive them. +- **Spark cannot trigger native spilling.** `NativeMemoryConsumer.spill()` returns `0`, so native + reservations are released only when a native operator decides to spill on its own failed + `try_grow`. JVM consumers in the same task can be starved behind them. + +[Issue #4576](https://github.com/apache/datafusion-comet/issues/4576) tracks work on the first two. + +## 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 | +| [Tracing](tracing.md#analyzing-memory-usage) | `jemalloc_allocated` vs summed pool reservations; the accounting gap | +| `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: + +1. Confirm which budget was exceeded. Exit code 137 / `OOMKilled` on the pod is the cgroup. Exit + code 52 with `java.lang.OutOfMemoryError` in the executor log is JVM heap exhaustion; Spark + treats it as fatal, so the executor is lost either way and the exit code is what distinguishes + them. A failed task with `SparkOutOfMemoryError` and a surviving executor is Spark's managed + memory pool, which is the only one of the three that is recoverable at task level. +2. Compare `jemalloc_allocated` against the summed pool reservations from 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. +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 + `spark.comet.exec.forceShuffledHashJoin=true` converts a spillable sort-merge join into one that + is not.