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/01/20 02:18:49 UTC
[incubator-celeborn] branch main updated: [CELEBORN-55][FEATURE] Split maxReqsInFlight limitation into level of target worker (#1102)
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 c5be79ee [CELEBORN-55][FEATURE] Split maxReqsInFlight limitation into level of target worker (#1102)
c5be79ee is described below
commit c5be79ee3d763603fca518eaed55da86c90f1828
Author: zy.jordan <82...@users.noreply.github.com>
AuthorDate: Fri Jan 20 10:18:45 2023 +0800
[CELEBORN-55][FEATURE] Split maxReqsInFlight limitation into level of target worker (#1102)
---
.../spark/shuffle/celeborn/SortBasedPusher.java | 1 +
.../shuffle/celeborn/HashBasedShuffleWriter.java | 1 +
.../shuffle/celeborn/RssShuffleWriterSuiteJ.java | 3 +-
.../shuffle/celeborn/HashBasedShuffleWriter.java | 1 +
.../shuffle/celeborn/RssShuffleWriterSuiteJ.java | 3 +-
.../org/apache/celeborn/client/ShuffleClient.java | 7 +
.../apache/celeborn/client/ShuffleClientImpl.java | 129 ++++++------
.../celeborn/client/write/DataPushQueue.java | 118 +++++++++++
.../apache/celeborn/client/write/DataPusher.java | 96 +++++----
.../client/write/InFlightRequestTracker.java | 220 +++++++++++++++++++++
.../apache/celeborn/client/write/PushState.java | 94 +++------
.../apache/celeborn/client/DummyShuffleClient.java | 45 ++++-
.../celeborn/client/write/DataPushQueueSuitJ.java | 174 ++++++++++++++++
.../org/apache/celeborn/common/CelebornConf.scala | 4 +-
docs/configuration/client.md | 2 +-
.../apache/celeborn/tests/spark/HugeDataTest.scala | 1 +
.../celeborn/tests/spark/PushdataTimeoutTest.scala | 1 +
.../celeborn/tests/spark/SkewJoinSuite.scala | 1 +
18 files changed, 729 insertions(+), 172 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 82dd6c9a..01e75384 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
@@ -109,6 +109,7 @@ public class SortBasedPusher extends MemoryConsumer {
rssShuffleClient,
afterPush,
mapStatusLengths);
+ dataPusher.startPushThread();
pushBufferMaxSize = conf.pushBufferMaxSize();
pushSortMemoryThreshold = conf.pushSortMemoryThreshold();
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 532d822d..24e03c29 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
@@ -153,6 +153,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
rssShuffleClient,
writeMetrics::incBytesWritten,
mapStatusLengths);
+ dataPusher.startPushThread();
}
@Override
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 59334a3b..6e7df3ce 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
@@ -233,7 +233,8 @@ public class RssShuffleWriterSuiteJ {
final File tempFile = new File(tempDir, UUID.randomUUID().toString());
final RssShuffleHandle<Integer, String, String> handle =
new RssShuffleHandle<>(appId, host, port, userIdentifier, shuffleId, numMaps, dependency);
- final ShuffleClient client = new DummyShuffleClient(tempFile);
+ final ShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ ((DummyShuffleClient) client).initReducePartitionMap(shuffleId, numPartitions, 1);
if (hashWriter) {
final HashBasedShuffleWriter<Integer, String, String> writer =
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 e54ba063..bd26433c 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
@@ -161,6 +161,7 @@ public class HashBasedShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
rssShuffleClient,
writeMetrics::incBytesWritten,
mapStatusLengths);
+ dataPusher.startPushThread();
if (conf.columnarShuffleEnabled()) {
this.schema = SparkUtils.getSchema(dep);
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 c7e59ce9..e5408bcb 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
@@ -210,7 +210,8 @@ public class RssShuffleWriterSuiteJ {
final File tempFile = new File(tempDir, UUID.randomUUID().toString());
final RssShuffleHandle<Integer, String, String> handle =
new RssShuffleHandle<>(appId, host, port, userIdentifier, shuffleId, numMaps, dependency);
- final ShuffleClient client = new DummyShuffleClient(tempFile);
+ final ShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ ((DummyShuffleClient) client).initReducePartitionMap(shuffleId, numPartitions, 1);
final HashBasedShuffleWriter<Integer, String, String> writer =
new HashBasedShuffleWriter<>(
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
index a3d3fa82..460705d5 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -19,6 +19,7 @@ package org.apache.celeborn.client;
import java.io.IOException;
import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BooleanSupplier;
import io.netty.buffer.ByteBuf;
@@ -28,6 +29,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.read.RssInputStream;
+import org.apache.celeborn.client.write.PushState;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.protocol.PartitionLocation;
@@ -235,4 +237,9 @@ public abstract class ShuffleClient {
public abstract PartitionLocation registerMapPartitionTask(
String appId, int shuffleId, int numMappers, int mapId, int attemptId);
+
+ public abstract ConcurrentHashMap<Integer, PartitionLocation> getPartitionLocation(
+ String applicationId, int shuffleId, int numMappers, int numPartitions);
+
+ public abstract PushState getPushState(String mapKey);
}
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 0676d086..063fdc40 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -195,7 +195,7 @@ public class ShuffleClientImpl extends ShuffleClient {
} else if (mapperEnded(shuffleId, mapId, attemptId)) {
logger.debug(
"Retrying push data, but the mapper(map {} attempt {}) has ended.", mapId, attemptId);
- pushState.removeBatch(batchId);
+ pushState.removeBatch(batchId, loc.hostAndPushPort());
} else {
PartitionLocation newLoc = reducePartitionMap.get(shuffleId).get(partitionId);
logger.info("Revive success, new location for reduce {} is {}.", partitionId, newLoc);
@@ -209,7 +209,7 @@ public class ShuffleClientImpl extends ShuffleClient {
PushData newPushData =
new PushData(MASTER_MODE, shuffleKey, newLoc.getUniqueId(), newBuffer);
ChannelFuture future = client.pushData(newPushData, callback);
- pushState.pushStarted(batchId, future, callback);
+ pushState.pushStarted(batchId, future, callback, loc.hostAndPushPort());
} else {
throw new RuntimeException(
"Mock push data submit retry failed. remainReviveTimes = " + remainReviveTimes + ".");
@@ -287,7 +287,7 @@ public class ShuffleClientImpl extends ShuffleClient {
remainReviveTimes);
}
if (reviveFailedBatchesMap.isEmpty()) {
- pushState.removeBatch(oldGroupedBatchId);
+ pushState.removeBatch(oldGroupedBatchId, batches.get(0).loc.hostAndPushPort());
} else {
pushDataRetryPool.submit(
() ->
@@ -350,6 +350,18 @@ public class ShuffleClientImpl extends ShuffleClient {
return partitionLocationMap.get(partitionId);
}
+ @Override
+ public ConcurrentHashMap<Integer, PartitionLocation> getPartitionLocation(
+ String applicationId, int shuffleId, int numMappers, int numPartitions) {
+ return reducePartitionMap.computeIfAbsent(
+ shuffleId, (id) -> registerShuffle(applicationId, shuffleId, numMappers, numPartitions));
+ }
+
+ @Override
+ public PushState getPushState(String mapKey) {
+ return pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
+ }
+
private ConcurrentHashMap<Integer, PartitionLocation> registerShuffleInternal(
int shuffleId,
int numMappers,
@@ -405,45 +417,21 @@ public class ShuffleClientImpl extends ShuffleClient {
return null;
}
- private void limitMaxInFlight(String mapKey, PushState pushState, int limit) throws IOException {
- if (pushState.exception.get() != null) {
- throw pushState.exception.get();
- }
-
- long timeoutMs = conf.pushLimitInFlightTimeoutMs();
- long delta = conf.pushLimitInFlightSleepDeltaMs();
- long times = timeoutMs / delta;
- try {
- while (times > 0) {
- if (pushState.inflightBatchCount() <= limit) {
- break;
- }
- if (pushState.exception.get() != null) {
- throw pushState.exception.get();
- }
+ private void limitMaxInFlight(
+ String mapKey, PushState pushState, int limit, String hostAndPushPort) throws IOException {
+ boolean reachLimit = pushState.limitMaxInFlight(hostAndPushPort, limit);
- pushState.failExpiredBatch();
-
- Thread.sleep(delta);
- times--;
- }
- } catch (InterruptedException e) {
- pushState.exception.set(new IOException(e));
+ if (reachLimit) {
+ throw new IOException("wait timeout for task " + mapKey, pushState.exception.get());
}
+ }
- if (times <= 0) {
- logger.error(
- "After waiting for {} ms, there are still {} batches in flight for map {}, "
- + "which exceeds the limit {}.",
- timeoutMs,
- pushState.inflightBatchCount(),
- mapKey,
- limit);
+ private void limitZeroInFlight(String mapKey, PushState pushState) throws IOException {
+ boolean reachLimit = pushState.limitZeroInFlight();
+
+ if (reachLimit) {
throw new IOException("wait timeout for task " + mapKey, pushState.exception.get());
}
- if (pushState.exception.get() != null) {
- throw pushState.exception.get();
- }
}
private boolean waitRevivedLocation(
@@ -564,9 +552,7 @@ public class ShuffleClientImpl extends ShuffleClient {
}
// register shuffle if not registered
final ConcurrentHashMap<Integer, PartitionLocation> map =
- reducePartitionMap.computeIfAbsent(
- shuffleId,
- (id) -> registerShuffle(applicationId, shuffleId, numMappers, numPartitions));
+ getPartitionLocation(applicationId, shuffleId, numMappers, numPartitions);
if (map == null) {
throw new IOException("Register shuffle failed for shuffle " + shuffleKey);
@@ -612,10 +598,10 @@ public class ShuffleClientImpl extends ShuffleClient {
+ " is NULL!");
}
- PushState pushState = pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
+ PushState pushState = getPushState(mapKey);
// increment batchId
- final int nextBatchId = pushState.batchId.addAndGet(1);
+ final int nextBatchId = pushState.nextBatchId();
// compress data
final Compressor compressor = compressorThreadLocal.get();
@@ -641,10 +627,10 @@ public class ShuffleClientImpl extends ShuffleClient {
partitionId,
nextBatchId);
// check limit
- limitMaxInFlight(mapKey, pushState, currentMaxReqsInFlight.get());
+ limitMaxInFlight(mapKey, pushState, currentMaxReqsInFlight.get(), loc.hostAndPushPort());
// add inFlight requests
- pushState.addBatch(nextBatchId);
+ pushState.addBatch(nextBatchId, loc.hostAndPushPort());
// build PushData request
NettyManagedBuffer buffer = new NettyManagedBuffer(Unpooled.wrappedBuffer(body));
@@ -655,7 +641,7 @@ public class ShuffleClientImpl extends ShuffleClient {
new RpcResponseCallback() {
@Override
public void onSuccess(ByteBuffer response) {
- pushState.removeBatch(nextBatchId);
+ pushState.removeBatch(nextBatchId, loc.hostAndPushPort());
// TODO Need to adjust maxReqsInFlight if server response is congested, see
// CELEBORN-62
if (response.remaining() > 0 && response.get() == StatusCode.STAGE_ENDED.getValue()) {
@@ -788,7 +774,7 @@ public class ShuffleClientImpl extends ShuffleClient {
getPushDataFailCause(e.getMessage()),
remainReviveTimes));
} else {
- pushState.removeBatch(nextBatchId);
+ pushState.removeBatch(nextBatchId, loc.hostAndPushPort());
logger.info(
"Mapper shuffleId:{} mapId:{} attempt:{} already ended, remove batchId:{}.",
shuffleId,
@@ -805,7 +791,7 @@ public class ShuffleClientImpl extends ShuffleClient {
TransportClient client =
dataClientFactory.createClient(loc.getHost(), loc.getPushPort(), partitionId);
ChannelFuture future = client.pushData(pushData, wrappedCallback);
- pushState.pushStarted(nextBatchId, future, wrappedCallback);
+ pushState.pushStarted(nextBatchId, future, wrappedCallback, loc.hostAndPushPort());
} else {
throw new RuntimeException("Mock push data first time failed.");
}
@@ -820,7 +806,7 @@ public class ShuffleClientImpl extends ShuffleClient {
String addressPair = genAddressPair(loc);
boolean shouldPush = pushState.addBatchData(addressPair, loc, nextBatchId, body);
if (shouldPush) {
- limitMaxInFlight(mapKey, pushState, maxInFlight);
+ limitMaxInFlight(mapKey, pushState, maxInFlight, loc.hostAndPushPort());
DataBatches dataBatches = pushState.takeDataBatches(addressPair);
doPushMergedData(
addressPair.split("-")[0],
@@ -897,7 +883,7 @@ public class ShuffleClientImpl extends ShuffleClient {
final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
PushState pushState = pushStates.get(mapKey);
if (pushState != null) {
- limitMaxInFlight(mapKey, pushState, 0);
+ limitZeroInFlight(mapKey, pushState);
}
}
@@ -938,13 +924,13 @@ public class ShuffleClientImpl extends ShuffleClient {
ArrayList<Map.Entry<String, DataBatches>> batchesArr =
new ArrayList<>(pushState.batchesMap.entrySet());
while (!batchesArr.isEmpty()) {
- limitMaxInFlight(mapKey, pushState, currentMaxReqsInFlight.get());
Map.Entry<String, DataBatches> entry = batchesArr.get(RND.nextInt(batchesArr.size()));
+ String[] tokens = entry.getKey().split("-");
+ limitMaxInFlight(mapKey, pushState, currentMaxReqsInFlight.get(), tokens[0]);
ArrayList<DataBatches.DataBatch> batches = entry.getValue().requireBatches(pushBufferMaxSize);
if (entry.getValue().getTotalSize() == 0) {
batchesArr.remove(entry);
}
- String[] tokens = entry.getKey().split("-");
doPushMergedData(
tokens[0],
applicationId,
@@ -970,8 +956,8 @@ public class ShuffleClientImpl extends ShuffleClient {
final String host = splits[0];
final int port = Integer.parseInt(splits[1]);
- int groupedBatchId = pushState.batchId.addAndGet(1);
- pushState.addBatch(groupedBatchId);
+ int groupedBatchId = pushState.nextBatchId();
+ pushState.addBatch(groupedBatchId, hostPort);
final int numBatches = batches.size();
final String[] partitionUniqueIds = new String[numBatches];
@@ -1001,7 +987,7 @@ public class ShuffleClientImpl extends ShuffleClient {
mapId,
attemptId,
groupedBatchId);
- pushState.removeBatch(groupedBatchId);
+ pushState.removeBatch(groupedBatchId, hostPort);
// TODO Need to adjust maxReqsInFlight if server response is congested, see CELEBORN-62
if (response.remaining() > 0 && response.get() == StatusCode.STAGE_ENDED.getValue()) {
mapperEndMap
@@ -1135,7 +1121,7 @@ public class ShuffleClientImpl extends ShuffleClient {
if (!testRetryRevive || remainReviveTimes < 1) {
TransportClient client = dataClientFactory.createClient(host, port);
ChannelFuture future = client.pushMergedData(mergedData, wrappedCallback);
- pushState.pushStarted(groupedBatchId, future, wrappedCallback);
+ pushState.pushStarted(groupedBatchId, future, wrappedCallback, hostPort);
} else {
throw new RuntimeException("Mock push merge data failed");
}
@@ -1173,10 +1159,10 @@ public class ShuffleClientImpl extends ShuffleClient {
Integer partitionId)
throws IOException {
final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
- PushState pushState = pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
+ PushState pushState = getPushState(mapKey);
try {
- limitMaxInFlight(mapKey, pushState, 0);
+ limitZeroInFlight(mapKey, pushState);
MapperEndResponse response =
driverRssMetaService.askSync(
@@ -1441,10 +1427,10 @@ public class ShuffleClientImpl extends ShuffleClient {
return 0;
}
- PushState pushState = pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
+ PushState pushState = getPushState(mapKey);
// increment batchId
- final int nextBatchId = pushState.batchId.addAndGet(1);
+ final int nextBatchId = pushState.nextBatchId();
int totalLength = data.readableBytes();
data.markWriterIndex();
data.writerIndex(0);
@@ -1462,10 +1448,10 @@ public class ShuffleClientImpl extends ShuffleClient {
partitionId,
nextBatchId);
// check limit
- limitMaxInFlight(mapKey, pushState, maxInFlight);
+ limitMaxInFlight(mapKey, pushState, maxInFlight, location.hostAndPushPort());
// add inFlight requests
- pushState.addBatch(nextBatchId);
+ pushState.addBatch(nextBatchId, location.hostAndPushPort());
// build PushData request
NettyManagedBuffer buffer = new NettyManagedBuffer(data);
@@ -1477,7 +1463,7 @@ public class ShuffleClientImpl extends ShuffleClient {
@Override
public void onSuccess(ByteBuffer response) {
closeCallBack.getAsBoolean();
- pushState.removeBatch(nextBatchId);
+ pushState.removeBatch(nextBatchId, location.hostAndPushPort());
if (response.remaining() > 0) {
byte reason = response.get();
if (reason == StatusCode.STAGE_ENDED.getValue()) {
@@ -1498,7 +1484,7 @@ public class ShuffleClientImpl extends ShuffleClient {
@Override
public void onFailure(Throwable e) {
closeCallBack.getAsBoolean();
- pushState.removeBatch(nextBatchId);
+ pushState.removeBatch(nextBatchId, location.hostAndPushPort());
if (pushState.exception.get() != null) {
return;
}
@@ -1527,7 +1513,7 @@ public class ShuffleClientImpl extends ShuffleClient {
try {
TransportClient client = createClientWaitingInFlightRequest(location, mapKey, pushState);
ChannelFuture future = client.pushData(pushData, callback);
- pushState.pushStarted(nextBatchId, future, callback);
+ pushState.pushStarted(nextBatchId, future, callback, location.hostAndPushPort());
} catch (Exception e) {
logger.warn("PushData byteBuf failed", e);
callback.onFailure(new Exception(getPushDataFailCause(e.getMessage()).toString(), e));
@@ -1545,7 +1531,7 @@ public class ShuffleClientImpl extends ShuffleClient {
// makesure that messages have been sent by old client, in order to keep receiving data
// orderly
if (currentClient != null) {
- limitMaxInFlight(mapKey, pushState, 0);
+ limitZeroInFlight(mapKey, pushState);
}
currentClient = client;
}
@@ -1568,6 +1554,7 @@ public class ShuffleClientImpl extends ShuffleClient {
shuffleId,
mapId,
attemptId,
+ location,
pushState,
() -> {
String shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId);
@@ -1607,6 +1594,7 @@ public class ShuffleClientImpl extends ShuffleClient {
shuffleId,
mapId,
attemptId,
+ location,
pushState,
() -> {
String shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId);
@@ -1676,6 +1664,7 @@ public class ShuffleClientImpl extends ShuffleClient {
shuffleId,
mapId,
attemptId,
+ location,
pushState,
() -> {
final String shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId);
@@ -1697,6 +1686,7 @@ public class ShuffleClientImpl extends ShuffleClient {
int shuffleId,
int mapId,
int attemptId,
+ PartitionLocation location,
PushState pushState,
ThrowingExceptionSupplier<R, Exception> supplier)
throws IOException {
@@ -1713,16 +1703,17 @@ public class ShuffleClientImpl extends ShuffleClient {
attemptId);
return null;
}
+ pushState = getPushState(mapKey);
// force data has been send
- limitMaxInFlight(mapKey, pushState, 0);
+ limitZeroInFlight(mapKey, pushState);
// add inFlight requests
- batchId = pushState.batchId.incrementAndGet();
- pushState.addBatch(batchId);
+ batchId = pushState.nextBatchId();
+ pushState.addBatch(batchId, location.hostAndPushPort());
return retrySendMessage(supplier);
} finally {
if (pushState != null) {
- pushState.removeBatch(batchId);
+ pushState.removeBatch(batchId, location.hostAndPushPort());
}
}
}
diff --git a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
new file mode 100644
index 00000000..6fb3def0
--- /dev/null
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
@@ -0,0 +1,118 @@
+/*
+ * 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.client.write;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.Utils;
+
+/*
+ * Queue for push data,
+ * it can take one PushTask whose worker inflight request is not reach limit,
+ * and it can add one PushTask.
+ *
+ * */
+public class DataPushQueue {
+ private static final Logger logger = LoggerFactory.getLogger(DataPushQueue.class);
+
+ private final long WAIT_TIME_NANOS = TimeUnit.MILLISECONDS.toNanos(500);
+
+ private final LinkedBlockingQueue<PushTask> workingQueue;
+ private final PushState pushState;
+ private final DataPusher dataPusher;
+ private final int maxInFlight;
+
+ private final String appId;
+ private final int shuffleId;
+ private final int numMappers;
+ private final int numPartitions;
+ private final ShuffleClient client;
+ private final Set<String> reachLimitWorker = new HashSet<>();
+
+ public DataPushQueue(
+ CelebornConf conf,
+ DataPusher dataPusher,
+ ShuffleClient client,
+ String appId,
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int numMappers,
+ int numPartitions) {
+ this.appId = appId;
+ this.shuffleId = shuffleId;
+ this.numMappers = numMappers;
+ this.numPartitions = numPartitions;
+ this.client = client;
+ this.dataPusher = dataPusher;
+ final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ this.pushState = client.getPushState(mapKey);
+ this.maxInFlight = conf.pushMaxReqsInFlight();
+ final int capacity = conf.pushQueueCapacity();
+ workingQueue = new LinkedBlockingQueue<>(capacity);
+ }
+
+ /*
+ * Now, `takePushTask` is only used by one thread,
+ * so it is not thread-safe.
+ * */
+ public PushTask takePushTask() throws IOException {
+ while (dataPusher.stillRunning()) {
+ reachLimitWorker.clear();
+ Iterator<PushTask> iterator = workingQueue.iterator();
+ while (iterator.hasNext()) {
+ PushTask task = iterator.next();
+ int partitionId = task.getPartitionId();
+ Map<Integer, PartitionLocation> partitionLocationMap =
+ client.getPartitionLocation(appId, shuffleId, numMappers, numPartitions);
+ PartitionLocation loc = partitionLocationMap.get(partitionId);
+
+ if (!reachLimitWorker.contains(loc.hostAndPushPort())) {
+ boolean reachLimit = pushState.reachLimit(loc.hostAndPushPort(), maxInFlight);
+ if (!reachLimit) {
+ iterator.remove();
+ return task;
+ } else {
+ reachLimitWorker.add(loc.hostAndPushPort());
+ }
+ }
+ }
+ }
+ return null;
+ }
+
+ public boolean addPushTask(PushTask pushTask) throws InterruptedException {
+ return workingQueue.offer(pushTask, WAIT_TIME_NANOS, TimeUnit.NANOSECONDS);
+ }
+
+ public void clear() {
+ workingQueue.clear();
+ }
+}
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 4e39d0ee..f49d01ad 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
@@ -18,6 +18,7 @@
package org.apache.celeborn.client.write;
import java.io.IOException;
+import java.util.Objects;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
@@ -33,8 +34,8 @@ public class DataPusher {
private final long WAIT_TIME_NANOS = TimeUnit.MILLISECONDS.toNanos(500);
private final LinkedBlockingQueue<PushTask> idleQueue;
- private final LinkedBlockingQueue<PushTask> workingQueue;
-
+ // partition -> PushTask Queue
+ private final DataPushQueue dataPushQueue;
private final ReentrantLock idleLock = new ReentrantLock();
private final Condition idleFull = idleLock.newCondition();
@@ -51,6 +52,7 @@ public class DataPusher {
private volatile boolean terminated;
private final LongAdder[] mapStatusLengths;
+ private Thread pushThread;
public DataPusher(
String appId,
@@ -69,7 +71,9 @@ public class DataPusher {
final int pushBufferMaxSize = conf.pushBufferMaxSize();
idleQueue = new LinkedBlockingQueue<>(pushQueueCapacity);
- workingQueue = new LinkedBlockingQueue<>(pushQueueCapacity);
+ dataPushQueue =
+ new DataPushQueue(
+ conf, this, client, appId, shuffleId, mapId, attemptId, numMappers, numPartitions);
for (int i = 0; i < pushQueueCapacity; i++) {
try {
@@ -90,40 +94,46 @@ public class DataPusher {
this.afterPush = afterPush;
this.mapStatusLengths = mapStatusLengths;
- new Thread("DataPusher-" + taskId) {
- private void reclaimTask(PushTask task) throws InterruptedException {
- idleLock.lockInterruptibly();
- try {
- idleQueue.put(task);
- if (idleQueue.remainingCapacity() == 0) {
- idleFull.signal();
+ pushThread =
+ new Thread("DataPusher-" + taskId) {
+ private void reclaimTask(PushTask task) throws InterruptedException {
+ idleLock.lockInterruptibly();
+ try {
+ idleQueue.put(task);
+ if (idleQueue.remainingCapacity() == 0) {
+ idleFull.signal();
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ exceptionRef.set(new IOException(e));
+ } finally {
+ idleLock.unlock();
+ }
}
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- exceptionRef.set(new IOException(e));
- } finally {
- idleLock.unlock();
- }
- }
- @Override
- public void run() {
- while (!terminated && exceptionRef.get() == null) {
- try {
- PushTask task = workingQueue.poll(WAIT_TIME_NANOS, TimeUnit.NANOSECONDS);
- if (task == null) {
- continue;
+ @Override
+ public void run() {
+ while (stillRunning()) {
+ try {
+ PushTask task = dataPushQueue.takePushTask();
+ if (Objects.isNull(task)) {
+ continue;
+ }
+ pushData(task);
+ reclaimTask(task);
+ } catch (InterruptedException e) {
+ exceptionRef.set(new IOException(e));
+ } catch (IOException e) {
+ exceptionRef.set(e);
+ }
}
- pushData(task);
- reclaimTask(task);
- } catch (InterruptedException e) {
- exceptionRef.set(new IOException(e));
- } catch (IOException e) {
- exceptionRef.set(e);
}
- }
- }
- }.start();
+ };
+ pushThread.setDaemon(true);
+ }
+
+ public void startPushThread() {
+ pushThread.start();
}
public void addTask(int partitionId, byte[] buffer, int size) throws IOException {
@@ -136,7 +146,7 @@ public class DataPusher {
task.setSize(size);
task.setPartitionId(partitionId);
System.arraycopy(buffer, 0, task.getBuffer(), 0, size);
- while (!workingQueue.offer(task, WAIT_TIME_NANOS, TimeUnit.NANOSECONDS)) {
+ while (!dataPushQueue.addPushTask(task)) {
checkException();
}
} catch (InterruptedException e) {
@@ -157,8 +167,12 @@ public class DataPusher {
}
terminated = true;
+ try {
+ pushThread.join();
+ } catch (InterruptedException ignored) {
+ }
idleQueue.clear();
- workingQueue.clear();
+ dataPushQueue.clear();
checkException();
}
@@ -197,4 +211,16 @@ public class DataPusher {
idleLock.unlock();
}
}
+
+ protected boolean stillRunning() {
+ return !terminated && !Objects.nonNull(exceptionRef.get());
+ }
+
+ public DataPushQueue getDataPushQueue() {
+ return dataPushQueue;
+ }
+
+ public Thread getPushThread() {
+ return pushThread;
+ }
}
diff --git a/client/src/main/java/org/apache/celeborn/client/write/InFlightRequestTracker.java b/client/src/main/java/org/apache/celeborn/client/write/InFlightRequestTracker.java
new file mode 100644
index 00000000..0a54f1e9
--- /dev/null
+++ b/client/src/main/java/org/apache/celeborn/client/write/InFlightRequestTracker.java
@@ -0,0 +1,220 @@
+/*
+ * 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.client.write;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import io.netty.channel.ChannelFuture;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.client.RpcResponseCallback;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+
+/*
+ * This class is for track in flight request and limit request.
+ * */
+public class InFlightRequestTracker {
+ private static final Logger logger = LoggerFactory.getLogger(InFlightRequestTracker.class);
+
+ private final long waitInflightTimeoutMs;
+ private final long pushTimeoutMs;
+ private final long delta;
+ private final PushState pushState;
+
+ private final AtomicInteger batchId = new AtomicInteger();
+ private final ConcurrentHashMap<String, ConcurrentHashMap<Integer, BatchInfo>>
+ inflightBatchesPerAddress = new ConcurrentHashMap<>();
+
+ public InFlightRequestTracker(CelebornConf conf, PushState pushState) {
+ this.waitInflightTimeoutMs = conf.pushLimitInFlightTimeoutMs();
+ this.pushTimeoutMs = conf.pushDataTimeoutMs();
+ this.delta = conf.pushLimitInFlightSleepDeltaMs();
+ this.pushState = pushState;
+ }
+
+ public void addFlightBatch(int batchId, String hostAndPushPort) {
+ ConcurrentHashMap<Integer, BatchInfo> batchIdSetPerPair =
+ inflightBatchesPerAddress.computeIfAbsent(hostAndPushPort, id -> new ConcurrentHashMap<>());
+ batchIdSetPerPair.computeIfAbsent(batchId, id -> new BatchInfo());
+ }
+
+ public void removeFlightBatch(int batchId, String hostAndPushPort) {
+ ConcurrentHashMap<Integer, BatchInfo> batchIdMap =
+ inflightBatchesPerAddress.get(hostAndPushPort);
+ BatchInfo info = batchIdMap.remove(batchId);
+ if (info != null && info.channelFuture != null) {
+ info.channelFuture.cancel(true);
+ }
+ if (batchIdMap.size() == 0) {
+ inflightBatchesPerAddress.remove(hostAndPushPort);
+ }
+ }
+
+ public ConcurrentHashMap<Integer, BatchInfo> getBatchIdSetByAddressPair(String hostAndPort) {
+ return inflightBatchesPerAddress.computeIfAbsent(
+ hostAndPort, pair -> new ConcurrentHashMap<>());
+ }
+
+ public boolean limitMaxInFlight(String hostAndPushPort, int maxInFlight) throws IOException {
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+
+ ConcurrentHashMap<Integer, BatchInfo> batchIdMap = getBatchIdSetByAddressPair(hostAndPushPort);
+ long times = waitInflightTimeoutMs / delta;
+ try {
+ while (times > 0) {
+ if (batchIdMap.size() <= maxInFlight) {
+ break;
+ }
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+ failExpiredBatch();
+ Thread.sleep(delta);
+ times--;
+ }
+ } catch (InterruptedException e) {
+ pushState.exception.set(new IOException(e));
+ }
+
+ if (times <= 0) {
+ logger.warn(
+ "After waiting for {} ms, "
+ + "there are still {} batches in flight "
+ + "for hostAndPushPort {}, "
+ + "which exceeds the limit {}.",
+ waitInflightTimeoutMs,
+ batchIdMap.size(),
+ hostAndPushPort,
+ maxInFlight);
+ }
+
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+
+ return times <= 0;
+ }
+
+ public boolean limitZeroInFlight() throws IOException {
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+ long times = waitInflightTimeoutMs / delta;
+
+ int inFlightSize = 0;
+ try {
+ while (times > 0) {
+ Optional<Integer> inFlightSizeOptional =
+ inflightBatchesPerAddress.values().stream().map(Map::size).reduce(Integer::sum);
+ inFlightSize = inFlightSizeOptional.orElse(0);
+ if (inFlightSize == 0) {
+ break;
+ }
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+ failExpiredBatch();
+ Thread.sleep(delta);
+ times--;
+ }
+ } catch (InterruptedException e) {
+ pushState.exception.set(new IOException(e));
+ }
+
+ if (times <= 0) {
+ logger.error(
+ "After waiting for {} ms, there are still {} batches in flight, expect 0 batches",
+ waitInflightTimeoutMs,
+ inFlightSize);
+ }
+
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+
+ return times <= 0;
+ }
+
+ public boolean reachLimit(String hostAndPushPort, int maxInFlight) throws IOException {
+ if (pushState.exception.get() != null) {
+ throw pushState.exception.get();
+ }
+
+ ConcurrentHashMap<Integer, BatchInfo> batchIdMap = getBatchIdSetByAddressPair(hostAndPushPort);
+ return batchIdMap.size() > maxInFlight;
+ }
+
+ protected int nextBatchId() {
+ return batchId.incrementAndGet();
+ }
+
+ public synchronized void failExpiredBatch() {
+ long currentTime = System.currentTimeMillis();
+ inflightBatchesPerAddress
+ .values()
+ .forEach(
+ batchIdMap ->
+ batchIdMap
+ .values()
+ .forEach(
+ info -> {
+ if (currentTime - info.pushTime > pushTimeoutMs) {
+ if (info.callback != null) {
+ info.channelFuture.cancel(true);
+ info.callback.onFailure(
+ new IOException(StatusCode.PUSH_DATA_TIMEOUT.getMessage()));
+ info.channelFuture = null;
+ info.callback = null;
+ }
+ }
+ }));
+ }
+
+ public void cleanup() {
+ if (!inflightBatchesPerAddress.isEmpty()) {
+ logger.debug("Cancel all {} futures.", inflightBatchesPerAddress.size());
+ inflightBatchesPerAddress
+ .values()
+ .forEach(
+ batchIdMap ->
+ batchIdMap
+ .values()
+ .forEach(
+ info -> {
+ if (info.channelFuture != null) {
+ info.channelFuture.cancel(true);
+ }
+ }));
+ inflightBatchesPerAddress.clear();
+ }
+ }
+
+ static class BatchInfo {
+ ChannelFuture channelFuture;
+ long pushTime;
+ RpcResponseCallback callback;
+ }
+}
diff --git a/client/src/main/java/org/apache/celeborn/client/write/PushState.java b/client/src/main/java/org/apache/celeborn/client/write/PushState.java
index a8f96240..6449a70d 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/PushState.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/PushState.java
@@ -19,7 +19,6 @@ package org.apache.celeborn.client.write;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import io.netty.channel.ChannelFuture;
@@ -29,65 +28,23 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.protocol.PartitionLocation;
-import org.apache.celeborn.common.protocol.message.StatusCode;
public class PushState {
- class BatchInfo {
- ChannelFuture channelFuture;
- long pushTime;
- RpcResponseCallback callback;
- }
-
private static final Logger logger = LoggerFactory.getLogger(PushState.class);
- private int pushBufferMaxSize;
- private long pushTimeout;
-
- public final AtomicInteger batchId = new AtomicInteger();
- private final ConcurrentHashMap<Integer, BatchInfo> inflightBatchInfos =
- new ConcurrentHashMap<>();
+ private final int pushBufferMaxSize;
public AtomicReference<IOException> exception = new AtomicReference<>();
+ private final InFlightRequestTracker inFlightRequestTracker;
public PushState(CelebornConf conf) {
pushBufferMaxSize = conf.pushBufferMaxSize();
- pushTimeout = conf.pushDataTimeoutMs();
+ inFlightRequestTracker = new InFlightRequestTracker(conf, this);
}
- public void addBatch(int batchId) {
- inflightBatchInfos.computeIfAbsent(batchId, id -> new BatchInfo());
- }
-
- public void removeBatch(int batchId) {
- BatchInfo info = inflightBatchInfos.remove(batchId);
- if (info != null && info.channelFuture != null) {
- info.channelFuture.cancel(true);
- }
- }
-
- public int inflightBatchCount() {
- return inflightBatchInfos.size();
- }
-
- public synchronized void failExpiredBatch() {
- long currentTime = System.currentTimeMillis();
- inflightBatchInfos
- .values()
- .forEach(
- info -> {
- if (currentTime - info.pushTime > pushTimeout) {
- if (info.callback != null) {
- info.channelFuture.cancel(true);
- info.callback.onFailure(
- new IOException(StatusCode.PUSH_DATA_TIMEOUT.getMessage()));
- info.channelFuture = null;
- info.callback = null;
- }
- }
- });
- }
-
- public void pushStarted(int batchId, ChannelFuture future, RpcResponseCallback callback) {
- BatchInfo info = inflightBatchInfos.get(batchId);
+ public void pushStarted(
+ int batchId, ChannelFuture future, RpcResponseCallback callback, String hostAndPort) {
+ InFlightRequestTracker.BatchInfo info =
+ inFlightRequestTracker.getBatchIdSetByAddressPair(hostAndPort).get(batchId);
// In rare cases info could be null. For example, a speculative task has one thread pushing,
// and other thread retry-pushing. At time 1 thread 1 find StageEnded, then it cleans up
// PushState, at the same time thread 2 pushes data and calles pushStarted,
@@ -100,18 +57,7 @@ public class PushState {
}
public void cleanup() {
- if (!inflightBatchInfos.isEmpty()) {
- logger.debug("Cancel all {} futures.", inflightBatchInfos.size());
- inflightBatchInfos
- .values()
- .forEach(
- entry -> {
- if (entry.channelFuture != null) {
- entry.channelFuture.cancel(true);
- }
- });
- inflightBatchInfos.clear();
- }
+ inFlightRequestTracker.cleanup();
}
// key: ${master addr}-${slave addr} value: list of data batch
@@ -135,4 +81,28 @@ public class PushState {
public DataBatches takeDataBatches(String addressPair) {
return batchesMap.remove(addressPair);
}
+
+ public int nextBatchId() {
+ return inFlightRequestTracker.nextBatchId();
+ }
+
+ public void addBatch(int batchId, String hostAndPushPort) {
+ inFlightRequestTracker.addFlightBatch(batchId, hostAndPushPort);
+ }
+
+ public void removeBatch(int batchId, String hostAndPushPort) {
+ inFlightRequestTracker.removeFlightBatch(batchId, hostAndPushPort);
+ }
+
+ public boolean limitMaxInFlight(String hostAndPushPort, int maxInFlight) throws IOException {
+ return inFlightRequestTracker.limitMaxInFlight(hostAndPushPort, maxInFlight);
+ }
+
+ public boolean limitZeroInFlight() throws IOException {
+ return inFlightRequestTracker.limitZeroInFlight();
+ }
+
+ public boolean reachLimit(String hostAndPushPort, int maxInFlight) throws IOException {
+ return inFlightRequestTracker.reachLimit(hostAndPushPort, maxInFlight);
+ }
}
diff --git a/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java b/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
index 07f03038..c88b1d30 100644
--- a/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -17,12 +17,19 @@
package org.apache.celeborn.client;
+import static org.apache.celeborn.common.protocol.PartitionLocation.Mode.MASTER;
+
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BooleanSupplier;
import io.netty.buffer.ByteBuf;
@@ -30,6 +37,8 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.read.RssInputStream;
+import org.apache.celeborn.client.write.PushState;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
@@ -38,9 +47,14 @@ public class DummyShuffleClient extends ShuffleClient {
private static final Logger LOG = LoggerFactory.getLogger(DummyShuffleClient.class);
private final OutputStream os;
+ private final CelebornConf conf;
+
+ private final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>> reducePartitionMap =
+ new HashMap<>();
- public DummyShuffleClient(File file) throws Exception {
+ public DummyShuffleClient(CelebornConf conf, File file) throws Exception {
this.os = new BufferedOutputStream(new FileOutputStream(file));
+ this.conf = conf;
}
@Override
@@ -184,4 +198,33 @@ public class DummyShuffleClient extends ShuffleClient {
String appId, int shuffleId, int numMappers, int mapId, int attemptId) {
return null;
}
+
+ @Override
+ public ConcurrentHashMap<Integer, PartitionLocation> getPartitionLocation(
+ String applicationId, int shuffleId, int numMappers, int numPartitions) {
+ return reducePartitionMap.get(shuffleId);
+ }
+
+ @Override
+ public PushState getPushState(String mapKey) {
+ return new PushState(conf);
+ }
+
+ public void initReducePartitionMap(int shuffleId, int numPartitions, int workerNum) {
+ ConcurrentHashMap<Integer, PartitionLocation> map = new ConcurrentHashMap<>();
+ String host = "host";
+ List<PartitionLocation> partitionLocationList = new ArrayList<>();
+ for (int i = 0; i < workerNum; i++) {
+ partitionLocationList.add(
+ new PartitionLocation(0, 0, host, 1000 + i, 2000 + i, 3000 + i, 4000 + i, MASTER));
+ }
+ for (int i = 0; i < numPartitions; i++) {
+ map.put(i, partitionLocationList.get(i % workerNum));
+ }
+ reducePartitionMap.put(shuffleId, map);
+ }
+
+ public Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>> getReducePartitionMap() {
+ return reducePartitionMap;
+ }
}
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
new file mode 100644
index 00000000..868b193c
--- /dev/null
+++ b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
@@ -0,0 +1,174 @@
+/*
+ * 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.client.write;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.LongAdder;
+
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.slf4j.Logger;
+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.Utils;
+
+public class DataPushQueueSuitJ {
+ private static final Logger LOG = LoggerFactory.getLogger(DataPushQueueSuitJ.class);
+ private static File tempDir = null;
+
+ private final int shuffleId = 0;
+ private final int numPartitions = 10;
+
+ @BeforeClass
+ public static void beforeAll() {
+ tempDir = Utils.createTempDir(System.getProperty("java.io.tmpdir"), "celeborn_test");
+ }
+
+ @AfterClass
+ public static void afterAll() {
+ try {
+ JavaUtils.deleteRecursively(tempDir);
+ } catch (IOException e) {
+ LOG.error("Failed to delete temporary directory.", e);
+ }
+ }
+
+ @Test
+ public void testDataPushQueue() throws Exception {
+ final int numWorker = 3;
+ List<List<Integer>> workerData = new ArrayList<>();
+ for (int i = 0; i < numWorker; i++) {
+ workerData.add(new ArrayList<>());
+ }
+ for (int i = 0; i < numPartitions; i++) {
+ workerData.get(i % numWorker).add(i);
+ }
+ List<List<Integer>> tarWorkerData = new ArrayList<>();
+ for (int i = 0; i < numWorker; i++) {
+ tarWorkerData.add(new ArrayList<>());
+ }
+
+ Map<Integer, Integer> partitionBatchIdMap = new HashMap<>();
+
+ CelebornConf conf = new CelebornConf();
+ conf.set("celeborn.push.maxReqsInFlight", "2");
+
+ String app = "APP-1";
+ int shuffleId = 0;
+ int mapId = 0;
+ int attemptId = 0;
+ int numMappers = 10;
+ int numPartitions = 10;
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(shuffleId, numPartitions, numWorker);
+
+ LongAdder[] mapStatusLengths = new LongAdder[numPartitions];
+ for (int i = 0; i < numPartitions; i++) {
+ mapStatusLengths[i] = new LongAdder();
+ }
+ DataPusher dataPusher =
+ new DataPusher(
+ app,
+ shuffleId,
+ mapId,
+ attemptId,
+ 0,
+ numMappers,
+ numPartitions,
+ conf,
+ client,
+ integer -> {},
+ mapStatusLengths);
+
+ final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ DataPushQueue dataPushQueue = dataPusher.getDataPushQueue();
+ PushState pushState = client.getPushState(mapKey);
+ Map<Integer, PartitionLocation> reducePartitionMap =
+ client.getReducePartitionMap().get(shuffleId);
+
+ for (int i = 0; i < numPartitions; i++) {
+ byte[] b = intToBytes(workerData.get(i % numWorker).get(i / numWorker));
+ dataPusher.addTask(i, b, b.length);
+ int batchId = pushState.nextBatchId();
+ pushState.addBatch(batchId, reducePartitionMap.get(i).hostAndPushPort());
+ partitionBatchIdMap.put(i, batchId);
+ }
+
+ AtomicBoolean running = new AtomicBoolean(true);
+ new Thread(
+ () -> {
+ while (running.get()) {
+ try {
+ PushTask task = dataPushQueue.takePushTask();
+ byte[] buffer = task.getBuffer();
+ int partitionId = task.getPartitionId();
+ tarWorkerData.get(partitionId % numWorker).add(bytesToInt(buffer));
+ pushState.removeBatch(
+ partitionBatchIdMap.get(partitionId),
+ reducePartitionMap.get(partitionId).hostAndPushPort());
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ })
+ .start();
+
+ Thread.sleep(15 * 1000);
+ running.set(false);
+
+ for (int i = 0; i < numWorker; i++) {
+ Assert.assertArrayEquals(workerData.get(i).toArray(), tarWorkerData.get(i).toArray());
+ }
+
+ client.shutdown();
+ }
+
+ public static byte[] intToBytes(int value) {
+ byte[] src = new byte[4];
+ src[0] = (byte) (value & 0xFF);
+ src[1] = (byte) ((value >> 8) & 0xFF);
+ src[2] = (byte) ((value >> 16) & 0xFF);
+ src[3] = (byte) ((value >> 24) & 0xFF);
+ return src;
+ }
+
+ public static int bytesToInt(byte[] src) {
+ int value;
+ value =
+ (int)
+ ((src[0] & 0xFF)
+ | ((src[1] & 0xFF) << 8)
+ | ((src[2] & 0xFF) << 16)
+ | ((src[3] & 0xFF) << 24));
+ return value;
+ }
+}
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 c73003d2..816a7679 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1272,11 +1272,11 @@ object CelebornConf extends Logging {
.withAlternative("rss.push.data.maxReqsInFlight")
.categories("client")
.version("0.2.0")
- .doc("Amount of Netty in-flight requests. The maximum memory is " +
+ .doc("Amount of Netty in-flight requests per worker. The maximum memory is " +
"`celeborn.push.maxReqsInFlight` * `celeborn.push.buffer.max.size` * " +
"compression ratio(1 in worst case), default: 64Kib * 32 = 2Mib")
.intConf
- .createWithDefault(32)
+ .createWithDefault(4)
val PUSH_MAX_REVIVE_TIMES: ConfigEntry[Int] =
buildConf("celeborn.push.revive.maxRetries")
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index 0ac32a45..e3cc0fe6 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -34,7 +34,7 @@ license: |
| celeborn.push.data.timeout | 120s | Timeout for a task to push data rpc message. | 0.2.0 |
| celeborn.push.limit.inFlight.sleepInterval | 50ms | Sleep interval when check netty in-flight requests to be done. | 0.2.0 |
| celeborn.push.limit.inFlight.timeout | 240s | Timeout for netty in-flight requests to be done. | 0.2.0 |
-| celeborn.push.maxReqsInFlight | 32 | Amount of Netty in-flight requests. The maximum memory is `celeborn.push.maxReqsInFlight` * `celeborn.push.buffer.max.size` * compression ratio(1 in worst case), default: 64Kib * 32 = 2Mib | 0.2.0 |
+| celeborn.push.maxReqsInFlight | 4 | Amount of Netty in-flight requests per worker. The maximum memory is `celeborn.push.maxReqsInFlight` * `celeborn.push.buffer.max.size` * compression ratio(1 in worst case), default: 64Kib * 32 = 2Mib | 0.2.0 |
| celeborn.push.queue.capacity | 512 | Push buffer queue size for a task. The maximum memory is `celeborn.push.buffer.max.size` * `celeborn.push.queue.capacity`, default: 64KiB * 512 = 32MiB | 0.2.0 |
| celeborn.push.replicate.enabled | true | When true, Celeborn worker will replicate shuffle data to another Celeborn worker asynchronously to ensure the pushed shuffle data won't be lost after the node failure. | 0.2.0 |
| celeborn.push.retry.threads | 8 | Thread number to process shuffle re-send push data requests. | 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 05accffe..e31b54d1 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
@@ -38,6 +38,7 @@ class HugeDataTest extends AnyFunSuite
test("celeborn spark integration test - huge data") {
val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[4]")
+ .set("spark.celeborn.shuffle.register.maxRetries", "10")
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 74eb2d60..4120c2ab 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
@@ -46,6 +46,7 @@ class PushdataTimeoutTest extends AnyFunSuite
test("celeborn spark integration test - pushdata timeout") {
val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[4]")
.set("spark.celeborn.push.data.timeout", "10s")
+ .set("spark.celeborn.shuffle.register.maxRetries", "10")
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/SkewJoinSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/SkewJoinSuite.scala
index 69f3768f..6a0244ad 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
@@ -58,6 +58,7 @@ class SkewJoinSuite extends AnyFunSuite
.set("spark.sql.autoBroadcastJoinThreshold", "-1")
.set("spark.sql.parquet.compression.codec", "gzip")
.set("spark.celeborn.shuffle.compression.codec", codec.name)
+ .set("spark.celeborn.shuffle.register.maxRetries", "10")
enableRss(sparkConf)