Conversation
|
@andygrove please resolve the conflicts, overall lgtm |
sunchao
left a comment
There was a problem hiding this comment.
Correctness
This removes the separate accounting used by Comet's opt-in on-heap testing mode. Native plans now use UnboundedMemoryPool, and JVM shuffle uses a task-owned Unsafe page table. The production off-heap pool selection, budget formula and Spark allocation path remain unchanged.
I checked the maintained Spark 3.5 and 4.0 sources. Their on-heap Tungsten allocator returns heap-backed pages, so keeping Unsafe pages for rows dereferenced by Rust is necessary. Both JVM writer entry points share one allocator with every component that decodes its page addresses. The JNI caller, declaration and native entry point remove the same argument. Page-table exhaustion is reported before allocating or changing usage, and the existing managed-OOM handlers can spill and retry. Constructor and writer cleanup remain responsible for these pages, since Spark's task cleanup does not own them.
[P2] Preserve the compatibility path for released on-heap settings. The inline comment covers the unconditional removal of configuration behavior shipped in 1.0.0 while this branch targets 1.1.0. The repository's configuration compatibility policy requires a deprecation/migration path. I found no additional P1/P2 issue in the allocator or JNI changes.
Validation
The head has 99 successful checks and 70 skips. Spark SQL jobs for 3.4.3, 3.5.9, 4.0.4 and 4.1.3 and Iceberg 1.8–1.11 workflows completed successfully. Sampled logs explicitly enable on-heap Comet. The shuffle job ran the renamed allocator suite, SpillSorterSuite and CometDiskBlockWriterSuite, with 491 Scala tests passing.
Those jobs checked out historical merge b2c2942440f8, combining this head f1c4e61045f7 with d62d3382ac5d. Every authored file matches this head at that tested merge. The current API merge is unavailable because of the conflict already discussed on the PR. This evidence does not qualify a later conflict resolution. Local validation was source and caller analysis plus git diff --check. The maintained Spark 3.4 and 4.1 branches were unavailable, so their runtime CI coverage is separate from the canonical-source comparison.
Performance
Removing the executor-wide on-heap allocator also removes its shared monitor contention, retention maps and blocking allocation protocol. Off-heap allocation behavior is unchanged. The tradeoff is higher possible memory use in the testing mode: native memory pressure no longer causes spilling, and the JVM sort path can retain map output until its row, pointer-array or page-table boundary. The bypass writer retains its row-count spill trigger. The completed CI runs establish test completion, not an RSS bound or a measured performance improvement. No benchmark was run.
Could you add a small before/after on-heap memory-pressure microbenchmark, recording peak RSS, spill counts and elapsed time for a multi-batch sort or aggregation? That would quantify the effect of removing the old bound; the green CI runs do not measure it.
Design
The two remaining allocation backends have a clear purpose: Spark arbitrates off-heap execution, while the testing mode supplies native-addressable pages without a second budget. Task ownership also makes the page namespace easier to reason about. The documentation explains the intentionally unbounded mode and loss of memory-pressure spill coverage. The compatibility transition for existing explicit settings needs to accompany that design change.
Abstraction & complexity
Deleting the six on-heap pool variants, the unused off-heap per-task limit, and allocateBlocking removes substantial special-case code. The remaining allocation trait still serves the writers' common allocate/free/address operations. Tests that depended on the removed blocking protocol are deleted, and the retained failure-cleanup tests use a deliberately constrained Spark memory manager. I found no additional abstraction or complexity issue worth raising.
| // changing Spark's memory configuration, and native memory cannot be charged to Spark's | ||
| // on-heap pool, so nothing is accounted. See the memory management contributor guide. | ||
| logDebug("on-heap mode: native memory is unbounded and unaccounted") | ||
| MemoryConfig(offHeapMode, memoryPoolType = "unbounded", memoryLimit = 0) |
There was a problem hiding this comment.
Correctness
[P2] Keep a compatibility path for the released on-heap memory controls
Could we preserve a deprecated path for the existing settings before making this unconditional in 1.1? For example, an on-heap test invocation with spark.comet.exec.onHeap.enabled=true, spark.comet.memoryOverhead=128m and spark.comet.exec.onHeap.memoryPool=greedy previously selected a finite native reservation pool. It now silently selects an unbounded pool, regardless of those explicit settings. The shuffle budget settings disappear too.
memoryOverhead, onHeap.memoryPool and shuffle.jvm.memoryFactor all shipped in 1.0.0. The versioning policy covers spark.comet.* configuration semantics, including these testing-category keys, and requires deprecation before removal in a major release. Keeping the old behavior available through the documented compatibility mechanism, with an upgrade-guide entry, would let this simplification land without silently removing an existing test harness's configured bound.
There was a problem hiding this comment.
@sunchao Hmm .. technically, this is correct, but this is a test-only config intended for running Spark SQL test suite and it is disabled by default. Our docs already say that no user should use on-heap mode. I'd like to push back against this review issue. WDYT?
On-heap mode exists so that Spark's own SQL suite and the Iceberg suites can run against Comet without changing Spark's memory configuration. It is off by default, sits in CATEGORY_TESTING, and the driver plugin disables Comet entirely unless a test opts in. Comet's own suites run off-heap. The accounting that mode performed did not protect anything. Native memory is not on the JVM heap, so there is no Spark pool it can honestly be charged to: NativeMemoryConsumer is hardcoded to MemoryMode.OFF_HEAP and Spark sizes that pool from spark.memory.offHeap.size alone, which is 0 when off-heap is disabled. The fixed-size DataFusion pool that stood in for one was sized from spark.comet.memoryOverhead, a config that predates the unified pool and that the driver plugin could not fold into the container on three of five supported Spark versions. On-heap mode now gets UnboundedMemoryPool on the native side and an unbounded per-task Unsafe page allocator on the JVM side, which removes: - six of the nine MemoryPoolType variants and memory_limit_per_task - spark.comet.memoryOverhead, spark.comet.exec.onHeap.memoryPool, spark.comet.shuffle.jvm.memoryFactor and spark.comet.shuffle.jvm.memoryWaitTimeout - the shared-pool blocking-allocation protocol in what is now CometUnboundedShuffleMemoryAllocator, along with allocateBlocking - the driver plugin's spark.executor.memoryOverhead mutation and the ShimCometDriverPlugin files that only it needed spark.comet.exec.onHeap.enabled stays: it is still the switch that keeps Comet off in on-heap mode unless a test opts in. Closes apache#6063
Four follow-ups from a self review of this PR. The renamed `CometUnboundedShuffleMemoryAllocatorSuite` was never renamed in `pr_build_linux.yml` or `pr_build_macos.yml`, so `check-suites` failed Preflight and the suite ran in neither job. `CometUnboundedShuffleMemoryAllocator` never refuses a page, which leaves its page table as the only limit it has. Exhausting it threw `IllegalStateException`, which nothing catches, so the task failed. Report `SparkOutOfMemoryError` instead, which `SpillWriter.acquireNewPageIfNecessary` and `CometShuffleExternalSorter.growPointerArrayIfNecessary` already answer by spilling and retrying. A page number indexes one allocator's own table, so an address encoded by one instance cannot be resolved by another. That was previously guaranteed by the executor-wide singleton and is now a requirement on callers, so say so on both the allocator and the factory. The JVM shuffle guide described a single allocator that spills when an allocation fails, which is true of the off-heap path only.
Scalafix `RemoveUnused` fails the Lint Java jobs on the 3.4, 3.5 and 4.0
profiles:
- private implicit val signaler: Signaler = ThreadSignaler
+ ThreadSignaler
This PR removed the five tests that covered the blocking-allocation protocol,
and they were the only callers of `failAfter`. With those gone the `TimeLimits`
mixin and the `Signaler` that served it are unreferenced, so remove both rather
than the val alone.
Lint did not catch this earlier in the PR because Preflight was failing on the
renamed allocator suite and gated every downstream job. Spark 4.1 is excluded
from the lint matrix, so formatting locally on the default profile does not run
scalafix at all.
…pool apache#6054 landed on main and its comment above warnIfExecutorMemoryOverheadUnset named spark.comet.memoryOverhead, which this branch removes, and said the share operators reserve is charged against a memory pool, which is now true in off-heap mode only.
f1c4e61 to
97217c9
Compare
sunchao
left a comment
There was a problem hiding this comment.
@andygrove I agree with the test-only scope in your reply. The flag is disabled by default, the removed controls are explicitly testing or internal, and the existing guide already excludes on-heap mode from production. I would drop my request to retain the old pools or add a legacy flag solely for this testing mode.
[P2] I would narrow the remaining compatibility concern to policy and merge order. The currently merged policy still covers these keys. The explicit testing/internal exemption in #6089 is a reasonable remedy. Could we land that exemption first and rebase this PR onto it? My review of #6089 objects to treating absence from configs.md as proof that a key is internal, not to the explicit testing/internal carve-out. Once that narrower exemption lands, this compatibility concern is resolved without restoring the old accounting machinery.
Re-reviewed 97217c9cbbf9 against 09b44ad6fa17, including the rebase. The production off-heap allocator and native pool implementations remain unchanged, the JNI parameter removal matches on both sides, and task-local allocator sharing and cleanup remain consistent. The rebase keeps #6054’s warnIfExecutorMemoryOverheadUnset. I found no new P1/P2 implementation issue.
Current CI has 85 successes, 10 skips and 2 failures. The shuffle job passed 495 tests, including the allocator, sorter and failure-cleanup cases. However, the Spark 4.0 SQL shard failed SQLAppStatusListenerMemoryLeakSuite’s “no memory leak” test at noLiveData(), which checks retained listener state. Required Checks consequently failed. This log does not establish an allocator leak, and I have not established that the failure is unrelated to the PR. Both sampled jobs checked out 33d8ceec9ee8, whose tree exactly matches this head.
There is still no new before/after RSS or spill measurement, so the passing jobs establish completion for their workloads, not a memory bound or performance improvement. Local checks were source/ABI analysis and git diff --check. Maintained Spark 3.5/4.0 sources were checked; maintained 3.4/4.1 sources remain unavailable. No local full suite or benchmark was run.
Which issue does this PR close?
Closes #6063.
Rationale for this change
On-heap mode is not a production configuration.
spark.comet.exec.onHeap.enableddefaults tofalse, lives inCATEGORY_TESTING, andCometDriverPlugin.initdisables Comet outright whenoff-heap is off and the flag is not set. Comet's own suites run off-heap (
CometTestBasesetsspark.memory.offHeap.enabled=truewith 2 GiB). The only things that reach the on-heap path areSpark's own SQL suite and the Iceberg suites, which set
ENABLE_COMET_ONHEAP=truebecause Spark'stest 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.
NativeMemoryConsumeris hardcoded toMemoryMode.OFF_HEAP, and Spark sizes the off-heap execution pool fromspark.memory.offHeap.sizealone (
MemoryManager.scala:61-66, byte-identical on 3.4.3, 3.5.9, 4.0.4 and 4.1.3), which is 0when off-heap is disabled. Registering as an
ON_HEAPconsumer 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
CometTaskMemoryManagerandthe 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_configreturnsUnboundedMemoryPoolwhenever off-heap is disabledand ignores the pool-type string there. Six of the nine
MemoryPoolTypevariants (greedy,fair_spill, and their_task_sharedand_globalpairings) and their arms inmod.rsare gone,as is
memory_limit_per_task, which no off-heap pool read.createPlanloses the correspondingJNI parameter.
JVM shuffle.
CometBoundedShuffleMemoryAllocatorbecomesCometUnboundedShuffleMemoryAllocator: a page table overUnsafeMemoryAllocatorwith no budget.The pages still have to be
Unsafe-allocated in either memory mode, becauseSpillWriterhandstheir addresses to
writeSortedFileNativefor Rust to dereference andTaskMemoryManager.allocatePagewould returnlong[]heap pages here. With no shared budget thereis 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
allocateBlockingcollapses intoallocateand leavesCometShuffleMemoryAllocatorTrait. The allocator is now created per task like the off-heap onerather than being an executor-wide singleton, and
getUsedreads anAtomicLonginstead of takingthe allocator's monitor, since
TaskMemoryManagercalls it while holding its own. Its page tableis the one limit it still has, and exhausting it now reports
SparkOutOfMemoryErrorrather thanIllegalStateException, so the callers that already answer a refused acquisition by spilling andretrying can make progress instead of failing the task. A page number indexes the allocator's own
table, so an address encoded by one instance cannot be resolved by another, and each caller has to
create one allocator and share it for the whole task. Both classes now say so.
Configs.
spark.comet.memoryOverhead,spark.comet.exec.onHeap.memoryPool,spark.comet.shuffle.jvm.memoryFactorandspark.comet.shuffle.jvm.memoryWaitTimeoutare removed,along with
getCometMemoryOverhead*,getCometShuffleMemorySizeandshouldOverrideMemoryConf.spark.comet.exec.onHeap.enabledstays: it is still the switch that keeps Comet off in on-heapmode unless a test opts in, and its doc now says the mode accounts for nothing.
Driver plugin.
CometDriverPlugin.initno longer adjustsspark.executor.memoryOverhead;there is nothing left to add. Both
ShimCometDriverPluginfiles are deleted, sincegetMemoryOverheadMinMibwas only needed by the removed calculation.CI and docs. Both
pr_build_linux.ymlandpr_build_macos.ymlname the renamed allocatorsuite. Without that
check-suitesfails Preflight and the suite runs in neither job. The JVMshuffle contributor guide's Memory Management section described a single allocator that spills when
an allocation fails, which is now true of the off-heap path only, so it describes both paths and
the row-count triggers instead.
Diffs. The four
dev/diffspatches drop.set("spark.comet.memoryOverhead", ...). They wereregenerated through the documented flow (clone at the tag, apply, edit,
git diff <tag>), nothand-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
indexline forParquetRowIndexSuite.scala, whose recordedpost-image hash did not match what
git applyactually produces. That line is inert — CI applieswith 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 spilling still triggers independently of any pool, but only the bypass writer really
keeps it.
CometDiskBlockWriterspills atmin(spark.comet.shuffle.jvm.spillThreshold, spark.comet.shuffle.jvm.batchSize), and the batchsize defaults to 8192, so that path stays bounded per writer.
CometShuffleExternalSortercomparesagainst
spark.comet.shuffle.jvm.spillThresholdalone, which defaults toInt.MaxValue, and itsremaining triggers were allocation failures that the unbounded allocator no longer raises. So on
the sort-based path a task in on-heap mode now buffers its whole map output until it closes. That
writer is selected when the partition count exceeds
spark.shuffle.sort.bypassMergeThresholdorpartitions times cores exceeds
spark.comet.shuffle.jvm.maxWritersPerExecutor.spark.comet.shuffle.native.maxBufferBytesstill triggers independently of any pool. If CImemory regresses, the cheap recovery is a single executor-wide
GreedyMemoryPoolwith a fixed cap,which is one arm in
parse_memory_pool_configrather 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.memoryOverheadmutation,
shouldOverrideMemoryConfand the shim files. This PR has to remove them too, because itdeletes the config they consumed. I will rebase onto
mainonce #6054 lands and keep itswarnIfExecutorMemoryOverheadUnsetaddition.How are these changes tested?
Existing tests, adjusted:
CometUnboundedShuffleMemoryAllocatorSuite(renamed) keeps thegetUsedaccounting andpage-table-exhaustion tests, drops the two that asserted budget exhaustion, and adds one
asserting that a 64 MiB request now succeeds. The page-table test asserts the
SparkOutOfMemoryErrorabove, and that freeing a page lets the next allocation through, which iswhat lets a spilling writer recover.
CometDiskBlockWriterSuiteloses the five tests that covered the deleted blocking protocol. Thethree that cover behavior this PR keeps — write() reclaiming buffered pages on a fatal error, the
SpillSorterconstructor not leaking on a failed allocation, and the unsafe writer allocatingnothing before
write()— are ported to off-heap, whereTestMemoryManager.limitsupplies thesame pressure and the assertions become
getUsed == 0/getMemoryConsumptionForThisTask == 0. That also moves them onto the production allocator.SpillSorterSuitenow builds one allocator per test inbeforeEachinstead of relying on theexecutor-wide singleton to hand the same instance to two call sites.
CometPluginsSuite's two override assertions become "the configured value stands" and "nothingis added when none was configured".
CometExecSuite"spill sort with (multiple) dictionaries" tests lose awithSQLConfwrapper for
spark.comet.memoryOverheadthat was already inert twice over: the suite runsoff-heap, and the config is read from
SparkConfrather thanSQLConf.Local verification so far: the four regenerated diffs apply cleanly to pristine
v3.4.3,v3.5.9,v4.0.4andv4.1.3trees;CometDiskBlockWriterSuiteandSpillSorterSuitepass. Broader suiteruns 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 suitesbefore it is queued.