This is an automated email from the ASF dual-hosted git repository.
angerszhuuuu 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 7adf1fca4 [CELEBORN-295] Optimize data push (#1232)
7adf1fca4 is described below
commit 7adf1fca41cd34eb3e9ddf5cc8cd24674ab585b4
Author: Keyong Zhou <[email protected]>
AuthorDate: Tue Feb 28 10:35:55 2023 +0800
[CELEBORN-295] Optimize data push (#1232)
* [CELEBORN-295] Add double buffer for sort pusher
---
.../spark/shuffle/celeborn/SortBasedPusher.java | 193 ++++++++++++++-------
.../spark/shuffle/celeborn/RssShuffleManager.java | 25 ++-
.../shuffle/celeborn/SortBasedShuffleWriter.java | 135 ++++++++++----
.../shuffle/celeborn/RssShuffleWriterSuiteJ.java | 5 +-
.../spark/shuffle/celeborn/RssShuffleManager.java | 32 +++-
.../shuffle/celeborn/SortBasedShuffleWriter.java | 136 ++++++++++++---
.../shuffle/celeborn/RssShuffleWriterSuiteJ.java | 2 +-
.../apache/celeborn/client/write/DataPusher.java | 6 +-
.../celeborn/client/write/DataPushQueueSuitJ.java | 2 +-
.../network/buffer/FileSegmentManagedBuffer.java | 2 +-
.../network/client/TransportClientFactory.java | 11 +-
.../common/network/server/TransportServer.java | 1 +
.../common/{network => }/util/JavaUtils.java | 33 +++-
.../org/apache/celeborn/common/CelebornConf.scala | 34 ++++
.../common/internal/config/ConfigBuilder.scala | 4 +-
.../org/apache/celeborn/common/util/Utils.scala | 7 +-
.../common/network/RpcIntegrationSuiteJ.java | 2 +-
.../celeborn/common/network/StreamTestHelper.java | 2 +-
.../network/TransportClientFactorySuiteJ.java | 2 +-
.../celeborn/common/util/JavaUtilsSuite.java | 34 ++++
docs/configuration/client.md | 2 +
docs/configuration/worker.md | 1 +
.../apache/celeborn/tests/spark/HugeDataTest.scala | 2 +-
.../celeborn/tests/spark/PushDataTimeoutTest.scala | 6 +-
.../tests/spark/RetryCommitFilesTest.scala | 2 +-
.../celeborn/tests/spark/RetryReviveTest.scala | 2 +-
.../tests/spark/RssHashCheckDiskSuite.scala | 2 +-
.../apache/celeborn/tests/spark/RssHashSuite.scala | 2 +-
...rtSuite.scala => RssNonPipelineSortSuite.scala} | 8 +-
...geDataTest.scala => RssPipelineSortSuite.scala} | 10 +-
.../tests/spark/ShuffleFallbackSuite.scala | 2 +-
.../celeborn/tests/spark/SkewJoinSuite.scala | 2 +-
.../celeborn/tests/spark/SparkTestBase.scala | 2 +-
.../service/deploy/worker/PushDataHandler.scala | 16 +-
.../celeborn/service/deploy/worker/Worker.scala | 1 -
.../deploy/worker/storage/FileWriterSuiteJ.java | 2 +-
.../storage/MapPartitionFileWriterSuiteJ.java | 2 +-
37 files changed, 559 insertions(+), 173 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 82dd6c9ab..089dbd6f9 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
@@ -19,6 +19,7 @@ package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
import java.util.LinkedList;
+import java.util.concurrent.ExecutorService;
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.common.CelebornConf;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
public class SortBasedPusher extends MemoryConsumer {
@@ -52,6 +54,7 @@ public class SortBasedPusher extends MemoryConsumer {
private final int pushBufferMaxSize;
private final long pushSortMemoryThreshold;
final int uaoSize = UnsafeAlignedOffset.getUaoSize();
+ static final long bytes8K = Utils.byteStringAsBytes("8k");
String appId;
int shuffleId;
@@ -63,6 +66,12 @@ public class SortBasedPusher extends MemoryConsumer {
CelebornConf conf;
Consumer<Integer> afterPush;
LongAdder[] mapStatusLengths;
+ // this lock is shared between different SortBasedPushers to synchronize
pushData
+ Object sharedPushLock;
+ volatile boolean asyncPushing = false;
+ int[] shuffledPartitions = null;
+ int[] inversedShuffledPartitions = null;
+ ExecutorService executorService;
public SortBasedPusher(
TaskMemoryManager memoryManager,
@@ -76,7 +85,10 @@ public class SortBasedPusher extends MemoryConsumer {
int numPartitions,
CelebornConf conf,
Consumer<Integer> afterPush,
- LongAdder[] mapStatusLengths)
+ LongAdder[] mapStatusLengths,
+ long pushSortMemoryThreshold,
+ Object sharedPushLock,
+ ExecutorService executorService)
throws IOException {
super(
memoryManager,
@@ -92,6 +104,13 @@ public class SortBasedPusher extends MemoryConsumer {
this.taskAttemptId = taskAttemptId;
this.numMappers = numMappers;
this.numPartitions = numPartitions;
+
+ if (conf.pushSortRandomizePartitionIdEnabled()) {
+ shuffledPartitions = new int[numPartitions];
+ inversedShuffledPartitions = new int[numPartitions];
+ JavaUtils.shuffleArray(shuffledPartitions, inversedShuffledPartitions);
+ }
+
this.conf = conf;
this.afterPush = afterPush;
this.mapStatusLengths = mapStatusLengths;
@@ -111,88 +130,95 @@ public class SortBasedPusher extends MemoryConsumer {
mapStatusLengths);
pushBufferMaxSize = conf.pushBufferMaxSize();
- pushSortMemoryThreshold = conf.pushSortMemoryThreshold();
+ this.pushSortMemoryThreshold = pushSortMemoryThreshold;
- inMemSorter = new ShuffleInMemorySorter(this, 4 * 1024 * 1024);
+ int initialSize = Math.min((int) pushSortMemoryThreshold / 8, 1024 * 1024);
+ inMemSorter = new ShuffleInMemorySorter(this, initialSize);
+ this.sharedPushLock = sharedPushLock;
+ this.executorService = executorService;
}
- /**
- * @return bytes of memory freed
- * @throws IOException
- */
public long pushData() throws IOException {
- final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords =
- inMemSorter.getSortedIterator();
-
- byte[] dataBuf = new byte[pushBufferMaxSize];
- int offSet = 0;
- int currentPartition = -1;
- while (sortedRecords.hasNext()) {
- sortedRecords.loadNext();
- final int partition = sortedRecords.packedRecordPointer.getPartitionId();
- assert (partition >= currentPartition);
- if (partition != currentPartition) {
- if (currentPartition == -1) {
- currentPartition = partition;
- } else {
- int bytesWritten =
- rssShuffleClient.mergeData(
- appId,
- shuffleId,
- mapId,
- attemptNumber,
- currentPartition,
- dataBuf,
- 0,
- offSet,
- numMappers,
- numPartitions);
- mapStatusLengths[currentPartition].add(bytesWritten);
- afterPush.accept(bytesWritten);
- currentPartition = partition;
+ // pushData should be synchronized between pushers
+ synchronized (sharedPushLock) {
+ final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords =
+ inMemSorter.getSortedIterator();
+
+ byte[] dataBuf = new byte[pushBufferMaxSize];
+ int offSet = 0;
+ int currentPartition = -1;
+ while (sortedRecords.hasNext()) {
+ sortedRecords.loadNext();
+ final int partition =
+ shuffledPartitions != null
+ ?
inversedShuffledPartitions[sortedRecords.packedRecordPointer.getPartitionId()]
+ : sortedRecords.packedRecordPointer.getPartitionId();
+ if (partition != currentPartition) {
+ if (currentPartition == -1) {
+ currentPartition = partition;
+ } else {
+ int bytesWritten =
+ rssShuffleClient.mergeData(
+ appId,
+ shuffleId,
+ mapId,
+ attemptNumber,
+ currentPartition,
+ dataBuf,
+ 0,
+ offSet,
+ numMappers,
+ numPartitions);
+ mapStatusLengths[currentPartition].add(bytesWritten);
+ afterPush.accept(bytesWritten);
+ currentPartition = partition;
+ offSet = 0;
+ }
+ }
+ final long recordPointer =
sortedRecords.packedRecordPointer.getRecordPointer();
+ final Object recordPage = taskMemoryManager.getPage(recordPointer);
+ final long recordOffsetInPage =
taskMemoryManager.getOffsetInPage(recordPointer);
+ int recordSize = UnsafeAlignedOffset.getSize(recordPage,
recordOffsetInPage);
+
+ if (offSet + recordSize > dataBuf.length) {
+ dataPusher.addTask(partition, dataBuf, offSet);
offSet = 0;
}
+
+ long recordReadPosition = recordOffsetInPage + uaoSize;
+ Platform.copyMemory(
+ recordPage,
+ recordReadPosition,
+ dataBuf,
+ Platform.BYTE_ARRAY_OFFSET + offSet,
+ recordSize);
+ offSet += recordSize;
}
- final long recordPointer =
sortedRecords.packedRecordPointer.getRecordPointer();
- final Object recordPage = taskMemoryManager.getPage(recordPointer);
- final long recordOffsetInPage =
taskMemoryManager.getOffsetInPage(recordPointer);
- int recordSize = UnsafeAlignedOffset.getSize(recordPage,
recordOffsetInPage);
-
- if (offSet + recordSize > dataBuf.length) {
- dataPusher.addTask(partition, dataBuf, offSet);
- offSet = 0;
+ if (offSet > 0) {
+ dataPusher.addTask(currentPartition, dataBuf, offSet);
}
- long recordReadPosition = recordOffsetInPage + uaoSize;
- Platform.copyMemory(
- recordPage, recordReadPosition, dataBuf, Platform.BYTE_ARRAY_OFFSET
+ offSet, recordSize);
- offSet += recordSize;
- }
- if (offSet > 0) {
- dataPusher.addTask(currentPartition, dataBuf, offSet);
- }
+ long freedBytes = freeMemory();
+ inMemSorter.freeMemory();
- long freedBytes = freeMemory();
- inMemSorter.freeMemory();
- return freedBytes;
+ return freedBytes;
+ }
}
- public void insertRecord(
+ public boolean insertRecord(
Object recordBase, long recordOffset, int recordSize, int partitionId,
boolean copySize)
throws IOException {
if (getUsed() > pushSortMemoryThreshold
- && pageCursor + Utils.byteStringAsBytes("8k")
- > currentPage.getBaseOffset() + currentPage.size()) {
+ && pageCursor + bytes8K > currentPage.getBaseOffset() +
currentPage.size()) {
logger.info(
- "Memory Used across threshold, trigger push. Memory: "
+ "Memory Used across threshold, need to trigger push. Memory: "
+ getUsed()
+ ", currentPage size: "
+ currentPage.size());
- pushData();
+ return false;
}
- final int uaoSize = UnsafeAlignedOffset.getUaoSize();
int required;
// Need 4 or 8 bytes to store the record recordSize.
if (copySize) {
@@ -218,7 +244,44 @@ public class SortBasedPusher extends MemoryConsumer {
Platform.copyMemory(recordBase, recordOffset, base, pageCursor,
recordSize);
pageCursor += recordSize;
}
- inMemSorter.insertRecord(recordAddress, partitionId);
+ if (shuffledPartitions != null) {
+ inMemSorter.insertRecord(recordAddress, shuffledPartitions[partitionId]);
+ } else {
+ inMemSorter.insertRecord(recordAddress, partitionId);
+ }
+
+ return true;
+ }
+
+ public void triggerPush() throws IOException {
+ asyncPushing = true;
+ dataPusher.checkException();
+ executorService.submit(
+ () -> {
+ try {
+ pushData();
+ asyncPushing = false;
+ } catch (IOException ie) {
+ dataPusher.setException(ie);
+ }
+ });
+ }
+
+ /**
+ * Since this method and pushData() are synchronized When this method
returns, it means pushData
+ * has released lock
+ *
+ * @throws IOException
+ */
+ public void waitPushFinish() throws IOException {
+ dataPusher.checkException();
+ while (asyncPushing) {
+ try {
+ Thread.sleep(50);
+ } catch (InterruptedException e) {
+ throw new IOException("Interrupted when waitPushFinish", e);
+ }
+ }
}
private void growPointerArrayIfNecessary() throws IOException {
@@ -316,11 +379,7 @@ public class SortBasedPusher extends MemoryConsumer {
@Override
public long spill(long l, MemoryConsumer memoryConsumer) throws IOException {
- logger.info("Pushdata in spill, memory used " + getUsed());
- if (getUsed() != 0) {
- logger.info("Pushdata is not empty , do push.");
- return pushData();
- }
+ logger.warn("SortBasedPusher not support spill yet");
return 0;
}
diff --git
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
index 00ed42e24..3c610d881 100644
---
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
+++
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
@@ -19,6 +19,8 @@ package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
import scala.Int;
@@ -33,6 +35,7 @@ import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.ShuffleMode;
+import org.apache.celeborn.common.util.ThreadUtils;
public class RssShuffleManager implements ShuffleManager {
@@ -53,11 +56,23 @@ public class RssShuffleManager implements ShuffleManager {
ConcurrentHashMap.newKeySet();
private final RssShuffleFallbackPolicyRunner fallbackPolicyRunner;
+ private final ExecutorService[] asyncPushers;
+ private AtomicInteger pusherIdx = new AtomicInteger(0);
+
public RssShuffleManager(SparkConf conf) {
this.conf = conf;
this.celebornConf = SparkUtils.fromSparkConf(conf);
this.cores = conf.getInt(SparkLauncher.EXECUTOR_CORES, 1);
this.fallbackPolicyRunner = new
RssShuffleFallbackPolicyRunner(celebornConf);
+ if (ShuffleMode.SORT.equals(celebornConf.shuffleWriterMode())
+ && celebornConf.pushSortPipelineEnabled()) {
+ asyncPushers = new ExecutorService[cores];
+ for (int i = 0; i < asyncPushers.length; i++) {
+ asyncPushers[i] =
ThreadUtils.newDaemonSingleThreadExecutor("async-pusher-" + i);
+ }
+ } else {
+ asyncPushers = null;
+ }
}
private boolean isDriver() {
@@ -162,8 +177,10 @@ public class RssShuffleManager implements ShuffleManager {
ShuffleClient.get(
h.rssMetaServiceHost(), h.rssMetaServicePort(), celebornConf,
h.userIdentifier());
if (ShuffleMode.SORT.equals(celebornConf.shuffleWriterMode())) {
+ ExecutorService pushThread =
+ celebornConf.pushSortPipelineEnabled() ? getPusherThread() :
null;
return new SortBasedShuffleWriter<>(
- h.dependency(), h.newAppId(), h.numMaps(), context,
celebornConf, client);
+ h.dependency(), h.newAppId(), h.numMaps(), context,
celebornConf, client, pushThread);
} else if (ShuffleMode.HASH.equals(celebornConf.shuffleWriterMode())) {
return new HashBasedShuffleWriter<>(
h, mapId, context, celebornConf, client,
SendBufferPool.get(cores));
@@ -190,4 +207,10 @@ public class RssShuffleManager implements ShuffleManager {
}
return _sortShuffleManager.getReader(handle, startPartition, endPartition,
context);
}
+
+ private ExecutorService getPusherThread() {
+ ExecutorService pusherThread = asyncPushers[pusherIdx.get() %
asyncPushers.length];
+ pusherIdx.incrementAndGet();
+ return pusherThread;
+ }
}
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 5818675a0..7a347c764 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
@@ -18,6 +18,7 @@
package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
+import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.LongAdder;
import scala.Option;
@@ -68,8 +69,10 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private final int numPartitions;
private final long pushBufferMaxSize;
- private SortBasedPusher sortBasedPusher;
-
+ private final Object globalPushLock = new Object();
+ private final boolean pipelined;
+ private SortBasedPusher[] pushers = new SortBasedPusher[2];
+ private SortBasedPusher currentPusher;
private long peakMemoryUsedBytes = 0;
private final OpenByteArrayOutputStream serBuffer;
@@ -94,7 +97,8 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
int numMappers,
TaskContext taskContext,
CelebornConf conf,
- ShuffleClient client)
+ ShuffleClient client,
+ ExecutorService executorService)
throws IOException {
this.mapId = taskContext.partitionId();
this.dep = dep;
@@ -119,21 +123,48 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
tmpRecords = new long[numPartitions];
pushBufferMaxSize = conf.pushBufferMaxSize();
-
- sortBasedPusher =
- new SortBasedPusher(
- taskContext.taskMemoryManager(),
- rssShuffleClient,
- appId,
- shuffleId,
- mapId,
- taskContext.attemptNumber(),
- taskContext.taskAttemptId(),
- numMappers,
- numPartitions,
- conf,
- writeMetrics::incBytesWritten,
- mapStatusLengths);
+ pipelined = conf.pushSortPipelineEnabled();
+
+ if (pipelined) {
+ for (int i = 0; i < pushers.length; i++) {
+ pushers[i] =
+ new SortBasedPusher(
+ taskContext.taskMemoryManager(),
+ rssShuffleClient,
+ appId,
+ shuffleId,
+ mapId,
+ taskContext.attemptNumber(),
+ taskContext.taskAttemptId(),
+ numMappers,
+ numPartitions,
+ conf,
+ writeMetrics::incBytesWritten,
+ mapStatusLengths,
+ conf.pushSortMemoryThreshold() / 2,
+ globalPushLock,
+ executorService);
+ }
+ currentPusher = pushers[0];
+ } else {
+ currentPusher =
+ new SortBasedPusher(
+ taskContext.taskMemoryManager(),
+ rssShuffleClient,
+ appId,
+ shuffleId,
+ mapId,
+ taskContext.attemptNumber(),
+ taskContext.taskAttemptId(),
+ numMappers,
+ numPartitions,
+ conf,
+ writeMetrics::incBytesWritten,
+ mapStatusLengths,
+ conf.pushSortMemoryThreshold(),
+ globalPushLock,
+ null);
+ }
}
@Override
@@ -187,14 +218,34 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
pushGiantRecord(partitionId, giantBuffer, serializedRecordSize);
} else {
long insertStartTime = System.nanoTime();
- sortBasedPusher.insertRecord(
- row.getBaseObject(), row.getBaseOffset(), rowSize, partitionId,
true);
+ boolean success =
+ currentPusher.insertRecord(
+ row.getBaseObject(), row.getBaseOffset(), rowSize,
partitionId, true);
+ if (!success) {
+ pushAndSwitch();
+ success =
+ currentPusher.insertRecord(
+ row.getBaseObject(), row.getBaseOffset(), rowSize,
partitionId, true);
+ if (!success) {
+ throw new IOException("Unable to push after switching pusher!");
+ }
+ }
writeMetrics.incWriteTime(System.nanoTime() - insertStartTime);
}
tmpRecords[partitionId] += 1;
}
}
+ private void pushAndSwitch() throws IOException {
+ if (pipelined) {
+ currentPusher.triggerPush();
+ currentPusher = (currentPusher == pushers[0] ? pushers[1] : pushers[0]);
+ currentPusher.waitPushFinish();
+ } else {
+ currentPusher.pushData();
+ }
+ }
+
private void write0(scala.collection.Iterator iterator) throws IOException {
final scala.collection.Iterator<Product2<K, ?>> records = iterator;
@@ -214,12 +265,26 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
pushGiantRecord(partitionId, serBuffer.getBuf(), serializedRecordSize);
} else {
long insertStartTime = System.nanoTime();
- sortBasedPusher.insertRecord(
- serBuffer.getBuf(),
- Platform.BYTE_ARRAY_OFFSET,
- serializedRecordSize,
- partitionId,
- false);
+ boolean success =
+ currentPusher.insertRecord(
+ serBuffer.getBuf(),
+ Platform.BYTE_ARRAY_OFFSET,
+ serializedRecordSize,
+ partitionId,
+ false);
+ if (!success) {
+ pushAndSwitch();
+ success =
+ currentPusher.insertRecord(
+ serBuffer.getBuf(),
+ Platform.BYTE_ARRAY_OFFSET,
+ serializedRecordSize,
+ partitionId,
+ false);
+ if (!success) {
+ throw new IOException("Unable to push after switching pusher!");
+ }
+ }
writeMetrics.incWriteTime(System.nanoTime() - insertStartTime);
}
tmpRecords[partitionId] += 1;
@@ -247,10 +312,22 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
private void close() throws IOException {
- logger.info("Pushdata in close, memory used " + sortBasedPusher.getUsed());
+ if (pipelined) {
+ logger.info("Memory used {}", (pushers[0].getUsed() +
pushers[1].getUsed()));
+ } else {
+ logger.info("Memory used {}", currentPusher.getUsed());
+ }
long pushStartTime = System.nanoTime();
- sortBasedPusher.pushData();
- sortBasedPusher.close();
+ if (pipelined) {
+ for (int i = 0; i < pushers.length; i++) {
+ pushers[i].waitPushFinish();
+ pushers[i].pushData();
+ pushers[i].close();
+ }
+ } else {
+ currentPusher.pushData();
+ currentPusher.close();
+ }
writeMetrics.incWriteTime(System.nanoTime() - pushStartTime);
rssShuffleClient.pushMergedData(appId, shuffleId, mapId,
taskContext.attemptNumber());
diff --git
a/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
b/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
index 6e7df3ce5..2bf6c3fb6 100644
---
a/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
+++
b/client-spark/spark-2/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
@@ -76,7 +76,7 @@ import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
-import org.apache.celeborn.common.network.util.JavaUtils;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
public class RssShuffleWriterSuiteJ {
@@ -304,7 +304,8 @@ public class RssShuffleWriterSuiteJ {
}
} else {
final SortBasedShuffleWriter<Integer, String, String> writer =
- new SortBasedShuffleWriter<>(dependency, appId, numPartitions,
taskContext, conf, client);
+ new SortBasedShuffleWriter<>(
+ dependency, appId, numPartitions, taskContext, conf, client,
null);
AtomicInteger total = new AtomicInteger(0);
Iterator iterator = getIterator(approximateSize, total, useUnsafe,
false);
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
index d3d2b0bc5..300ac6267 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/RssShuffleManager.java
@@ -19,6 +19,8 @@ package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
import org.apache.spark.*;
import org.apache.spark.launcher.SparkLauncher;
@@ -31,6 +33,7 @@ import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.ShuffleMode;
+import org.apache.celeborn.common.util.ThreadUtils;
public class RssShuffleManager implements ShuffleManager {
@@ -51,11 +54,23 @@ public class RssShuffleManager implements ShuffleManager {
ConcurrentHashMap.newKeySet();
private final RssShuffleFallbackPolicyRunner fallbackPolicyRunner;
+ private final ExecutorService[] asyncPushers;
+ private AtomicInteger pusherIdx = new AtomicInteger(0);
+
public RssShuffleManager(SparkConf conf) {
this.conf = conf;
this.celebornConf = SparkUtils.fromSparkConf(conf);
this.cores = conf.getInt(SparkLauncher.EXECUTOR_CORES, 1);
this.fallbackPolicyRunner = new
RssShuffleFallbackPolicyRunner(celebornConf);
+ if (ShuffleMode.SORT.equals(celebornConf.shuffleWriterMode())
+ && celebornConf.pushSortPipelineEnabled()) {
+ asyncPushers = new ExecutorService[cores];
+ for (int i = 0; i < asyncPushers.length; i++) {
+ asyncPushers[i] =
ThreadUtils.newDaemonSingleThreadExecutor("async-pusher-" + i);
+ }
+ } else {
+ asyncPushers = null;
+ }
}
private boolean isDriver() {
@@ -160,8 +175,17 @@ public class RssShuffleManager implements ShuffleManager {
ShuffleClient.get(
h.rssMetaServiceHost(), h.rssMetaServicePort(), celebornConf,
h.userIdentifier());
if (ShuffleMode.SORT.equals(celebornConf.shuffleWriterMode())) {
+ ExecutorService pushThread =
+ celebornConf.pushSortPipelineEnabled() ? getPusherThread() :
null;
return new SortBasedShuffleWriter<>(
- h.dependency(), h.newAppId(), h.numMappers(), context,
celebornConf, client, metrics);
+ h.dependency(),
+ h.newAppId(),
+ h.numMappers(),
+ context,
+ celebornConf,
+ client,
+ metrics,
+ pushThread);
} else if (ShuffleMode.HASH.equals(celebornConf.shuffleWriterMode())) {
return new HashBasedShuffleWriter<>(
h, context, celebornConf, client, metrics,
SendBufferPool.get(cores));
@@ -267,4 +291,10 @@ public class RssShuffleManager implements ShuffleManager {
context,
metrics);
}
+
+ private ExecutorService getPusherThread() {
+ ExecutorService pusherThread = asyncPushers[pusherIdx.get() %
asyncPushers.length];
+ pusherIdx.incrementAndGet();
+ return pusherThread;
+ }
}
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 ce9c1f85d..4a04e41c0 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
@@ -18,6 +18,7 @@
package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
+import java.util.concurrent.ExecutorService;
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.common.CelebornConf;
+import org.apache.celeborn.common.exception.CelebornIOException;
@Private
public class SortBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
@@ -70,7 +72,11 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private final int numPartitions;
private final long pushBufferMaxSize;
- private SortBasedPusher sortBasedPusher;
+ // this lock is shared between different SortBasedPushers to synchronize
pushData
+ private final Object sharedPushLock = new Object();
+ private final boolean pipelined;
+ private SortBasedPusher[] pushers = new SortBasedPusher[2];
+ private SortBasedPusher currentPusher;
@Nullable private long peakMemoryUsedBytes = 0;
@@ -97,7 +103,8 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
TaskContext taskContext,
CelebornConf conf,
ShuffleClient client,
- ShuffleWriteMetricsReporter metrics)
+ ShuffleWriteMetricsReporter metrics,
+ ExecutorService executorService)
throws IOException {
this.mapId = taskContext.partitionId();
this.dep = dep;
@@ -122,21 +129,48 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
tmpRecords = new long[numPartitions];
pushBufferMaxSize = conf.pushBufferMaxSize();
-
- sortBasedPusher =
- new SortBasedPusher(
- taskContext.taskMemoryManager(),
- rssShuffleClient,
- appId,
- shuffleId,
- mapId,
- taskContext.attemptNumber(),
- taskContext.taskAttemptId(),
- numMappers,
- numPartitions,
- conf,
- writeMetrics::incBytesWritten,
- mapStatusLengths);
+ pipelined = conf.pushSortPipelineEnabled();
+
+ if (pipelined) {
+ for (int i = 0; i < pushers.length; i++) {
+ pushers[i] =
+ new SortBasedPusher(
+ taskContext.taskMemoryManager(),
+ rssShuffleClient,
+ appId,
+ shuffleId,
+ mapId,
+ taskContext.attemptNumber(),
+ taskContext.taskAttemptId(),
+ numMappers,
+ numPartitions,
+ conf,
+ writeMetrics::incBytesWritten,
+ mapStatusLengths,
+ conf.pushSortMemoryThreshold() / 2,
+ sharedPushLock,
+ executorService);
+ }
+ currentPusher = pushers[0];
+ } else {
+ currentPusher =
+ new SortBasedPusher(
+ taskContext.taskMemoryManager(),
+ rssShuffleClient,
+ appId,
+ shuffleId,
+ mapId,
+ taskContext.attemptNumber(),
+ taskContext.taskAttemptId(),
+ numMappers,
+ numPartitions,
+ conf,
+ writeMetrics::incBytesWritten,
+ mapStatusLengths,
+ conf.pushSortMemoryThreshold(),
+ sharedPushLock,
+ null);
+ }
}
@Override
@@ -189,14 +223,34 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
pushGiantRecord(partitionId, giantBuffer, serializedRecordSize);
} else {
long insertStartTime = System.nanoTime();
- sortBasedPusher.insertRecord(
- row.getBaseObject(), row.getBaseOffset(), rowSize, partitionId,
true);
+ boolean success =
+ currentPusher.insertRecord(
+ row.getBaseObject(), row.getBaseOffset(), rowSize,
partitionId, true);
+ if (!success) {
+ pushAndSwitch();
+ success =
+ currentPusher.insertRecord(
+ row.getBaseObject(), row.getBaseOffset(), rowSize,
partitionId, true);
+ if (!success) {
+ throw new CelebornIOException("Unable to push after switching
pusher!");
+ }
+ }
writeMetrics.incWriteTime(System.nanoTime() - insertStartTime);
}
tmpRecords[partitionId] += 1;
}
}
+ private void pushAndSwitch() throws IOException {
+ if (pipelined) {
+ currentPusher.triggerPush();
+ currentPusher = (currentPusher == pushers[0] ? pushers[1] : pushers[0]);
+ currentPusher.waitPushFinish();
+ } else {
+ currentPusher.pushData();
+ }
+ }
+
private void write0(scala.collection.Iterator iterator) throws IOException {
final scala.collection.Iterator<Product2<K, ?>> records = iterator;
@@ -216,12 +270,26 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
pushGiantRecord(partitionId, serBuffer.getBuf(), serializedRecordSize);
} else {
long insertStartTime = System.nanoTime();
- sortBasedPusher.insertRecord(
- serBuffer.getBuf(),
- Platform.BYTE_ARRAY_OFFSET,
- serializedRecordSize,
- partitionId,
- false);
+ boolean success =
+ currentPusher.insertRecord(
+ serBuffer.getBuf(),
+ Platform.BYTE_ARRAY_OFFSET,
+ serializedRecordSize,
+ partitionId,
+ false);
+ if (!success) {
+ pushAndSwitch();
+ success =
+ currentPusher.insertRecord(
+ serBuffer.getBuf(),
+ Platform.BYTE_ARRAY_OFFSET,
+ serializedRecordSize,
+ partitionId,
+ false);
+ if (!success) {
+ throw new IOException("Unable to push after switching pusher!");
+ }
+ }
writeMetrics.incWriteTime(System.nanoTime() - insertStartTime);
}
tmpRecords[partitionId] += 1;
@@ -249,10 +317,22 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
}
private void close() throws IOException {
- logger.info("Pushdata in close, memory used " + sortBasedPusher.getUsed());
+ if (pipelined) {
+ logger.info("Memory used {}", (pushers[0].getUsed() +
pushers[1].getUsed()));
+ } else {
+ logger.info("Memory used {}", currentPusher.getUsed());
+ }
long pushStartTime = System.nanoTime();
- sortBasedPusher.pushData();
- sortBasedPusher.close();
+ if (pipelined) {
+ for (int i = 0; i < pushers.length; i++) {
+ pushers[i].waitPushFinish();
+ pushers[i].pushData();
+ pushers[i].close();
+ }
+ } else {
+ currentPusher.pushData();
+ currentPusher.close();
+ }
writeMetrics.incWriteTime(System.nanoTime() - pushStartTime);
long pushMergedDataTime = System.nanoTime();
diff --git
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
index e5408bcb3..f147f0ec2 100644
---
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
+++
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/RssShuffleWriterSuiteJ.java
@@ -74,7 +74,7 @@ import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
-import org.apache.celeborn.common.network.util.JavaUtils;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
public class RssShuffleWriterSuiteJ {
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 a75b84121..48b0a2163 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
@@ -180,7 +180,11 @@ public class DataPusher {
checkException();
}
- private void checkException() throws IOException {
+ public void setException(IOException ie) {
+ exceptionRef.compareAndSet(null, ie);
+ }
+
+ public void checkException() throws IOException {
if (exceptionRef.get() != null) {
throw exceptionRef.get();
}
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 3a8044081..9311a3440 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
@@ -36,8 +36,8 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.PushState;
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/buffer/FileSegmentManagedBuffer.java
b/common/src/main/java/org/apache/celeborn/common/network/buffer/FileSegmentManagedBuffer.java
index debf8ce16..dbd9d71ad 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/buffer/FileSegmentManagedBuffer.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/buffer/FileSegmentManagedBuffer.java
@@ -26,9 +26,9 @@ import com.google.common.base.Objects;
import com.google.common.io.ByteStreams;
import io.netty.channel.DefaultFileRegion;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.network.util.LimitedInputStream;
import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.util.JavaUtils;
/** A {@link ManagedBuffer} backed by a segment in a file. */
public final class FileSegmentManagedBuffer extends ManagedBuffer {
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
index de8e61de8..0043d729c 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
@@ -36,6 +36,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.network.TransportContext;
import org.apache.celeborn.common.network.server.TransportChannelHandler;
import org.apache.celeborn.common.network.util.*;
+import org.apache.celeborn.common.util.JavaUtils;
/**
* Factory for creating {@link TransportClient}s by using createClient.
@@ -97,12 +98,10 @@ public class TransportClientFactory implements Closeable {
/**
* Create a {@link TransportClient} connecting to the given remote host /
port.
*
- * <p>We maintains an array of clients (size determined by
spark.shuffle.io.numConnectionsPerPeer)
- * and randomly picks one to use. If no client was previously created in the
randomly selected
- * spot, this function creates a new client and places it there.
- *
- * <p>Prior to the creation of a new TransportClient, we will execute all
{@link
- * TransportClientBootstrap}s that are registered with this factory.
+ * <p>We maintains an array of clients (size determined by
+ * celeborn.$module.io.numConnectionsPerPeer) and randomly picks one to use.
If no client was
+ * previously created in the randomly selected spot, this function creates a
new client and places
+ * it there.
*
* <p>This blocks until a connection is successfully established and fully
bootstrapped.
*
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
b/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
index 278b12079..b26566896 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
@@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.network.TransportContext;
import org.apache.celeborn.common.network.util.*;
+import org.apache.celeborn.common.util.JavaUtils;
/** Server for the efficient, low-level streaming service. */
public class TransportServer implements Closeable {
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/util/JavaUtils.java
b/common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java
similarity index 93%
rename from
common/src/main/java/org/apache/celeborn/common/network/util/JavaUtils.java
rename to common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java
index 9606d4b71..ef284e3a6 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/util/JavaUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java
@@ -15,13 +15,14 @@
* limitations under the License.
*/
-package org.apache.celeborn.common.network.util;
+package org.apache.celeborn.common.util;
import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.nio.charset.StandardCharsets;
import java.util.Locale;
+import java.util.Random;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@@ -33,6 +34,8 @@ import org.apache.commons.lang3.SystemUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.common.network.util.ByteUnit;
+
/**
* General utilities available in the network package. Many of these are
sourced from Spark's own
* Utils, just accessible within this package.
@@ -387,4 +390,32 @@ public class JavaUtils {
}
return sb.toString();
}
+
+ /**
+ * Shuffle the elements of an array into a random order, modifying the
original array. Returns the
+ * original array.
+ */
+ public static void randomizeInPlace(int[] arr) {
+ Random rand = new Random();
+ for (int i = arr.length - 1; i >= 1; i--) {
+ int j = rand.nextInt(i + 1);
+ int tmp = arr[j];
+ arr[j] = arr[i];
+ arr[i] = tmp;
+ }
+ }
+
+ public static void shuffleArray(int[] arr, int[] inversed) {
+ if (arr == null || inversed == null || arr.length != inversed.length) {
+ return;
+ } else {
+ for (int i = 0; i < arr.length; i++) {
+ arr[i] = i;
+ }
+ JavaUtils.randomizeInPlace(arr);
+ for (int i = 0; i < arr.length; i++) {
+ inversed[arr[i]] = i;
+ }
+ }
+ }
}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index a3febb8b4..473ccf13f 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -497,6 +497,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def workerWorkingDir: String = get(WORKER_WORKING_DIR)
def workerCloseIdleConnections: Boolean = get(WORKER_CLOSE_IDLE_CONNECTIONS)
def workerReplicateFastFailDuration: Long =
get(WORKER_REPLICATE_FAST_FAIL_DURATION)
+ def workerReplicateRandomConnectionEnabled: Boolean =
+ get(WORKER_REPLICATE_RANDOM_CONNECTION_ENABLED)
def workerDeviceStatusCheckTimeout: Long =
get(WORKER_DEVICE_STATUS_CHECK_TIMEOUT)
def workerCheckFileCleanMaxRetries: Int =
get(WORKER_CHECK_FILE_CLEAN_MAX_RETRIES)
def workerCheckFileCleanTimeout: Long = get(WORKER_CHECK_FILE_CLEAN_TIMEOUT)
@@ -659,6 +661,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def pushMaxReqsInFlight: Int = get(PUSH_MAX_REQS_IN_FLIGHT)
def pushMaxReviveTimes: Int = get(PUSH_MAX_REVIVE_TIMES)
def pushSortMemoryThreshold: Long = get(PUSH_SORT_MEMORY_THRESHOLD)
+ def pushSortPipelineEnabled: Boolean = get(PUSH_SORT_PIPELINE_ENABLED)
+ def pushSortRandomizePartitionIdEnabled: Boolean =
get(PUSH_SORT_RANDOMIZE_PARITION_ENABLED)
def pushRetryThreads: Int = get(PUSH_RETRY_THREADS)
def pushStageEndTimeout: Long =
get(PUSH_STAGE_END_TIMEOUT).getOrElse(get(RPC_ASK_TIMEOUT) *
(requestCommitFilesMaxRetries + 1))
@@ -1857,6 +1861,17 @@ object CelebornConf extends Logging {
.timeConf(TimeUnit.MILLISECONDS)
.createWithDefaultString("60s")
+ val WORKER_REPLICATE_RANDOM_CONNECTION_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.worker.replicate.randomConnection.enabled")
+ .categories("worker")
+ .doc("Whether worker will create random connection to peer when
replicate data. When false, worker tend to " +
+ "reuse the same cached TransportClient to a specific replicate worker;
when true, worker tend to use " +
+ "different cached TransportClient. Netty will use the same thread to
serve the same connection, so " +
+ "with more connections replicate server can leverage more netty
threads")
+ .version("0.2.1")
+ .booleanConf
+ .createWithDefault(true)
+
val WORKER_DEVICE_STATUS_CHECK_TIMEOUT: ConfigEntry[Long] =
buildConf("celeborn.worker.disk.check.timeout")
.withAlternative("rss.worker.status.check.timeout")
@@ -2234,6 +2249,25 @@ object CelebornConf extends Logging {
.bytesConf(ByteUnit.BYTE)
.createWithDefaultString("64m")
+ val PUSH_SORT_PIPELINE_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.push.sort.pipeline.enabled")
+ .categories("client")
+ .doc("Whether to enable pipelining for sort based shuffle writer. If
true, double buffering" +
+ " will be used to pipeline push")
+ .version("0.2.1")
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_SORT_RANDOMIZE_PARITION_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.push.sort.randomizePartitionId.enabled")
+ .categories("client")
+ .doc(
+ "Whether to randomize partitionId in push sorter. If true, partitionId
will be randomized " +
+ "when sort data to avoid skew when push to worker")
+ .version("0.2.1")
+ .booleanConf
+ .createWithDefault(false)
+
val PUSH_RETRY_THREADS: ConfigEntry[Int] =
buildConf("celeborn.push.retry.threads")
.withAlternative("rss.pushdata.retry.thread.num")
diff --git
a/common/src/main/scala/org/apache/celeborn/common/internal/config/ConfigBuilder.scala
b/common/src/main/scala/org/apache/celeborn/common/internal/config/ConfigBuilder.scala
index 7e2882ad4..f972643c4 100644
---
a/common/src/main/scala/org/apache/celeborn/common/internal/config/ConfigBuilder.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/internal/config/ConfigBuilder.scala
@@ -23,8 +23,8 @@ import java.util.regex.PatternSyntaxException
import scala.util.matching.Regex
import org.apache.celeborn.common.CelebornConf
-import org.apache.celeborn.common.network.util.{ByteUnit, JavaUtils}
-import org.apache.celeborn.common.util.Utils
+import org.apache.celeborn.common.network.util.ByteUnit
+import org.apache.celeborn.common.util.{JavaUtils, Utils}
private object ConfigHelpers {
diff --git a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
index 951afa8d2..5e697e225 100644
--- a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
@@ -45,8 +45,8 @@ import org.apache.celeborn.common.exception.CelebornException
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.meta.{DiskStatus, FileInfo, WorkerInfo}
import org.apache.celeborn.common.network.protocol.TransportMessage
-import org.apache.celeborn.common.network.util.{JavaUtils, TransportConf}
-import org.apache.celeborn.common.protocol.{PartitionLocation,
PartitionSplitMode, PartitionType}
+import org.apache.celeborn.common.network.util.TransportConf
+import org.apache.celeborn.common.protocol.{PartitionLocation,
PartitionSplitMode, PartitionType, TransportModuleConstants}
import org.apache.celeborn.common.protocol.message.{ControlMessages, Message,
StatusCode}
import
org.apache.celeborn.common.protocol.message.ControlMessages.WorkerResource
@@ -499,6 +499,9 @@ object Utils extends Logging {
val numThreads = defaultNumThreads(numUsableCores)
conf.setIfMissing(s"celeborn.$module.io.serverThreads",
numThreads.toString)
conf.setIfMissing(s"celeborn.$module.io.clientThreads",
numThreads.toString)
+ if (TransportModuleConstants.PUSH_MODULE == module) {
+ conf.setIfMissing(s"celeborn.$module.io.numConnectionsPerPeer",
numThreads.toString)
+ }
// TODO remove after releasing 0.2.0
conf.setIfMissing(s"rss.$module.io.serverThreads", numThreads.toString)
conf.setIfMissing(s"rss.$module.io.clientThreads", numThreads.toString)
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/RpcIntegrationSuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/network/RpcIntegrationSuiteJ.java
index d54fe0cf3..2cabc41d7 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/RpcIntegrationSuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/RpcIntegrationSuiteJ.java
@@ -40,8 +40,8 @@ import
org.apache.celeborn.common.network.client.TransportClientFactory;
import org.apache.celeborn.common.network.protocol.*;
import org.apache.celeborn.common.network.server.BaseMessageHandler;
import org.apache.celeborn.common.network.server.TransportServer;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.util.JavaUtils;
public class RpcIntegrationSuiteJ {
static TransportConf conf;
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/StreamTestHelper.java
b/common/src/test/java/org/apache/celeborn/common/network/StreamTestHelper.java
index ed1281d66..04f85d216 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/StreamTestHelper.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/StreamTestHelper.java
@@ -28,8 +28,8 @@ import com.google.common.io.Files;
import org.apache.celeborn.common.network.buffer.FileSegmentManagedBuffer;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.util.JavaUtils;
class StreamTestHelper {
static final String[] STREAMS = {"largeBuffer", "smallBuffer",
"emptyBuffer", "file"};
diff --git
a/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java
b/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java
index fa8faedfb..b98c60e7c 100644
---
a/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java
+++
b/common/src/test/java/org/apache/celeborn/common/network/TransportClientFactorySuiteJ.java
@@ -32,8 +32,8 @@ import
org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.client.TransportClientFactory;
import org.apache.celeborn.common.network.server.BaseMessageHandler;
import org.apache.celeborn.common.network.server.TransportServer;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.util.JavaUtils;
public class TransportClientFactorySuiteJ {
private TransportContext context;
diff --git
a/common/src/test/java/org/apache/celeborn/common/util/JavaUtilsSuite.java
b/common/src/test/java/org/apache/celeborn/common/util/JavaUtilsSuite.java
new file mode 100644
index 000000000..95d48166b
--- /dev/null
+++ b/common/src/test/java/org/apache/celeborn/common/util/JavaUtilsSuite.java
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.util;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class JavaUtilsSuite {
+
+ @Test
+ public void testNormalPackedPartitionId() {
+ int[] arr = new int[1000];
+ int[] inverse = new int[1000];
+ JavaUtils.shuffleArray(arr, inverse);
+ for (int i = 0; i < arr.length; i++) {
+ Assert.assertTrue(i == inverse[arr[i]]);
+ }
+ }
+}
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index 58ec93a49..34f95b1bb 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -41,6 +41,8 @@ license: |
| celeborn.push.revive.maxRetries | 5 | Max retry times for reviving when
celeborn push data failed. | 0.3.0 |
| celeborn.push.slowStart.initialSleepTime | 500ms | The initial sleep time if
the current max in flight requests is 0 | 0.3.0 |
| celeborn.push.slowStart.maxSleepTime | 2s | If celeborn.push.limit.strategy
is set to SLOWSTART, push side will take a sleep strategy for each batch of
requests, this controls the max sleep time if the max in flight requests limit
is 1 for a long time | 0.3.0 |
+| celeborn.push.sort.pipeline.enabled | false | Whether to enable pipelining
for sort based shuffle writer. If true, double buffering will be used to
pipeline push | 0.2.1 |
+| celeborn.push.sort.randomizePartitionId.enabled | false | Whether to
randomize partitionId in push sorter. If true, partitionId will be randomized
when sort data to avoid skew when push to worker | 0.2.1 |
| celeborn.push.sortMemory.threshold | 64m | When SortBasedPusher use memory
over the threshold, will trigger push data. | 0.2.0 |
| celeborn.push.splitPartition.threads | 8 | Thread number to process shuffle
split request in shuffle client. | 0.2.0 |
| celeborn.push.stageEnd.timeout | <undefined> | Timeout for waiting
StageEnd. Default value should be `celeborn.rpc.askTimeout *
(celeborn.rpc.requestCommitFiles.maxRetries + 1)`. | 0.2.0 |
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index 2fdafe8f1..db9133198 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -87,6 +87,7 @@ license: |
| celeborn.worker.replicate.fastFail.duration | 60s | If a replicate request
not replied during the duration, worker will mark the replicate data request as
failed. | 0.2.0 |
| celeborn.worker.replicate.io.threads | <undefined> | Netty IO thread
number of worker to replicate shuffle data. The default threads number is the
number of flush thread. | 0.2.0 |
| celeborn.worker.replicate.port | 0 | Server port for Worker to receive
replicate data request from other Workers. | 0.2.0 |
+| celeborn.worker.replicate.randomConnection.enabled | true | Whether worker
will create random connection to peer when replicate data. When false, worker
tend to reuse the same cached TransportClient to a specific replicate worker;
when true, worker tend to use different cached TransportClient. Netty will use
the same thread to serve the same connection, so with more connections
replicate server can leverage more netty threads | 0.2.1 |
| celeborn.worker.replicate.threads | 64 | Thread number of worker to
replicate shuffle data. | 0.2.0 |
| celeborn.worker.rpc.port | 0 | Server port for Worker to receive RPC
request. | 0.2.0 |
| celeborn.worker.shuffle.commit.timeout | <value of
celeborn.rpc.askTimeout> | Timeout for a Celeborn worker to commit files of
a shuffle. | 0.2.0 |
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
index 05accffe2..2918a3288 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
@@ -37,7 +37,7 @@ class HugeDataTest extends AnyFunSuite
}
test("celeborn spark integration test - huge data") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]")
val ss = SparkSession.builder().config(updateSparkConf(sparkConf,
false)).getOrCreate()
val value = Range(1, 10000).mkString(",")
val tuples = ss.sparkContext.parallelize(1 to 10000, 2)
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
index 7238f5019..bc18a91ca 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
@@ -33,7 +33,7 @@ class PushDataTimeoutTest extends AnyFunSuite
val workerConf = Map(
"celeborn.test.pushMasterDataTimeout" -> "true",
"celeborn.test.pushSlaveDataTimeout" -> "true",
- "celeborn.push.timeoutCheck.interval" -> "2s")
+ "celeborn.push.timeoutCheck.interval" -> "1s")
setUpMiniCluster(masterConfs = null, workerConfs = workerConf)
}
@@ -47,7 +47,7 @@ class PushDataTimeoutTest extends AnyFunSuite
test("celeborn spark integration test - pushdata timeout") {
Seq("false", "true").foreach { enabled =>
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]")
.set("spark.celeborn.push.data.timeout", "5s")
.set("spark.celeborn.push.timeoutCheck.interval", "2s")
.set("spark.celeborn.push.replicate.enabled", enabled)
@@ -80,7 +80,7 @@ class PushDataTimeoutTest extends AnyFunSuite
}
test("celeborn spark integration test - pushdata timeout will add to
balcklist") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]")
.set("spark.celeborn.push.data.timeout", "5s")
.set("spark.celeborn.push.replicate.enabled", "true")
.set("spark.celeborn.client.blacklistSlave.enabled", "true")
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryCommitFilesTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryCommitFilesTest.scala
index 379e67265..9c4d71b9a 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryCommitFilesTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryCommitFilesTest.scala
@@ -46,7 +46,7 @@ class RetryCommitFilesTest extends AnyFunSuite
test("celeborn spark integration test - retry commit files") {
val sparkConf = new SparkConf()
.set("spark.celeborn.test.retryCommitFiles", "true")
- .setAppName("rss-demo").setMaster("local[4]")
+ .setAppName("rss-demo").setMaster("local[2]")
val ss = SparkSession.builder().config(updateSparkConf(sparkConf,
false)).getOrCreate()
ss.sparkContext.parallelize(1 to 1000, 2)
.map { i => (i, Range(1, 1000).mkString(",")) }.groupByKey(16).collect()
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
index 611bc0fe6..145247d9f 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
@@ -44,7 +44,7 @@ class RetryReviveTest extends AnyFunSuite
test("celeborn spark integration test - retry revive as configured times") {
val sparkConf = new SparkConf()
.set("spark.celeborn.test.retryRevive", "true")
- .setAppName("rss-demo").setMaster("local[4]")
+ .setAppName("rss-demo").setMaster("local[2]")
val ss = SparkSession.builder().config(updateSparkConf(sparkConf,
false)).getOrCreate()
ss.sparkContext.parallelize(1 to 1000, 2)
.map { i => (i, Range(1, 1000).mkString(",")) }.groupByKey(16).collect()
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala
index 24a91fdf3..02d507d4b 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashCheckDiskSuite.scala
@@ -49,7 +49,7 @@ class RssHashCheckDiskSuite extends AnyFunSuite
}
test("celeborn spark integration test - hash-checkDiskFull") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]").set(
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]").set(
"spark.celeborn.shuffle.expired.checkInterval",
"5s")
val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashSuite.scala
index f7fa360a2..d3fa885d3 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssHashSuite.scala
@@ -37,7 +37,7 @@ class RssHashSuite extends AnyFunSuite
}
test("celeborn spark integration test - hash") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]")
val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
val combineResult = combine(sparkSession)
val groupbyResult = groupBy(sparkSession)
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssSortSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssNonPipelineSortSuite.scala
similarity index 89%
rename from
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssSortSuite.scala
rename to
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssNonPipelineSortSuite.scala
index 159f37131..1e24d968c 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssSortSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssNonPipelineSortSuite.scala
@@ -24,7 +24,7 @@ import org.scalatest.funsuite.AnyFunSuite
import org.apache.celeborn.client.ShuffleClient
-class RssSortSuite extends AnyFunSuite
+class RssNonPipelineSortSuite extends AnyFunSuite
with SparkTestBase
with BeforeAndAfterEach {
@@ -36,8 +36,10 @@ class RssSortSuite extends AnyFunSuite
System.gc()
}
- test("celeborn spark integration test - sort") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
+ test("celeborn spark integration test - non pipeline sort") {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]")
+ .set("spark.celeborn.push.sort.pipeline.enabled", "false")
+ .set("spark.celeborn.push.sort.randomizePartitionId.enabled", "false")
val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
val combineResult = combine(sparkSession)
val groupbyResult = groupBy(sparkSession)
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssPipelineSortSuite.scala
similarity index 85%
copy from
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
copy to
tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssPipelineSortSuite.scala
index 05accffe2..59d909271 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/HugeDataTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RssPipelineSortSuite.scala
@@ -24,7 +24,7 @@ import org.scalatest.funsuite.AnyFunSuite
import org.apache.celeborn.client.ShuffleClient
-class HugeDataTest extends AnyFunSuite
+class RssPipelineSortSuite extends AnyFunSuite
with SparkTestBase
with BeforeAndAfterEach {
@@ -36,9 +36,11 @@ class HugeDataTest extends AnyFunSuite
System.gc()
}
- test("celeborn spark integration test - huge data") {
- val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[4]")
- val ss = SparkSession.builder().config(updateSparkConf(sparkConf,
false)).getOrCreate()
+ test("celeborn spark integration test - pipeline sort") {
+ val sparkConf = new
SparkConf().setAppName("rss-demo").setMaster("local[2]")
+ .set("spark.celeborn.push.sort.pipeline.enabled", "true")
+ .set("spark.celeborn.push.sort.randomizePartitionId.enabled", "true")
+ val ss = SparkSession.builder().config(updateSparkConf(sparkConf,
true)).getOrCreate()
val value = Range(1, 10000).mkString(",")
val tuples = ss.sparkContext.parallelize(1 to 10000, 2)
.map { i => (i, value) }.groupByKey(16).collect()
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/ShuffleFallbackSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/ShuffleFallbackSuite.scala
index 92b036fa9..fd3227d57 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/ShuffleFallbackSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/ShuffleFallbackSuite.scala
@@ -46,7 +46,7 @@ class ShuffleFallbackSuite extends AnyFunSuite
test(s"celeborn spark integration test - fallback") {
val sparkConf = new SparkConf().setAppName("rss-demo")
- .setMaster("local[4]")
+ .setMaster("local[2]")
.set("spark.celeborn.shuffle.forceFallback.enabled", "true")
enableRss(sparkConf)
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SkewJoinSuite.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SkewJoinSuite.scala
index d74cdef15..40e9b09f3 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SkewJoinSuite.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SkewJoinSuite.scala
@@ -48,7 +48,7 @@ class SkewJoinSuite extends AnyFunSuite
CompressionCodec.values.foreach { codec =>
test(s"celeborn spark integration test - skew join - $codec") {
val sparkConf = new SparkConf().setAppName("rss-demo")
- .setMaster("local[4]")
+ .setMaster("local[2]")
.set("spark.sql.adaptive.enabled", "true")
.set("spark.sql.adaptive.skewJoin.enabled", "true")
.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes",
"16MB")
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
index 08dd18fa9..2c1cc3a51 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SparkTestBase.scala
@@ -54,7 +54,7 @@ trait SparkTestBase extends AnyFunSuite
sparkConf.set("spark.sql.adaptive.localShuffleReader.enabled", "false")
sparkConf.set("spark.celeborn.master.endpoints",
masterInfo._1.rpcEnv.address.toString)
if (sort) {
- sparkConf.set("spark.celeborn.shuffle.writer.mode", "sort")
+ sparkConf.set("spark.celeborn.shuffle.writer", "sort")
}
sparkConf
}
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
index 6ca9c6afe..f2f37da23 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@@ -315,8 +315,7 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
}
}
try {
- val client =
- pushClientFactory.createClient(peer.getHost,
peer.getReplicatePort, location.getId)
+ val client = getClient(peer.getHost, peer.getReplicatePort,
location.getId)
val newPushData = new PushData(
PartitionLocation.Mode.SLAVE.mode(),
shuffleKey,
@@ -569,10 +568,7 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
}
try {
- val client = pushClientFactory.createClient(
- peer.getHost,
- peer.getReplicatePort,
- location.getId)
+ val client = getClient(peer.getHost, peer.getReplicatePort,
location.getId)
val newPushMergedData = new PushMergedData(
PartitionLocation.Mode.SLAVE.mode(),
shuffleKey,
@@ -1078,6 +1074,14 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
val id = partitionUniqueId.split("-")(0).toInt
(PackedPartitionId.getRawPartitionId(id),
PackedPartitionId.getAttemptId(id))
}
+
+ private def getClient(host: String, port: Int, partitionId: Int):
TransportClient = {
+ if (conf.workerReplicateRandomConnectionEnabled) {
+ pushClientFactory.createClient(host, port)
+ } else {
+ pushClientFactory.createClient(host, port, partitionId)
+ }
+ }
}
object PushDataHandler {
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index 68825aa4b..8d375d2f0 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -17,7 +17,6 @@
package org.apache.celeborn.service.deploy.worker
-import java.io.IOException
import java.lang.{Long => JLong}
import java.util.{HashMap => JHashMap, HashSet => JHashSet}
import java.util.concurrent._
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/FileWriterSuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/FileWriterSuiteJ.java
index 6ef563c22..6f425fa3e 100644
---
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/FileWriterSuiteJ.java
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/FileWriterSuiteJ.java
@@ -69,11 +69,11 @@ import
org.apache.celeborn.common.network.protocol.OpenStream;
import org.apache.celeborn.common.network.protocol.StreamHandle;
import org.apache.celeborn.common.network.server.TransportServer;
import org.apache.celeborn.common.network.server.memory.MemoryManager;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.PartitionSplitMode;
import org.apache.celeborn.common.protocol.PartitionType;
import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.FetchHandler;
diff --git
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriterSuiteJ.java
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriterSuiteJ.java
index c2cd82434..3667da801 100644
---
a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriterSuiteJ.java
+++
b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionFileWriterSuiteJ.java
@@ -44,9 +44,9 @@ import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.network.server.memory.MemoryManager;
-import org.apache.celeborn.common.network.util.JavaUtils;
import org.apache.celeborn.common.protocol.PartitionSplitMode;
import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
