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/03/16 06:07:15 UTC
[incubator-celeborn] branch main updated: [CELEBORN-432][FLINK] Refine RemoteShuffleOutputGate and RemoteShuffleInputGateFactory (#1355)
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 66cc72086 [CELEBORN-432][FLINK] Refine RemoteShuffleOutputGate and RemoteShuffleInputGateFactory (#1355)
66cc72086 is described below
commit 66cc720861a9d6f8210a6c7a39b87b68502ab9c5
Author: Keyong Zhou <zh...@apache.org>
AuthorDate: Thu Mar 16 14:07:11 2023 +0800
[CELEBORN-432][FLINK] Refine RemoteShuffleOutputGate and RemoteShuffleInputGateFactory (#1355)
---
.../flink/RemoteShuffleInputGateFactory.java | 7 ++--
.../plugin/flink/RemoteShuffleOutputGate.java | 39 ++++++++++------------
.../flink/RemoteShuffleOutputGateSuiteJ.java | 12 +++----
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 2 +-
4 files changed, 30 insertions(+), 30 deletions(-)
diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
index a3485721f..ad3620518 100644
--- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
@@ -70,7 +70,10 @@ public class RemoteShuffleInputGateFactory {
long configuredMemorySize =
org.apache.celeborn.common.util.Utils.byteStringAsBytes(
PluginConf.getValue(flinkConf, PluginConf.MEMORY_PER_INPUT_GATE));
- if (configuredMemorySize < MIN_BUFFERS_PER_GATE) {
+ long minConfiguredMemorySize =
+ org.apache.celeborn.common.util.Utils.byteStringAsBytes(
+ PluginConf.getValue(flinkConf, PluginConf.MIN_MEMORY_PER_GATE));
+ if (configuredMemorySize < minConfiguredMemorySize) {
throw new IllegalArgumentException(
String.format(
"Insufficient network memory per input gate, please increase %s to at " + "least %s.",
@@ -118,7 +121,7 @@ public class RemoteShuffleInputGateFactory {
SupplierWithException<BufferPool, IOException> bufferPoolFactory,
BufferDecompressor bufferDecompressor) {
return new RemoteShuffleInputGate(
- this.celebornConf, owningTaskName, gateIndex, igdd, bufferPoolFactory, bufferDecompressor);
+ celebornConf, owningTaskName, gateIndex, igdd, bufferPoolFactory, bufferDecompressor);
}
private SupplierWithException<BufferPool, IOException> createBufferPoolFactory(
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 d0aa990d9..85c504a7c 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
@@ -58,7 +58,7 @@ public class RemoteShuffleOutputGate {
private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleOutputGate.class);
private final RemoteShuffleDescriptor shuffleDesc;
protected final int numSubs;
- protected FlinkShuffleClientImpl shuffleWriteClient;
+ protected FlinkShuffleClientImpl flinkShuffleClient;
protected final SupplierWithException<BufferPool, IOException> bufferPoolFactory;
protected BufferPool bufferPool;
private CelebornConf celebornConf;
@@ -104,11 +104,9 @@ public class RemoteShuffleOutputGate {
this.mapId = shuffleDesc.getShuffleResource().getMapPartitionShuffleDescriptor().getMapId();
this.attemptId =
shuffleDesc.getShuffleResource().getMapPartitionShuffleDescriptor().getAttemptId();
- this.rssMetaServiceHost =
- ((RemoteShuffleResource) shuffleDesc.getShuffleResource()).getRssMetaServiceHost();
- this.rssMetaServicePort =
- ((RemoteShuffleResource) shuffleDesc.getShuffleResource()).getRssMetaServicePort();
- this.shuffleWriteClient = createWriteClient();
+ this.rssMetaServiceHost = shuffleDesc.getShuffleResource().getRssMetaServiceHost();
+ this.rssMetaServicePort = shuffleDesc.getShuffleResource().getRssMetaServicePort();
+ this.flinkShuffleClient = getShuffleClient();
}
/** Initialize transportation gate. */
@@ -139,16 +137,16 @@ public class RemoteShuffleOutputGate {
* @param isBroadcast Whether it's a broadcast region.
*/
public void regionStart(boolean isBroadcast) {
- Optional<PartitionLocation> newPartitionLoc = null;
+ Optional<PartitionLocation> newPartitionLoc;
try {
if (isFirstHandShake) {
- handshake(isFirstHandShake);
+ handshake(true);
isFirstHandShake = false;
LOG.info("send firstHandShake:" + isBroadcast);
}
newPartitionLoc =
- shuffleWriteClient.regionStart(
+ flinkShuffleClient.regionStart(
applicationId,
shuffleId,
mapId,
@@ -162,7 +160,7 @@ public class RemoteShuffleOutputGate {
// send handshake again
handshake(false);
// send regionstart again
- shuffleWriteClient.regionStart(
+ flinkShuffleClient.regionStart(
applicationId,
shuffleId,
mapId,
@@ -183,7 +181,7 @@ public class RemoteShuffleOutputGate {
public void regionFinish() throws InterruptedException {
bufferPacker.drain();
try {
- shuffleWriteClient.regionFinish(
+ flinkShuffleClient.regionFinish(
applicationId, shuffleId, mapId, attemptId, partitionLocation);
currentRegionIndex++;
} catch (IOException e) {
@@ -193,7 +191,7 @@ public class RemoteShuffleOutputGate {
/** Indicates the writing/spilling is finished. */
public void finish() throws InterruptedException, IOException {
- shuffleWriteClient.mapPartitionMapperEnd(
+ flinkShuffleClient.mapPartitionMapperEnd(
applicationId, shuffleId, mapId, attemptId, numMappers, partitionLocation.getId());
}
@@ -203,7 +201,7 @@ public class RemoteShuffleOutputGate {
bufferPool.lazyDestroy();
}
bufferPacker.close();
- shuffleWriteClient.cleanup(applicationId, shuffleId, mapId, attemptId);
+ flinkShuffleClient.cleanup(applicationId, shuffleId, mapId, attemptId);
}
/** Returns shuffle descriptor. */
@@ -212,15 +210,15 @@ public class RemoteShuffleOutputGate {
}
@VisibleForTesting
- FlinkShuffleClientImpl createWriteClient() {
+ FlinkShuffleClientImpl getShuffleClient() {
return FlinkShuffleClientImpl.get(
rssMetaServiceHost, rssMetaServicePort, celebornConf, userIdentifier);
}
/** Writes a piece of data to a subpartition. */
- public void write(ByteBuf byteBuf, int subIdx) throws InterruptedException {
+ public void write(ByteBuf byteBuf, int subIdx) {
try {
- shuffleWriteClient.pushDataToLocation(
+ flinkShuffleClient.pushDataToLocation(
applicationId,
shuffleId,
mapId,
@@ -235,17 +233,16 @@ public class RemoteShuffleOutputGate {
}
public void handshake(boolean isFirstHandShake) throws IOException {
- if (partitionLocation == null) {
+ if (isFirstHandShake) {
partitionLocation =
- shuffleWriteClient.registerMapPartitionTask(
+ flinkShuffleClient.registerMapPartitionTask(
applicationId, shuffleId, numMappers, mapId, attemptId);
Utils.checkNotNull(partitionLocation);
- }
- if (isFirstHandShake) {
+
currentRegionIndex = 0;
}
try {
- shuffleWriteClient.pushDataHandShake(
+ flinkShuffleClient.pushDataHandShake(
applicationId, shuffleId, mapId, attemptId, numSubs, bufferSize, partitionLocation);
} catch (IOException e) {
Utils.rethrowAsRuntimeException(e);
diff --git a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGateSuiteJ.java b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGateSuiteJ.java
index 3dd9ec64f..1b72b9c09 100644
--- a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGateSuiteJ.java
+++ b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleOutputGateSuiteJ.java
@@ -48,7 +48,7 @@ public class RemoteShuffleOutputGateSuiteJ {
@Before
public void setup() throws IOException {
- remoteShuffleOutputGate.shuffleWriteClient = shuffleClient;
+ remoteShuffleOutputGate.flinkShuffleClient = shuffleClient;
networkBufferPool = new NetworkBufferPool(10, BUFFER_SIZE);
bufferPool = networkBufferPool.createBufferPool(10, 10);
}
@@ -61,12 +61,12 @@ public class RemoteShuffleOutputGateSuiteJ {
when(shuffleClient.registerMapPartitionTask(any(), anyInt(), anyInt(), anyInt(), anyInt()))
.thenAnswer(t -> partitionLocation);
doNothing()
- .when(remoteShuffleOutputGate.shuffleWriteClient)
+ .when(remoteShuffleOutputGate.flinkShuffleClient)
.pushDataHandShake(anyString(), anyInt(), anyInt(), anyInt(), anyInt(), anyInt(), any());
remoteShuffleOutputGate.handshake(true);
- when(remoteShuffleOutputGate.shuffleWriteClient.regionStart(
+ when(remoteShuffleOutputGate.flinkShuffleClient.regionStart(
any(), anyInt(), anyInt(), anyInt(), any(), anyInt(), anyBoolean()))
.thenAnswer(t -> Optional.empty());
remoteShuffleOutputGate.regionStart(false);
@@ -74,16 +74,16 @@ public class RemoteShuffleOutputGateSuiteJ {
remoteShuffleOutputGate.write(bufferPool.requestBuffer(), 0);
doNothing()
- .when(remoteShuffleOutputGate.shuffleWriteClient)
+ .when(remoteShuffleOutputGate.flinkShuffleClient)
.regionFinish(any(), anyInt(), anyInt(), anyInt(), any());
remoteShuffleOutputGate.regionFinish();
doNothing()
- .when(remoteShuffleOutputGate.shuffleWriteClient)
+ .when(remoteShuffleOutputGate.flinkShuffleClient)
.mapperEnd(any(), anyInt(), anyInt(), anyInt(), anyInt());
remoteShuffleOutputGate.finish();
- doNothing().when(remoteShuffleOutputGate.shuffleWriteClient).shutdown();
+ doNothing().when(remoteShuffleOutputGate.flinkShuffleClient).shutdown();
remoteShuffleOutputGate.close();
}
diff --git a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
index aa9a66070..b05fccb7c 100644
--- a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
+++ b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
@@ -455,7 +455,7 @@ public class RemoteShuffleResultPartitionSuiteJ {
}
@Override
- FlinkShuffleClientImpl createWriteClient() {
+ FlinkShuffleClientImpl getShuffleClient() {
FlinkShuffleClientImpl client = mock(FlinkShuffleClientImpl.class);
doNothing().when(client).cleanup(anyString(), anyInt(), anyInt(), anyInt());
return client;