Skip to content

Remove memory accounting from Comet's on-heap mode #6063

Description

@andygrove

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 fix: make CometDiskBlockWriter spill registry per-task instead of executor-global #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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions