viirya commented on code in PR #7096:
URL: https://github.com/apache/datafusion-comet/pull/7096#discussion_r4321652975
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -153,22 +209,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 {
+ let (task_share, tasks) = self.task_share();
+ // `num_spill` counts this consumer.
+ let limit = task_share.saturating_sub(state.unspillable) /
state.num_spill.max(1);
Review Comment:
This ceiling, `(task_share - unspillable) / num_spill`, can be well below
what Spark would grant the task. DataFusion 55.2's `ExternalSorter` fails
instead of spilling in a few spots:
- `reserve_memory_for_batch_and_maybe_spill` when `in_mem_batches` is empty
- the same function's retry after a spill
- the `try_resize` in `sort_batch_stream`
Each returns "Not enough memory to continue external sort". The description
mentions hitting exactly this in a suite ("five tasks left a sort a share
smaller than one of its batches"). Keeping the fraction as a per-task cap fixed
the suite, but the same thing can happen in production. Take 16 cores, 1 GiB of
off-heap and three or four spillable operators in a task. Each one gets tens of
MiB at most, and one batch of wide string rows can be bigger than that. On
main, that sort could use the task's whole Spark grant.
Is this a trade-off you have weighed? If it is accepted, could the 1.2.0
migration note say that a sort can now fail when a single batch exceeds its
share, and point to `greedy_unified` or a larger off-heap size as the way out?
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -851,8 +851,7 @@ impl ExecutionPlan for IcebergWriteExec {
context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
let metrics = WriteMetrics::new(&self.metrics, partition);
- // One consumer for the whole task, however many files it opens: a
consumer per file
- // would shrink every other consumer's share of the fair pool as a
fanout write widened.
+ // One consumer for the whole task, however many files it opens.
Review Comment:
The fanout path does respond to a refusal. `reserve()` writes out and closes
the partitions holding the most until the resize fits. With this change an
unspillable consumer is held to no share, so a fanout write grows until Spark
refuses, and every byte it holds is subtracted from the shares of the task's
spillable operators. In a plan that sorts and writes in the same task, the sort
ends up spilling where the writer could have closed files instead.
Should the fanout writer register `with_can_spill(true)`? The unpartitioned
and clustered writers would stay unspillable, since a refusal is fatal for
them. Either way, the new sentence at `memory_management.md:454` ("They cannot
spill") contradicts the paragraph right above it, which describes the fanout
writer closing files on refusal.
##########
spark/src/main/scala/org/apache/spark/Plugins.scala:
##########
@@ -322,7 +322,7 @@ object CometDriverPlugin extends Logging {
// spark.comet.exec.memoryPool.fraction was documented as holding back part
of the off-heap pool
// for the native memory that Comet does not reserve. It cannot: Spark hands
out the whole pool
- // to the tasks that ask for it, and the fraction only caps each task's
consumers under
+ // to the tasks that ask for it, and the fraction only limits the operators
that can spill under
Review Comment:
"the fraction only limits the operators that can spill" is not quite
accurate. `pool_size` also caps the task's total for unspillable consumers,
through the `else if self.pool_size < used.saturating_add(additional)` branch
in `try_grow`. Something like "the fraction only caps each task's total under
fair_unified" would match the code.
##########
docs/source/contributor-guide/memory_management.md:
##########
@@ -311,34 +312,50 @@ JNI, which goes through Spark's ordinary
`TaskMemoryManager`. That means:
is refused unless Spark can cover both, so operators spill until the debt is
repaid. Both pools'
`Display` output and their `try_grow` errors report the current overcommit.
-`CometFairMemoryPool` additionally applies two local checks before it asks
Spark, and refuses the
+`CometFairMemoryPool` exists because Spark cannot ask a native operator to
spill. Spark limits each
+task to an even share of its off-heap pool among the tasks that hold memory
from it. Once one
+operator holds its task's share, Spark refuses the task's other operators, and
one that cannot
+spill, such as a sort's merge or the build side of a hash join, fails the
task. The pool keeps
+memory for those operators by limiting the ones that can spill to a share of
the task's, as
+DataFusion's `FairSpillPool` does over its whole pool
+([issue #7088](https://github.com/apache/datafusion-comet/issues/7088)).
+
+The task's share is `off_heap_size` divided by the number of tasks on the
executor that are running
+native plans, that is, that hold a task-shared pool of either type, or
`pool_size` if that is
+smaller, which it is only while the deprecated fraction is below `1 / N`. Like
Spark's own divisor,
+the number follows tasks as they start and finish, so a task that runs alone
has the whole pool.
+Spark's count differs in two ways: it includes tasks that run no native plan,
and it counts a task
Review Comment:
Could this also say what each difference does? Spark's `ExecutionMemoryPool`
counts a task from its first acquire until its balance returns to zero
(`memoryForTask.remove`).
- When tasks that run no native plan hold off-heap memory, Comet divides by
fewer tasks than Spark does. The spillable shares then add up to more than
Spark will grant, so #7088 can still happen in mixed JVM and native stages.
- When tasks hold a native pool but no memory, for example a native scan
feeding JVM shuffle, Comet divides by more. Sorts then spill earlier than Spark
requires.
That would tell the next reader which way to look when a mixed workload
misbehaves.
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -31,14 +32,31 @@ use parking_lot::Mutex;
/// A DataFusion fair `MemoryPool` implementation for Comet. Internally this is
/// implemented via delegating calls to
[`crate::jvm_bridge::CometTaskMemoryManager`].
+///
+/// Spark limits each task to an even share of its off-heap pool among the
tasks that hold memory
+/// from it, and it cannot ask a native consumer to spill. So, as DataFusion's
`FairSpillPool` does
+/// over its whole pool, this pool limits each consumer that can spill to an
even share of what the
Review Comment:
Could this doc comment, and the matching paragraph in the guide, state the
limit as a formula as well? For example:
```text
task_share = min(off_heap_size / running_tasks, pool_size)
limit = (task_share - unspillable) / num_spill
```
The prose version nests "an even share of what ... leave of the task's
share, which is an even share of ..." and takes a few reads to parse. The guide
already uses a formula block for `memory_limit`.
##########
native/core/src/execution/delta_dv.rs:
##########
@@ -1837,21 +1808,16 @@ mod tests {
}
/// One `attach_access_plans` call must add exactly one consumer to the
task's pool no
- /// matter how many DV'd files it attaches. `CometFairMemoryPool` divides
the pool by the
- /// number of registered consumers, and the reservations live in the
returned files until
- /// the task ends, so one consumer per file would lower every other native
operator's fair
- /// limit for the whole task even when the DV bytes themselves are tiny.
Three DV'd files
- /// must leave a later consumer (a hash join build, say) its half of the
pool. Each file's
- /// bytes must still return to the pool when that file alone drops, while
the shared
- /// registration stays until the last file is gone.
+ /// matter how many DV'd files it attaches. Each file's bytes must still
return to the pool
Review Comment:
Now that the fair-share reason is gone, is there still a reason to pin one
consumer per call rather than one per file? Keeping `TrackConsumersPool`'s
top-10 list readable, or keeping the pool's per-consumer map small, might be
one. If so, could this comment and the one on `attach_access_plans` say so?
Otherwise the test asserts "not one per file" without anything explaining why
that matters.
--
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]
