andygrove opened a new issue, #6196: URL: https://github.com/apache/datafusion-comet/issues/6196
### What is the problem the feature request solves? The local native shuffle writer allocates buffers of `spark.comet.shuffle.native.writeBufferSize` bytes that no `MemoryReservation` covers. Because of #6183 the default has been a 1-byte buffer, so these have been close to empty. #6191 fixes that, and with the 1 MiB default in effect, a multi-partition shuffle task that has spilled holds four of them: - the data file `BufWriter`, allocated when `LocalPartitionWriter` is created (`native/shuffle/src/writers/local/local_partition_writer.rs`) - the spill file `BufWriter`, allocated on the first spill and kept until the writer is dropped (`ensure_spill_file_created` in `native/shuffle/src/writers/local/spill.rs`) - the buffer `finish_partition` reads spilled ranges through, allocated for the first partition with spilled data and also kept (`spill_reader` in `local_partition_writer.rs`) - the recycled scratch that blocks are encoded into, which `BufBatchWriter::flush` shrinks back to the write buffer size rather than freeing (`native/shuffle/src/writers/buf_batch_writer.rs`). Between flushes it can reach the write buffer size plus one block. DataFusion's default spill writer has no buffer of its own, so that is all of them: about 4 MiB per task, or 2 MiB for a task that does not spill. The single-partition writer holds only the scratch, but `SinglePartitionShufflePartitioner` has no reservation to charge it to. The remote shuffle writer does not use this setting. An executor running 16 tasks at once can hold 64 MiB that neither Comet's pool nor Spark's `TaskMemoryManager` sees, and raising the setting to `8m` for write performance makes that 512 MiB. Like the rest of Comet's untracked memory, it has to fit in `spark.executor.memoryOverhead`, and when it does not, the cluster manager kills the executor instead of the writer spilling. ### Describe the potential solution As suggested in the #6191 review, charge the buffers to the `ShuffleRepartitioner[N]` consumer that `MultiPartitionShuffleRepartitioner` already registers. A few things constrain how: - Not in the existing reservation. `spill()` releases it with `reservation.free()` and reports what it freed as `memory_spilled_bytes`, and `spark.comet.shuffle.native.maxBufferBytes` is compared against its size. Buffers charged there would be released by the first spill while still allocated, counted as spilled, and would make the fixed limit trigger early. `reservation.new_empty()` gives a second reservation on the same consumer. - Not through a new consumer. `fair_unified` divides the pool by the number of registered consumers (#5961), so another registration would lower the limit for every consumer in the task. - `LocalPartitionWriter` allocates the data file buffer before the repartitioner registers its consumer, so the reservation has to be created first and passed in, or the buffer allocated lazily. - The spill file buffer is allocated inside `spill()`, usually just after the pool refused to grow, and the copy buffer while the last round of batches is still reserved, so growing a reservation for either can fail at that point. They could be reserved up front, which charges every multi-partition task for buffers it may never use, or fall back to a small buffer when the grow is refused. The writer ran with 1-byte buffers until #6191, and `copy_spill_range` already handles ranges larger than its buffer, so falling back costs speed, not correctness, and is better than failing the task. A test can spill through a `GreedyMemoryPool` and check that `memory_pool.reserved()` includes the buffers after a spill and returns to zero once the writer is dropped. ### Additional context Raised in https://github.com/apache/datafusion-comet/pull/6191#discussion_r4095031420. #6115 is the same problem for the Parquet and Iceberg writers. The zstd context that `ShuffleCodecContext` keeps between blocks is also untracked, about 1.3 MiB at level 1 according to the table in `native/shuffle/src/codec_context.rs`, and could be charged the same way when zstd is configured. If this is not done, the list of untracked memory under "Configuring Comet Memory" in the tuning guide should name these buffers. It currently mentions the shuffle writer only for its buffered partitions, which are reserved. -- 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]
