andygrove commented on code in PR #5466:
URL: https://github.com/apache/datafusion-comet/pull/5466#discussion_r3905935196
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -187,4 +191,29 @@ impl MemoryPool for CometFairMemoryPool {
fn reserved(&self) -> usize {
self.state.lock().used
}
+
+ fn memory_limit(&self) -> MemoryLimit {
+ MemoryLimit::Finite(self.pool_size)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion::execution::memory_pool::UnboundedMemoryPool;
+
+ #[test]
+ fn fair_share_uses_requesting_reservation_and_reports_pool_limit() {
+ let backing: Arc<dyn MemoryPool> =
Arc::new(UnboundedMemoryPool::default());
+ let other = MemoryConsumer::new("other").register(&backing);
+ let requesting = MemoryConsumer::new("requesting").register(&backing);
+ other.grow(10);
+ requesting.grow(6);
+
+ assert_eq!(fair_limit_exceeded(32, 2, &requesting, 10), None);
+ assert_eq!(fair_limit_exceeded(32, 2, &requesting, 11), Some((6, 16)));
+
+ let pool = CometFairMemoryPool::new(Arc::new(Global::null()), 32);
Review Comment:
This test only reaches the free function with a hardcoded `num` of 2, so
nothing exercises `try_grow` itself. If someone reverted `try_grow` back to
passing `state.used`, this would still pass.
The `other` reservation also has no bearing on either assertion, since
`fair_limit_exceeded` never looks at other consumers. The test would behave
identically without it, so it does not really reproduce the 10 plus 6 case in
its name.
Would it be possible to register both consumers on the `CometFairMemoryPool`
so that `num` comes from the pool's own bookkeeping rather than a literal? I
realize `try_grow` is awkward to reach because `acquire` needs a live JVM. If
putting the JNI calls behind a small trait is more than you want to take on
here, then a comment explaining why the test stops at the helper would at least
make the gap visible to the next person reading it.
On the `memory_limit()` half, I checked and `Global::null()` is safe to drop
without a JVM in jni 0.22.4 because `Drop` short-circuits on a null reference,
so no concern there.
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -142,21 +152,15 @@ impl MemoryPool for CometFairMemoryPool {
fn try_grow(
&self,
- _reservation: &MemoryReservation,
+ reservation: &MemoryReservation,
additional: usize,
) -> Result<(), DataFusionError> {
if additional > 0 {
let mut state = self.state.lock();
let num = state.num;
- let limit = self
- .pool_size
- .checked_div(num)
- .expect("overflow in checked_div");
- // We use state.used instead of reservation.size() because
DataFusion 53+
- // calls pool.try_grow() before incrementing the reservation's
atomic size,
- // so reservation.size() would not include prior grows.
- let used = state.used;
- if limit < used + additional {
+ if let Some((used, limit)) =
+ fair_limit_exceeded(self.pool_size, num, reservation,
additional)
+ {
return resources_err!(
"Failed to acquire {additional} bytes where {used} bytes
already reserved and the fair limit is {limit} bytes, {num} registered"
);
Review Comment:
I want to back up @sunchao's point about losing the aggregate bound, with
two specifics that I think make it more concrete.
The first is that there is a documentation claim this breaks.
`docs/source/user-guide/latest/tuning.md` lines 66 to 67 say the shared pool
"ensures that the combined memory usage stays within the per-task limit". After
this change nothing in the fair pool enforces that. `state.used` becomes pure
bookkeeping and the only real ceiling left is Spark's `TaskMemoryManager`,
which applies the full off-heap size and knows nothing about
`spark.comet.exec.memoryPool.fraction`.
The second is that the sibling-reservation case is not hypothetical for the
operators we actually run. `new_empty()` and `split()` show up in
`sorts/sort.rs:644`, `sorts/streaming_merge.rs:248`, `sorts/stream.rs:187` and
`:253`, and `aggregates/row_hash.rs:1368`. Each sibling shares one consumer
registration but carries its own size counter, so each one independently gets a
full `pool_size / num` allowance. A single sort can hold several multiples of
its intended share.
In fairness, upstream `FairSpillPool` has the same late-registration
overshoot property, so this is not a hazard we are inventing. But our pool
takes a `pool_size` and after this change uses it only as a divisor. Could you
keep the corrected per-reservation check and add `state.used + additional >
pool_size` next to it? If you would rather not, then I think `tuning.md` lines
66 to 70 need updating in this PR, since the `pool_size / num_reservations`
wording on line 70 also gets less accurate once the denominator is
registrations and the numerator is per-reservation.
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -44,6 +43,17 @@ struct CometFairPoolState {
num: usize,
}
+fn fair_limit_exceeded(
+ pool_size: usize,
+ num: usize,
+ reservation: &MemoryReservation,
+ additional: usize,
+) -> Option<(usize, usize)> {
+ let used = reservation.size();
+ let limit = pool_size.checked_div(num).expect("overflow in checked_div");
Review Comment:
Since this is moving into a helper anyway, could it use
`.unwrap_or(pool_size)` the way `FairSpillPool` does with
`.checked_div(state.num_spill).unwrap_or(spill_available)`? The current message
describes the wrong failure, because `checked_div` returns `None` on division
by zero rather than overflow, and a panic in here aborts the process instead of
surfacing as a Spark error. The helper is now a free function that is directly
callable with `num = 0`, so it is a bit easier to hit than it was before.
--
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]
