This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 16642f4597 [CELEBORN-2502] Swap DataPusher buffers instead of copying
pushed data
16642f4597 is described below
commit 16642f45970979baab3965c60cbc87e8dbb6de38
Author: ForVic <[email protected]>
AuthorDate: Mon Oct 26 14:17:43 2026 +0800
[CELEBORN-2502] Swap DataPusher buffers instead of copying pushed data
### What changes were proposed in this pull request?
`DataPusher.addTask` copies the caller's full push buffer into an idle
`PushTask` buffer before queueing it. The Spark writers only call it with a
buffer they are about to reuse, so they can give the buffer away instead of
copying it.
- **client**: `DataPusher.addTaskAndSwapBuffer` puts the caller's buffer
into the idle `PushTask` and returns the task's previous buffer for the caller
to continue with. `addTask` keeps copying and is still used by the Tez writer.
- **spark**: `HashBasedShuffleWriter` and `SortBasedPusher` use
`addTaskAndSwapBuffer`. Both only flush a full buffer of
`celeborn.client.push.buffer.max.size` bytes, the size of the `PushTask`
buffers, so the buffer the writer gets back is the same size.
`HashBasedShuffleWriter` now caps `celeborn.client.push.buffer.initial.size` at
the max size, so a send buffer always reaches exactly that size before it is
flushed.
`PushTask` queues are reused across writers through `SendBufferPool`, so an
idle task can hold a buffer of another writer's push buffer size. In that case
the writer gets a newly allocated buffer of its own length instead, so it never
writes past the end of a smaller buffer. Within one application all writers
share the push buffer size, so this does not happen outside tests.
Every buffer is owned by either the writer or a single `PushTask`, so the
push thread never reads a buffer the writer is still filling. The task is
handed over through the same queues as before, which also publish the swapped
buffer to the push thread.
### Why are the changes needed?
Spark shuffle writes copy every pushed byte once more on the task thread.
That copy goes away.
Writer-side cost of handing one full buffer to `DataPusher`, the step this
PR changes, measured with a no-op push and one writer thread (median of the
warm rounds):
| Buffer size | `addTask` (copy) | `addTaskAndSwapBuffer` |
|---|---|---|
| 64 KiB (default) | 4.1 us | 1.6 us |
| 256 KiB | 25.0 us | 1.9 us |
The CPU this saves in a real shuffle, on one m6id.8xlarge (Celeborn master
and worker on the local NVMe disk, Spark 3.5.8 `local[16]`). Each
`spark-submit` caches 100M rows and repartitions them into 400 partitions three
times. Main and this PR swap only the Spark client jar and alternate order
every round. Summed CPU time of the map tasks, median of the warm iterations,
10 per side, p from a permutation test:
| Writer, codec | Shuffle per iteration | Map task CPU, main -> PR |
|---|---|---|
| hash, lz4 | 8.8 GiB | 109.2 s -> 107.9 s (-1.3 s, p=0.01) |
| hash, none | 13.0 GiB | 97.6 s -> 96.1 s (-1.5 s, p<0.001) |
That is about the cost of copying 9 to 13 GiB, and an earlier run of the
first commit of this PR gave similar savings (-1.1 s and -1.5 s). The sort
writer pushes only the full buffers of a partition through `DataPusher`, so its
saving is smaller and was below the run-to-run noise of this benchmark.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
- New `DataPushQueueSuiteJ#testAddTaskAndSwapBuffer`: pushes 1000 buffers,
each filled with its own index, and checks that every returned buffer is a
different array of the same size. It also checks that the push thread sees each
payload intact and in order while the writer keeps filling the returned buffers.
- New
`DataPushQueueSuiteJ#testAddTaskAndSwapBufferWithPooledTasksOfOtherSize`:
pooled tasks with smaller buffers must not shrink the writer's buffer.
- Existing `DataPushQueueSuiteJ`, `HashBasedShuffleWriterSuiteJ`,
`SortBasedShuffleWriterSuiteJ` and `SparkShuffleManagerSuite`.
- The benchmarks above. Both variants wrote the same amount of shuffle data
in every iteration.
Closes #3892 from ForVic/ForVic/datapusher-buffer-swap.
Lead-authored-by: ForVic <[email protected]>
Co-authored-by: zhengtao <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../spark/shuffle/celeborn/SortBasedPusher.java | 4 +-
.../shuffle/celeborn/HashBasedShuffleWriter.java | 12 ++-
.../apache/celeborn/client/write/DataPusher.java | 33 ++++++-
.../org/apache/celeborn/client/write/PushTask.java | 4 +
.../celeborn/client/write/DataPushQueueSuiteJ.java | 110 +++++++++++++++++++++
5 files changed, 155 insertions(+), 8 deletions(-)
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
index 6e5a7d48a4..a4a1985b5b 100644
---
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
@@ -250,7 +250,7 @@ public class SortBasedPusher extends MemoryConsumer {
if (offSet + recordSize > dataBuf.length) {
try {
- dataPusher.addTask(partition, dataBuf, offSet);
+ dataBuf = dataPusher.addTaskAndSwapBuffer(partition, dataBuf,
offSet);
uncompressedBytesWritten += offSet;
memoryThresholdManager.updateStats(offSet, true);
} catch (InterruptedException e) {
@@ -266,7 +266,7 @@ public class SortBasedPusher extends MemoryConsumer {
}
if (offSet > 0) {
try {
- dataPusher.addTask(currentPartition, dataBuf, offSet);
+ dataPusher.addTaskAndSwapBuffer(currentPartition, dataBuf, offSet);
uncompressedBytesWritten += offSet;
memoryThresholdManager.updateStats(offSet, offSet ==
pushBufferMaxSize);
} catch (InterruptedException e) {
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index d1e47e66f0..f9c281df9c 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -135,8 +135,10 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
mapStatusLengths[i] = new LongAdder();
}
- PUSH_BUFFER_INIT_SIZE = conf.clientPushBufferInitialSize();
PUSH_BUFFER_MAX_SIZE = conf.clientPushBufferMaxSize();
+ // A full send buffer is swapped with a DataPusher buffer of the max size,
so a send buffer
+ // must not start out larger than that.
+ PUSH_BUFFER_INIT_SIZE = Math.min(conf.clientPushBufferInitialSize(),
PUSH_BUFFER_MAX_SIZE);
this.sendBufferPool = sendBufferPool;
sendBuffers = sendBufferPool.acquireBuffer(numPartitions);
@@ -325,20 +327,22 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
if ((buffer.length - offset) < serializedRecordSize) {
- flushSendBuffer(partitionId, buffer, offset);
+ sendBuffers[partitionId] = flushSendBuffer(partitionId, buffer, offset);
updateRecordsWrittenMetrics();
offset = 0;
}
return offset;
}
- private void flushSendBuffer(int partitionId, byte[] buffer, int size)
+ /** Hands {@code buffer} to the data pusher and returns the buffer to
continue writing into. */
+ private byte[] flushSendBuffer(int partitionId, byte[] buffer, int size)
throws IOException, InterruptedException {
long start = System.nanoTime();
if (logger.isDebugEnabled()) logger.debug("Flush buffer, size {}.",
Utils.bytesToString(size));
- dataPusher.addTask(partitionId, buffer, size);
+ byte[] newBuffer = dataPusher.addTaskAndSwapBuffer(partitionId, buffer,
size);
writeMetrics.incWriteTime(System.nanoTime() - start);
uncompressedBytesWritten += size;
+ return newBuffer;
}
protected void closeWrite() throws IOException {
diff --git
a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
index 0677502ab2..a87a928010 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
@@ -153,6 +153,22 @@ public class DataPusher {
public void addTask(int partitionId, byte[] buffer, int size)
throws IOException, InterruptedException {
+ addTask(partitionId, buffer, size, false);
+ }
+
+ /**
+ * Like {@link #addTask(int, byte[], int)}, but hands {@code buffer} to the
pusher instead of
+ * copying it. The caller must not touch {@code buffer} afterwards and
should continue with the
+ * returned buffer, which has the same length as {@code buffer}: an idle
buffer taken from the
+ * pusher, or a new one if the idle buffer has another length.
+ */
+ public byte[] addTaskAndSwapBuffer(int partitionId, byte[] buffer, int size)
+ throws IOException, InterruptedException {
+ return addTask(partitionId, buffer, size, true);
+ }
+
+ private byte[] addTask(int partitionId, byte[] buffer, int size, boolean
swapBuffer)
+ throws IOException, InterruptedException {
client.computeBatchCRC(shuffleId, mapId, attemptId, partitionId, buffer,
0, size);
try {
PushTask task = null;
@@ -160,12 +176,25 @@ public class DataPusher {
checkException();
task = idleQueue.poll(WAIT_TIME_NANOS, TimeUnit.NANOSECONDS);
}
- task.setSize(size);
+ byte[] idleBuffer = null;
+ if (swapBuffer) {
+ idleBuffer = task.getBuffer();
+ // Tasks reused through SendBufferPool may come from a writer with
another push buffer
+ // size, so do not hand back a buffer of a different length than the
caller's.
+ if (idleBuffer.length != buffer.length) {
+ idleBuffer = new byte[buffer.length];
+ }
+ task.setBuffer(buffer);
+ task.setSize(size);
+ } else {
+ task.setSize(size);
+ System.arraycopy(buffer, 0, task.getBuffer(), 0, size);
+ }
task.setPartitionId(partitionId);
- System.arraycopy(buffer, 0, task.getBuffer(), 0, size);
while (!dataPushQueue.addPushTask(task)) {
checkException();
}
+ return idleBuffer;
} catch (InterruptedException e) {
logger.error("DataPusher thread interrupted while adding push task.");
pushThread.interrupt();
diff --git
a/client/src/main/java/org/apache/celeborn/client/write/PushTask.java
b/client/src/main/java/org/apache/celeborn/client/write/PushTask.java
index f4dbd61d4d..382292c827 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/PushTask.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/PushTask.java
@@ -49,4 +49,8 @@ public class PushTask {
public byte[] getBuffer() {
return buffer;
}
+
+ public void setBuffer(byte[] buffer) {
+ this.buffer = buffer;
+ }
}
diff --git
a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
index b789875ac2..d9d561e060 100644
---
a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
+++
b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuiteJ.java
@@ -20,6 +20,7 @@ package org.apache.celeborn.client.write;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
@@ -189,6 +190,115 @@ public class DataPushQueueSuiteJ {
client.shutdown();
}
+ @Test
+ public void testAddTaskAndSwapBuffer() throws Exception {
+ int shuffleId = 0;
+ int mapId = 0;
+ int attemptId = 0;
+ int numMappers = 10;
+ int numPartitions = 10;
+ int numTasks = 1000;
+ CelebornConf conf = new CelebornConf();
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ // Initialize partition locations to exercise the normal worker-selection
path.
+ client.initReducePartitionMap(shuffleId, numPartitions, 1);
+ LongAdder[] mapStatusLengths = new LongAdder[numPartitions];
+ for (int i = 0; i < numPartitions; i++) {
+ mapStatusLengths[i] = new LongAdder();
+ }
+ List<Integer> pushedPartitions = new ArrayList<>();
+ List<Integer> pushedValues = new ArrayList<>();
+ DataPusher dataPusher =
+ new DataPusher(
+ shuffleId,
+ mapId,
+ attemptId,
+ 0,
+ numMappers,
+ numPartitions,
+ conf,
+ client,
+ null,
+ integer -> {},
+ mapStatusLengths) {
+ @Override
+ protected void pushData(PushTask task) throws IOException {
+ byte[] buffer = task.getBuffer();
+ int value = bytesToInt(buffer);
+ for (int i = 4; i < task.getSize(); i += 4) {
+ Assert.assertEquals(value, bytesToInt(Arrays.copyOfRange(buffer,
i, i + 4)));
+ }
+ pushedPartitions.add(task.getPartitionId());
+ pushedValues.add(value);
+ }
+ };
+
+ int size = conf.clientPushBufferMaxSize();
+ byte[] buffer = new byte[size];
+ for (int i = 0; i < numTasks; i++) {
+ byte[] value = intToBytes(i);
+ for (int j = 0; j < size; j += 4) {
+ System.arraycopy(value, 0, buffer, j, 4);
+ }
+ byte[] pushed = buffer;
+ buffer = dataPusher.addTaskAndSwapBuffer(i % numPartitions, pushed,
size);
+ Assert.assertNotSame(pushed, buffer);
+ Assert.assertEquals(size, buffer.length);
+ }
+
+ dataPusher.waitOnTermination();
+
+ Assert.assertEquals(numTasks, pushedValues.size());
+ for (int i = 0; i < numTasks; i++) {
+ Assert.assertEquals(i % numPartitions, (int) pushedPartitions.get(i));
+ Assert.assertEquals(i, (int) pushedValues.get(i));
+ }
+ client.shutdown();
+ }
+
+ @Test
+ public void testAddTaskAndSwapBufferWithPooledTasksOfOtherSize() throws
Exception {
+ int numPartitions = 1;
+ CelebornConf conf = new CelebornConf();
+ int size = conf.clientPushBufferMaxSize();
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(0, numPartitions, 1);
+ LongAdder[] mapStatusLengths = new LongAdder[] {new LongAdder()};
+ // Tasks left in a pooled queue by a writer with a smaller push buffer
size.
+ LinkedBlockingQueue<PushTask> pushTasks =
+ new LinkedBlockingQueue<>(conf.clientPushQueueCapacity());
+ for (int i = 0; i < conf.clientPushQueueCapacity(); i++) {
+ pushTasks.put(new PushTask(size / 2));
+ }
+ DataPusher dataPusher =
+ new DataPusher(
+ 0,
+ 0,
+ 0,
+ 0,
+ 10,
+ numPartitions,
+ conf,
+ client,
+ pushTasks,
+ integer -> {},
+ mapStatusLengths) {
+ @Override
+ protected void pushData(PushTask task) {}
+ };
+
+ byte[] buffer = new byte[size];
+ for (int i = 0; i < 2 * conf.clientPushQueueCapacity(); i++) {
+ buffer = dataPusher.addTaskAndSwapBuffer(0, buffer, size);
+ Assert.assertEquals(size, buffer.length);
+ }
+
+ dataPusher.waitOnTermination();
+ client.shutdown();
+ }
+
@Test
public void testTaskAddedWhileTakingAllTasksIsNotLost() throws Exception {
CelebornConf conf = new CelebornConf();
