andygrove commented on code in PR #7096:
URL: https://github.com/apache/datafusion-comet/pull/7096#discussion_r4345560907
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -153,22 +219,32 @@ impl MemoryPool for CometFairMemoryPool {
) -> Result<(), DataFusionError> {
if additional > 0 {
let mut state = self.state.lock();
- let num = state.consumers.len();
- let limit = self
- .pool_size
- .checked_div(num)
- .expect("overflow in checked_div");
- let consumer_used = *state.consumer_used(reservation);
- if limit < consumer_used.saturating_add(additional) {
- return resources_err!(
- "Failed to acquire {additional} bytes where this consumer
already holds {consumer_used} bytes and the fair limit is {limit} bytes, {num}
registered ({} bytes overcommitted)",
- self.spark.overcommit()
- );
- }
- // The shares alone do not bound the pool's total, because a
consumer keeps what it
- // reserved before another consumer registered.
+ let can_spill = reservation.consumer().can_spill();
let used = state.used;
- if self.pool_size < used.saturating_add(additional) {
+ if can_spill {
Review Comment:
Good catch. I checked DataFusion 55.2 and it is as you describe.
`OrderedFinalAggregateStream` registers with `with_can_spill(true)` as the
#17334 workaround, builds a spill context only for `PartiallySorted` input, and
returns the refusal as the error otherwise.
The pool now treats it as a consumer that cannot spill.
`limited_as_spillable` in `fair_pool.rs` is what `register`, `unregister`,
`grow`, `shrink` and `try_grow` now ask instead of `can_spill()`. It returns
false for `OrderedFinalAggregateStream[`, so that stream takes Spark's grant,
and its bytes come out of the shares of the consumers that can spill. The
partially sorted stream still spills when Spark refuses it.
`an_ordered_final_aggregate_is_limited_as_a_consumer_that_cannot_spill` has
a sort and the aggregate in one task. The aggregate takes 60 of the task's 100
bytes, and the sort is then limited to the 40 left. With `limited_as_spillable`
reduced to `can_spill()`, the aggregate is refused at 50 and the test fails.
The contributor guide names the exception next to the paragraph on consumers
that cannot spill. (d79cb838b)
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -203,15 +282,213 @@ impl MemoryPool for CometFairMemoryPool {
#[cfg(test)]
mod tests {
use super::super::spark_memory::fake::FakeSpark;
+ use super::super::task_shared::RunningTask;
use super::*;
use std::sync::Arc;
+ /// A pool of `pool_size` bytes, all of the off-heap pool as with the
default fraction, which
+ /// `tasks` share, connected to `fake`.
+ fn shared_pool(
+ fake: &Arc<FakeSpark>,
+ pool_size: usize,
+ tasks: &RunningTasks,
+ ) -> Arc<dyn MemoryPool> {
+ Arc::new(CometFairMemoryPool::with_spark(
+ fake.memory(),
+ pool_size,
+ pool_size,
+ tasks.clone(),
+ ))
+ }
+
+ /// A pool of `pool_size` bytes for a task that runs alone, connected to
`fake`.
+ fn pool(fake: &Arc<FakeSpark>, pool_size: usize) -> Arc<dyn MemoryPool> {
+ shared_pool(fake, pool_size, &RunningTasks::default())
+ }
+
+ /// Starts `n` tasks, which run until the returned guards drop.
+ fn start(tasks: &RunningTasks, n: usize) -> Vec<RunningTask> {
+ (0..n).map(|_| tasks.start()).collect()
+ }
+
+ /// A consumer that can spill, like a sort's buffer.
+ fn spillable(name: &str, pool: &Arc<dyn MemoryPool>) -> MemoryReservation {
+ MemoryConsumer::new(name)
+ .with_can_spill(true)
+ .register(pool)
+ }
+
+ #[test]
+ fn consumers_that_can_spill_share_their_tasks_share_of_the_pool() {
+ // Eight running tasks share a 1040-byte pool, and Spark grants this
one its 130 bytes.
+ let fake = FakeSpark::with(130);
+ let tasks = RunningTasks::default();
+ let _running = start(&tasks, 8);
+ let pool = shared_pool(&fake, 1040, &tasks);
+ let sort = spillable("ExternalSorter[0]", &pool);
+ let _other_sort = spillable("ExternalSorter[1]", &pool);
+ let other_merge =
MemoryConsumer::new("ExternalSorterMerge[1]").register(&pool);
+
+ // The two sorts share the task's 130 bytes, so the first is refused,
and spills, at 65.
+ sort.try_grow(65).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ // That leaves the rest of what Spark grants the task for the other
sort's merge, which
+ // cannot spill.
+ other_merge.try_grow(65).unwrap();
+ assert_eq!(fake.held(), 130);
+ }
+
+ #[test]
+ fn a_consumer_that_cannot_spill_is_limited_only_by_the_pool() {
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let _running = start(&tasks, 4);
+ let pool = shared_pool(&fake, 400, &tasks);
+ let _sort = spillable("ExternalSorter[0]", &pool);
+ let _other_sort = spillable("ExternalSorter[1]", &pool);
+ let merge =
MemoryConsumer::new("ExternalSorterMerge[0]").register(&pool);
+
+ // Neither a share of the task's 100 bytes nor the task's share itself
limits the merge.
+ merge.try_grow(400).unwrap();
+ // The pool does, but only binds before Spark does when the deprecated
+ // `spark.comet.exec.memoryPool.fraction` is below 1.0.
+ assert!(merge.try_grow(1).is_err());
+ assert_eq!(fake.held(), 400);
+ }
+
+ #[test]
+ fn consumers_that_can_spill_share_what_the_others_leave() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool = pool(&fake, 100);
+ let first = spillable("ExternalSorter[0]", &pool);
+ let second = spillable("ExternalSorter[1]", &pool);
+ let merge =
MemoryConsumer::new("ExternalSorterMerge[0]").register(&pool);
+ merge.try_grow(40).unwrap();
+
+ // The two sorts share the 60 bytes that the merge leaves.
+ first.try_grow(30).unwrap();
+ assert!(first.try_grow(1).is_err());
+
+ // Their shares grow as the merge releases memory.
+ merge.shrink(40);
+ first.try_grow(20).unwrap();
+ assert!(first.try_grow(1).is_err());
+ second.try_grow(50).unwrap();
+ assert!(second.try_grow(1).is_err());
+ assert_eq!(pool.reserved(), 100);
+ }
+
+ #[test]
+ fn
memory_recorded_for_a_consumer_that_cannot_spill_comes_out_of_the_shares() {
+ let fake = FakeSpark::with(usize::MAX);
+ let pool = pool(&fake, 100);
+ let sort = spillable("ExternalSorter[0]", &pool);
+ let _other_sort = spillable("ExternalSorter[1]", &pool);
+ // Like a batch read back from a spill file, which already exists.
+ let reader = MemoryConsumer::new("spill reader").register(&pool);
+ reader.grow(40);
+
+ // The two sorts share the 60 bytes left.
+ sort.try_grow(30).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ #[test]
+ fn a_tasks_share_follows_the_number_of_running_tasks() {
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let mut running = start(&tasks, 4);
+ let pool = shared_pool(&fake, 400, &tasks);
+ let sort = spillable("ExternalSorter[0]", &pool);
+
+ // Four running tasks share the 400-byte pool.
+ sort.try_grow(100).unwrap();
+ assert!(sort.try_grow(1).is_err());
+
+ // Once two of them finish, the other two share it.
+ running.truncate(2);
+ sort.try_grow(100).unwrap();
+ assert!(sort.try_grow(1).is_err());
+
+ // A task that starts takes part of the share back, so the sort spills
before it grows.
+ running.push(tasks.start());
+ assert!(sort.try_grow(1).is_err());
+ sort.shrink(150);
+ sort.try_grow(83).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ #[test]
+ fn the_tasks_total_is_checked_against_the_tasks_share_not_the_pool() {
+ // One task has the whole 200-byte pool, and its sort takes half.
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let _running = tasks.start();
+ let pool = shared_pool(&fake, 200, &tasks);
+ let first = spillable("ExternalSorter[0]", &pool);
+ first.try_grow(100).unwrap();
+
+ // Another task starts, which halves the task's share to the 100 bytes
that the sort holds.
+ // A second sort is under its own limit of 50 bytes, but the task's
total is at its share.
+ let _other_task = tasks.start();
+ let second = spillable("ExternalSorter[1]", &pool);
+ let error = second.try_grow(50).unwrap_err().to_string();
+ assert!(error.contains("the task's share is 100 bytes"), "{error}");
+
+ // A consumer that cannot spill is limited by the pool, not by the
task's share.
+ let merge =
MemoryConsumer::new("ExternalSorterMerge[1]").register(&pool);
+ merge.try_grow(100).unwrap();
+ assert_eq!(pool.reserved(), 200);
+ }
+
+ /// The trade-off that the migration guide describes: a consumer that
cannot spill can take
+ /// what Spark grants, which can be the task's whole share, and then a
consumer that can spill
+ /// is refused its first request, so a sort fails where it used to run.
+ #[test]
+ fn
a_consumer_that_can_spill_is_refused_once_the_others_hold_the_tasks_share() {
+ let fake = FakeSpark::with(usize::MAX);
+ let tasks = RunningTasks::default();
+ let _running = start(&tasks, 4);
+ let pool = shared_pool(&fake, 400, &tasks);
+ // Like a shuffled hash join's build side, which Spark lets grow past
`off_heap_size / 4`
+ // while the other tasks hold nothing yet.
+ let build = MemoryConsumer::new("HashJoinInput[0]").register(&pool);
+ build.try_grow(100).unwrap();
+ let sort = spillable("ExternalSorter[0]", &pool);
+
+ let error = sort.try_grow(1).unwrap_err().to_string();
+ assert!(error.contains("the fair limit is 0 bytes"), "{error}");
+
+ // The sort can reserve again once the build side gives memory back.
+ build.shrink(60);
+ sort.try_grow(60).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ #[test]
+ fn a_tasks_share_is_at_most_the_pool() {
Review Comment:
It is, thanks. Dropped. The `mod.rs` test covers the same case through
`create_pool`. (d79cb838b)
--
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]
