andygrove opened a new issue, #6063:
URL: https://github.com/apache/datafusion-comet/issues/6063
### What is the problem the feature request solves?
On-heap mode is not a production configuration.
`spark.comet.exec.onHeap.enabled` defaults to
`false`, lives in `CATEGORY_TESTING`, and `CometDriverPlugin.init` disables
Comet outright when
off-heap is off and the flag is not set. Comet's own suites run off-heap:
`CometTestBase` sets
`spark.memory.offHeap.enabled=true` with 2 GiB
(`CometTestBase.scala:85-86`), which makes the
`spark.comet.exec.onHeap.enabled` and `spark.comet.memoryOverhead` it also
sets inert. The only
things that reach the on-heap path are Spark's own SQL suite and the Iceberg
suites, which set
`ENABLE_COMET_ONHEAP=true` because Spark's test harness does not configure
off-heap memory.
That path nevertheless carries a complete second memory-accounting
implementation, and the
accounting it performs does not protect anything.
**Comet cannot honestly join Spark's ledger in on-heap mode.**
`CometTaskMemoryManager`'s
`NativeMemoryConsumer` is hardcoded to `MemoryMode.OFF_HEAP`
(`CometTaskMemoryManager.java:106`).
Spark sizes the off-heap execution pool from `spark.memory.offHeap.size`
alone
(`MemoryManager.scala:61-66`, byte-identical on 3.4.3, 3.5.9, 4.0.4 and
4.1.3), and that is 0 when
off-heap is disabled, so `acquireExecutionMemory` grants nothing and every
native `try_grow` would
fail. Registering as an `ON_HEAP` consumer instead would be a category
error: Comet's bytes are in
the Rust heap, so Spark would evict cached blocks and spill its own sorters
to make room for memory
that is not on the heap, while doing nothing about the native RSS that
actually gets an executor
OOM-killed. `parse_memory_pool_config` therefore rejects the unified pools
in on-heap mode
(`config.rs:63-67`) and falls back to DataFusion's own pools, sized from
`spark.comet.memoryOverhead`.
**That budget is a leftover, not a design.** `spark.comet.memoryOverhead`
was in the initial commit
(PR #1, February 2024), a month before `CometTaskMemoryManager` and the
unified pool existed
(PR #83). It survived because deleting it broke the tests: PR #1062 required
off-heap and removed
the on-heap branch, and PR #1177 restored it a month later with the
rationale "after #1062 we have
not been running Spark tests for native execution". PR #2554 then made
on-heap opt-in and
testing-only, closing #2342 ("Remove on-heap memory pools") without removing
them.
**And it does not bound what a container cares about.** Per the analysis in
#6054, the driver
plugin's attempt to fold `spark.comet.memoryOverhead` into
`spark.executor.memoryOverhead` is inert
on 3.4, 3.5 and 4.0, because `ResourceProfileManager` has already
materialized and cached the
default profile by the time plugins initialize. Separately, on-heap sizes
the native pool *and* the
JVM shuffle pool from that same figure —
`spark.comet.shuffle.jvm.memoryFactor` defaults to `1.0` —
so the two together can allocate roughly twice what is declared.
What the accounting does buy is complexity, all of it on a path no user runs:
- Six of the nine `MemoryPoolType` variants (`greedy`, `fair_spill`, and
their `_task_shared` and
`_global` pairings) plus `memory_limit_per_task`, which no off-heap pool
reads.
- Four configs: `spark.comet.memoryOverhead`,
`spark.comet.exec.onHeap.memoryPool`,
`spark.comet.shuffle.jvm.memoryFactor`,
`spark.comet.shuffle.jvm.memoryWaitTimeout`.
- Most of `CometBoundedShuffleMemoryAllocator` (352 lines). Because on-heap
mode shares one
executor-wide bounded allocator across all tasks, and #5493 removed the
cross-task force-spill
that used to paper over contention on it, that class now carries a
blocking-allocation protocol:
per-thread retention accounting, two fail-fast liveness predicates, a
5-minute timeout,
30-second progress logging, and cooperative task-kill polling. Eight tests
in
`CometDiskBlockWriterSuite` exist solely to cover it.
### Describe the potential solution
Stop accounting in on-heap mode rather than accounting badly.
1. `parse_memory_pool_config` returns `MemoryPoolType::Unbounded` whenever
off-heap is disabled.
Delete the six on-heap pool variants and their arms in
`memory_pools/mod.rs`.
2. Delete `memory_limit_per_task` from `MemoryConfig`, `Native.createPlan`
and `jni_api.rs`. Nothing
else reads it.
3. Reduce `CometBoundedShuffleMemoryAllocator` to an unbounded page table
over
`UnsafeMemoryAllocator` — the pages must stay `Unsafe`-allocated in
either memory mode because
`SpillWriter` hands their addresses to `writeSortedFileNative` for Rust
to dereference — and
rename it accordingly. With no shared budget there is nothing to wait
for, so `allocateBlocking`
collapses into `allocate` and leaves `CometShuffleMemoryAllocatorTrait`,
and the allocator
becomes per-task like the off-heap one instead of an executor-wide
singleton.
4. Delete the four configs above and the now-unused
`getCometMemoryOverhead*` and
`getCometShuffleMemorySize` helpers.
5. Drop `.set("spark.comet.memoryOverhead", ...)` from the four `dev/diffs`
patches, regenerating
them through the documented flow rather than editing them by hand.
6. Update the tuning guide, `configs.md`, and the memory management page,
which currently has to
explain the on-heap pool types in order to exclude them.
`spark.comet.exec.onHeap.enabled` stays: it is still the switch that keeps
Comet off in on-heap mode
unless a test opts in.
### Additional context
**What this gives up.** The on-heap pool is the only bound the Spark SQL
suite runs under, so any
memory-pressure-driven native spill it triggers today stops happening, and
nothing caps Comet's RSS
in those jobs. Given the suite's data sizes the spill coverage is probably
near zero already, but
that is an assumption, not a measurement. Row-count-driven JVM shuffle
spilling
(`spark.comet.shuffle.jvm.spillThreshold`) and
`spark.comet.shuffle.native.maxBufferBytes` both
still trigger independently of any pool, and the off-heap "memory pressure
spills only writers of
the requesting task" test in `CometDiskBlockWriterSuite` keeps covering the
spill policy itself. If
CI memory does regress, the cheap recovery is a single executor-wide
`GreedyMemoryPool` with a fixed
cap, which is one arm in `parse_memory_pool_config` rather than the whole
subsystem.
**Two alternatives were considered and rejected.**
*Set `spark.memory.offHeap.size` in the Spark SQL tests while leaving
`spark.memory.offHeap.enabled`
false.* This is legal and unvalidated on every supported version:
`MemoryManager` sizes the off-heap
execution pool from the size alone, `tungstenMemoryMode` is the only thing
gated on the flag, and
nothing in `SparkConf.validateSettings` objects. Comet's native pool would
then run the real unified
path under the Spark SQL suite, which is a genuine coverage gain. But it
does not unify the JVM
shuffle allocator: `TaskMemoryManager.allocatePage` asserts
`consumer.getMode() == tungstenMemoryMode`
and returns `long[]` heap pages in on-heap mode, which cannot be handed to
native code, so that
allocator still forks on page provenance. It also means regenerating the
diffs against a different
memory model and accepting that `Utils.checkOffHeapEnabled` returns 0, so
the size never reaches the
container — a mechanism that works only because it is a test harness.
*Set `spark.memory.offHeap.enabled=true` in `SharedSparkSessionBase`.* This
is the only option that
truly leaves one code path, but it flips Spark's own Tungsten to off-heap
for the entire Comet test
run, so Spark's operators change behavior and the diffs absorb the fallout
across four versions.
This is what PR #1177 backed out of in December 2024.
**Sequencing.** #6054 removes the driver plugin's
`spark.executor.memoryOverhead` mutation and
`shouldOverrideMemoryConf`, which this change also touches. That should land
first.
Related: #5212, #2342, #6050.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
