diff --git a/.ai/skills/review-comet-ffi-pr/SKILL.md b/.ai/skills/review-comet-ffi-pr/SKILL.md index f57ef627367..facf9ba926a 100644 --- a/.ai/skills/review-comet-ffi-pr/SKILL.md +++ b/.ai/skills/review-comet-ffi-pr/SKILL.md @@ -135,9 +135,13 @@ a new subclass needs no new case as long as `getValueVector` returns an Arrow ve ## 5. Memory Accounting An FFI change is usually also a memory change, and the two are easy to review separately and miss -the interaction. Imported JVM buffers come from `CometArrowAllocator`, which is an unbounded -`RootAllocator(Long.MaxValue)` that no budget sees. Exported native batches are usually no longer -pool-reserved by the time the JVM receives them, but they stay resident until the JVM closes them. +the interaction. The JVM buffers native imports belong to a child of `CometArrowAllocator`, an +unbounded `RootAllocator(Long.MaxValue)` that no budget sees, and a native operator that retains +them reserves them itself. Batches the JVM owns within a task come from `CometTaskArrowAllocator`, +whose listener charges them to Spark and refuses what Spark cannot cover, so a new path that exports +them without going through `CometArrowStream`, which moves the charge off the task, charges them on +both sides. Exported native batches are usually no longer pool-reserved by the time the JVM receives +them, but they stay resident until the JVM closes them. If the PR changes how long either side holds a batch, use `review-comet-memory-pr` as well. diff --git a/.ai/skills/review-comet-memory-pr/SKILL.md b/.ai/skills/review-comet-memory-pr/SKILL.md index fbd4610b20a..9ba6d2ec6f1 100644 --- a/.ai/skills/review-comet-memory-pr/SKILL.md +++ b/.ai/skills/review-comet-memory-pr/SKILL.md @@ -158,15 +158,27 @@ memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction `spark.executor.memoryOverhead` is the only real slack in the container, and JVM non-heap usage already consumes much of it. -## 6. Unaccounted Allocators +## 6. JVM-Side Allocators Two allocators sit outside Comet's pool entirely, and they are easy to confuse: - **`CometArrowAllocator`** (`spark/src/main/scala/org/apache/comet/package.scala`) is a - process-wide `RootAllocator(Long.MaxValue)`. It is unbounded and no budget sees it. Child - allocators are cut from it for FFI stream export, broadcast coalescing, and - `CometSparkToColumnarExec`. A PR that adds a child allocator, or makes an existing one hold more, - is adding uncounted container RSS. Say so even if the volume is small. + process-wide `RootAllocator(Long.MaxValue)` with no allocation listener. It is unbounded. Buffers + allocated from it, or from a child cut with the three-argument `newChildAllocator`, are seen by no + budget: FFI stream export, `NativeUtil`, the JVM UDF result, and the import path's + `CometArrowImportAllocator`. JVM-owned allocations inside a task go through + `CometTaskArrowAllocator.forCurrentTask()` instead, a per-task child whose + `CometArrowAllocationListener` charges what it owns to that task's `TaskMemoryManager` and refuses + an allocation Spark cannot cover, so it fails with Arrow's `OutOfMemoryException`. A PR that + allocates from the root inside a task, or makes a root-backed allocator hold more, is adding + uncounted container RSS. Say so even if the volume is small. A PR that allocates something native + will retain from the task allocator is charging it twice, unless the batch reaches native through + `CometArrowStream`, whose reader moves the charge off the task allocator. A new call site on the + task allocator also has to cope with that refusal: nothing spills these buffers, so it fails the + task. Check that the listener stays without a lock in `getUsed` and `spill`, that only + `onPreAllocation` throws, and that it sizes the reservation from what the allocator owns rather + than from a sum of its callbacks, because an ownership transfer between allocators calls no + listener. - **JVM shuffle pages** go through `CometShuffleMemoryAllocator.getInstance`, which returns `CometUnifiedShuffleMemoryAllocator`, an ordinary Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`. These **are** arbitrated by Spark against its other consumers. Do diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index ea8a23d204c..68e646f6d5b 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -543,6 +543,7 @@ jobs: org.apache.spark.CometPluginsSuite org.apache.spark.CometRuntimeShutdownSuite org.apache.spark.CometTaskMemoryManagerSuite + org.apache.spark.comet.CometArrowAllocationListenerSuite org.apache.spark.CometExecIteratorLifecycleSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsMemoryOverheadWarningSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 3a351d1e6f9..451b58249e3 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -191,6 +191,7 @@ jobs: org.apache.spark.CometPluginsSuite org.apache.spark.CometRuntimeShutdownSuite org.apache.spark.CometTaskMemoryManagerSuite + org.apache.spark.comet.CometArrowAllocationListenerSuite org.apache.spark.CometExecIteratorLifecycleSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsMemoryOverheadWarningSuite diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 9768698b0d8..cb7270ea0fd 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -76,13 +76,14 @@ Enabling Comet does not add one new memory consumer, it adds several, and they a 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 | 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 | +| 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 | Reservations only | +| Comet JVM Arrow owned by a task | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | +| Comet JVM Arrow passed to or from native | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | +| Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | Several observations follow. @@ -107,13 +108,69 @@ pool over JNI. A native allocation therefore consumes the same accounting budget 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 -`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. +**Comet's JVM-side Arrow memory is charged to the task that owns it, and bounded by the pool.** +`CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single +process-wide `new RootAllocator(Long.MaxValue)` with no allocation listener. JVM-owned allocations +do not use it directly: `CometTaskArrowAllocator.forCurrentTask()` hands each task a child of it +carrying its own `CometArrowAllocationListener`, a Spark `MemoryConsumer` for that task, which +reserves what the child owns from the off-heap execution pool in whole 1 MiB blocks. Those bytes +therefore appear in `showMemoryUsage` and are arbitrated against Spark's other off-heap consumers, +Comet's native pool among them. The reservation is made before Arrow allocates, in +`onPreAllocation`, the one `AllocationListener` callback Arrow allows to throw. If Spark cannot +cover the allocation, any partial grant is handed back and the allocation is refused with Arrow's +`OutOfMemoryException`. That is `try_grow` rather than `grow`: nothing is allocated that is not +covered, and there is no overcommit. Nothing can spill these buffers, so a refusal fails the +allocating task, the way a Spark consumer fails when it cannot acquire memory. The deprecated +`spark.comet.legacy.unboundedJvmArrowMemory=true` restores the earlier behavior of charging +nothing, in which case `forCurrentTask()` returns the root. It also returns the root in on-heap mode, and off a task, as +for broadcast coalescing on the driver or for a reader built on a native thread that is pulling a +stream. + +**The owner is the allocator, not the calling thread.** Arrow hands `AllocationListener` nothing +but a size, and it reports a release on whichever thread drops the last reference, which for +anything exported over the C Data Interface is a Tokio worker with no task context installed. +Attributing by `TaskContext` would therefore drop those releases and leave a task charged for +memory it had already freed. Binding one listener to one task's allocator is what makes the release +land on the task that allocated. Children cut from a task allocator inherit its listener and roll +their bytes up into it, so the paths that make their own children are covered without knowing +about any of this. + +**What is charged is what the allocator owns, not a running total of the callbacks.** When two +allocators hold references to one buffer and the owner releases first, Arrow moves the charge to +the other with `transferBalance`, which calls no listener, and the final `onRelease` goes to the new +owner. `ColumnarBatchArrowReader`, which feeds every JVM batch that reaches native through +`CometArrowStream.wrapColumnarBatchRDD`, does exactly this to each batch: it retains the buffers +into the stream's allocator and closes the source. A listener that summed its callbacks would +charge the task for every batch it ever streamed, until the task ended. The listener therefore +treats a callback only as a prompt to re-read what its allocator owns, from Arrow's own accountant, +which Arrow keeps correct across transfers. `ColumnarBatchArrowReader` also has the source's task +allocator reconcile as soon as it closes the source, so the task stops paying for a batch the moment +native has it, even when that is the last batch the task reads, as it is for a broadcast coalesced +into one batch. Any other transfer is picked up at the next allocation or release on the task's +allocator, and whatever is still reserved is returned when the task completes. + +**Memory passed to or from native is not charged on the JVM side.** Native memory is native's to +account for: a native operator that retains a batch reserves its buffers through Comet's pool, +`ExternalSorter` through `get_reserved_bytes_for_record_batch` and the hash join build side through +`get_record_batch_memory_size`, and both read sizes off the `ArrayData`, so they count imported +buffers too. Charging those buffers on the JVM side as well would reserve the same memory twice. +So `NativeUtil`, `CometUdfBridge` together with the codegen output vector it exports, and +`CometArrowStream.stream` allocate from the root, and imports go through +`CometArrowImportAllocator`, which must stay listener-less as well, because Arrow's +`wrapForeignAllocation` reports an imported buffer to the importing allocator's listener at full +capacity. Allocations that stay in the JVM, which are `StreamReader`, the cached batch serializer, +`CometArrowStream.readerBatchIter` and the Python runner, use the task allocator. A batch read that +way and then streamed to native needs no special handling, because the stream's reader moves its +charge off the task as described above. What remains is a buffer exported column by column with +`Data.exportVector`, which stays owned by, and charged to, the task allocator while native reads it. +`NativeColumnarToRowConverter` is the one caller that does this with a batch the JVM may own, and +native's columnar-to-row conversion reserves nothing, so the bytes are not charged twice. + +**A task allocator is closed only once it has been drained.** The process-wide allocator exists +because Arrow buffers can outlive the task that created them, and Arrow treats closing an allocator +that still owns bytes as a leak. At task completion the Spark reservation is dropped and the +allocator is closed if it is empty. Otherwise it is parked and closed by a later task once the +stragglers are released, since `BaseAllocator` keeps every child in a map until it closes. 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 @@ -136,36 +193,43 @@ flowchart LR 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"] + NASS["CometArrowStream.stream
batches exported to native"] + BC["Utils broadcast-coalesce
on the driver"] + SR["StreamReader
broadcast and IPC reads"] + NASR["CometArrowStream.readerBatchIter
batches read in the JVM"] 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"] + TASK["CometTaskArrowAllocator
a child of the root per task
CometArrowAllocationListener"] + subgraph OTHER["other consumers of Spark's off-heap pool"] + SHUF["Comet JVM shuffle pages
CometUnifiedShuffleMemoryAllocator"] + NPOOL["Comet native memory pool
declared reservations only"] + end 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 + NASS --> ROOT BC --> ROOT + SR --> TASK + NASR --> TASK + CACHE --> TASK + PY --> TASK ROOT --> NOBODY + TASK -->|"what it owns, in 1 MiB blocks"| TMM 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. +`CometTaskMemoryManager` and `CometArrowAllocationListener` are the two places where Comet code acts +as a Spark `MemoryConsumer`, and the rules below are why they are written the way they are. 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 @@ -186,19 +250,24 @@ reaches this method from native over JNI, so whatever escapes crosses the JNI bo **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. +overrides `getUsed` to report Comet's own tally instead, and `CometArrowAllocationListener` reports +what its allocator owns: without that, 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. +`CometArrowAllocationListener` measures it as the change in that figure across the call, with both +reads and the call inside the `TaskMemoryManager` monitor so that no other consumer in the task can +acquire in between, and releases it before refusing the allocation. **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. +`NativeMemoryConsumer.spill` returns `0`, as does `CometArrowAllocationListener.spill`, so Spark can +select either 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 @@ -356,14 +425,17 @@ accounting (if any) saw the bytes appear. Whether the buffer is _reserved_ in Co 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`. +**JVM → native (`ScanExec`).** The JVM exports the whole per-partition iterator once as an +`ArrowArrayStream`, and every batch it hands over belongs to the stream's allocator, a child of the +listener-less root. A batch built for the stream, such as a row-to-Arrow conversion, is allocated +there. A batch the JVM had already read into a task allocator is retained there by +`ColumnarBatchArrowReader` and its source closed, which moves the charge off the task allocator. `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 +reference survives. Spark's `TaskMemoryManager` does not see them on the JVM side, deliberately, +because native decides whether to charge them. 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 @@ -456,7 +528,7 @@ hard ceiling on the sum of everything in the container. That cgroup counts, amon - 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`), +- Comet's JVM-side Arrow buffers (`CometArrowAllocator` and the per-task allocators cut from it), - page cache charged to the cgroup by the container's file I/O, including spill files. Everything the cgroup counts, and who accounts for each part: @@ -469,10 +541,11 @@ flowchart TB 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"] + ARROWT["Comet JVM Arrow owned by a task
CometTaskArrowAllocator, refused when the pool is short"] 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"] + ARROWR["Comet JVM Arrow passed to or from native
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"] @@ -480,12 +553,13 @@ flowchart TB 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. +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, +and a task's Arrow allocator reserves what an allocation needs before Arrow makes it, refusing the +allocation when Spark cannot cover it. 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 @@ -519,7 +593,9 @@ 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 bounded only for what a task owns.** A task's allocator is refused + what Spark cannot cover, but Arrow memory passed to or from native, and anything allocated off a + task, is charged nowhere on the JVM side, and the root itself has no limit. - **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. diff --git a/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md index 59d5c711faf..16267761d60 100644 --- a/docs/source/user-guide/latest/configs.md +++ b/docs/source/user-guide/latest/configs.md @@ -72,6 +72,16 @@ shuffle, including applications that rely only on remote reliable storage or dec +## Legacy Settings + +Each of these settings restores a behavior that an earlier release had, after a release changed it. +They are deprecated from the moment they are added, and may be removed in any future major release, +at which point the newer behavior becomes unconditional. The +[Upgrade Guide](migration-guide.md) describes each change and the setting that reverts it. + + + + ## Development & Testing Settings These settings exist for Comet's own test suites and for debugging. They are **not covered by the diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index 1de5f884672..8248b61586b 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -53,6 +53,25 @@ Treat setting one of these keys as a temporary measure. If you find you cannot s legacy behavior, please open an issue describing your use case so it can be considered before the key is removed. +## Upgrading to Comet 1.2.0 + +### JVM-side Arrow memory is charged to Spark's off-heap memory pool + +The Arrow memory that Comet allocates on the JVM side for a task, for example to read a cached table +or to exchange batches with a Python worker, is now charged to `spark.memory.offHeap.size` as a +memory consumer of that task. An allocation the pool cannot cover fails the task with an +`OutOfMemoryException` whose message begins `Unable to reserve`, the way a Spark operator fails when +it cannot acquire memory. The charge also leaves Comet's native operators and Spark's own consumers +less of the pool, so queries can spill sooner than before. Earlier releases charged this memory to no +pool, so it never failed an allocation and had to fit in `spark.executor.memoryOverhead` instead. + +Arrow memory that is handed to Comet's native code, or received from it, is not charged on the JVM +side, because native accounts for what it retains. + +To restore the previous behavior, set `spark.comet.legacy.unboundedJvmArrowMemory=true` when the +application is submitted. It is read on the executors, so setting it in a running session has no +effect. + ## Upgrading to Comet 1.1.0 Comet `1.1.0` makes no behavior changes that need a `spark.comet.legacy.*` key. diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 0fb4288abac..431be35e868 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -50,7 +50,7 @@ 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, +- Arrow buffers passed between the JVM and native code, unless a native operator reserves them, - allocator overhead: buffer padding, size-class rounding, fragmentation, and pages the allocator retains after a free rather than returning to the operating system. @@ -64,6 +64,16 @@ while Comet runs; see [Sizing the Overhead from the Memory Usage Log]. fraction to each task separately, where Spark's own limit of an even share of the pool per running task is tighter whenever more than one task is running, and the `greedy_unified` pool ignores it. +Arrow memory that Comet allocates on the JVM side for a task, for example to read a cached table or to exchange +batches with a Python worker, is charged to the same pool as a memory consumer of that task. An allocation the pool +cannot cover fails the task that makes it, the way a Spark operator fails when it cannot acquire memory, and the charge +leaves Comet's native operators and Spark's own consumers less of the pool, so they can spill sooner. Comet 1.1.0 and +earlier did not charge this memory to the pool, so if you sized `spark.memory.offHeap.size` against one of those +releases, check how often queries spill, and raise it if tasks fail with `Unable to reserve ... for a JVM Arrow +allocation`. Setting the deprecated `spark.comet.legacy.unboundedJvmArrowMemory=true` restores the earlier behavior +while you adjust, in which case this memory has to fit in `spark.executor.memoryOverhead` instead; see the +[Upgrade Guide](../migration-guide.md). + For more details about Spark off-heap memory mode, please refer to [Spark documentation]. [Spark documentation]: https://spark.apache.org/docs/latest/configuration.html @@ -104,7 +114,8 @@ container to include it, so the memory that Comet's operators explicitly reserve not have room is everything Comet allocates without reserving it — the untracked categories listed under [Configuring Comet Memory]. Those allocations are made by the Rust global allocator and live in the native heap, outside the JVM heap and outside Spark's off-heap allocations, and nothing in -the container sizing accounts for them. The same applies to Comet's JVM-side Arrow buffers. +the container sizing accounts for them. The same applies to Comet's JVM-side Arrow buffers on their +way to or from native code. `spark.executor.memoryOverhead` is the only slack the container has for this, and the JVM's own non-heap usage — metaspace, code cache, thread stacks, GC structures — is already drawing on it. @@ -169,7 +180,7 @@ JVM's own non-heap memory. To size the overhead from it: 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. + zstd, or Comet's JVM-side Arrow buffers on their way to or from native code. 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 diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index d2881bfe449..118dbc17158 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -210,6 +210,9 @@ private static void evaluateInternal( }); assert udf != null : "reflective instantiation returned null for " + udfClassName; + // Neither allocator reports to Spark, and neither should: the inputs wrap memory native owns, + // and the result is exported straight back to native, which accounts for what it retains. See + // CometTaskArrowAllocator. BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator(); // See CometArrowImportAllocator: inputs are imported against that allocator, so that tracing // can report the import path's charges apart from the rest of Comet's Arrow memory. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index eeb5cf8045d..c3552b83124 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -68,6 +68,7 @@ object CometConf extends ShimCometConf { private val CATEGORY_SHUFFLE = "shuffle" private val CATEGORY_TUNING = "tuning" private val CATEGORY_TESTING = "testing" + private val CATEGORY_LEGACY = "legacy" def register(conf: ConfigEntry[_]): Unit = { assert(conf.category.nonEmpty, s"${conf.key} does not have a category defined") @@ -374,6 +375,21 @@ object CometConf extends ShimCometConf { .checkValue(_ >= 0, "The memory usage log interval must not be negative") .createWithDefault(TimeUnit.SECONDS.toMillis(10)) + val COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY: ConfigEntry[Boolean] = + conf("spark.comet.legacy.unboundedJvmArrowMemory") + .category(CATEGORY_LEGACY) + .doc( + "When true, restores the pre-1.2.0 behavior of the Arrow memory that Comet " + + "allocates on the JVM side for a task, which was charged to no memory pool and so " + + "could never fail an allocation. Since 1.2.0 it is charged to Spark's off-heap " + + "memory pool as a memory consumer of the task, and an allocation the pool cannot " + + "cover fails the task. Arrow memory handed to Comet's native code, or received from " + + "it, is not affected. This is an executor setting, read when an executor first " + + "allocates Arrow memory for a task, so it must be set when the application is " + + "submitted. This config is deprecated and will be removed in a future major release.") + .booleanConf + .createWithDefault(false) + val COMET_SHUFFLE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.enabled") diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 33e6c0c0355..bff1a363ba9 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -87,6 +87,8 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { * Closes the vector on any failure so a partially-initialized tree doesn't leak buffers. */ def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int): FieldVector = { + // The listener-less root rather than the task allocator: `CometUdfBridge` exports this vector + // straight to native, which accounts for what it retains. See `CometTaskArrowAllocator`. val vec: FieldVector = field.getType match { case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList => val v = new RenamedListVector(field, CometArrowAllocator) diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala index 702359e335e..0007f581cb7 100644 --- a/spark/src/main/scala/org/apache/comet/package.scala +++ b/spark/src/main/scala/org/apache/comet/package.scala @@ -32,6 +32,12 @@ package object comet { * Until the reference count is zero, the memory will not be released. If the consumer side is * finished later than the close of the allocator, the allocator will think the memory is * leaked. To avoid this, we use a single allocator for the whole execution process. + * + * It carries no allocation listener, so neither it nor a child cut with the three-argument + * `newChildAllocator` charges anything to Spark. That suits buffers handed to native, which + * accounts for what it retains. JVM-owned allocations within a task go through + * `CometTaskArrowAllocator.forCurrentTask()` instead, which describes which allocations go + * where. */ val CometArrowAllocator = new RootAllocator(Long.MaxValue) @@ -55,6 +61,10 @@ package object comet { * * So read this and the root's total as allocator charges. Their difference is not a bound on * the Arrow memory the JVM allocated itself, and neither is a count of unique physical bytes. + * + * It must stay without an allocation listener, as the root is. Arrow's `wrapForeignAllocation` + * reports an imported buffer to the importing allocator's listener at full capacity, as though + * the JVM had allocated it, so a listener here would charge Spark for native memory. */ val CometArrowImportAllocator: BufferAllocator = CometArrowAllocator.newChildAllocator("comet-ffi-imports", 0, Long.MaxValue) diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 5ef24321733..373793e2d08 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -47,7 +47,9 @@ import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator} class NativeUtil extends AutoCloseable { import Utils._ - /** Use the global allocator */ + // The listener-less root, not the task allocator: what this allocates, the FFI structs and + // materialised constant vectors, is handed to native, which accounts for what it retains. See + // CometTaskArrowAllocator. private val allocator = CometArrowAllocator /** ArrowImporter does not hold any state and does not need to be closed */ diff --git a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala index 805eae988e8..8a47fe44557 100644 --- a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala +++ b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala @@ -26,17 +26,19 @@ import scala.util.control.NonFatal import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel} import org.apache.arrow.vector.ipc.message.MessageChannelReader +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.CometArrowAllocator - /** * A reader that consumes Arrow data from an input channel, and produces Comet batches. */ case class StreamReader(channel: ReadableByteChannel, source: String) extends AutoCloseable { + // Decoded into JVM-owned buffers, so it is accounted to the task doing the read. Held in a val + // because the reader and the channel reader must share one allocator. + private val allocator = CometTaskArrowAllocator.forCurrentTask() private val channelReader = - new MessageChannelReader(new ReadChannel(channel), CometArrowAllocator) - private var arrowReader = new ArrowStreamReader(channelReader, CometArrowAllocator) + new MessageChannelReader(new ReadChannel(channel), allocator) + private var arrowReader = new ArrowStreamReader(channelReader, allocator) // Reading the schema allocates the root's vectors, so it can fail with buffers already taken. // No caller holds this reader until its constructor returns, so close it here or nothing will. diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala new file mode 100644 index 00000000000..abc44af7d24 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala @@ -0,0 +1,325 @@ +/* + * 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.atomic.AtomicBoolean +import java.util.function.LongSupplier + +import scala.util.control.NonFatal + +import org.apache.arrow.memory.{AllocationListener, AllocationOutcome, OutOfMemoryException} +import org.apache.spark.internal.Logging +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager} + +import org.apache.comet.CometConf + +/** + * Charges one task's JVM-side Arrow allocations to Spark's off-heap execution pool, and refuses + * an allocation the pool cannot cover. + * + * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, but the bytes a task's + * allocator hands out are resident in the container all the same. Charged here, they appear in + * `TaskMemoryManager.showMemoryUsage`, are arbitrated against Spark's other off-heap consumers, + * and are bounded by the same pool. + * + * '''Reserved before it is allocated.''' Arrow calls [[onPreAllocation]] before it allocates a + * buffer, and documents it as the one callback that may throw to terminate the allocation. The + * reservation is taken there, in whole blocks, and if Spark cannot cover the allocation whatever + * it did grant is handed straight back and the allocation is refused with Arrow's own + * `OutOfMemoryException`, the exception Arrow's vectors already clean up after. That is + * `try_grow` rather than `grow`: nothing is allocated that is not covered, and there is no + * overcommit. Nothing can spill these buffers either, since the JVM holds them until the batch is + * done with, so a refusal fails the allocating task, the way a Spark consumer fails when it + * cannot acquire memory. + * + * '''Ownership.''' One instance is created per task and attached to that task's Arrow allocator + * by [[CometTaskArrowAllocator]]. Arrow calls the listener of the allocator that makes a buffer + * on allocation, and the listener of the allocator that '''owns''' it once the last reference + * drops, on whichever thread that happens to be. `AllocationListener` is handed nothing but a + * size, so binding the listener to an allocator is the only way to attribute a release. Reading + * `TaskContext` inside the callbacks would get it wrong: a buffer exported over the C Data + * Interface is released by whichever native thread drops it last, and that thread has no task + * context installed. + * + * '''What is charged is what the allocator owns, not a running total of the callbacks.''' The two + * come apart because ownership moves between allocators without either listener hearing about it. + * When two allocators hold references to one buffer and the owner lets go first, Arrow hands the + * charge to the other through `transferBalance`, which calls no listener, and the final + * `onRelease` goes to the new owner. That is the everyday path rather than a corner: + * `ColumnarBatchArrowReader` retains every batch it streams to native into its own allocator, a + * child of the unaccounted root, and then closes the source. Summing the callbacks would charge + * the task for each such batch until the task ended. So the callbacks only say when to look, and + * the size comes from the task allocator's accountant, which Arrow keeps right across transfers. + * Every callback trims the reservation to what is owned and pending, so a transfer out is + * reflected at the next allocation or release on the task's allocator, and the reader calls + * [[reconcile]] as soon as it closes a source, so that a task stops paying for a batch once + * native has it. A transfer in is charged at the next allocation, which is refused if the pool + * cannot cover both. + * + * '''Only [[onPreAllocation]] may throw.''' Arrow documents that the other callbacks cannot, and + * `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so + * throwing from any of them loses a buffer Arrow has already created. They only ever settle or + * trim the reservation, and contain anything the memory manager throws. How the acquisition in + * [[onPreAllocation]] can fail is described at [[acquire]]. + * + * '''Lock order.''' This listener's monitor is taken before Spark's and never the other way + * round: [[onPreAllocation]] holds ours across `acquireExecutionMemory`, and [[acquire]] + * additionally takes the `TaskMemoryManager` monitor itself, so that the two usage snapshots + * either side of that call cannot be split by another consumer. [[getUsed]] and [[spill]] must + * therefore stay lock-free, because Spark calls both while holding its own monitor: were either + * to take ours, a native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio + * thread could hold Spark's monitor and wait for ours while an Arrow allocation on the same task + * held ours and waited for Spark's. For the same reason, nothing reachable from a `spill` + * callback may allocate JVM Arrow memory. + */ +private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) + with AllocationListener { + + import CometArrowAllocationListener._ + + /** + * Bytes the task's allocator currently owns, read from Arrow's own accountant, which is an + * atomic, so [[getUsed]] can read it without taking this listener's monitor; see the lock order + * note above. Set by [[bind]], because the allocator is built with this listener and so cannot + * exist before it. + */ + @volatile private var owned: LongSupplier = NothingOwned + + /** Bytes currently reserved with Spark. Guarded by this listener's monitor. */ + private var reserved = 0L + + /** + * Bytes admitted by [[onPreAllocation]] whose allocation Arrow has not reported yet. They count + * against the reservation, so that two allocations on different threads cannot both be admitted + * against the same headroom. Guarded by this listener's monitor. + */ + private var pending = 0L + + /** Set once the owning task has finished. Volatile so [[getUsed]] can read it lock-free. */ + @volatile private var completed = false + + /** Attaches the listener to what it accounts for. Called once, before the allocator is used. */ + private[comet] def bind(owned: LongSupplier): Unit = this.owned = owned + + /** + * Reserves what the allocation needs before Arrow makes it, or refuses it. Once the task has + * finished there is nothing left to charge, so a straggling allocation goes through. + */ + override def onPreAllocation(size: Long): Unit = synchronized { + if (!completed) { + val needed = owned.getAsLong + pending + size + if (reserved < needed) { + reserved += acquireAtLeast(needed - reserved, roundUpToBlock(needed) - reserved, size) + } + pending += size + } + } + + override def onAllocation(size: Long): Unit = settle(size) + + /** Arrow's own limit refused the allocation after it was admitted here, so hand it back. */ + override def onFailedAllocation(size: Long, outcome: AllocationOutcome): Boolean = { + settle(size) + false + } + + override def onRelease(size: Long): Unit = reconcile() + + /** Trims the reservation after buffers moved out of the allocator without a callback. */ + private[comet] def reconcile(): Unit = settle(0L) + + /** + * Reports what the task's allocator owns. 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`; the reservation is made + * on Arrow's behalf, not through the consumer's own page accounting. + * + * Reports zero once the task has finished, so that buffers deliberately allowed to outlive + * their task are not reported by `cleanUpAllAllocatedMemory` as a Spark memory leak. + */ + override def getUsed: Long = if (completed) 0L else owned.getAsLong + + /** The JVM holds these buffers until the batch is done with, so there is nothing to spill. */ + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + + /** + * Drops the whole reservation and stops accounting. + * + * Called from the owning task's completion listener. Anything still alive afterwards is a + * buffer that outlives its task, which the process-wide allocator exists to allow; those + * releases are ignored rather than charged to whichever task happens to be running by then. + */ + private[comet] def taskCompleted(): Unit = { + try { + synchronized { + completed = true + releaseDownTo(0L) + } + } catch { + case NonFatal(e) => warnOnReleaseFailure(e) + } + } + + /** Bytes currently reserved with Spark on this task's behalf. Visible for testing. */ + private[comet] def reservedBytes: Long = synchronized(reserved) + + /** + * Moves `size` bytes out of pending, since Arrow has put them in the allocator's accountant by + * the time it calls back, and trims the reservation to what is owned and pending, rounded up to + * a block. The rounding keeps a change within a block from reaching Spark, although a change + * that crosses a block boundary still does every time. Releasing cannot fail in practice, but + * this is reached from callbacks that may not throw, so anything unforeseen is contained. + */ + private def settle(size: Long): Unit = { + try { + synchronized { + pending = math.max(0L, pending - size) + if (!completed) { + releaseDownTo(roundUpToBlock(owned.getAsLong + pending)) + } + } + } catch { + case NonFatal(e) => warnOnReleaseFailure(e) + } + } + + /** + * Returns everything reserved above `keep`, in one call rather than one per block, because + * `releaseExecutionMemory` synchronizes on the executor-wide pool. Called with this listener's + * monitor held. + */ + private def releaseDownTo(keep: Long): Unit = { + if (reserved > keep) { + taskMemoryManager.releaseExecutionMemory(reserved - keep, this) + reserved = keep + } + } + + /** + * Asks Spark for `request` bytes, which takes the reservation up to a block boundary so that + * the next small allocation does not go straight back into Spark's lock, and returns what it + * granted. A grant short of `request` is accepted as long as it covers `shortfall`, because the + * rounding is there to save lock traffic and must not be the reason an allocation that fits is + * refused. Anything less is handed back and the allocation refused. + */ + private def acquireAtLeast(shortfall: Long, request: Long, size: Long): Long = { + val granted = acquire(request, size) + if (granted < shortfall) { + if (granted > 0L) { + taskMemoryManager.releaseExecutionMemory(granted, this) + } + throw refusal(size, shortfall, granted, null) + } + granted + } + + /** + * Asks Spark for `request` bytes and returns what it granted, or refuses the allocation if the + * acquisition fails. + * + * It is fallible in three ways, and only the first is caught by `NonFatal`: it runs other + * consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an I/O failure + * into a `SparkOutOfMemoryError`, and the execution pool itself parks in `lock.wait()`, so + * killing a task can raise a plain `InterruptedException`. Each becomes the refusal's cause, + * and an interrupt also re-arms the thread's flag. + * + * `acquireExecutionMemory` takes its first grant from the pool and only then asks other + * consumers to spill, so when a spill throws it has already charged the task for bytes it never + * reports back, and nothing would release them before `cleanUpAllAllocatedMemory` at the very + * end of the task. They are measured as the change in what the pool says this task holds, and + * released before the refusal is thrown. + * + * Spark reports that figure per task rather than per consumer, so it only measures '''our''' + * grant if nothing else in the task can move it while we are looking. Both snapshots and the + * acquisition therefore run as one transaction under the `TaskMemoryManager` monitor. That is + * the same monitor `acquireExecutionMemory` takes and holds for its whole duration, spills + * included, and it is reentrant, so taking it here only widens that window to cover the two + * reads. Every acquisition in the task funnels through that method, so with it held no other + * consumer can take memory between a snapshot and the call and have it released here. + * + * What the monitor does not cover is a release, which reaches the pool without it. Another + * consumer returning memory, or a spill that frees some bytes before throwing, makes the figure + * too small, which is the safe direction: too small leaves bytes charged to the task until it + * ends, which is what would have happened anyway. `request` bounds it from above. + */ + private def acquire(request: Long, size: Long): Long = taskMemoryManager.synchronized { + val heldBefore = taskMemoryManager.getMemoryConsumptionForThisTask + try { + taskMemoryManager.acquireExecutionMemory(request, this) + } catch { + case e @ (_: InterruptedException | _: SparkOutOfMemoryError | NonFatal(_)) => + if (e.isInstanceOf[InterruptedException]) { + // Refusing is the only way out of an Arrow callback, but the cancellation must not be + // lost with it: re-arm the flag so the task still observes it. + Thread.currentThread().interrupt() + } + val orphaned = math.max( + 0L, + math.min(taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore, request)) + if (orphaned > 0L) { + taskMemoryManager.releaseExecutionMemory(orphaned, this) + } + throw refusal(size, request, 0L, e) + } + } + + private def refusal( + size: Long, + needed: Long, + granted: Long, + cause: Throwable): OutOfMemoryException = { + // As Spark's own consumers do before failing, so the task's consumers are in the log. + taskMemoryManager.showMemoryUsage() + new OutOfMemoryException( + s"Unable to reserve $needed bytes of Spark off-heap execution memory for a JVM Arrow " + + s"allocation of $size bytes, got $granted. Increase spark.memory.offHeap.size, or set " + + s"${CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key}=true to restore the previous " + + "behavior of not charging JVM Arrow allocations to Spark.", + cause) + } +} + +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[comet] val BLOCK_SIZE = 1024L * 1024L + + private val NothingOwned: LongSupplier = () => 0L + + private val releaseFailureLogged = new AtomicBoolean(false) + + private def roundUpToBlock(bytes: Long): Long = + ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE + + private def warnOnReleaseFailure(e: Throwable): Unit = { + if (releaseFailureLogged.compareAndSet(false, true)) { + logWarning( + "Failed to return JVM Arrow memory to Spark's memory manager. The task keeps it " + + "reserved until it ends, when Spark releases everything the task still holds.", + e) + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala new file mode 100644 index 00000000000..7405872d76b --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala @@ -0,0 +1,213 @@ +/* + * 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, ConcurrentLinkedQueue} + +import scala.util.control.NonFatal + +import org.apache.arrow.memory.BufferAllocator +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.memory.MemoryMode +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.vector.CometVector + +/** + * Hands out the Arrow allocator that JVM-owned allocations should use, one per Spark task. + * + * Each task gets a child of `CometArrowAllocator` carrying its own + * [[CometArrowAllocationListener]], so what that allocator owns is charged to the task's + * `TaskMemoryManager`, and an allocation Spark cannot cover fails with an `OutOfMemoryException` + * from `buffer`. The allocator, not the calling thread, is what identifies the owner: Arrow's + * `AllocationListener` is given only a size, and a buffer is released on whichever thread drops + * the last reference, which for anything that reaches native over the C Data Interface is a Comet + * Tokio worker with no task context installed. Child allocators cut from the returned allocator + * inherit its listener and roll their bytes up into it, so the paths that make their own children + * are covered too. + * + * '''This is for buffers the JVM owns.''' Buffers allocated to be handed straight to native -- by + * `NativeUtil`, for the JVM UDF result, and in `CometArrowStream.stream` -- come from the + * listener-less root instead, and imports from native come from `CometArrowImportAllocator`, a + * listener-less child of it. Native memory is native's to account for: a native operator that + * retains a batch reserves its buffers through Comet's pool, which charges the same Spark task, + * so charging them here as well would reserve the same memory twice. A batch read into this + * allocator and then streamed to native stops being charged here once the stream has it; see + * [[closeAndReconcile]]. The root is also what callers with no task get, such as broadcast + * coalescing on the driver or a reader built while native pulls a stream, and what Comet's + * on-heap mode gets, where charging an off-heap consumer would be wrong. + * + * '''Lifetime.''' The task allocator cannot simply be closed when the task ends. The process-wide + * allocator exists precisely because Arrow buffers can outlive the task that created them, and + * Arrow treats closing an allocator that still owns bytes as a leak. So at task completion the + * Spark reservation is dropped and the allocator is closed only if it has been drained; otherwise + * it is parked and closed by a later task once the stragglers are released. Leaving it open + * indefinitely is not an option: `BaseAllocator` keeps every child in a map until it closes. + */ +object CometTaskArrowAllocator extends Logging { + + private class TaskAllocator( + val allocator: BufferAllocator, + val listener: CometArrowAllocationListener) + + private val perTask = new ConcurrentHashMap[Long, TaskAllocator]() + + /** Allocators whose task has ended but which still own bytes. Guarded by [[closeLock]]. */ + private val awaitingClose = new ConcurrentLinkedQueue[BufferAllocator]() + + private val closeLock = new Object + + /** + * Whether tasks get an accounted allocator, which the deprecated + * `spark.comet.legacy.unboundedJvmArrowMemory` turns off. Read from the `SparkConf` rather than + * `SQLConf`, because this is reached from executor threads where `SQLConf` does not carry + * Comet's settings, and when a task first needs its allocator rather than once per JVM, so that + * the value belongs to whichever `SparkContext` is running. Tests that install a task context + * by hand have no `SparkEnv`, and account as the default does. + */ + private def accountingEnabled: Boolean = !Option(SparkEnv.get).exists { env => + env.conf.getBoolean( + CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key, + CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.defaultValue.get) + } + + /** + * The allocator to use for JVM-owned Arrow buffers on the calling thread. Never null, and never + * an allocator belonging to a task other than the current one. + */ + def forCurrentTask(): BufferAllocator = { + // Cheapest check first, and the one that eliminates the most callers: the driver, and native + // threads pulling a stream, allocate with no task in scope. + val taskContext = TaskContext.get() + if (taskContext == null) { + CometArrowAllocator + } else { + val existing = perTask.get(taskContext.taskAttemptId()) + if (existing != null) existing.allocator else create(taskContext) + } + } + + private def create(taskContext: TaskContext): BufferAllocator = { + val taskMemoryManager = taskContext.taskMemoryManager() + // Comet's on-heap mode accounts for nothing and exists so the Spark SQL suite can run without + // off-heap memory configured, so it gets the root: the same switch that + // `CometShuffleMemoryAllocator.getInstance` makes for shuffle pages. + if (taskMemoryManager == null || + taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP || !accountingEnabled) { + return CometArrowAllocator + } + + val taskAttemptId = taskContext.taskAttemptId() + closeDrained() + val listener = new CometArrowAllocationListener(taskMemoryManager) + val allocator = CometArrowAllocator + .newChildAllocator(s"comet-task-$taskAttemptId", listener, 0L, Long.MaxValue) + listener.bind(() => allocator.getAllocatedMemory) + val created = new TaskAllocator(allocator, listener) + val previous = perTask.putIfAbsent(taskAttemptId, created) + if (previous != null) { + // Lost the race with another thread in the same task; the loser's allocator is untouched. + closeQuietly(allocator) + previous.allocator + } else { + // Deliberately outside the map operation: `addTaskCompletionListener` runs the callback + // inline if the task has already completed, and that callback removes from this same map, + // which would be a recursive update inside a mapping function. + taskContext.addTaskCompletionListener[Unit](_ => taskCompleted(taskAttemptId)) + // ...and if it did run inline, the allocator just created is already closed, so hand back + // the root rather than something the caller cannot allocate from. + if (perTask.containsKey(taskAttemptId)) allocator else CometArrowAllocator + } + } + + private def taskCompleted(taskAttemptId: Long): Unit = { + val finished = perTask.remove(taskAttemptId) + if (finished != null) { + finished.listener.taskCompleted() + closeLock.synchronized(awaitingClose.add(finished.allocator)) + } + closeDrained() + } + + /** Closes every parked allocator that has been drained, including one parked just now. */ + private def closeDrained(): Unit = { + if (!awaitingClose.isEmpty) { + closeLock.synchronized { + val parked = awaitingClose.iterator() + while (parked.hasNext) { + val allocator = parked.next() + if (allocator.getAllocatedMemory == 0L) { + closeQuietly(allocator) + parked.remove() + } + } + } + } + } + + private def closeQuietly(allocator: BufferAllocator): Unit = { + try { + allocator.close() + } catch { + case NonFatal(e) => + // Closing is bookkeeping: the bytes are already gone and the Spark reservation is already + // released, so a failure here must not propagate into a task completion listener. + logWarning(s"Failed to close Arrow allocator ${allocator.getName}", e) + } + } + + /** + * Closes a batch whose buffers another allocator has just retained, then trims the reservation + * of the task allocator they came from. Closing hands ownership to the retaining allocator + * without calling any listener, so without this the task would go on paying for them until its + * next allocation or release. A batch is built within one task, so it has at most one task + * allocator, and a `NullVector` owns no buffers and has no allocator at all. + */ + private[spark] def closeAndReconcile(batch: ColumnarBatch): Unit = { + var listener: CometArrowAllocationListener = null + var i = 0 + while (i < batch.numCols()) { + batch.column(i) match { + case vector: CometVector => + val allocator = vector.getValueVector.getAllocator + if (allocator != null) { + allocator.getListener match { + case l: CometArrowAllocationListener => listener = l + case _ => + } + } + case _ => + } + i += 1 + } + batch.close() + if (listener != null) { + listener.reconcile() + } + } + + /** The listener accounting for the given task, if it has one. Visible for testing. */ + private[comet] def listenerForTask( + taskAttemptId: Long): Option[CometArrowAllocationListener] = { + Option(perTask.get(taskAttemptId)).map(_.listener) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 47691556848..0414683351e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -24,6 +24,7 @@ import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float = import scala.collection.JavaConverters._ import org.apache.spark.TaskContext +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, StartsWith, UnsafeProjection} @@ -38,7 +39,7 @@ import org.apache.spark.storage.StorageLevel import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.io.ChunkedByteBuffer -import org.apache.comet.{CometArrowAllocator, CometConf, DataTypeSupport} +import org.apache.comet.{CometConf, DataTypeSupport} import org.apache.comet.vector.NativeUtil /** @@ -398,6 +399,9 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val readerFields = arrowSchema.getFields.asScala.toIndexedSeq val codec = CachedBatchIpc.compressionCodec(settings.codecName, settings.zstdLevel) val orderings = boundsOrderings(attrs) + // Everything allocated here is transient and closed before the next batch, so it is charged to + // the task that is writing the cache. + val allocator = CometTaskArrowAllocator.forCurrentTask() batches.map { batch => // Bounds and null counts are read from the input batch before it is serialized, and the row @@ -413,11 +417,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { Utils.isArrowBacked(batch) && CachedBatchIpc.matchesReaderLayout(batch, readerFields) val (bytes, columnSizes) = if (writeDirectly) { - CachedBatchIpc.serialize(batch, codec, CometArrowAllocator, settings.chunkSize) + CachedBatchIpc.serialize(batch, codec, allocator, settings.chunkSize) } else { val arrowBatch = - CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) - try CachedBatchIpc.serialize(arrowBatch, codec, CometArrowAllocator, settings.chunkSize) + CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, allocator) + try CachedBatchIpc.serialize(arrowBatch, codec, allocator, settings.chunkSize) finally arrowBatch.close() } @@ -605,7 +609,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Decoding happens during construction, so `batches` below can hand out the root directly. // `load` releases everything it allocated if it throws, so there is nothing to unwind here. - private val root = projection.load(cached.bytes, CometArrowAllocator) + private val root = projection.load(cached.bytes, CometTaskArrowAllocator.forCurrentTask()) private var closed = false // A cached batch's columns all cover the same rows. Check rather than trust: a mismatch would @@ -674,7 +678,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // of session timezone, so no values are converted. It also matches Comet's native // schema, avoiding a cast at the native boundary. CometArrowStream.NATIVE_TIMEZONE, - CometArrowAllocator) + CometTaskArrowAllocator.forCurrentTask()) encodeBatches(iter, schema, settings) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala index 4a15355d269..58f1b2835a8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ColumnarBatchArrowReader.scala @@ -25,6 +25,7 @@ import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.{VectorSchemaRoot, VectorUnloader} import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.Schema +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.vector.{CometVector, CometVectorUtils} @@ -69,7 +70,9 @@ private[comet] class ColumnarBatchArrowReader( // Do not close this borrowed root: src and the helper own its vectors. } } finally { - src.close() + // Closing leaves this reader's allocator owning the buffers it retained, so the task + // allocator they came from stops paying for them here. + CometTaskArrowAllocator.closeAndReconcile(src) } true } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index a60e2811ecb..f709cb30739 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -27,6 +27,7 @@ import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema} import org.apache.spark.TaskContext +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.comet.util.Utils @@ -230,6 +231,8 @@ object CometArrowStream extends Logging { name: String, readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = { val context = TaskContext.get() + // A child of the listener-less root rather than of the task allocator, since everything loaded + // here is exported to native, which accounts for what it retains. See CometTaskArrowAllocator. val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) var reader: ArrowReader = null var arrowStream: ArrowArrayStream = null @@ -274,7 +277,10 @@ object CometArrowStream extends Logging { name: String, readerFactory: BufferAllocator => ArrowReader): Iterator[ColumnarBatch] = { val context = TaskContext.get() - val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) + // Accounted, unlike `stream`: these batches are consumed in the JVM, so nothing on the native + // side reserves them. + val allocator = + CometTaskArrowAllocator.forCurrentTask().newChildAllocator(name, 0, Long.MaxValue) val reader = try readerFactory(allocator) catch { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index eb5947f30c0..4351b32dcc3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -340,7 +340,8 @@ object Utils extends CometTypeShim with Logging { return (Array.empty, 0L, 0L) } - val allocator = org.apache.comet.CometArrowAllocator + val allocator = org.apache.spark.comet.CometTaskArrowAllocator + .forCurrentTask() .newChildAllocator("broadcast-coalesce", 0, Long.MaxValue) try { var targetRoot: VectorSchemaRoot = null @@ -548,7 +549,7 @@ object Utils extends CometTypeShim with Logging { cv.dataType(), rows, s"_const_$index", - org.apache.comet.CometArrowAllocator, + org.apache.spark.comet.CometTaskArrowAllocator.forCurrentTask(), "UTC") (materialized, None) diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index fecf781baac..3229f11d7c3 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -34,13 +34,13 @@ import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch, Me import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType} import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths} +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.{CometDecodedVector, CometDictionaryVector, CometVector, CometVectorUtils} /** @@ -124,7 +124,9 @@ private[python] trait CometArrowPythonRunnerBase new Writer(env, worker, inputIterator, partitionIndex, context) { private val allocator = - CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue) + CometTaskArrowAllocator + .forCurrentTask() + .newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue) private var batches = inputIterator.flatten // Upstream owns this batch. Even hasNext may close it and reuse its native buffers, so // leave the upstream iterators untouched until all ranges have been serialized. @@ -264,7 +266,9 @@ private[python] trait CometArrowPythonRunnerBase new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) { private val allocator = - CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue) + CometTaskArrowAllocator + .forCurrentTask() + .newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue) private var reader: ArrowStreamReader = _ private var root: VectorSchemaRoot = _ private var batchLoaded = true diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala new file mode 100644 index 00000000000..9043d7f8430 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala @@ -0,0 +1,236 @@ +/* + * 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.atomic.{AtomicBoolean, AtomicLong} + +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator} +import org.apache.spark.{CometTaskMemoryManager, TaskContext} +import org.apache.spark.benchmark.{Benchmark, BenchmarkBase} +import org.apache.spark.memory.{MemoryConsumer, MemoryManager, TaskMemoryManager} + +/** + * Measures what charging JVM Arrow allocations to Spark costs, since it is on by default. + * + * Each case runs the same allocate/release loop twice: once against a plain `RootAllocator`, + * which is what Comet did before [[CometArrowAllocationListener]] existed and what + * `spark.comet.legacy.unboundedJvmArrowMemory=true` restores, and once against a task allocator + * carrying the listener. Reservation call counts are printed under each table, because the + * interesting variable is not the per-buffer bookkeeping but how often a buffer size crosses a + * block boundary and has to go into `TaskMemoryManager` at all. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.comet.CometArrowAllocationListenerBenchmark + * }}} + */ +object CometArrowAllocationListenerBenchmark extends BenchmarkBase { + + private val blockSize = CometArrowAllocationListener.BLOCK_SIZE + private val poolBytes = 1024L * 1024L * 1024L + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + runBenchmark("JVM Arrow allocations charged to Spark") { + // Many sub-block buffers, the shape a codegen output vector's validity and offset buffers + // take. None of them comes close to a block, so the listener should never reach Spark after + // the first one. + allocateAndRelease("small buffers", bufferSize = 128L, buffersPerIteration = 512) + // A wide batch: many medium buffers alive at once, one block boundary crossed per few + // buffers on the way up and the same on the way down. + allocateAndRelease( + "wide batch buffers", + bufferSize = 64L * 1024L, + buffersPerIteration = 128) + // Worst case for the block batching: every allocation crosses a boundary, so every one of + // them takes the executor-wide lock in `acquireExecutionMemory`. + allocateAndRelease("block-sized buffers", bufferSize = blockSize, buffersPerIteration = 8) + // Same, with a second thread reserving from the same pool the way Comet's native side does. + // The pool has room for the whole set plus that thread's block, because a smaller one would + // refuse the allocations rather than slow them down, so what this measures is contention for + // the pool's lock rather than for its capacity. In the unaccounted arm only the pressure + // thread touches the pool: the delta is what the Arrow side adds by competing. + allocateAndRelease( + "block-sized buffers under native pressure", + bufferSize = blockSize, + buffersPerIteration = 8, + pool = blockSize * 9, + nativePressure = true) + // Per call site rather than per buffer, but it is the cost the root allocator `val` did not + // have: a TaskContext lookup and a concurrent map read. + allocatorLookup() + } + } + + private def allocateAndRelease( + name: String, + bufferSize: Long, + buffersPerIteration: Int, + pool: Long = poolBytes, + nativePressure: Boolean = false): Unit = { + val benchmark = + new Benchmark( + s"$name (${buffersPerIteration}x$bufferSize)", + buffersPerIteration, + output = output) + + // Both allocators are built once, outside the timed body, so what is measured is the + // steady-state cost of allocating and releasing rather than the cost of standing a task up. + val root = new RootAllocator(Long.MaxValue) + try { + withTaskAllocator(pool) { (accounted, memory) => + withNativePressure(nativePressure) { + benchmark.addCase("not accounted") { _ => + churn(root, bufferSize, buffersPerIteration) + } + benchmark.addCase("accounted") { _ => + churn(accounted, bufferSize, buffersPerIteration) + } + benchmark.run() + + // One more round with the counters zeroed, to report how often a single iteration + // reaches the memory manager. That, rather than the per-buffer bookkeeping, is the cost + // that scales with buffer size. + memory.reset() + churn(accounted, bufferSize, buffersPerIteration) + writeLine(s" accounted: ${memory.summary(buffersPerIteration)} per iteration") + } + } + } finally { + root.close() + } + } + + private def allocatorLookup(): Unit = { + val lookupsPerIteration = 100000 + val benchmark = + new Benchmark("allocator lookup", lookupsPerIteration.toLong, output = output) + + // The "before" shape: call sites read a package-object `val`, which the JIT folds away + // entirely. The interesting number is therefore the absolute cost of the second case. + benchmark.addCase("process-wide val") { _ => + var i = 0 + var sink = 0 + while (i < lookupsPerIteration) { + sink += System.identityHashCode(org.apache.comet.CometArrowAllocator) + i += 1 + } + assert(sink != Int.MinValue) + } + benchmark.addCase("forCurrentTask()") { _ => + withTaskAllocator() { (_, _) => + var i = 0 + var sink = 0 + while (i < lookupsPerIteration) { + sink += System.identityHashCode(CometTaskArrowAllocator.forCurrentTask()) + i += 1 + } + assert(sink != Int.MinValue) + } + } + + benchmark.run() + } + + /** + * Allocates the whole set, then releases it, so the peak is what the reservation has to cover. + */ + private def churn(allocator: BufferAllocator, bufferSize: Long, count: Int): Unit = { + val buffers = new Array[ArrowBuf](count) + var i = 0 + while (i < count) { + buffers(i) = allocator.buffer(bufferSize) + i += 1 + } + i = 0 + while (i < count) { + buffers(i).close() + i += 1 + } + } + + /** + * Runs the body against a task allocator, torn down afterwards, so that repeated iterations do + * not accumulate reservations or allocators. + */ + private def withTaskAllocator[T](pool: Long = poolBytes)( + f: (BufferAllocator, CountingTaskMemoryManager) => T): T = { + val (context, memory) = TestTasks.newTask(pool)(new CountingTaskMemoryManager(_, _)) + TestTasks.withInstalled(context)(f(CometTaskArrowAllocator.forCurrentTask(), memory)) + } + + /** + * Hammers the task's pool from another thread for the duration of the body, when asked to, the + * way native reservations do. Built on the task's thread, which `CometTaskMemoryManager` needs. + */ + private def withNativePressure[T](enabled: Boolean)(f: => T): T = { + if (!enabled) { + f + } else { + val stop = new AtomicBoolean(false) + val native = new CometTaskMemoryManager(0L, TaskContext.get().taskAttemptId()) + val thread = new Thread(() => { + while (!stop.get()) { + native.releaseMemory(native.acquireMemory(blockSize)) + } + }) + thread.setDaemon(true) + thread.setName("native-reservations") + thread.start() + try f + finally { + stop.set(true) + thread.join() + } + } + } + + private def writeLine(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes("UTF-8"))) + } + + private class CountingTaskMemoryManager(memoryManager: MemoryManager, taskAttemptId: Long) + extends TaskMemoryManager(memoryManager, taskAttemptId) { + private val acquires = new AtomicLong(0L) + private val releases = new AtomicLong(0L) + + // Only the Arrow listener's own calls are counted, so the pressure thread's traffic does not + // land in the reported figure. + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + if (consumer.isInstanceOf[CometArrowAllocationListener]) acquires.incrementAndGet() + super.acquireExecutionMemory(required, consumer) + } + + override def releaseExecutionMemory(size: Long, consumer: MemoryConsumer): Unit = { + if (consumer.isInstanceOf[CometArrowAllocationListener]) releases.incrementAndGet() + super.releaseExecutionMemory(size, consumer) + } + + def reset(): Unit = { + acquires.set(0L) + releases.set(0L) + } + + def summary(buffers: Int): String = + s"${acquires.get()} acquire and ${releases.get()} release calls for $buffers buffers" + } +} diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala new file mode 100644 index 00000000000..0f6f1670cb8 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala @@ -0,0 +1,761 @@ +/* + * 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.io.{InterruptedIOException, IOException} +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} + +import scala.jdk.CollectionConverters._ + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.arrow.c.Data +import org.apache.arrow.memory.OutOfMemoryException +import org.apache.arrow.vector.{FieldVector, IntVector, VectorSchemaRoot} +import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkContext, TaskContext, TaskContextImpl} +import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager} +import org.apache.spark.sql.comet.execution.arrow.CometArrowStream +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.ThreadUtils + +import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometConf} +import org.apache.comet.vector.NativeUtil + +/** + * Tests that JVM Arrow allocations are charged to Spark, that they are charged to the task that + * made them rather than whichever task happens to be on the releasing thread, that what is + * charged follows Arrow's ownership as buffers move between allocators, and that an allocation + * Spark cannot cover is refused before Arrow makes it, leaving no buffer, reservation or grant + * behind. + */ +class CometArrowAllocationListenerSuite extends AnyFunSuite { + + private val blockSize = CometArrowAllocationListener.BLOCK_SIZE + private val poolBytes = 64L * 1024 * 1024 + + // --------------------------------------------------------------------------------------------- + // Reservation arithmetic. Driven through the listener directly, with a stand-in for the + // allocator's accountant, since Arrow's rounding policy would otherwise decide the sizes under + // test. + // --------------------------------------------------------------------------------------------- + + test("allocations are charged to the current task in whole blocks") { + withTask() { task => + val owner = new StubOwner(task.taskMemoryManager) + val listener = owner.listener + + // Far smaller than a block, so the reservation should round up to exactly one block. + owner.allocate(128L) + assert(listener.reservedBytes == blockSize) + + // Still inside the first block, so Spark is not asked again. + owner.allocate(1024L) + assert(listener.reservedBytes == blockSize) + + listener.taskCompleted() + } + } + + test("a request larger than a block rounds up to a block multiple") { + withTask() { task => + val owner = new StubOwner(task.taskMemoryManager) + owner.allocate(blockSize * 3 + 7L) + val listener = owner.listener + // Rounded up rather than sized to the exact deficit, so growth leaves headroom and the next + // small allocation does not go straight back into Spark. + assert(listener.reservedBytes == blockSize * 4) + listener.taskCompleted() + } + } + + test("releasing returns whole blocks to Spark") { + withTask() { task => + val owner = new StubOwner(task.taskMemoryManager) + val listener = owner.listener + owner.allocate(blockSize * 2) + assert(listener.reservedBytes == blockSize * 2) + + owner.release(blockSize * 2) + assert(listener.reservedBytes == 0L) + listener.taskCompleted() + } + } + + // --------------------------------------------------------------------------------------------- + // Which allocator is handed out, and what it charges. + // --------------------------------------------------------------------------------------------- + + test("a real Arrow allocation from the task allocator is charged to that task") { + withTask() { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + assert(allocator ne CometArrowAllocator) + val buf = allocator.buffer(blockSize) + try { + assert(reservedFor(task) == blockSize) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + } finally { + buf.close() + } + assert(reservedFor(task) == 0L) + } + } + + test("a child of the task allocator is charged to the same task") { + withTask() { task => + // The Python runner and the native Arrow source cut their own children. Arrow passes the + // parent's listener down, so they are accounted without knowing anything about it. + val child = + CometTaskArrowAllocator.forCurrentTask().newChildAllocator("probe", 0L, Long.MaxValue) + try { + val buf = child.buffer(blockSize) + try { + assert(reservedFor(task) == blockSize) + } finally { + buf.close() + } + } finally { + child.close() + } + assert(reservedFor(task) == 0L) + } + } + + test("the allocators on either side of the FFI boundary are not accounted") { + withTask() { task => + // Establish the task allocator first, so this asserts "not charged" rather than "no task". + CometTaskArrowAllocator.forCurrentTask() + // Buffers allocated for export come from the listener-less root, and imports from a + // listener-less child of it, because native accounts for what it retains. Arrow reports an + // imported buffer to the importing allocator's listener at full capacity, so giving the + // import allocator a listener would charge Spark for native memory. + for (allocator <- Seq(CometArrowAllocator, CometArrowImportAllocator)) { + val buf = allocator.buffer(blockSize) + try { + assert(reservedFor(task) == 0L, s"${allocator.getName} was charged to the task") + } finally { + buf.close() + } + } + } + } + + test("no active task uses the unaccounted root allocator") { + TaskContext.unset() + // Broadcast coalescing on the driver, and native threads pulling a stream, have no task. + assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator) + } + + test("on-heap mode uses the unaccounted root allocator") { + withTask(offHeap = false) { task => + // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory. + // Charging an off-heap consumer there would be wrong. + assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator) + assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty) + } + } + + test("the legacy setting hands tasks the unaccounted root, and is off by default") { + // Read from the executors' SparkConf, so each arm runs a task in a context of its own. + def taskGetsRoot(legacy: Option[String]): Boolean = { + val conf = new SparkConf() + .setMaster("local[1]") + .setAppName(getClass.getSimpleName) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", poolBytes.toString) + legacy.foreach(conf.set(CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key, _)) + val sc = new SparkContext(conf) + try { + sc.parallelize(Seq(0), 1) + .map(_ => CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator) + .collect() + .head + } finally { + sc.stop() + } + } + assert(!taskGetsRoot(None), "a task was not accounted by default") + assert(taskGetsRoot(Some("true")), "the legacy setting did not restore the unaccounted root") + } + + // --------------------------------------------------------------------------------------------- + // Ownership: a release is attributed to the task that allocated, not to the releasing thread. + // --------------------------------------------------------------------------------------------- + + test("a release on a thread with no task context is charged to the allocating task") { + withTask() { task => + val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + assert(reservedFor(task) == blockSize) + + // This is what happens when a shuffle-read batch is handed on to a native operator: native + // pins it and drops it later from a Tokio worker with no task context installed. Reading + // TaskContext in onRelease would ignore this release and leave the task charged for memory + // it had already freed. + ThreadUtils.runInNewThread("detached-release")(buf.close()) + + assert(reservedFor(task) == 0L) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + + test("a release under a different task does not touch that task's accounting") { + withTask() { taskA => + val bufA = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + withTask() { taskB => + val bufB = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + assert(reservedFor(taskB) == blockSize) + + // A's buffer, released while B's context is on the thread. B must not pay for it. + bufA.close() + + assert(reservedFor(taskB) == blockSize) + assert(reservedFor(taskA) == 0L) + bufB.close() + } + } + } + + // --------------------------------------------------------------------------------------------- + // What is charged follows ownership. Arrow moves a buffer's charge between allocators with + // `transferBalance`, which calls no listener, so a tally of the callbacks drifts from what the + // task's allocator actually owns. + // --------------------------------------------------------------------------------------------- + + test("a buffer whose ownership moves to another allocator stops being charged to the task") { + withTask() { task => + val stream = CometArrowAllocator.newChildAllocator("stream", 0L, Long.MaxValue) + try { + val allocator = CometTaskArrowAllocator.forCurrentTask() + val source = allocator.buffer(blockSize * 4) + assert(reservedFor(task) == blockSize * 4) + + // What ColumnarBatchArrowReader does to every batch it streams to native: retain the + // buffers into the stream's allocator, then close the source. Arrow makes the stream's + // allocator the owner, and neither listener is told. + val retained = source.getReferenceManager.retain(source, stream) + try { + source.close() + assert(allocator.getAllocatedMemory == 0L) + assert(listenerFor(task).getUsed == 0L) + } finally { + // Native drops it later, and the release goes to the stream allocator's listener, which + // is the root's no-op. Nothing on the task's allocator hears about it. + retained.close() + } + + // With no reconcile, the reservation catches up at the next callback on the task's + // allocator: one block for the new buffer, not four for a buffer the task no longer owns. + val next = allocator.buffer(128L) + assert(reservedFor(task) == blockSize) + next.close() + assert(reservedFor(task) == 0L) + } finally { + stream.close() + } + } + } + + test("batches streamed to native stop being charged to the task that read them") { + withTask() { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + val importer = CometArrowAllocator.newChildAllocator("importer", 0L, Long.MaxValue) + val numBatches = 16 + // Half a block of values per batch, which Arrow allocates as one block-sized buffer. + val rowsPerBatch = (blockSize / 2 / IntVector.TYPE_WIDTH).toInt + val batches = Iterator.tabulate(numBatches) { b => + val vector = new IntVector("i", allocator) + vector.allocateNew(rowsPerBatch) + (0 until rowsPerBatch).foreach(i => vector.set(i, b + i)) + vector.setValueCount(rowsPerBatch) + val root = new VectorSchemaRoot(java.util.Arrays.asList[FieldVector](vector)) + root.setRowCount(rowsPerBatch) + NativeUtil.rootAsBatch(root): ColumnarBatch + } + try { + // The same export a native operator's JVM input takes, drained here by a JVM importer + // standing in for native's ScanExec. + val stream = CometArrowStream.fromColumnarBatchIter( + batches, + StructType(Seq(StructField("i", IntegerType))), + CometArrowStream.NATIVE_TIMEZONE, + "listener-suite") + val reader = Data.importArrayStream(importer, stream) + var imported = 0 + try { + while (reader.loadNextBatch()) { + imported += 1 + // Native holds this batch now, and the task no longer owns any of it, so it should not + // be paying for it either. Summing the callbacks would have kept a block for every + // batch read so far, and waiting for the task's next allocation would have kept one + // for this batch, which for a broadcast coalesced into one batch is all of it. + assert(allocator.getAllocatedMemory == 0L) + assert( + reservedFor(task) == 0L, + s"reserved ${reservedFor(task)} bytes after handing over $imported batches") + } + } finally { + reader.close() + } + assert(imported == numBatches) + } finally { + importer.close() + } + } + } + + test("a buffer adopted from another allocator is charged while the task owns it") { + withTask() { task => + val other = CometArrowAllocator.newChildAllocator("other", 0L, Long.MaxValue) + try { + val allocator = CometTaskArrowAllocator.forCurrentTask() + val foreign = other.buffer(blockSize) + // The mirror image: the task allocator takes a reference and the original owner lets go, + // so Arrow makes the task allocator the owner without telling either listener. + val adopted = foreign.getReferenceManager.retain(foreign, allocator) + try { + foreign.close() + assert(allocator.getAllocatedMemory == blockSize) + assert(listenerFor(task).getUsed == blockSize) + } finally { + // The final release then reaches this task's listener for bytes it never saw allocated. + adopted.close() + } + // That must not leave the task's figure short, so the next allocation is charged in full. + val buf = allocator.buffer(blockSize) + try { + assert(reservedFor(task) == blockSize) + } finally { + buf.close() + } + assert(reservedFor(task) == 0L) + } finally { + other.close() + } + } + } + + // --------------------------------------------------------------------------------------------- + // Task completion, and buffers that outlive their task. + // --------------------------------------------------------------------------------------------- + + test("task completion releases the whole reservation and closes the allocator") { + val allocatorName = withTask() { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + buf.close() + + task.context.markTaskCompleted(None) + + assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + allocator.getName + } + assert(!rootChildNames().contains(allocatorName)) + } + + test("a buffer outliving its task parks the allocator until it is released") { + withTask() { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize) + + task.context.markTaskCompleted(None) + + // The reservation goes back to Spark even though the buffer is still alive: the task is + // over, and leaving it charged would be reported as a Spark memory leak. + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty) + // Arrow treats closing an allocator that still owns bytes as a leak, so it has to stay open. + assert(rootChildNames().contains(allocator.getName)) + + // A late release is ignored rather than charged to whoever is running by then... + buf.close() + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + + // ...and the drained allocator is reaped by the next task, so the root does not accumulate + // one child per task attempt. + withTask() { _ => CometTaskArrowAllocator.forCurrentTask() } + assert(!rootChildNames().contains(allocator.getName)) + } + } + + // --------------------------------------------------------------------------------------------- + // Refusal. An allocation Spark cannot cover is refused in onPreAllocation, before Arrow has made + // it, and nothing is left behind: no buffer, no reservation, and no grant charged to the task. + // --------------------------------------------------------------------------------------------- + + test("an allocation the pool cannot cover is refused, and one that fits still succeeds") { + // Two blocks of budget, one of them held by a consumer that cannot give anything back. + withTask(pool = blockSize * 2) { task => + val other = new OtherConsumer(task.taskMemoryManager) + assert(other.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize * 2)) + assert(refused.getMessage.contains(CometConf.COMET_LEGACY_UNBOUNDED_JVM_ARROW_MEMORY.key)) + // Spark granted the block it had before coming up short. That partial grant is handed back + // rather than kept, and Arrow never allocated anything. + assert(allocator.getAllocatedMemory == 0L) + assert(reservedFor(task) == 0L) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + + val buf = allocator.buffer(blockSize) + try { + assert(reservedFor(task) == blockSize) + } finally { + buf.close() + } + } + } + + test("a grant short of a block is accepted when it covers the allocation") { + // Reservations are requested in whole blocks to save lock traffic, but a pool with less than a + // block left must still admit an allocation that fits in what is left. + val left = 64L * 1024 + withTask(pool = blockSize) { task => + val other = new OtherConsumer(task.taskMemoryManager) + assert(other.take(blockSize - left) == blockSize - left) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(1024L) + try { + assert(reservedFor(task) == left) + // What does not fit in what is left is still refused. + intercept[OutOfMemoryException](allocator.buffer(left)) + assert(reservedFor(task) == left) + } finally { + buf.close() + } + // A reservation that is not a whole block is returned once nothing needs it. + assert(reservedFor(task) == 0L) + } + } + + test("an allocation still in flight counts against the reservation") { + withTask(pool = blockSize) { task => + val owner = new StubOwner(task.taskMemoryManager) + val listener = owner.listener + + // Admitted, but Arrow has not reported it yet... + listener.onPreAllocation(blockSize) + assert(listener.reservedBytes == blockSize) + // ...so a second allocation racing it on another thread cannot be admitted against the same + // block. + intercept[OutOfMemoryException](listener.onPreAllocation(blockSize)) + assert(listener.reservedBytes == blockSize) + + owner.made(blockSize) + owner.release(blockSize) + assert(listener.reservedBytes == 0L) + listener.taskCompleted() + } + } + + for ((label, failure) <- Seq( + "an I/O failure" -> new IOException("spill failed"), + "an interrupted spill" -> new InterruptedIOException("task killed"))) { + test(s"$label while acquiring refuses the allocation and leaks nothing") { + // Exactly one block of budget, already taken by a consumer whose spill throws, so the + // acquisition has to go through Spark's spill path and comes back throwing. + withTask(pool = blockSize) { task => + val hostile = new OtherConsumer(task.taskMemoryManager, Some(failure)) + assert(hostile.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize)) + // The memory manager's failure is the cause, which is what says the acquisition really + // went down the spill path rather than being refused for want of memory alone. + assert(refused.getCause != null) + assert(allocator.getAllocatedMemory == 0L) + assert(reservedFor(task) == 0L) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + } + } + } + + test("an interrupt while acquiring refuses the allocation and leaves the flag set") { + // Spark's execution pool parks in `lock.wait()` when a task is below its fair share, so a task + // kill raises a plain InterruptedException out of `acquireExecutionMemory`. TestMemoryManager + // never parks, so the interrupt is injected through a failing spill instead. + withTask(pool = blockSize) { task => + val hostile = + new OtherConsumer(task.taskMemoryManager, Some(new InterruptedException("task killed"))) + assert(hostile.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val refused = intercept[OutOfMemoryException](allocator.buffer(blockSize)) + assert(refused.getCause.isInstanceOf[InterruptedException]) + assert(allocator.getAllocatedMemory == 0L) + assert(reservedFor(task) == 0L) + // Cleared here as well as asserted, so the flag does not leak into the next test. + assert(Thread.interrupted(), "the interrupt was swallowed instead of being re-armed") + } + } + + test("a partial grant lost to a failing spill is released rather than stranded") { + // One block already taken, one still in the pool, and a two-block request: Spark hands over the + // block it has and only then asks the other consumer to spill, which throws. It never reports + // the block it already took, so nothing would release it before the task ended. + withTask(pool = blockSize * 2) { task => + val hostile = + new OtherConsumer(task.taskMemoryManager, Some(new IOException("spill failed"))) + assert(hostile.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + intercept[OutOfMemoryException](allocator.buffer(blockSize * 2)) + assert(reservedFor(task) == 0L) + // Only the other consumer's block is left. Without releasing the orphan this would be two + // blocks, one of them charged to the task and owned by nobody. + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + } + } + + test("a concurrent acquisition is not released as this listener's lost grant") { + // The orphan is measured as a change in the task's total consumption, which Spark reports per + // task rather than per consumer. If another consumer could acquire between the snapshot taken + // before the acquisition and the acquisition itself, its bytes would be counted as the orphan + // and released here, leaving that consumer holding bytes the pool no longer charges for. Both + // snapshots and the call therefore run as one transaction under the + // TaskMemoryManager monitor. This forces the interleaving that transaction exists to exclude. + val spare = 1024L + withTask(pool = blockSize * 2 + spare) { task => + val hostile = + new OtherConsumer(task.taskMemoryManager, Some(new IOException("spill failed"))) + assert(hostile.take(blockSize) == blockSize) + + val interloper = new OtherConsumer(task.taskMemoryManager) + // Once the acquisition has taken what was left, the pool is empty and Spark answers the + // interloper by asking the hostile consumer to spill, which throws. That is a legitimate + // outcome for the interloper and not what is under test here; what it ends up holding is. + val interloperThread = daemonThread("interloper") { + try interloper.take(spare) + catch { case _: SparkOutOfMemoryError => } + } + + // Fires once, in place of the snapshot taken before the acquisition: exactly the window the + // transaction has to close. The interloper either runs to completion here, which is the bug, + // or blocks on the monitor the acquisition is holding, which is the fix. + task.snapshotHook.set(() => { + interloperThread.start() + awaitBlockedOrFinished(interloperThread) + }) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + intercept[OutOfMemoryException](allocator.buffer(blockSize * 2)) + interloperThread.join(30000L) + assert( + !interloperThread.isAlive, + "the interloper never finished; the transaction deadlocked") + // Every byte the task is charged for belongs to someone. Without the transaction the + // listener counts the interloper's bytes as its own orphan and releases them, and the task + // comes out charged for less than the other two consumers hold. + assert( + listenerFor(task).reservedBytes + hostile.getUsed + interloper.getUsed == + task.taskMemoryManager.getMemoryConsumptionForThisTask, + "the listener released bytes belonging to another consumer") + } + } + + // --------------------------------------------------------------------------------------------- + // Lock order. Spark calls getUsed and spill while holding the TaskMemoryManager monitor, and the + // listener holds its own monitor while waiting for that one. + // --------------------------------------------------------------------------------------------- + + test("the usage snapshot does not take the reservation monitor") { + withTask() { task => + val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + try { + val listener = listenerFor(task) + val used = new AtomicLong(-1L) + val spilled = new AtomicLong(-1L) + listener.synchronized { + // A native reservation arriving through CometTaskMemoryManager on a Tokio thread holds + // Spark's monitor here. If either call waited on this one, it would deadlock against an + // Arrow allocation on the same task that already holds this monitor and wants Spark's. + val probe = daemonThread("lock-order-probe") { + used.set(listener.getUsed) + spilled.set(listener.spill(blockSize, listener)) + } + probe.start() + probe.join(30000L) + assert(!probe.isAlive, "getUsed or spill blocked on the reservation monitor") + } + assert(used.get == blockSize) + assert(spilled.get == 0L) + } finally { + buf.close() + } + } + } + + test("concurrent Arrow and native reservations make progress") { + // Two blocks of budget, and the native side asks for both, so whenever the Arrow side holds + // one the native grant comes up short and Spark walks its consumer list, calling getUsed and + // spill on the Arrow listener while holding its own monitor. Either side can come up short, + // and an Arrow allocation refused for it is a legitimate outcome: finishing is what is tested. + withTask(pool = blockSize * 2) { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + val native = new CometTaskMemoryManager(0L, task.taskAttemptId) + val failure = new AtomicReference[Throwable]() + + val arrowThread = loopingThread("arrow-allocations", failure) { + try allocator.buffer(blockSize).close() + catch { case _: OutOfMemoryException => } + } + val nativeThread = loopingThread("native-reservations", failure) { + native.releaseMemory(native.acquireMemory(blockSize * 2)) + } + + Seq(arrowThread, nativeThread).foreach(_.start()) + Seq(arrowThread, nativeThread).foreach { t => + t.join(60000L) + assert(!t.isAlive, s"${t.getName} did not finish; concurrent reservations deadlocked") + } + + Option(failure.get).foreach(e => fail("a worker failed", e)) + assert(allocator.getAllocatedMemory == 0L) + assert(native.getUsed == 0L) + } + } + + // --------------------------------------------------------------------------------------------- + // Fixtures. + // --------------------------------------------------------------------------------------------- + + /** + * A listener with a stand-in for the task allocator's accountant, calling the listener the way + * Arrow would. + */ + private class StubOwner(tmm: TaskMemoryManager) { + val listener = new CometArrowAllocationListener(tmm) + private val owned = new AtomicLong(0L) + listener.bind(() => owned.get) + + /** What `BaseAllocator.buffer` does around a successful allocation. */ + def allocate(bytes: Long): Unit = { + listener.onPreAllocation(bytes) + made(bytes) + } + + /** The second half of an allocation that `onPreAllocation` has already admitted. */ + def made(bytes: Long): Unit = { + owned.addAndGet(bytes) + listener.onAllocation(bytes) + } + + def release(bytes: Long): Unit = { + owned.addAndGet(-bytes) + listener.onRelease(bytes) + } + } + + /** + * Any other consumer in the task: holds memory and cannot give it back. With a `spillFailure`, + * its spill throws that instead, so `trySpillAndAcquire` throws on its behalf. + */ + private class OtherConsumer(tmm: TaskMemoryManager, spillFailure: Option[Exception] = None) + extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) { + def take(bytes: Long): Long = acquireMemory(bytes) + override def spill(size: Long, trigger: MemoryConsumer): Long = spillFailure match { + case Some(failure) => throw failure + case None => 0L + } + } + + /** + * Runs one action immediately after the usage snapshot the listener takes before asking Spark + * for memory, so a test can drive what happens in the window between that snapshot and the + * acquisition. The real value is read first, which is what makes the window the one under test: + * running the action before the read would fold whatever it does into the snapshot itself. + */ + private class HookedTaskMemoryManager( + memoryManager: MemoryManager, + taskAttemptId: Long, + val hook: AtomicReference[Runnable] = new AtomicReference[Runnable]()) + extends TaskMemoryManager(memoryManager, taskAttemptId) { + override def getMemoryConsumptionForThisTask(): Long = { + val held = super.getMemoryConsumptionForThisTask() + val pending = hook.getAndSet(null) + if (pending != null) pending.run() + held + } + } + + private case class TaskFixture( + context: TaskContextImpl, + taskMemoryManager: HookedTaskMemoryManager) { + def taskAttemptId: Long = context.taskAttemptId + def snapshotHook: AtomicReference[Runnable] = taskMemoryManager.hook + } + + private def reservedFor(task: TaskFixture): Long = + CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).map(_.reservedBytes).getOrElse(0L) + + private def listenerFor(task: TaskFixture): CometArrowAllocationListener = + CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get + + /** Runs the body in a task of its own, completed afterwards, and returns what the body does. */ + private def withTask[T](offHeap: Boolean = true, pool: Long = poolBytes)( + f: TaskFixture => T): T = { + val (context, taskMemoryManager) = + TestTasks.newTask(pool, offHeap)(new HookedTaskMemoryManager(_, _)) + val task = TaskFixture(context, taskMemoryManager) + TestTasks.withInstalled(context)(f(task)) + } + + /** An unstarted daemon thread, so a test can choose the moment it runs. */ + private def daemonThread(name: String)(body: => Unit): Thread = { + val thread = new Thread(() => body) + thread.setDaemon(true) + thread.setName(name) + thread + } + + /** + * Waits until the thread is either blocked on a monitor or finished, whichever happens first, + * so that a test can tell the two orderings apart without depending on timing. + */ + private def awaitBlockedOrFinished(thread: Thread): Unit = { + val deadline = System.currentTimeMillis() + 30000L + var state = thread.getState + while (state != Thread.State.BLOCKED && state != Thread.State.TERMINATED && + System.currentTimeMillis() < deadline) { + Thread.sleep(1L) + state = thread.getState + } + } + + /** Runs the body 500 times on an unstarted daemon thread, recording the first failure. */ + private def loopingThread(name: String, failure: AtomicReference[Throwable])( + body: => Unit): Thread = daemonThread(name) { + try { + var i = 0 + while (i < 500) { + body + i += 1 + } + } catch { + case t: Throwable => failure.compareAndSet(null, t) + } + } + + private def rootChildNames(): Set[String] = + CometArrowAllocator.getChildAllocators.asScala.map(_.getName).toSet +} diff --git a/spark/src/test/scala/org/apache/spark/comet/TestTasks.scala b/spark/src/test/scala/org/apache/spark/comet/TestTasks.scala new file mode 100644 index 00000000000..b29f9dadb1d --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/comet/TestTasks.scala @@ -0,0 +1,89 @@ +/* + * 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.Properties +import java.util.concurrent.atomic.AtomicLong + +import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryManager, TaskMemoryManager, TestMemoryManager} + +/** + * Spark tasks built by hand, each over a memory pool of its own, for driving the JVM Arrow + * accounting without a `SparkContext`. + */ +private[comet] object TestTasks { + + /** Task attempt ids are keys in a process-wide map, so no two tasks may share one. */ + private val nextTaskAttemptId = new AtomicLong(1000L) + + /** + * A task over a pool of `pool` bytes, off-heap unless `offHeap` is false, whose + * `TaskMemoryManager` comes from `newTaskMemoryManager` so that a caller can subclass it. + */ + def newTask[M <: TaskMemoryManager](pool: Long, offHeap: Boolean = true)( + newTaskMemoryManager: (MemoryManager, Long) => M): (TaskContextImpl, M) = { + val conf = new SparkConf(false) + if (offHeap) { + conf + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", pool.toString) + } + val memoryManager = new TestMemoryManager(conf) + memoryManager.limit(pool) + val taskAttemptId = nextTaskAttemptId.getAndIncrement() + val taskMemoryManager = newTaskMemoryManager(memoryManager, taskAttemptId) + val context = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = taskAttemptId, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + (context, taskMemoryManager) + } + + /** + * Runs the body with the task installed on this thread, then completes the task, which fires + * the completion listener that drops the Arrow reservation and is harmless if it already ran, + * and restores whatever was installed before. + */ + def withInstalled[T](context: TaskContextImpl)(f: => T): T = { + val previous = TaskContext.get() + TaskContext.setTaskContext(context) + try { + f + } finally { + try { + context.markTaskCompleted(None) + context.taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous) + } + } + } +}