andygrove opened a new pull request, #5983: URL: https://github.com/apache/datafusion-comet/pull/5983
## 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. #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 in the pools themselves. Before asking Spark for memory, both off-heap pools refuse a reservation when real native usage plus the request would exceed Comet's off-heap allotment. The operators that can spill then spill; the ones that cannot fail the task rather than the executor. Unlike #5959 this is not a new pool type. `fair_unified` and `greedy_unified` both get the check, and it is on by default, because a protection that has to be opted into does not help the users who hit this. ## What changes are included in this PR? - New `CheckedMemoryPool<P: MemoryPool>` (`native/core/src/execution/memory_pools/checked_pool.rs`). `try_grow` returns `ResourcesExhausted` when `native_allocated + additional > budget`, naming the request, the bytes in use, the budget and the reserved total, and otherwise delegates. Everything else delegates; `memory_limit` reports the budget. The inner pool is generic so the gate is unit-testable without a JVM. - `create_memory_pool` wraps 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. - New `spark.comet.exec.memoryPool.checkNativeUsage` (default `true`) turns the check off. It rides the existing config proto rather than the `createPlan` JNI signature, so no signature change was needed. - The budget is `spark.memory.offHeap.size`, deliberately **not** the pool's own limit. That limit is the off-heap size times `spark.comet.exec.memoryPool.fraction`, and lowering the fraction is how operators force spilling. Deriving the budget from it as well would turn a small fraction into denied reservations instead of the spills it was set to cause, which is exactly what `CometTaskMetricsSuite` does deliberately. - `spark.comet.exec.memoryPool.fraction` now defaults to `0.8` rather than `1.0`, so some of the pool absorbs allocations Comet makes without reserving them. - `alloc-accounting` becomes a default cargo feature. All `cfg(feature)` guards remain, so `--no-default-features` still compiles; such a build reports zero bytes in use, which leaves the check a passthrough. - Docs: the tuning guide, the memory management contributor guide, and the tracing guide. Semantics worth stating up front: the balance and the budget are both process-wide. When any task pushes real usage to the budget, every task's next non-zero reservation is denied. There is no per-task attribution, and allocations themselves are never refused; this is a reservation gate, not a hard limit. Per-task attribution is a follow-up. ## How are these changes tested? Rust unit tests in `checked_pool.rs`: - `denies_when_real_bytes_plus_request_exceed_the_budget`: over `UnboundedMemoryPool` with a budget of the current balance plus 64 MiB, holding a 256 MiB block the pool never heard about makes a 1-byte `try_grow` fail with `ResourcesExhausted` and reserve nothing; dropping the block makes the same call succeed. - Zero-byte grows never fail, `memory_limit` reports the budget, and 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 small `spark.memory.offHeap.size` rather than a `withSQLConf` override: - Under both `fair_unified` and `greedy_unified`, a sort's reservation is refused with the check's message. - With `spark.comet.exec.memoryPool.checkNativeUsage=false`, the same query under the same off-heap size is no longer held back in Comet, so the reservation reaches Spark's ledger and fails there instead. The assertion is on which component refuses it, which is what isolates the check. `CometTaskMetricsSuite` is the regression this design is shaped around: it sets `memoryPool.fraction` to `0.002` to force spilling, and an earlier revision that derived the budget from the fraction turned those spills into hard failures. -- 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]
