andygrove opened a new issue, #5383:
URL: https://github.com/apache/datafusion-comet/issues/5383
## What / Why
Every DataFusion `try_grow` / `shrink` against `CometUnifiedMemoryPool` or
`CometFairMemoryPool` becomes a JNI round-trip into `CometTaskMemoryManager`,
and there is no batching, no slack, and no hysteresis: memory is acquired in
exactly the requested amount and released the moment a reservation shrinks,
however small the amount.
This breaks out positions 17 and 18 of #5212 into a standalone issue so the
work can be sequenced and benchmarked on its own.
## Anatomy of one grow/shrink
`native/core/src/execution/memory_pools/unified_pool.rs:64-78`
```rust
fn acquire_from_spark(&self, additional: usize) -> CometResult<i64> {
let handle = self.task_memory_manager_handle.as_obj();
JVMClasses::with_env(|env| unsafe {
jni_call!(env,
comet_task_memory_manager(handle).acquire_memory(additional as
i64) -> i64)
})
}
```
Each call pays:
1. `JVMClasses::with_env` (`native/jni-bridge/src/lib.rs:339-368`) —
`attach_current_thread_guard` plus `with_local_frame`, i.e. a
PushLocalFrame/PopLocalFrame pair, even though `acquire_memory(long) -> long`
and `release_memory(long) -> void` create no local references at all.
2. `call_method_unchecked` plus `check_exception`
(`native/jni-bridge/src/lib.rs:78-100`).
3. On the JVM side, `internal.acquireExecutionMemory`
(`spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java:62`), which
takes Spark's executor-wide `synchronized` `MemoryManager` lock. This is the
part that scales badly — every concurrent task's grow and shrink contends on it.
Callers hit this per batch: DataFusion's aggregate (`try_resize` per input
batch), sort accumulation and repartition, plus Comet's own native shuffle
writer, which does `try_grow` per batch and `reservation.free()` per flush
(`native/shuffle/src/partitioners/multi_partition.rs:463`, `:528`) — a full
release-then-reacquire cycle on every spill.
## Options
### A. Chunked acquire with retained slack
Track two numbers instead of one: `used` (what DataFusion reservations hold)
and `granted` (what Spark has handed us), with the invariant `granted >= used`.
- `try_grow(n)`: if `used + n <= granted`, bump `used` and return — no JNI
at all. Otherwise ask Spark for `max(need, chunk)`.
- `shrink(n)`: decrement `used`, and call Spark only when `granted - used`
exceeds the slack cap, releasing chunk-aligned amounts.
- Release all remaining `granted` in `Drop`. Note that today's `Drop` only
warns (`unified_pool.rs:81-91`), so this has to be added regardless.
One knob, something like `spark.comet.exec.memoryPool.acquireChunkSize`,
defaulting to a few MB with `0` disabling the behaviour; slack cap equal to the
chunk size. Steady-state churn within a chunk becomes free, which is exactly
the shuffle writer's access pattern.
Two sharp edges:
- `acquireExecutionMemory` will spill *other* consumers to satisfy an
oversized request, so only round small requests up; pass large ones through at
their exact size.
- On a short grant, keep it when `grant >= need`, otherwise release it and
return `resources_err` as today, so the DataFusion-side spill still triggers.
Hoarding memory Spark is actively contending for is worse than the extra
round-trip.
Implementation note: the `granted`/`used` bookkeeping and the JNI handle
should live in one shared type used by both `unified_pool.rs` and
`fair_pool.rs`, so the fair pool gets the same benefit and position 17 of #5212
(mutex held across the JNI round-trip) can be fixed in the same place.
### B. Make the retained slack reclaimable
`NativeMemoryConsumer.spill()` returns 0 unconditionally
(`CometTaskMemoryManager.java:110-113`), so Spark can never reclaim anything
from Comet. Holding slack makes that worse.
Fix: have `spill(size, trigger)` call back into native — a new `@native def
releasePoolSlack(taskAttemptId: Long): Long` resolved through the existing
`TASK_SHARED_MEMORY_POOLS` registry (`memory_pools/task_shared.rs:25`) — and
return the number of bytes of *free slack* dropped. No data is spilled, so it
is cheap and always safe.
This is what makes hysteresis defensible under memory pressure rather than
just shifting pain onto other tasks in the executor. It also partly addresses
position 5 of #5212, since Comet's consumer would stop being invisible to
Spark's spill-victim ordering.
I would not ship A enabled by default without B.
### C. Cheaper per-call path
Add a `with_env` variant that skips `with_local_frame` for calls that create
no local references. Two fewer JNI transitions per acquire and release, no
behavioural change. Worth doing whether or not A lands.
### D. Pre-reserve the whole per-task budget
Chunk size equal to `memoryLimitPerTask`, then run a native
`GreedyMemoryPool` over it and touch Spark only on exhaustion. Nearly
eliminates the JNI traffic, but hoards memory from tasks that need little and
defeats Spark's dynamic sharing. Reasonable as an opt-in config, not as a
default.
### E. Reduce call frequency at the source
Round Comet's own reservations up so growth is stepwise rather than per
batch; the native shuffle writer is the main offender. Narrower than A,
complementary, and needs no new config.
## Suggested sequencing
1. **Instrumentation.** Count acquire/release calls and bytes, and measure
cumulative time spent in the JNI round-trip, then report per task. Without this
there is no way to show a win — or to tell whether the win is worth the added
complexity. This should be its own PR.
2. **C** — small, unconditional, independent of everything else.
3. **A and B together**, behind a config, with numbers from position 1 above.
4. **E** if the shuffle writer still shows up in the counters afterwards.
D only if a workload turns up that genuinely wants it.
## Related
- #5212 — parent EPIC; this is positions 17 and 18, and B overlaps position
5.
- #4576 — adopting DataFusion's allocator-level accounting would change what
triggers these calls, but not what each one costs.
--
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]
