andygrove opened a new pull request, #6066:
URL: https://github.com/apache/datafusion-comet/pull/6066
## Which issue does this PR close?
Closes #6063.
## Rationale for this change
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). 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 carried a complete second memory-accounting
implementation, and the
accounting it performed did not protect anything.
Comet cannot honestly join Spark's ledger in on-heap mode.
`NativeMemoryConsumer` is hardcoded to
`MemoryMode.OFF_HEAP`, and 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), which is 0
when off-heap is disabled. 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.
The fixed-size DataFusion pool that stood in for one was sized from
`spark.comet.memoryOverhead`,
which was in the initial commit (#1, February 2024), a month before
`CometTaskMemoryManager` and
the unified pool existed (#83). It survived because deleting it broke the
tests: #1062 required
off-heap and removed the on-heap branch, and #1177 restored it with the
rationale "after #1062 we
have not been running Spark tests for native execution". Per #6054, the
driver plugin could not
even fold that figure into the container on three of the five supported
Spark versions.
## What changes are included in this PR?
**Native.** `parse_memory_pool_config` returns `UnboundedMemoryPool`
whenever off-heap is disabled
and ignores the pool-type string there. Six of the nine `MemoryPoolType`
variants (`greedy`,
`fair_spill`, and their `_task_shared` and `_global` pairings) and their
arms in `mod.rs` are gone,
as is `memory_limit_per_task`, which no off-heap pool read. `createPlan`
loses the corresponding
JNI parameter.
**JVM shuffle.** `CometBoundedShuffleMemoryAllocator` becomes
`CometUnboundedShuffleMemoryAllocator`: a page table over
`UnsafeMemoryAllocator` with no budget.
The pages still have to be `Unsafe`-allocated in either memory mode, because
`SpillWriter` hands
their addresses to `writeSortedFileNative` for Rust to dereference and
`TaskMemoryManager.allocatePage` would return `long[]` heap pages here. With
no shared budget there
is nothing to wait for, so the blocking-allocation protocol added in #5493
goes away — per-thread
retention accounting, the two fail-fast liveness predicates, the timeout,
the progress logging and
the task-kill polling — and `allocateBlocking` collapses into `allocate` and
leaves
`CometShuffleMemoryAllocatorTrait`. The allocator is now created per task
like the off-heap one
rather than being an executor-wide singleton, and `getUsed` reads an
`AtomicLong` instead of taking
the allocator's monitor, since `TaskMemoryManager` calls it while holding
its own.
**Configs.** `spark.comet.memoryOverhead`,
`spark.comet.exec.onHeap.memoryPool`,
`spark.comet.shuffle.jvm.memoryFactor` and
`spark.comet.shuffle.jvm.memoryWaitTimeout` are removed,
along with `getCometMemoryOverhead*`, `getCometShuffleMemorySize` and
`shouldOverrideMemoryConf`.
`spark.comet.exec.onHeap.enabled` stays: it is still the switch that keeps
Comet off in on-heap
mode unless a test opts in, and its doc now says the mode accounts for
nothing.
**Driver plugin.** `CometDriverPlugin.init` no longer adjusts
`spark.executor.memoryOverhead`;
there is nothing left to add. Both `ShimCometDriverPlugin` files are
deleted, since
`getMemoryOverheadMinMib` was only needed by the removed calculation.
**Diffs.** The four `dev/diffs` patches drop
`.set("spark.comet.memoryOverhead", ...)`. They were
regenerated through the documented flow (clone at the tag, apply, edit, `git
diff <tag>`), not
hand-edited, and each was verified to apply cleanly to a pristine tree
afterwards. Regenerating
4.0.4 and 4.1.3 also corrects a stale `index` line for
`ParquetRowIndexSuite.scala`, whose recorded
post-image hash did not match what `git apply` actually produces. That line
is inert — CI applies
with plain `git apply` — but it is now consistent with the hunks.
## Behavior change
Nothing bounds Comet's allocations in on-heap mode any more. The on-heap
pool was the only bound
the Spark SQL suite ran 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 rather
than 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. If CI
memory regresses, 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.
Off-heap mode is untouched.
## Overlap with #6054
#6054 is in the merge queue and also removes the driver plugin's
`spark.executor.memoryOverhead`
mutation, `shouldOverrideMemoryConf` and the shim files. This PR has to
remove them too, because it
deletes the config they consumed. I will rebase onto `main` once #6054 lands
and keep its
`warnIfExecutorMemoryOverheadUnset` addition.
## How are these changes tested?
Existing tests, adjusted:
- `CometUnboundedShuffleMemoryAllocatorSuite` (renamed) keeps the `getUsed`
accounting and
page-table-exhaustion tests, drops the two that asserted budget
exhaustion, and adds one
asserting that a 64 MiB request now succeeds.
- `CometDiskBlockWriterSuite` loses the five tests that covered the deleted
blocking protocol. The
three that cover behavior this PR keeps — write() reclaiming buffered
pages on a fatal error, the
`SpillSorter` constructor not leaking on a failed allocation, and the
unsafe writer allocating
nothing before `write()` — are ported to off-heap, where
`TestMemoryManager.limit` supplies the
same pressure and the assertions become `getUsed == 0` /
`getMemoryConsumptionForThisTask == 0`. That also moves them onto the
production allocator.
- `SpillSorterSuite` now builds one allocator per test in `beforeEach`
instead of relying on the
executor-wide singleton to hand the same instance to two call sites.
- `CometPluginsSuite`'s two override assertions become "the configured value
stands" and "nothing
is added when none was configured".
- The two `CometExecSuite` "spill sort with (multiple) dictionaries" tests
lose a `withSQLConf`
wrapper for `spark.comet.memoryOverhead` that was already inert twice
over: the suite runs
off-heap, and the config is read from `SparkConf` rather than `SQLConf`.
Local verification so far: the four regenerated diffs apply cleanly to
pristine `v3.4.3`, `v3.5.9`,
`v4.0.4` and `v4.1.3` trees; `CometDiskBlockWriterSuite` and
`SpillSorterSuite` pass. Broader suite
runs are still in progress and I will report them here. This change touches
the serde-adjacent JNI
signature, a native operator path and `dev/diffs`, so it needs the Spark SQL
and Iceberg suites
before it is queued.
--
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]
