Conversation
Comet's memory pool counts declared reservations. Plenty of real allocation never goes through it -- Arrow builders, expression kernels, decompression buffers, Parquet metadata, object_store buffers, tokio itself -- so pool reservations are a lower bound on Comet's footprint and the size of the gap is currently unmeasurable at runtime. Diagnosing an OOM means guessing at it, and spark.comet.exec.memoryPool.fraction asks operators to hand-tune a haircut for a quantity nobody can see. AccountingAllocator wraps the selected global allocator (jemalloc, mimalloc, or system) and maintains one signed process-wide byte balance. executePlan reports it as the native_allocated tracing metric, next to the per-thread pool reservations it should be compared against. This is observability only: it never rejects an allocation, never panics, and does not touch the memory pool. Because it cannot fail an allocation, realloc can account after delegating rather than before, which avoids over-counting a failed realloc. Per-thread deltas are batched and flushed into the shared balance at 64 KiB, so the common path is a thread-local add-and-compare rather than an atomic RMW. ThreadDrift's destructor settles the remainder when a thread exits, which matters because the blocking pool churns on tokio's idle timeout and would otherwise bias the balance on a long-lived executor. Touching that destructor-bearing thread-local can itself allocate on first use, so track() keeps a destructor-free re-entrancy flag and settles re-entrant calls straight into the shared balance. Off by default; a build without the feature has no wrapper and no per-allocation work. Verified clippy -D warnings and the native test suite across default, alloc-accounting, jemalloc+alloc-accounting, and mimalloc+alloc-accounting. The thread-exit test was mutation-checked: neutering the destructor fails it.
Answers the question the feature has to answer before anyone proposes enabling it by default. The liveness assertion is the point of the harness as much as the timings are: without it a 'with the feature' run can silently be a second baseline, which is exactly what happened on the first attempt here.
`dealloc` accounted after calling the inner allocator, mirroring `alloc` and `realloc`. A free cannot fail, so that ordering bought nothing, and it opened a window: jemalloc decrements `stats.allocated` at the start of a large free and then, for blocks above its 8 MiB oversize threshold, unmaps the pages eagerly, which takes milliseconds for a block of a few hundred megabytes. For that whole window the balance still carried a block the allocator had already given back, so `native_allocated` read above `jemalloc_allocated` by the size of the block in flight. On TPC-H SF100 about 2% of trace samples showed the excess, up to 160 MB, during task teardown in Q10, Q17 and Q18. Subtract before delegating. A new test wraps a recording inner allocator and asserts the balance has already dropped by the time the inner free is called; it fails with "inner dealloc saw balance 67182807, expected at most 33628375" under the previous ordering.
…nation Select the allocator backend once, in a `backend` module whose three cfgs partition every feature combination, and let the single accounting `#[global_allocator]` refer to whatever that resolved to. Previously each backend predicate was repeated per arm and the accounting fallback to the system allocator excluded `mimalloc`, so `jemalloc,mimalloc,alloc-accounting` on a non-MSVC target matched no arm at all: the process ran on the unwrapped default allocator while `native_allocated` stayed enabled and read zero. A build without the feature is unchanged: the unwrapped allocator lives in the backend module that owns it, and no explicit allocator is installed when the selection is the system allocator. Make the accounting tests immune to parallel test noise. `BALANCE` is process-wide, so with the wrapper installed the crate's other tests move it concurrently and the margin-based assertions failed intermittently. The tests that observe the balance now serialize against each other, the real-allocation probe uses an untouched 256 MiB block that nothing else in the crate can mask, and the thread-exit test injects its drift directly so it no longer needs the feature and runs in the default CI build. Add `threshold_churn` to the `alloc_overhead` benchmark: alloc/free loops at 32 KiB (never flushes) and 64 KiB (flushes on every call), single threaded and from every core at once, so shared-counter contention can be quantified rather than inferred.
…ed allocator An `--extern` crate that nothing names is dropped from the crate graph, and the accounting-off run of this benchmark named nothing in `comet`, so the `#[global_allocator]` in `lib.rs` never reached the binary: the "jemalloc" baseline was measuring glibc malloc. The accounting-on run names `comet::alloc_accounting`, so it did link the crate, and every off/on comparison so far was glibc against jemalloc plus the wrapper. Add `extern crate comet` so the crate is always linked, and a jemalloc liveness assertion next to the existing accounting one, so a run against the wrong allocator fails instead of producing a plausible number.
The alloc_overhead benchmark asserted on jemalloc's counters whenever the jemalloc feature was on, but lib.rs selects jemalloc only when mimalloc was not also requested and the target is not MSVC. With jemalloc,mimalloc the library falls back to the system allocator and the bench aborted before measuring anything. Export the selection lib.rs made as comet::ALLOCATOR_BACKEND, one NAME per backend module, and have the bench announce it and run the jemalloc assertion only when the library reports jemalloc. The predicate stays written once, so the guard cannot drift from the selection it verifies.
…ing feature in CI Address the third review round on the alloc-accounting feature. - analyze_trace only matched jemalloc_allocated, so a trace from a build with alloc-accounting and no jemalloc reported zero allocated and no excess. It now analyzes native_allocated when present, falls back to jemalloc_allocated otherwise, names the counter in its output, and rejects a trace with neither. tracing.md describes the selection. - Nothing in CI compiled the feature. The rust-test action now lints the jemalloc,alloc-accounting build with --all-targets, runs the accounting tests with the wrapper installed, and checks the system-allocator arm. - lib.rs: the three backend modules only name the type; the two global_allocator statics sit together, so the default build now installs System explicitly. ALLOCATOR_BACKEND is doc(hidden). - alloc_accounting.rs: inline the clamp, table-drive the settle tests, add a realloc test that pins the size-difference accounting and the after-delegating order, and confine the thread-exit test to the default build, where a missing destructor is actually caught. Trim the prose. - alloc_overhead bench: drop the redundant extern crate, run both liveness guards from a single Once that every bench function calls, merge the two alloc/free loops, and use iter instead of iter_batched. Benchmark IDs are unchanged.
…format tracing.md A zero pool reservation is a real sample that allocation can exceed, so the analyzer now defers the comparison only until it has seen at least one pool sample, and says so when it never does instead of reporting that allocation never exceeded reservations. Reflow the tracing.md label table with prettier, which the Preflight check enforces.
Comet's off-heap memory pools count only what operators explicitly reserve. Native memory that never goes through a pool, such as scratch buffers inside kernels or intermediate Arrow arrays, stays invisible until the executor exceeds its container limit and is killed. Wrap both fair_unified and greedy_unified in a check that compares the bytes the native allocator has actually handed out against Comet's off-heap budget, and refuse a reservation that would take real usage past it. Operators that can spill then spill, and the ones that cannot fail the task instead of the executor. The budget is spark.memory.offHeap.size times spark.comet.exec.memoryPool.fraction, the value the pools already receive as their limit, so the check adds no sizing config of its own. The fraction now defaults to 0.8 so a margin of the pool absorbs unreserved allocations, and the check can be turned off with spark.comet.exec.memoryPool.checkNativeUsage=false. The alloc-accounting cargo feature that supplies the byte count becomes a default feature. A build without it reports zero bytes in use, which leaves the check a passthrough. The balance and the budget are both process-wide, so there is no per-task attribution, and allocations themselves are never refused: this gates reservations only.
…ive-usage-check # Conflicts: # docs/source/contributor-guide/tracing.md # native/core/Cargo.toml # native/core/src/alloc_accounting.rs
…fraction Deriving the check's budget from the pool's own limit overloaded spark.comet.exec.memoryPool.fraction with a second meaning. The fraction bounds what Comet may reserve, and lowering it is how operators force spilling, so a small fraction also dropped the ceiling on real usage below Comet's baseline footprint and denied reservations outright instead of provoking the spills it was set to cause. CometTaskMetricsSuite does exactly this, and one of its tests failed with the reservation refused for a consumer that cannot spill. Take the budget from spark.memory.offHeap.size instead, so the fraction keeps its single meaning and the two ceilings move independently. The value reaches native through the existing config proto rather than the createPlan JNI signature. Move the pool tests into their own suite, since the budget is fixed when the Spark session starts and cannot be lowered with withSQLConf.
github-actions
Bot
added
enhancement
andygrove
changed the title
Enforcing on the first iteration was too aggressive for three reasons, none of which are measured yet: - The false positive rate on real workloads is unknown. The TPC-H numbers from the accounting PR show moments of more than 1 GB of native usage against near zero reservations, which is exactly the regime this fires in. - Spilling releases reserved bytes, so a denial provoked by untracked allocations may not relieve the pressure it reports. The likely outcome there is spill thrash and a failed task, which is worse than not denying at all. - The balance is process-wide, so the reservation refused belongs to whichever operator asks next rather than to whoever caused the overshoot. So the pools now always compare real usage against the off-heap size and log a crossing once per pool, and refuse the reservation only when spark.comet.exec.memoryPool.enforceNativeUsage is set. The config is renamed from checkNativeUsage to say what enabling it now does, and defaults to false. memory_limit reports the budget only while enforcing, since advertising a limit that nothing rejects would mislead callers. Also revert the spark.comet.exec.memoryPool.fraction default to 1.0. Tightening that crude manual proxy in the same change that adds the measured version of it was contradictory, and it changed behaviour for every off-heap user for reasons unrelated to this feature. Document what the check is not: the budget is shared with Spark's own off-heap allocations and Comet's JVM-side shuffle pages while the balance counts only Comet's native allocations, so it is a loose backstop against losing the executor rather than a bound on total off-heap usage.
andygrove
changed the title
The warning claimed real usage had passed spark.memory.offHeap.size when the condition that fired was usage plus the pending request. It reported 340742 bytes as having passed a 2097152 byte budget, which is false and would misdirect anyone tuning from the log line. Name both the request and the usage, since either can be what crosses the budget.
CI's Lint job reported a rustfmt diff in jni_api.rs, where the spark_config import list no longer fitted the wrapping rustfmt wanted after the enforce flag was added. The syntactic scalafix check reported two s-prefixed strings in CometMemoryPoolNativeUsageSuite that interpolate nothing.
CheckedMemoryPool returned MemoryLimit::Finite(budget) while enforcing. TrackConsumersPool forwards memory_limit straight through, and AggregateExec::should_use_partial_reduce_hash_stream bails out whenever the pool reports Finite, with a TODO saying its memory-limited path is unimplemented. So enabling the check silently disabled a DataFusion aggregation fast path even when it never refused a single reservation. This surfaced in a TPC-H SF100 comparison at a 2 GB off-heap budget: observing logged 44 crossings while enforcing logged none, with zero refusals on either side, so the gate itself could not have been what made them differ. Always delegate to the inner pool instead. The budget is a process-wide backstop compared against process-wide allocator usage, not this pool's reservable limit, so reporting it as one was wrong on its own terms as well, and a memory guard should not change which operator strategy DataFusion picks.
A refusal is returned as a DataFusionError rather than logged, and DataFusion normally answers it by spilling and retrying, which swallows the error. An enforcing run could therefore refuse a reservation on every task and leave no trace at all, in executor stderr or in the driver log. That is not hypothetical. A TPC-H SF100 sweep at a 2 GB off-heap budget came back with zero refusals recorded for the enforcing run, which read as "the check never fired" when it may well have been firing constantly and holding real usage down. An operator running with enforcement on had no way to tell whether it was doing anything. Log the first refusal per pool, reusing the flag that already rate limits the observe-mode warning, so one line per pool means one line per task that reached the budget rather than one per reservation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Relates to #4576. An alternative to #5959, which added the same check as a separate opt-in pool type rather than applying it to the pools that already exist.
Rationale for this change
Comet's off-heap memory pools count only what operators explicitly reserve. Native memory that never goes through a pool, such as scratch buffers inside kernels, intermediate Arrow arrays, and allocations made inside DataFusion operators that are not reserved, is invisible to the pool until the executor exceeds its container limit and is killed. That failure is materially worse than a task failure: every task on the executor dies, its cached blocks are lost, and its shuffle files become unavailable.
#5934 made the bytes the Rust allocator has actually handed out available at runtime, but only as a tracing metric. This PR consumes that number inside the pools themselves, so that for the first time a Comet memory pool sees measured usage rather than only declared intent.
Both
fair_unifiedandgreedy_unifiedget the check. It is observe only by default: a crossing is logged once per task, and nothing is refused.spark.comet.exec.memoryPool.enforceNativeUsage=trueturns it into a real gate, where operators that can spill do so and those that cannot fail the task rather than the executor.Why enforcement is off by default
Deliberately, and I would rather land it this way than default it on:
try_growby spilling and retrying, but spilling releases reserved bytes. If the overshoot lives in untracked allocations, the retry is denied again, and the likely outcome is spill thrash followed by a failed task, which is worse than not denying at all.Observing first lets these be measured on real workloads before anyone's job starts failing on them.
What changes are included in this PR?
CheckedMemoryPool<P: MemoryPool>(native/core/src/execution/memory_pools/checked_pool.rs). On a crossing it logs once per pool in either mode, then returnsResourcesExhaustedwhen enforcing or delegates when observing. Both modes log deliberately: a refusal is an error return, and DataFusion normally answers it by spilling and retrying, which swallows the error, so without the log an enforcing run can refuse a reservation on every task and leave no trace at all that the check ran. Everything else delegates, includingmemory_limit, which always reports the inner pool's limit and never the budget. That matters:TrackConsumersPoolforwardsmemory_limitthrough, andAggregateExec::should_use_partial_reduce_hash_streamdisables a fast path whenever the pool reportsMemoryLimit::Finite, so reporting the backstop there would silently change the aggregation strategy as a side effect of enabling the check. The budget is a process-wide backstop, not this pool's reservable limit. The inner pool is generic so the gate is unit-testable without a JVM.create_memory_poolwraps both off-heap pools in it, inside the existing task-shared and consumer-tracking wrappers, so a denial is still annotated with the largest consumers.spark.comet.exec.memoryPool.enforceNativeUsage(defaultfalse). It rides the existing config proto rather than thecreatePlanJNI signature, so no signature change was needed.spark.memory.offHeap.size, deliberately not the pool's own limit. That limit is the off-heap size timesspark.comet.exec.memoryPool.fraction, and lowering the fraction is how operators force spilling. Deriving the budget from it as well turned a small fraction into denied reservations instead of the spills it was set to cause, which brokeCometTaskMetricsSuitein an earlier revision of this PR.alloc-accountingbecomes a default cargo feature. Allcfg(feature)guards remain, so--no-default-featuresstill compiles; such a build reports zero bytes in use, which leaves the wrapper a passthrough.What this check is not
Worth stating plainly, and documented alongside the code:
spark.memory.offHeap.sizeis shared with Spark's own Tungsten off-heap allocations and with Comet's JVM-side shuffle pages, while the balance counts only Comet's Rust allocations.CometArrowAllocatoris in neither. So this is a loose backstop against losing the executor, not a bound on total off-heap usage.How are these changes tested?
Rust unit tests in
checked_pool.rs, overUnboundedMemoryPoolwith a budget of the current balance plus 64 MiB, holding a 256 MiB block the pool never heard about:try_growfails withResourcesExhaustedand reserves nothing; dropping the block makes the same call succeed.memory_limitreports the budget when enforcing and defers to the inner pool when not; successful grows and shrinks reach the inner pool.New
CometMemoryPoolNativeUsageSuite. The budget is fixed when the Spark session starts, so the suite runs with a deliberately smallspark.memory.offHeap.sizerather than awithSQLConfoverride. At that size the query cannot complete either way, so the assertion is on which component refuses the reservation, which is what isolates the check:fair_unifiedandgreedy_unified, the reservation is refused by Comet with the check's message.CometTaskMetricsSuitepasses; it setsmemoryPool.fraction=0.002to force spilling and is the regression that shaped the budget decision above.TPC-H SF100
Measured on a 32-core box, 2 executors x 8 cores, Spark 4.1.1, asking two things: does the check fire on a healthy workload, and is there a budget where enabling it is the difference between exceeding
spark.memory.offHeap.sizeand staying under it.At the benchmark runner's default 16 GB budget it never fires. 22/22 queries in both modes, identical result hashes, zero crossings, zero refusals. (Those two runs predate the
memory_limitfix in this branch, where enforcing made the pool reportFiniteand disabled a DataFusion aggregation fast path, so their timings are not a clean measure of the check's cost. The query outcomes and hashes stand.)A 2 to 12 GB sweep in 2 GB steps shows it only engages at 2 GB. From 4 GB up: zero crossings in both modes, no difference between them, 22/22 every run.
At 2 GB, three repeats per mode, enforcement does not prevent the overshoot:
Enforcing refused 30 to 44 reservations per run, and real native usage still passed the limit in every run, exactly as it did when only observing. Every refusal was absorbed by a spill and retry: no task failed, no refusal surfaced as an error, and the timings are indistinguishable.
Why, and it is structural rather than a tuning problem. The gate refuses reservations, but the overshoot lives in allocations that never reserve. Making a reserving operator spill releases reserved bytes and does nothing about the untracked allocations that actually carry usage past the limit. The operators that play by the rules absorb the penalty for the ones that do not.
What that means for this PR. On TPC-H the check is harmless, costs nothing measurable and breaks nothing, but it does not achieve the goal of bounding real native usage. That is the case for shipping it observe-only: the logging makes the accounting gap visible at runtime for the first time, which is worth having on its own, while the gate has not earned the right to act by default. Anyone proposing to enforce by default needs evidence that refusing reservations actually relieves the overshoot on their workload, which this does not provide.
One measurement caveat: standalone Spark enforces no container limit, so nothing was killed here. "Exceeded
spark.memory.offHeap.size" is a proxy for what would get an executor killed under YARN or Kubernetes, not a demonstration of a prevented kill.