Conversation
|
Tested on Before PR the Q67 failed on After PR the query passed |
andygrove
added
enhancement
Collapse the two ways of enabling the real-usage gate into one. The `real_usage` memory pool type and `spark.comet.exec.memoryGuard.enabled` both armed the breaker and built a `RealUsageMemoryPool` ceilinged at the off-heap budget, differing only in the inner pool and in a `fair_share` value already derivable from the pool type. Drop the pool type, make `unbounded` selectable in off-heap mode, and have `memoryGuard.enabled` force it there. Layering the gate over a unified pool left both ceilings equal to the off-heap budget, so Spark's per-task accounting always rejected first and the gate never fired. Fix a hook that the rebase onto main silently clobbered. Tokio's `on_thread_start` is a setter rather than a list, so the `attach_thread_as_daemon` hook added upstream replaced the guard's `stamp_current_thread`. No worker thread was ever stamped, which left the breaker unable to fire on the threads doing the allocating. Both now run from a single `on_worker_thread_start`. Resolve `spark.comet.exec.memoryGuard.size` to a byte count before it crosses JNI. `bytesConf` only converts on read through the `ConfigEntry`, so the native side saw the raw `4g` string, failed to parse it, and silently fell back to the off-heap budget. Cleanups: five `#[global_allocator]` blocks collapse to two, the duplicate native reader for `spark.comet.exec.memoryPool` is gone, the gate skips `reserved()` and the task registry lock when no fair share is configured, `serializeCometSQLConfs` makes one pass with one `SQLConf.get`, `RealUsagePool` is renamed to `RealUsageMemoryPool` to match its siblings, and six repeated test fixtures fold into one.
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
Thanks for picking this up @comphead. I spent some time with this and with the allocator accounting in #5934, and my overall take is that this is a large and complex PR and I think we need to step back and build this incrementally. Perhaps we could start with adding the accounting allocator as proposed in #5934, and then layer the gate and any enforcement on top of that as separate PRs? That would give us one balance and one wrapper to reason about, and it would let us settle a few design questions below one at a time rather than all at once.
Detailed notes follow. Some of them come directly from what I learned validating #5934 on TPC-H SF100, so I have tried to include the numbers.
The breaker panics inside the allocator
oom_guard::track calls panic_any from inside GlobalAlloc::alloc, alloc_zeroed and realloc. The GlobalAlloc docs in core/src/alloc/global.rs say:
It's undefined behavior if global allocators unwind. This restriction may be lifted in the future, but currently a panic from any of these functions may lead to memory unsafety.
The catch_unwind scaffolding in jni_api.rs and the UNWINDING flag make the unwind land somewhere sensible when it works, but the "failed to initiate panic, exit 133" abort described in the comment above the ARMED CAS looks like this contract being hit. The CAS reduces how often it happens rather than making it sound. I think the enforcement point has to move out of the allocator. The gate in try_grow is already a safe enforcement point, and the executePlan poll loop already has an every-100-polls hook where a balance check could return a ResourcesExhausted error for the current task without unwinding through arbitrary code. That loses the "any allocation anywhere trips it" property, but it keeps the process alive, which is the whole point.
What the gate can and cannot see
From the SF100 TPC-H traces on #5934, pool reservations summed to near zero at points where the balance was over 1 GB. Q8 was 2011 MB real against 428 MB reserved, Q16 was 1192 MB against 0, Q17 was 1292 MB against 34 MB. Most native bytes at peak are scan buffers, shuffle and kernel intermediates that never call try_grow. Since RealUsageMemoryPool can only reject reserving operators, a task pushed over budget by unreserved allocation gets its sort or aggregate rejected, spills, retries, and is rejected again because spilling freed little. That ends in ResourcesExhausted for a task that was not the one consuming the memory, and the breaker does the same on an even more arbitrary thread. The Q67 result is encouraging, but it is one query that happens to have a large spillable sort. I think the design needs to say what the expected outcome is in the common case where the over-budget bytes are not behind a reservation.
Dealloc ordering
dealloc in oom_guard.rs calls inner.dealloc first and subtracts afterwards. jemalloc drops its stats.allocated at the start of a large free and then eagerly unmaps anything over its 8 MiB oversize threshold, which takes milliseconds for a 100+ MB block, so the balance carries a block the allocator has already released for that whole window. This was the root cause of the native_allocated > jemalloc_allocated samples on #5934 (up to 160 MB high, about 2% of samples, all during task teardown) and the same code is here. On #5934 it was cosmetic. Here the balance is enforced, so another thread's try_grow can be rejected, or the breaker can trip, on a value that is 100+ MB high. A free cannot fail, so the subtraction can move before the delegation, which is what ce9b13a did on #5934. The reasoning in the realloc comment about accounting before delegating is right for realloc and does not apply to dealloc.
Thread-exit drift
LOCAL_DRIFT is a bare Cell<isize>, so up to 64 KiB per dying thread is discarded, in either direction. Tokio's blocking pool churns on its idle timeout, so on a long-lived executor the enforced balance acquires a slow bias. #5934 handles this with a ThreadDrift destructor plus the re-entrancy flag and try_with needed to make a destructor-bearing thread-local safe inside an allocator. That carries over directly if this builds on that module.
Overlap with #5934 and the default feature
Both PRs install a wrapper around the same allocator selection in lib.rs with different feature names and different cfgs, and they will conflict on merge. #5934's alloc_accounting module is the same accounting with the two fixes above already applied, and its current_balance() is exactly what BalanceSource::Live needs. This PR also adds oom-guard to the default feature, so every build carries the wrapper with a relaxed load per allocation when idle. That cost is probably small but it is unmeasured, and the #5934 branch has an alloc_overhead bench that could answer it. For reference, with tracking fully on the cost on TPC-H SF100 was 1.4% overall and about 3% on Q21. #5934 is opt-in, this PR is always-on, and I would rather make that decision once.
Ceiling semantics
The gate's ceiling and the breaker's default limit are both memory_limit, which in off-heap mode is spark.memory.offHeap.size scaled by spark.comet.exec.memoryPool.fraction. The memoryGuard.size doc says it defaults to the off-heap size, which is only true when the fraction is 1.0. Separately, the balance excludes everything the JVM allocates off-heap through Unsafe, and with unbounded Spark's TaskMemoryManager no longer sees Comet's usage. Both consumers can then each approach the full off-heap size independently, so the executor's real off-heap footprint is bounded by roughly twice the budget rather than once. The tuning.md paragraph hints at the second half of this but not the consequence. Is the intent that operators set memoryGuard.size to a container limit? If so the docs should say that.
Fair share
has_per_task_budget returns false for Unbounded, and in off-heap mode the guard always forces Unbounded, so fair_share is always None and the gate is first-come. The comments in createPlan and try_grow describe fair share as if it applies in that path. It only applies in on-heap mode, where the divisor falls back to executor_cores because the greedy per-task pools keep no registry. Could the comments be scoped accordingly, or is there a plan to make it apply?
Smaller items
- Exposing
unboundedas a user-selectable off-heap pool type intuning.mdand thememoryPooldoc means a user can remove all limits by setting one config without the guard. If it exists only so the guard can force it, it does not need to be a documented choice. armruns on everycreatePlanand overwrites the process-wideLIMIT. After a trip,ARMEDstays false until the nextcreatePlananywhere, so the breaker is off for every other task in that window.stamp_current_threadcovers tokio workers and the JNI caller thread. Allocations on the blocking pool are tracked but never enforced. The prototype listed that as a follow-up and it is worth stating here too.- The repro in the comment above uses
spark.comet.exec.memoryPool=real_usage, which is not a value this diff accepts. The tested configuration in the description should match the code.
Tests
It might be worth adding a JVM-level test that enables the guard with a small memoryGuard.size and runs a sort or aggregate to confirm the spill-and-succeed path, and one with an impossible budget to confirm a retriable failure rather than a hang or a process abort. A unit test for the 4g size resolution added to serializeCometSQLConfs would also help. The unit tests that are here cover the pure helpers well.
2 participants
Which issue does this PR close?
Experiment for #4576 .
Adopted from #4582
Rationale for this change
What changes are included in this PR?
How are these changes tested?