This is an automated email from the ASF dual-hosted git repository.
ethanfeng 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 b5dc106a [CELEBORN-291] optimize shuffleclientimpl creating client and
pushdata for mappartition (#1224)
b5dc106a is described below
commit b5dc106af8f406a72856d33b09cb318214429ae2
Author: zhongqiangchen <[email protected]>
AuthorDate: Fri Feb 17 19:07:19 2023 +0800
[CELEBORN-291] optimize shuffleclientimpl creating client and pushdata for
mappartition (#1224)
---
.../plugin/flink/RemoteShuffleOutputGate.java | 1 +
.../org/apache/celeborn/client/ShuffleClient.java | 3 +--
.../apache/celeborn/client/ShuffleClientImpl.java | 31 ++++++++++++----------
.../apache/celeborn/client/DummyShuffleClient.java | 3 +--
.../celeborn/client/ShuffleClientImplSuiteJ.java | 10 +++----
.../common/network/client/TransportClient.java | 26 ++++++++++++++++--
6 files changed, 49 insertions(+), 25 deletions(-)
diff --git
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGate.java
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGate.java
index ff4143ef..3d7adef3 100644
---
a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGate.java
+++
b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGate.java
@@ -203,6 +203,7 @@ public class RemoteShuffleOutputGate {
bufferPool.lazyDestroy();
}
bufferPacker.close();
+ shuffleWriteClient.cleanup(applicationId, shuffleId, mapId, attemptId);
}
/** Returns shuffle descriptor. */
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 184400b3..4d3d65c6 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -20,7 +20,6 @@ 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;
import org.apache.hadoop.conf.Configuration;
@@ -208,7 +207,7 @@ public abstract class ShuffleClient {
int partitionId,
ByteBuf data,
PartitionLocation location,
- BooleanSupplier closeCallBack)
+ Runnable closeCallBack)
throws IOException;
public abstract Optional<PartitionLocation> regionStart(
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 d0dc3428..8cdaa486 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -25,7 +25,6 @@ import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
-import java.util.function.BooleanSupplier;
import scala.reflect.ClassTag$;
@@ -130,7 +129,7 @@ public class ShuffleClientImpl extends ShuffleClient {
// key: shuffleId
private final Map<Integer, ReduceFileGroups> reduceFileGroupsMap = new
ConcurrentHashMap<>();
- private TransportClient currentClient;
+ private ConcurrentHashMap<String, TransportClient> currentClient = new
ConcurrentHashMap<>();
public ShuffleClientImpl(CelebornConf conf, UserIdentifier userIdentifier) {
super();
@@ -1209,6 +1208,9 @@ public class ShuffleClientImpl extends ShuffleClient {
pushState.exception.compareAndSet(null, new IOException("Cleaned Up"));
pushState.cleanup();
}
+ if (currentClient != null) {
+ currentClient.remove(mapKey);
+ }
}
@Override
@@ -1404,7 +1406,7 @@ public class ShuffleClientImpl extends ShuffleClient {
int partitionId,
ByteBuf data,
PartitionLocation location,
- BooleanSupplier closeCallBack)
+ Runnable closeCallBack)
throws IOException {
// mapKey
final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
@@ -1431,13 +1433,14 @@ public class ShuffleClientImpl extends ShuffleClient {
int totalLength = data.readableBytes();
data.markWriterIndex();
data.writerIndex(0);
- data.writeInt(mapId);
+ data.writeInt(partitionId);
data.writeInt(attemptId);
data.writeInt(nextBatchId);
data.writeInt(totalLength - BATCH_HEADER_SIZE);
data.resetWriterIndex();
logger.debug(
- "Do push data byteBuf for app {} shuffle {} map {} attempt {} reduce
{} batch {}.",
+ "Do push data byteBuf size {} for app {} shuffle {} map {} attempt {}
reduce {} batch {}.",
+ totalLength,
applicationId,
shuffleId,
mapId,
@@ -1459,7 +1462,6 @@ public class ShuffleClientImpl extends ShuffleClient {
new RpcResponseCallback() {
@Override
public void onSuccess(ByteBuffer response) {
- closeCallBack.getAsBoolean();
pushState.removeBatch(nextBatchId, location.hostAndPushPort());
if (response.remaining() > 0) {
byte reason = response.get();
@@ -1480,7 +1482,6 @@ public class ShuffleClientImpl extends ShuffleClient {
@Override
public void onFailure(Throwable e) {
- closeCallBack.getAsBoolean();
pushState.removeBatch(nextBatchId, location.hostAndPushPort());
if (pushState.exception.get() != null) {
return;
@@ -1509,7 +1510,7 @@ public class ShuffleClientImpl extends ShuffleClient {
// do push data
try {
TransportClient client = createClientWaitingInFlightRequest(location,
mapKey, pushState);
- client.pushData(pushData, pushDataTimeout, callback);
+ client.pushData(pushData, pushDataTimeout, callback, closeCallBack);
} catch (Exception e) {
logger.warn("PushData byteBuf failed", e);
callback.onFailure(
@@ -1529,15 +1530,15 @@ public class ShuffleClientImpl extends ShuffleClient {
TransportClient client =
dataClientFactory.createClient(
location.getHost(), location.getPushPort(), location.getId());
- if (currentClient != client) {
+ if (currentClient.get(mapKey) != client) {
// makesure that messages have been sent by old client, in order to keep
receiving data
// orderly
- if (currentClient != null) {
+ if (currentClient.get(mapKey) != null) {
limitZeroInFlight(mapKey, pushState);
}
- currentClient = client;
+ currentClient.put(mapKey, client);
}
- return currentClient;
+ return currentClient.get(mapKey);
}
@Override
@@ -1601,7 +1602,8 @@ public class ShuffleClientImpl extends ShuffleClient {
() -> {
String shuffleKey = Utils.makeShuffleKey(applicationId, shuffleId);
logger.info(
- "regionStart shuffleKey:{}, attemptId:{}, locationId:{}",
+ "regionStart regionId:{}, shuffleKey:{}, attemptId:{},
locationId:{}",
+ currentRegionIdx,
shuffleKey,
attemptId,
location.getUniqueId());
@@ -1735,7 +1737,7 @@ public class ShuffleClientImpl extends ShuffleClient {
while (!Thread.currentThread().isInterrupted()
&& !isSuccess
&& retryTimes <
conf.networkIoMaxRetries(TransportModuleConstants.PUSH_MODULE)) {
- logger.info("retrySendMessage times: {}", retryTimes);
+ logger.debug("retrySendMessage times: {}", retryTimes);
try {
result = supplier.get();
isSuccess = true;
@@ -1771,6 +1773,7 @@ public class ShuffleClientImpl extends ShuffleClient {
|| (e.getCause() != null && e.getCause() instanceof
TimeoutException)
|| (e.getCause() != null && e.getCause() instanceof IOException)
|| (e instanceof RuntimeException
+ && e.getMessage() != null
&& e.getMessage().startsWith(IOException.class.getName()));
return isIOException;
}
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 0c5c1568..37b72be0 100644
--- a/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -30,7 +30,6 @@ 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;
import org.slf4j.Logger;
@@ -160,7 +159,7 @@ public class DummyShuffleClient extends ShuffleClient {
int partitionId,
ByteBuf data,
PartitionLocation location,
- BooleanSupplier closeCallBack) {
+ Runnable closeCallBack) {
return 0;
}
diff --git
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
index 33b0c2af..d3b22784 100644
---
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
+++
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientImplSuiteJ.java
@@ -76,7 +76,7 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
TEST_REDUCRE_ID,
byteBuf,
masterLocation,
- () -> true);
+ () -> {});
Assert.assertEquals(BufferSize, pushDataLen);
}
@@ -102,13 +102,13 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
TEST_REDUCRE_ID,
byteBuf,
masterLocation,
- () -> true);
+ () -> {});
}
@Test
public void testPushDataByteBufFail() throws IOException {
ByteBuf byteBuf = Unpooled.wrappedBuffer(TEST_BUF1);
- when(client.pushData(any(), anyLong(), any()))
+ when(client.pushData(any(), anyLong(), any(), any()))
.thenAnswer(
t -> {
RpcResponseCallback rpcResponseCallback =
@@ -125,7 +125,7 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
TEST_REDUCRE_ID,
byteBuf,
masterLocation,
- () -> true);
+ () -> {});
boolean isFailed = false;
// second push will throw exception
@@ -138,7 +138,7 @@ public class ShuffleClientImplSuiteJ extends
ShuffleClientBaseSuiteJ {
TEST_REDUCRE_ID,
byteBuf,
masterLocation,
- () -> true);
+ () -> {});
} catch (IOException e) {
isFailed = true;
} finally {
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
index 5631518b..050f7f4c 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
@@ -159,6 +159,14 @@ public class TransportClient implements Closeable {
public ChannelFuture pushData(
PushData pushData, long pushDataTimeout, RpcResponseCallback callback) {
+ return pushData(pushData, pushDataTimeout, callback, null);
+ }
+
+ public ChannelFuture pushData(
+ PushData pushData,
+ long pushDataTimeout,
+ RpcResponseCallback callback,
+ Runnable rpcSendoutCallback) {
if (logger.isTraceEnabled()) {
logger.trace("Pushing data to {}", NettyUtils.getRemoteAddress(channel));
}
@@ -168,8 +176,7 @@ public class TransportClient implements Closeable {
PushRequestInfo info = new PushRequestInfo(dueTime, callback);
handler.addPushRequest(requestId, info);
pushData.requestId = requestId;
-
- PushChannelListener listener = new PushChannelListener(requestId,
callback);
+ PushChannelListener listener = new PushChannelListener(requestId,
callback, rpcSendoutCallback);
ChannelFuture channelFuture =
channel.writeAndFlush(pushData).addListener(listener);
info.setChannelFuture(channelFuture);
return channelFuture;
@@ -334,11 +341,26 @@ public class TransportClient implements Closeable {
private class PushChannelListener extends StdChannelListener {
final long pushRequestId;
final RpcResponseCallback callback;
+ Runnable rpcSendOutCallback;
PushChannelListener(long pushRequestId, RpcResponseCallback callback) {
+ this(pushRequestId, callback, null);
+ }
+
+ PushChannelListener(
+ long pushRequestId, RpcResponseCallback callback, Runnable
rpcSendOutCallback) {
super("PUSH " + pushRequestId);
this.pushRequestId = pushRequestId;
this.callback = callback;
+ this.rpcSendOutCallback = rpcSendOutCallback;
+ }
+
+ @Override
+ public void operationComplete(Future<? super Void> future) throws
Exception {
+ super.operationComplete(future);
+ if (rpcSendOutCallback != null) {
+ rpcSendOutCallback.run();
+ }
}
@Override
