You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2023/07/17 08:35:23 UTC
[incubator-celeborn] branch main updated: [CELEBORN-802] Reuse DataPusher#idleQueue by pooling to avoid too many byte[] objects
This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 10a1def51 [CELEBORN-802] Reuse DataPusher#idleQueue by pooling to avoid too many byte[] objects
10a1def51 is described below
commit 10a1def512f6054fa7d742ac90735f19314af433
Author: zky.zhoukeyong <zk...@alibaba-inc.com>
AuthorDate: Mon Jul 17 16:35:14 2023 +0800
[CELEBORN-802] Reuse DataPusher#idleQueue by pooling to avoid too many byte[] objects
### What changes were proposed in this pull request?
Reuse ```DataPusher#idleQueue``` by pooling in ```SendBufferPool``` to avoid too many ```byte[]```
objects in ```PushTask```.
### Why are the changes needed?
I'm testing 3T TPCDS. Before this PR, I encountered Container killed because of OOM, GC is about 9.6h. For alive Executors, I dumped the memory and see number of PushTask object is 2w, and the number of ```64k``` byte[] is 23356, total around 1.7G:
After this PR, no container is killed because of OOM, GC is about 8.6h. I also dumped Executor and found number
of PushTask object is 3584, and the number of ```64K``` byte[] objects is 5783, total around 361M:
Also, before this PR, total execution time is ```3313.8s```, after this PR, total execution time is ```3229.5s```.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Passes GA and Manual test.
Closes #1722 from waitinfuture/802.
Authored-by: zky.zhoukeyong <zk...@alibaba-inc.com>
Signed-off-by: zky.zhoukeyong <zk...@alibaba-inc.com>
---
.../apache/spark/shuffle/celeborn/SendBufferPool.java | 19 +++++++++++++++++++
.../spark/shuffle/celeborn/SortBasedPusher.java | 11 ++++++++++-
.../spark/shuffle/celeborn/SortBasedPusherSuiteJ.java | 3 ++-
.../shuffle/celeborn/HashBasedShuffleWriter.java | 5 +++++
.../shuffle/celeborn/SortBasedShuffleWriter.java | 9 ++++++---
.../spark/shuffle/celeborn/SparkShuffleManager.java | 8 +++++++-
.../celeborn/SortBasedShuffleWriterSuiteJ.java | 2 +-
.../shuffle/celeborn/HashBasedShuffleWriter.java | 5 +++++
.../shuffle/celeborn/SortBasedShuffleWriter.java | 15 ++++++++++-----
.../spark/shuffle/celeborn/SparkShuffleManager.java | 9 ++++++++-
.../celeborn/SortBasedShuffleWriterSuiteJ.java | 2 +-
.../org/apache/celeborn/client/write/DataPusher.java | 14 +++++++++++---
.../celeborn/client/write/DataPushQueueSuitJ.java | 1 +
13 files changed, 86 insertions(+), 17 deletions(-)
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SendBufferPool.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SendBufferPool.java
index f568cc28c..30989d75c 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SendBufferPool.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SendBufferPool.java
@@ -19,6 +19,9 @@ package org.apache.spark.shuffle.celeborn;
import java.util.Iterator;
import java.util.LinkedList;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import org.apache.celeborn.client.write.PushTask;
public class SendBufferPool {
private static volatile SendBufferPool _instance;
@@ -38,11 +41,13 @@ public class SendBufferPool {
// numPartitions -> buffers
private final LinkedList<byte[][]> buffers;
+ private final LinkedList<LinkedBlockingQueue<PushTask>> pushTaskQueues;
private SendBufferPool(int capacity) {
assert capacity > 0;
this.capacity = capacity;
buffers = new LinkedList<>();
+ pushTaskQueues = new LinkedList<>();
}
public synchronized byte[][] acquireBuffer(int numPartitions) {
@@ -60,10 +65,24 @@ public class SendBufferPool {
return new byte[numPartitions][];
}
+ public synchronized LinkedBlockingQueue<PushTask> acquirePushTaskQueue() {
+ if (!pushTaskQueues.isEmpty()) {
+ return pushTaskQueues.removeFirst();
+ }
+ return null;
+ }
+
public synchronized void returnBuffer(byte[][] buffer) {
if (buffers.size() == capacity) {
buffers.removeFirst();
}
buffers.addLast(buffer);
}
+
+ public synchronized void returnPushTaskQueue(LinkedBlockingQueue<PushTask> pushTaskQueue) {
+ if (pushTaskQueues.size() == capacity) {
+ pushTaskQueues.removeFirst();
+ }
+ pushTaskQueues.addLast(pushTaskQueue);
+ }
}
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 5e5fc635d..305f83bc2 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
@@ -20,6 +20,7 @@ package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
import java.util.LinkedList;
import java.util.concurrent.ExecutorService;
+import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Consumer;
@@ -36,6 +37,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.write.DataPusher;
+import org.apache.celeborn.client.write.PushTask;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
@@ -71,6 +73,7 @@ public class SortBasedPusher extends MemoryConsumer {
private int[] shuffledPartitions = null;
private int[] inversedShuffledPartitions = null;
private final ExecutorService executorService;
+ private final SendBufferPool sendBufferPool;
public SortBasedPusher(
TaskMemoryManager memoryManager,
@@ -86,7 +89,8 @@ public class SortBasedPusher extends MemoryConsumer {
LongAdder[] mapStatusLengths,
long pushSortMemoryThreshold,
Object sharedPushLock,
- ExecutorService executorService) {
+ ExecutorService executorService,
+ SendBufferPool sendBufferPool) {
super(
memoryManager,
(int) Math.min(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES, memoryManager.pageSizeBytes()),
@@ -109,7 +113,10 @@ public class SortBasedPusher extends MemoryConsumer {
this.afterPush = afterPush;
this.mapStatusLengths = mapStatusLengths;
+ this.sendBufferPool = sendBufferPool;
+
try {
+ LinkedBlockingQueue<PushTask> pushTaskQueue = sendBufferPool.acquirePushTaskQueue();
dataPusher =
new DataPusher(
shuffleId,
@@ -120,6 +127,7 @@ public class SortBasedPusher extends MemoryConsumer {
numPartitions,
conf,
shuffleClient,
+ pushTaskQueue,
afterPush,
mapStatusLengths);
} catch (InterruptedException e) {
@@ -437,6 +445,7 @@ public class SortBasedPusher extends MemoryConsumer {
cleanupResources();
try {
dataPusher.waitOnTermination();
+ sendBufferPool.returnPushTaskQueue(dataPusher.getIdleQueue());
} catch (InterruptedException e) {
TaskInterruptedHelper.throwTaskKillException();
}
diff --git a/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedPusherSuiteJ.java b/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedPusherSuiteJ.java
index 3ab993937..58e8ad053 100644
--- a/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedPusherSuiteJ.java
+++ b/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedPusherSuiteJ.java
@@ -93,7 +93,8 @@ public class SortBasedPusherSuiteJ {
/*mapStatusLengths=*/ null,
/*pushSortMemoryThreshold=*/ Utils.byteStringAsBytes("1m"),
/*sharedPushLock=*/ null,
- /*executorService=*/ null);
+ /*executorService=*/ null,
+ SendBufferPool.get(4));
// default page size == 2 MiB
assertEquals(unifiedMemoryManager.pageSizeBytes(), Utils.byteStringAsBytes("2m"));
diff --git a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index e79524fb6..89826be05 100644
--- a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++ b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -18,6 +18,7 @@
package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
+import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.LongAdder;
import javax.annotation.Nullable;
@@ -49,6 +50,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.write.DataPusher;
+import org.apache.celeborn.client.write.PushTask;
import org.apache.celeborn.common.CelebornConf;
@Private
@@ -138,6 +140,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
sendOffsets = new int[numPartitions];
try {
+ LinkedBlockingQueue<PushTask> pushTaskQueue = sendBufferPool.acquirePushTaskQueue();
dataPusher =
new DataPusher(
shuffleId,
@@ -148,6 +151,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
numPartitions,
conf,
shuffleClient,
+ pushTaskQueue,
writeMetrics::incBytesWritten,
mapStatusLengths);
} catch (InterruptedException e) {
@@ -312,6 +316,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
private void close() throws IOException, InterruptedException {
// here we wait for all the in-flight batches to return which sent by dataPusher thread
dataPusher.waitOnTermination();
+ sendBufferPool.returnPushTaskQueue(dataPusher.getIdleQueue());
shuffleClient.prepareForMergeData(shuffleId, mapId, taskContext.attemptNumber());
// merge and push residual data to reduce network traffic
diff --git a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index 03f6f8626..62c05858d 100644
--- a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++ b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -99,7 +99,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
TaskContext taskContext,
CelebornConf conf,
ShuffleClient client,
- ExecutorService executorService)
+ ExecutorService executorService,
+ SendBufferPool sendBufferPool)
throws IOException {
this.mapId = taskContext.partitionId();
this.dep = dep;
@@ -143,7 +144,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
mapStatusLengths,
conf.clientPushSortMemoryThreshold() / 2,
globalPushLock,
- executorService);
+ executorService,
+ sendBufferPool);
}
currentPusher = pushers[0];
} else {
@@ -162,7 +164,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
mapStatusLengths,
conf.clientPushSortMemoryThreshold(),
globalPushLock,
- null);
+ null,
+ sendBufferPool);
}
}
diff --git a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index 8b924b793..95641a6fc 100644
--- a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++ b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -191,7 +191,13 @@ public class SparkShuffleManager implements ShuffleManager {
ExecutorService pushThread =
celebornConf.clientPushSortPipelineEnabled() ? getPusherThread() : null;
return new SortBasedShuffleWriter<>(
- h.dependency(), h.numMaps(), context, celebornConf, client, pushThread);
+ h.dependency(),
+ h.numMaps(),
+ context,
+ celebornConf,
+ client,
+ pushThread,
+ SendBufferPool.get(cores));
} else if (ShuffleMode.HASH.equals(celebornConf.shuffleWriterMode())) {
return new HashBasedShuffleWriter<>(
h, mapId, context, celebornConf, client, SendBufferPool.get(cores));
diff --git a/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java b/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java
index 2fbd5dcf4..a3638f4eb 100644
--- a/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java
+++ b/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java
@@ -34,6 +34,6 @@ public class SortBasedShuffleWriterSuiteJ extends CelebornShuffleWriterSuiteBase
CelebornShuffleHandle handle, TaskContext context, CelebornConf conf, ShuffleClient client)
throws IOException {
return new SortBasedShuffleWriter<Integer, String, String>(
- handle.dependency(), numPartitions, context, conf, client, null);
+ handle.dependency(), numPartitions, context, conf, client, null, SendBufferPool.get(4));
}
}
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 349629edf..0684c72f8 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
@@ -18,6 +18,7 @@
package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
+import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.LongAdder;
import javax.annotation.Nullable;
@@ -52,6 +53,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.write.DataPusher;
+import org.apache.celeborn.client.write.PushTask;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.Utils;
@@ -158,6 +160,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
sendOffsets = new int[numPartitions];
try {
+ LinkedBlockingQueue<PushTask> pushTaskQueue = sendBufferPool.acquirePushTaskQueue();
dataPusher =
new DataPusher(
shuffleId,
@@ -168,6 +171,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
numPartitions,
conf,
shuffleClient,
+ pushTaskQueue,
writeMetrics::incBytesWritten,
mapStatusLengths);
} catch (InterruptedException e) {
@@ -447,6 +451,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
// here we wait for all the in-flight batches to return which sent by dataPusher thread
long pushMergedDataTime = System.nanoTime();
dataPusher.waitOnTermination();
+ sendBufferPool.returnPushTaskQueue(dataPusher.getIdleQueue());
shuffleClient.prepareForMergeData(shuffleId, mapId, taskContext.attemptNumber());
if (isColumnarShuffle) {
closeColumnarWrite();
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index a6d4ff5e4..3766c623d 100644
--- a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++ b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -100,7 +100,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
CelebornConf conf,
ShuffleClient client,
ShuffleWriteMetricsReporter metrics,
- ExecutorService executorService)
+ ExecutorService executorService,
+ SendBufferPool sendBufferPool)
throws IOException {
this.mapId = taskContext.partitionId();
this.dep = dep;
@@ -143,7 +144,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
mapStatusLengths,
conf.clientPushSortMemoryThreshold() / 2,
sharedPushLock,
- executorService);
+ executorService,
+ sendBufferPool);
}
currentPusher = pushers[0];
} else {
@@ -162,7 +164,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
mapStatusLengths,
conf.clientPushSortMemoryThreshold(),
sharedPushLock,
- null);
+ null,
+ sendBufferPool);
}
}
@@ -172,7 +175,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
CelebornConf conf,
ShuffleClient client,
ShuffleWriteMetricsReporter metrics,
- ExecutorService executorService)
+ ExecutorService executorService,
+ SendBufferPool sendBufferPool)
throws IOException {
this(
handle.dependency(),
@@ -181,7 +185,8 @@ public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
conf,
client,
metrics,
- executorService);
+ executorService,
+ sendBufferPool);
}
private void updatePeakMemoryUsed() {
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index 3b6c2ea65..f9ff30d04 100644
--- a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++ b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -193,7 +193,14 @@ public class SparkShuffleManager implements ShuffleManager {
ExecutorService pushThread =
celebornConf.clientPushSortPipelineEnabled() ? getPusherThread() : null;
return new SortBasedShuffleWriter<>(
- h.dependency(), h.numMappers(), context, celebornConf, client, metrics, pushThread);
+ h.dependency(),
+ h.numMappers(),
+ context,
+ celebornConf,
+ client,
+ metrics,
+ pushThread,
+ SendBufferPool.get(cores));
} else if (ShuffleMode.HASH.equals(celebornConf.shuffleWriterMode())) {
return new HashBasedShuffleWriter<>(
h, context, celebornConf, client, metrics, SendBufferPool.get(cores));
diff --git a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java
index 50c455ab8..db94d8441 100644
--- a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java
+++ b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriterSuiteJ.java
@@ -36,6 +36,6 @@ public class SortBasedShuffleWriterSuiteJ extends CelebornShuffleWriterSuiteBase
ShuffleWriteMetricsReporter metrics)
throws IOException {
return new SortBasedShuffleWriter<Integer, String, String>(
- handle, context, conf, client, metrics, null);
+ handle, context, conf, client, metrics, null, SendBufferPool.get(4));
}
}
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 bc26dab69..5444e31f4 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
@@ -69,18 +69,23 @@ public class DataPusher {
int numPartitions,
CelebornConf conf,
ShuffleClient client,
+ LinkedBlockingQueue<PushTask> pushTasks,
Consumer<Integer> afterPush,
LongAdder[] mapStatusLengths)
throws InterruptedException {
final int pushQueueCapacity = conf.clientPushQueueCapacity();
final int pushBufferMaxSize = conf.clientPushBufferMaxSize();
- idleQueue = new LinkedBlockingQueue<>(pushQueueCapacity);
+ if (pushTasks == null) {
+ idleQueue = new LinkedBlockingQueue<>(pushQueueCapacity);
+ } else {
+ idleQueue = pushTasks;
+ }
dataPushQueue =
new DataPushQueue(
conf, this, client, shuffleId, mapId, attemptId, numMappers, numPartitions);
- for (int i = 0; i < pushQueueCapacity; i++) {
+ for (int i = idleQueue.size(); i < pushQueueCapacity; i++) {
idleQueue.put(new PushTask(pushBufferMaxSize));
}
@@ -174,7 +179,6 @@ public class DataPusher {
Thread.currentThread().interrupt();
throw e;
}
- idleQueue.clear();
dataPushQueue.clear();
checkException();
}
@@ -225,4 +229,8 @@ public class DataPusher {
public DataPushQueue getDataPushQueue() {
return dataPushQueue;
}
+
+ public LinkedBlockingQueue<PushTask> getIdleQueue() {
+ return idleQueue;
+ }
}
diff --git a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
index 5d5d2563b..bbd432bc8 100644
--- a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
+++ b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
@@ -105,6 +105,7 @@ public class DataPushQueueSuitJ {
numPartitions,
conf,
client,
+ null,
integer -> {},
mapStatusLengths);