You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by et...@apache.org on 2023/05/05 15:00:55 UTC
[incubator-celeborn] branch main updated: [CELEBORN-548][FLINK] Support flink 1.17. (#1472)
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 114b1b4d6 [CELEBORN-548][FLINK] Support flink 1.17. (#1472)
114b1b4d6 is described below
commit 114b1b4d62337ff59fa58ec8148bf1415af14c49
Author: Ethan Feng <et...@apache.org>
AuthorDate: Fri May 5 23:00:49 2023 +0800
[CELEBORN-548][FLINK] Support flink 1.17. (#1472)
---
.github/workflows/maven.yml | 2 +
...bstractRemoteShuffleResultPartitionFactory.java | 4 +-
.../celeborn/plugin/flink/RemoteShuffleMaster.java | 9 +-
.../plugin/flink/ResultPartitionAdapter.java | 24 +
.../flink/RemoteShuffleResultPartitionFactory.java | 2 +-
.../plugin/flink/RemoteShuffleServiceFactory.java | 2 +-
.../plugin/flink/SimpleResultPartitionAdapter.java | 28 +
.../plugin/flink/RemoteShuffleMasterTest.java | 3 +-
.../flink/RemoteShuffleResultPartitionFactory.java | 2 +-
.../plugin/flink/RemoteShuffleServiceFactory.java | 2 +-
.../plugin/flink/SimpleResultPartitionAdapter.java | 28 +
.../plugin/flink/RemoteShuffleMasterTest.java | 3 +-
client-flink/flink-1.17-shaded/pom.xml | 128 +++++
client-flink/flink-1.17/pom.xml | 74 +++
.../plugin/flink/RemoteShuffleEnvironment.java | 82 +++
.../plugin/flink/RemoteShuffleInputGate.java | 295 ++++++++++
.../flink/RemoteShuffleInputGateFactory.java | 59 ++
.../plugin/flink/RemoteShuffleResultPartition.java | 210 +++++++
.../flink/RemoteShuffleResultPartitionFactory.java | 2 +-
.../plugin/flink/RemoteShuffleServiceFactory.java | 2 +-
.../plugin/flink/SimpleResultPartitionAdapter.java | 27 +
.../plugin/flink/RemoteShuffleMasterTest.java | 37 +-
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 618 +++++++++++++++++++++
.../flink/RemoteShuffleServiceFactorySuitJ.java | 58 ++
dev/reformat | 1 +
pom.xml | 22 +
.../celeborn/tests/flink/WordCountTest.scala | 7 +-
27 files changed, 1710 insertions(+), 21 deletions(-)
diff --git a/.github/workflows/maven.yml b/.github/workflows/maven.yml
index 1078a4276..567b85b82 100644
--- a/.github/workflows/maven.yml
+++ b/.github/workflows/maven.yml
@@ -126,6 +126,8 @@ jobs:
- 11
flink:
- '1.14'
+ - '1.15'
+ - '1.17'
steps:
- uses: actions/checkout@v2
- name: Setup JDK ${{ matrix.java }}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartitionFactory.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartitionFactory.java
index f486591f5..4b2f65a25 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartitionFactory.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartitionFactory.java
@@ -147,7 +147,7 @@ public abstract class AbstractRemoteShuffleResultPartitionFactory {
new BufferCompressor(networkBufferSize, compressionCodec);
RemoteShuffleDescriptor rsd = (RemoteShuffleDescriptor) shuffleDescriptor;
ResultPartition partition =
- createResmoteShuffleResultPartitionInternal(
+ createRemoteShuffleResultPartitionInternal(
taskNameWithSubtaskAndId,
partitionIndex,
id,
@@ -163,7 +163,7 @@ public abstract class AbstractRemoteShuffleResultPartitionFactory {
return partition;
}
- abstract ResultPartition createResmoteShuffleResultPartitionInternal(
+ abstract ResultPartition createRemoteShuffleResultPartitionInternal(
String taskNameWithSubtaskAndId,
int partitionIndex,
ResultPartitionID id,
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
index 9e86f5168..b196f9b0f 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
@@ -59,9 +59,12 @@ public class RemoteShuffleMaster implements ShuffleMaster<RemoteShuffleDescripto
1,
ThreadUtils.createFactoryWithDefaultExceptionHandler(
"remote-shuffle-master-executor", LOG));
+ private final ResultPartitionAdapter resultPartitionDelegation;
- public RemoteShuffleMaster(ShuffleMasterContext shuffleMasterContext) {
+ public RemoteShuffleMaster(
+ ShuffleMasterContext shuffleMasterContext, ResultPartitionAdapter resultPartitionDelegation) {
this.shuffleMasterContext = shuffleMasterContext;
+ this.resultPartitionDelegation = resultPartitionDelegation;
}
@Override
@@ -198,7 +201,9 @@ public class RemoteShuffleMaster implements ShuffleMaster<RemoteShuffleDescripto
TaskInputsOutputsDescriptor taskInputsOutputsDescriptor) {
for (ResultPartitionType partitionType :
taskInputsOutputsDescriptor.getPartitionTypes().values()) {
- if (!partitionType.isBlocking()) {
+ boolean isBlockingShuffle =
+ resultPartitionDelegation.isBlockingResultPartition(partitionType);
+ if (!isBlockingShuffle) {
throw new RuntimeException(
"Blocking result partition type expected but found " + partitionType);
}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ResultPartitionAdapter.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ResultPartitionAdapter.java
new file mode 100644
index 000000000..6ee9daf62
--- /dev/null
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ResultPartitionAdapter.java
@@ -0,0 +1,24 @@
+/*
+ * 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.plugin.flink;
+
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+
+public interface ResultPartitionAdapter {
+ boolean isBlockingResultPartition(ResultPartitionType partitionType);
+}
diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
index 7d18d5c09..59011c51a 100644
--- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
@@ -46,7 +46,7 @@ public class RemoteShuffleResultPartitionFactory
}
@Override
- public ResultPartition createResmoteShuffleResultPartitionInternal(
+ public ResultPartition createRemoteShuffleResultPartitionInternal(
String taskNameWithSubtaskAndId,
int partitionIndex,
ResultPartitionID id,
diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
index 74deffa40..7776572cc 100644
--- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
@@ -32,7 +32,7 @@ public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFac
@Override
public ShuffleMaster<RemoteShuffleDescriptor> createShuffleMaster(
ShuffleMasterContext shuffleMasterContext) {
- return new RemoteShuffleMaster(shuffleMasterContext);
+ return new RemoteShuffleMaster(shuffleMasterContext, new SimpleResultPartitionAdapter());
}
@Override
diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java
new file mode 100644
index 000000000..e1be22299
--- /dev/null
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java
@@ -0,0 +1,28 @@
+/*
+ * 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.plugin.flink;
+
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+
+public class SimpleResultPartitionAdapter implements ResultPartitionAdapter {
+ @Override
+ public boolean isBlockingResultPartition(ResultPartitionType partitionType) {
+ return partitionType.isBlocking();
+ }
+}
diff --git a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
index dc98b9d07..a4904e516 100644
--- a/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
+++ b/client-flink/flink-1.14/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
@@ -226,7 +226,8 @@ public class RemoteShuffleMasterTest {
public void onFatalError(Throwable throwable) {
System.exit(-1);
}
- });
+ },
+ new SimpleResultPartitionAdapter());
return remoteShuffleMaster;
}
diff --git a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
index 02eb9a806..50e349154 100644
--- a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ b/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
@@ -47,7 +47,7 @@ public class RemoteShuffleResultPartitionFactory
}
@Override
- public ResultPartition createResmoteShuffleResultPartitionInternal(
+ public ResultPartition createRemoteShuffleResultPartitionInternal(
String taskNameWithSubtaskAndId,
int partitionIndex,
ResultPartitionID id,
diff --git a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
index 74deffa40..7776572cc 100644
--- a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ b/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
@@ -32,7 +32,7 @@ public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFac
@Override
public ShuffleMaster<RemoteShuffleDescriptor> createShuffleMaster(
ShuffleMasterContext shuffleMasterContext) {
- return new RemoteShuffleMaster(shuffleMasterContext);
+ return new RemoteShuffleMaster(shuffleMasterContext, new SimpleResultPartitionAdapter());
}
@Override
diff --git a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java b/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java
new file mode 100644
index 000000000..e1be22299
--- /dev/null
+++ b/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java
@@ -0,0 +1,28 @@
+/*
+ * 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.plugin.flink;
+
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+
+public class SimpleResultPartitionAdapter implements ResultPartitionAdapter {
+ @Override
+ public boolean isBlockingResultPartition(ResultPartitionType partitionType) {
+ return partitionType.isBlocking();
+ }
+}
diff --git a/client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java b/client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
index b634d5a5c..a9b80b0c5 100644
--- a/client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
+++ b/client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
@@ -226,7 +226,8 @@ public class RemoteShuffleMasterTest {
public void onFatalError(Throwable throwable) {
System.exit(-1);
}
- });
+ },
+ new SimpleResultPartitionAdapter());
return remoteShuffleMaster;
}
diff --git a/client-flink/flink-1.17-shaded/pom.xml b/client-flink/flink-1.17-shaded/pom.xml
new file mode 100644
index 000000000..6e2da4b02
--- /dev/null
+++ b/client-flink/flink-1.17-shaded/pom.xml
@@ -0,0 +1,128 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+~ 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.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+ <parent>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-parent_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <relativePath>../../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>celeborn-client-flink-1.17-shaded_${scala.binary.version}</artifactId>
+ <packaging>jar</packaging>
+ <name>Celeborn Shaded Client for Flink 1.17</name>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-1.17_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ </dependencies>
+
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-shade-plugin</artifactId>
+ <configuration>
+ <relocations>
+ <relocation>
+ <pattern>com.google.protobuf</pattern>
+ <shadedPattern>${shading.prefix}.com.google.protobuf</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>com.google.common</pattern>
+ <shadedPattern>${shading.prefix}.com.google.common</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>io.netty</pattern>
+ <shadedPattern>${shading.prefix}.io.netty</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.apache.commons</pattern>
+ <shadedPattern>${shading.prefix}.org.apache.commons</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.roaringbitmap</pattern>
+ <shadedPattern>${shading.prefix}.org.roaringbitmap</shadedPattern>
+ </relocation>
+ </relocations>
+ <artifactSet>
+ <includes>
+ <include>org.apache.celeborn:*</include>
+ <include>com.google.protobuf:protobuf-java</include>
+ <include>com.google.guava:guava</include>
+ <include>io.netty:*</include>
+ <include>org.apache.commons:commons-lang3</include>
+ <include>org.roaringbitmap:RoaringBitmap</include>
+ </includes>
+ </artifactSet>
+ <filters>
+ <filter>
+ <artifact>*:*</artifact>
+ <excludes>
+ <exclude>META-INF/*.SF</exclude>
+ <exclude>META-INF/*.DSA</exclude>
+ <exclude>META-INF/*.RSA</exclude>
+ <exclude>**/log4j.properties</exclude>
+ </excludes>
+ </filter>
+ </filters>
+ </configuration>
+ </plugin>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-antrun-plugin</artifactId>
+ <version>${maven.plugin.antrun.version}</version>
+ <executions>
+ <execution>
+ <id>rename-native-library</id>
+ <goals>
+ <goal>run</goal>
+ </goals>
+ <phase>package</phase>
+ <configuration>
+ <target>
+ <echo message="unpacking netty jar"></echo>
+ <unzip dest="${project.build.directory}/unpacked/" src="${project.build.directory}/${artifactId}-${version}.jar"></unzip>
+ <echo message="renaming native epoll library"></echo>
+ <move includeemptydirs="false" todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_epoll_x86_64.so" to="liborg_apache_celeborn_shaded_netty_transport_native_epoll_x86_64.so" type="glob"></mapper>
+ </move>
+ <move includeemptydirs="false" todir="${project.build.directory}/unpacked/META-INF/native">
+ <fileset dir="${project.build.directory}/unpacked/META-INF/native"></fileset>
+ <mapper from="libnetty_transport_native_epoll_aarch_64.so" to="liborg_apache_celeborn_shaded_netty_transport_native_epoll_aarch_64.so.so" type="glob"></mapper>
+ </move>
+ <echo message="deleting native kqueue library"></echo>
+ <delete file="${project.build.directory}/unpacked/META-INF/native/libnetty_transport_native_kqueue_x86_64.jnilib"></delete>
+ <delete file="${project.build.directory}/unpacked/META-INF/native/libnetty_transport_native_kqueue_aarch_64.jnilib"></delete>
+ <delete file="${project.build.directory}/unpacked/META-INF/native/libnetty_resolver_dns_native_macos_aarch_64.jnilib"></delete>
+ <delete file="${project.build.directory}/unpacked/META-INF/native/libnetty_resolver_dns_native_macos_x86_64.jnilib"></delete>
+ <echo message="repackaging netty jar"></echo>
+ <jar basedir="${project.build.directory}/unpacked" destfile="${project.build.directory}/${artifactId}-${version}.jar"></jar>
+ </target>
+ </configuration>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+</project>
diff --git a/client-flink/flink-1.17/pom.xml b/client-flink/flink-1.17/pom.xml
new file mode 100644
index 000000000..9aa586bce
--- /dev/null
+++ b/client-flink/flink-1.17/pom.xml
@@ -0,0 +1,74 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+~ 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.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+ <parent>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-parent_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <relativePath>../../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>celeborn-client-flink-1.17_${scala.binary.version}</artifactId>
+ <packaging>jar</packaging>
+ <name>Celeborn Client for Flink 1.17</name>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-runtime</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <!-- Test dependencies -->
+ <dependency>
+ <groupId>junit</groupId>
+ <artifactId>junit</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.mockito</groupId>
+ <artifactId>mockito-core</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-slf4j-impl</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.logging.log4j</groupId>
+ <artifactId>log4j-1.2-api</artifactId>
+ <scope>test</scope>
+ </dependency>
+ </dependencies>
+</project>
diff --git a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
new file mode 100644
index 000000000..01740f1f3
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
@@ -0,0 +1,82 @@
+/*
+ * 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.plugin.flink;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor;
+import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
+import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate;
+import org.apache.flink.runtime.shuffle.ShuffleEnvironment;
+import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
+
+import org.apache.celeborn.common.CelebornConf;
+
+/**
+ * The implementation of {@link ShuffleEnvironment} based on the remote shuffle service, providing
+ * shuffle environment on flink TM side.
+ */
+public class RemoteShuffleEnvironment extends AbstractRemoteShuffleEnvironment
+ implements ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> {
+
+ /** Factory class to create {@link RemoteShuffleResultPartition}. */
+ private final RemoteShuffleResultPartitionFactory resultPartitionFactory;
+
+ private final RemoteShuffleInputGateFactory inputGateFactory;
+
+ /**
+ * @param networkBufferPool Network buffer pool for shuffle read and shuffle write.
+ * @param resultPartitionManager A trivial {@link ResultPartitionManager}.
+ * @param resultPartitionFactory Factory class to create {@link RemoteShuffleResultPartition}. //
+ * * @param inputGateFactory Factory class to create {@link RemoteShuffleInputGate}.
+ */
+ public RemoteShuffleEnvironment(
+ NetworkBufferPool networkBufferPool,
+ ResultPartitionManager resultPartitionManager,
+ RemoteShuffleResultPartitionFactory resultPartitionFactory,
+ RemoteShuffleInputGateFactory inputGateFactory,
+ CelebornConf conf) {
+
+ super(networkBufferPool, resultPartitionManager, conf);
+ this.resultPartitionFactory = resultPartitionFactory;
+ this.inputGateFactory = inputGateFactory;
+ }
+
+ @Override
+ public ResultPartitionWriter createResultPartitionWriterInternal(
+ ShuffleIOOwnerContext ownerContext,
+ int index,
+ ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor,
+ CelebornConf conf) {
+ return resultPartitionFactory.create(
+ ownerContext.getOwnerName(), index, resultPartitionDeploymentDescriptor, conf);
+ }
+
+ @Override
+ IndexedInputGate createInputGateInternal(
+ ShuffleIOOwnerContext ownerContext, int gateIndex, InputGateDeploymentDescriptor igdd) {
+ return inputGateFactory.create(ownerContext.getOwnerName(), gateIndex, igdd);
+ }
+
+ @VisibleForTesting
+ RemoteShuffleResultPartitionFactory getResultPartitionFactory() {
+ return resultPartitionFactory;
+ }
+}
diff --git a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
new file mode 100644
index 000000000..7437162a4
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
@@ -0,0 +1,295 @@
+/*
+ * 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.plugin.flink;
+
+import java.io.IOException;
+import java.net.InetAddress;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+
+import org.apache.flink.core.memory.MemorySegment;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+import org.apache.flink.runtime.event.TaskEvent;
+import org.apache.flink.runtime.executiongraph.IndexRange;
+import org.apache.flink.runtime.io.network.ConnectionID;
+import org.apache.flink.runtime.io.network.LocalConnectionManager;
+import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+import org.apache.flink.runtime.io.network.partition.consumer.BufferOrEvent;
+import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate;
+import org.apache.flink.runtime.io.network.partition.consumer.InputChannel;
+import org.apache.flink.runtime.io.network.partition.consumer.RemoteInputChannel;
+import org.apache.flink.runtime.io.network.partition.consumer.SingleInputGate;
+import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.runtime.throughput.ThroughputCalculator;
+import org.apache.flink.util.CloseableIterator;
+import org.apache.flink.util.FlinkRuntimeException;
+import org.apache.flink.util.clock.SystemClock;
+import org.apache.flink.util.function.SupplierWithException;
+
+import org.apache.celeborn.common.CelebornConf;
+
+/** A {@link IndexedInputGate} which ingest data from remote shuffle workers. */
+public class RemoteShuffleInputGate extends IndexedInputGate {
+
+ private final RemoteShuffleInputGateDelegation inputGateDelegation;
+
+ public RemoteShuffleInputGate(
+ CelebornConf celebornConf,
+ String taskName,
+ int gateIndex,
+ InputGateDeploymentDescriptor gateDescriptor,
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory,
+ BufferDecompressor bufferDecompressor,
+ int numConcurrentReading) {
+
+ inputGateDelegation =
+ new RemoteShuffleInputGateDelegation(
+ celebornConf,
+ taskName,
+ gateIndex,
+ gateDescriptor,
+ bufferPoolFactory,
+ bufferDecompressor,
+ numConcurrentReading,
+ availabilityHelper,
+ gateDescriptor.getConsumedSubpartitionIndexRange().getStartIndex(),
+ gateDescriptor.getConsumedSubpartitionIndexRange().getEndIndex());
+ }
+
+ /** Setup gate and build network connections. */
+ @Override
+ public void setup() throws IOException {
+ inputGateDelegation.setup();
+ }
+
+ /** Index of the gate of the corresponding computing task. */
+ @Override
+ public int getGateIndex() {
+ return inputGateDelegation.getGateIndex();
+ }
+
+ /** Get number of input channels. A channel is a data flow from one shuffle worker. */
+ @Override
+ public int getNumberOfInputChannels() {
+ return inputGateDelegation.getBufferReaders().size();
+ }
+
+ /** Whether reading is finished -- all channels are finished and cached buffers are drained. */
+ @Override
+ public boolean isFinished() {
+ return inputGateDelegation.isFinished();
+ }
+
+ @Override
+ public Optional<BufferOrEvent> getNext() {
+ throw new UnsupportedOperationException("Not implemented (DataSet API is not supported).");
+ }
+
+ /** Poll a received {@link BufferOrEvent}. */
+ @Override
+ public Optional<BufferOrEvent> pollNext() throws IOException {
+ return inputGateDelegation.pollNext();
+ }
+
+ /** Close all reading channels inside this {@link RemoteShuffleInputGate}. */
+ @Override
+ public void close() throws Exception {
+ inputGateDelegation.close();
+ }
+
+ /** Get {@link InputChannelInfo}s of this {@link RemoteShuffleInputGate}. */
+ @Override
+ public List<InputChannelInfo> getChannelInfos() {
+ return inputGateDelegation.getChannelsInfo();
+ }
+
+ @Override
+ public void requestPartitions() {
+ // do-nothing
+ }
+
+ @Override
+ public void checkpointStarted(CheckpointBarrier barrier) {
+ // do-nothing.
+ }
+
+ @Override
+ public void checkpointStopped(long cancelledCheckpointId) {
+ // do-nothing.
+ }
+
+ @Override
+ public void triggerDebloating() {
+ // do-nothing.
+ }
+
+ @Override
+ public List<InputChannelInfo> getUnfinishedChannels() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public EndOfDataStatus hasReceivedEndOfData() {
+ if (inputGateDelegation.getPendingEndOfDataEvents() > 0) {
+ return EndOfDataStatus.NOT_END_OF_DATA;
+ } else {
+ // Keep compatibility with streaming mode.
+ return EndOfDataStatus.DRAINED;
+ }
+ }
+
+ @Override
+ public void finishReadRecoveredState() {
+ // do-nothing.
+ }
+
+ @Override
+ public InputChannel getChannel(int channelIndex) {
+ return new FakedRemoteInputChannel(channelIndex);
+ }
+
+ @Override
+ public void sendTaskEvent(TaskEvent event) {
+ throw new FlinkRuntimeException("Method should not be called.");
+ }
+
+ @Override
+ public void resumeConsumption(InputChannelInfo channelInfo) {
+ throw new FlinkRuntimeException("Method should not be called.");
+ }
+
+ @Override
+ public void acknowledgeAllRecordsProcessed(InputChannelInfo inputChannelInfo) {}
+
+ @Override
+ public CompletableFuture<Void> getStateConsumedFuture() {
+ return CompletableFuture.completedFuture(null);
+ }
+
+ @Override
+ public String toString() {
+ return String.format(
+ "ReadGate [owning task: %s, gate index: %d, descriptor: %s]",
+ inputGateDelegation.getTaskName(),
+ inputGateDelegation.getGateIndex(),
+ inputGateDelegation.getGateDescriptor().toString());
+ }
+
+ /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
+ private class FakedRemoteInputChannel extends RemoteInputChannel {
+ FakedRemoteInputChannel(int channelIndex) {
+ super(
+ new SingleInputGate(
+ "",
+ inputGateDelegation.getGateIndex(),
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ new IndexRange(0, 0),
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ new FakedMemorySegmentProvider(),
+ 0,
+ new ThroughputCalculator(SystemClock.getInstance()),
+ null),
+ channelIndex,
+ new ResultPartitionID(),
+ 0,
+ new ConnectionID(
+ new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1),
+ 0),
+ new LocalConnectionManager(),
+ 0,
+ 0,
+ 0,
+ new SimpleCounter(),
+ new SimpleCounter(),
+ new FakedChannelStateWriter());
+ }
+ }
+
+ /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
+ private static class FakedMemorySegmentProvider implements MemorySegmentProvider {
+
+ @Override
+ public Collection<MemorySegment> requestUnpooledMemorySegments(int i) throws IOException {
+ return null;
+ }
+
+ @Override
+ public void recycleUnpooledMemorySegments(Collection<MemorySegment> collection)
+ throws IOException {}
+ }
+
+ /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
+ private static class FakedChannelStateWriter implements ChannelStateWriter {
+
+ @Override
+ public void start(long cpId, CheckpointOptions checkpointOptions) {}
+
+ @Override
+ public void addInputData(
+ long cpId, InputChannelInfo info, int startSeqNum, CloseableIterator<Buffer> data) {}
+
+ @Override
+ public void addOutputData(
+ long cpId, ResultSubpartitionInfo info, int startSeqNum, Buffer... data) {}
+
+ @Override
+ public void addOutputDataFuture(
+ long l,
+ ResultSubpartitionInfo resultSubpartitionInfo,
+ int i,
+ CompletableFuture<List<Buffer>> completableFuture)
+ throws IllegalArgumentException {}
+
+ @Override
+ public void finishInput(long checkpointId) {}
+
+ @Override
+ public void finishOutput(long checkpointId) {}
+
+ @Override
+ public void abort(long checkpointId, Throwable cause, boolean cleanup) {}
+
+ @Override
+ public ChannelStateWriteResult getAndRemoveWriteResult(long checkpointId) {
+ return null;
+ }
+
+ @Override
+ public void close() {}
+ }
+}
diff --git a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
new file mode 100644
index 000000000..89eba2b1a
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
@@ -0,0 +1,59 @@
+/*
+ * 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.plugin.flink;
+
+import java.io.IOException;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.util.function.SupplierWithException;
+
+import org.apache.celeborn.common.CelebornConf;
+
+/** Factory class to create {@link RemoteShuffleInputGate}. */
+public class RemoteShuffleInputGateFactory extends AbstractRemoteShuffleInputGateFactory {
+
+ public RemoteShuffleInputGateFactory(
+ Configuration flinkConf,
+ CelebornConf conf,
+ NetworkBufferPool networkBufferPool,
+ int networkBufferSize) {
+ super(flinkConf, conf, networkBufferPool, networkBufferSize);
+ }
+
+ // For testing.
+ protected RemoteShuffleInputGate createInputGate(
+ String owningTaskName,
+ int gateIndex,
+ InputGateDeploymentDescriptor igdd,
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory,
+ BufferDecompressor bufferDecompressor) {
+ return new RemoteShuffleInputGate(
+ this.celebornConf,
+ owningTaskName,
+ gateIndex,
+ igdd,
+ bufferPoolFactory,
+ bufferDecompressor,
+ numConcurrentReading);
+ }
+}
diff --git a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
new file mode 100644
index 000000000..0f79d1ebf
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
@@ -0,0 +1,210 @@
+/*
+ * 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.plugin.flink;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
+
+import javax.annotation.Nullable;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.runtime.checkpoint.CheckpointException;
+import org.apache.flink.runtime.event.AbstractEvent;
+import org.apache.flink.runtime.io.network.api.EndOfData;
+import org.apache.flink.runtime.io.network.api.EndOfPartitionEvent;
+import org.apache.flink.runtime.io.network.api.StopMode;
+import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.Buffer.DataType;
+import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.partition.*;
+import org.apache.flink.util.function.SupplierWithException;
+
+import org.apache.celeborn.plugin.flink.buffer.SortBuffer;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
+import org.apache.celeborn.plugin.flink.utils.Utils;
+
+/**
+ * A {@link ResultPartition} which appends records and events to {@link SortBuffer} and after the
+ * {@link SortBuffer} is full, all data in the {@link SortBuffer} will be copied and spilled to the
+ * remote shuffle service in subpartition index order sequentially. Large records that can not be
+ * appended to an empty {@link SortBuffer} will be spilled directly.
+ */
+public class RemoteShuffleResultPartition extends ResultPartition {
+
+ private final RemoteShuffleResultPartitionDelegation delegation;
+
+ public RemoteShuffleResultPartition(
+ String owningTaskName,
+ int partitionIndex,
+ ResultPartitionID partitionId,
+ ResultPartitionType partitionType,
+ int numSubpartitions,
+ int numTargetKeyGroups,
+ int networkBufferSize,
+ ResultPartitionManager partitionManager,
+ @Nullable BufferCompressor bufferCompressor,
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory,
+ RemoteShuffleOutputGate outputGate) {
+
+ super(
+ owningTaskName,
+ partitionIndex,
+ partitionId,
+ partitionType,
+ numSubpartitions,
+ numTargetKeyGroups,
+ partitionManager,
+ bufferCompressor,
+ bufferPoolFactory);
+
+ delegation =
+ new RemoteShuffleResultPartitionDelegation(
+ networkBufferSize, outputGate, numBuffersOut, numBytesOut, numSubpartitions);
+ }
+
+ @Override
+ public void setup() throws IOException {
+ super.setup();
+ BufferUtils.reserveNumRequiredBuffers(bufferPool, 1);
+ delegation.setup(
+ bufferPool,
+ bufferCompressor,
+ buffer -> canBeCompressed(buffer),
+ () -> checkInProduceState());
+ }
+
+ @Override
+ protected void setupInternal() {
+ // do not need to implement
+ }
+
+ @Override
+ public void emitRecord(ByteBuffer record, int targetSubpartition) throws IOException {
+ delegation.emit(record, targetSubpartition, DataType.DATA_BUFFER, false);
+ }
+
+ @Override
+ public void broadcastRecord(ByteBuffer record) throws IOException {
+ delegation.broadcast(record, DataType.DATA_BUFFER);
+ }
+
+ @Override
+ public void broadcastEvent(AbstractEvent event, boolean isPriorityEvent) throws IOException {
+ Buffer buffer = EventSerializer.toBuffer(event, isPriorityEvent);
+ try {
+ ByteBuffer serializedEvent = buffer.getNioBufferReadable();
+ delegation.broadcast(serializedEvent, buffer.getDataType());
+ } finally {
+ buffer.recycleBuffer();
+ }
+ }
+
+ @Override
+ public void alignedBarrierTimeout(long l) {}
+
+ @Override
+ public void abortCheckpoint(long l, CheckpointException e) {}
+
+ @Override
+ public void finish() throws IOException {
+ Utils.checkState(!isReleased(), "Result partition is already released.");
+ broadcastEvent(EndOfPartitionEvent.INSTANCE, false);
+ delegation.finish();
+ super.finish();
+ }
+
+ @Override
+ public synchronized void close() {
+ delegation.close(() -> super.close());
+ }
+
+ @Override
+ protected void releaseInternal() {
+ // no-op
+ }
+
+ @Override
+ public void flushAll() {
+ delegation.flushAll();
+ }
+
+ @Override
+ public void flush(int subpartitionIndex) {
+ flushAll();
+ }
+
+ @Override
+ public CompletableFuture<?> getAvailableFuture() {
+ return AVAILABLE;
+ }
+
+ @Override
+ public int getNumberOfQueuedBuffers() {
+ return 0;
+ }
+
+ @Override
+ public long getSizeOfQueuedBuffersUnsafe() {
+ return 0;
+ }
+
+ @Override
+ public int getNumberOfQueuedBuffers(int targetSubpartition) {
+ return 0;
+ }
+
+ @Override
+ public ResultSubpartitionView createSubpartitionView(
+ int index, BufferAvailabilityListener availabilityListener) {
+ throw new UnsupportedOperationException("Not supported.");
+ }
+
+ @Override
+ public void notifyEndOfData(StopMode mode) throws IOException {
+ if (!delegation.isEndOfDataNotified()) {
+ broadcastEvent(new EndOfData(mode), false);
+ delegation.setEndOfDataNotified(true);
+ }
+ }
+
+ @Override
+ public CompletableFuture<Void> getAllDataProcessedFuture() {
+ return CompletableFuture.completedFuture(null);
+ }
+
+ @Override
+ public String toString() {
+ return "ResultPartition "
+ + partitionId.toString()
+ + " ["
+ + partitionType
+ + ", "
+ + numSubpartitions
+ + " subpartitions, shuffle-descriptor: "
+ + delegation.getOutputGate().getShuffleDesc()
+ + "]";
+ }
+
+ @VisibleForTesting
+ public RemoteShuffleResultPartitionDelegation getDelegation() {
+ return delegation;
+ }
+}
diff --git a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
similarity index 97%
copy from client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
copy to client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
index 02eb9a806..b5584b9f3 100644
--- a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
@@ -47,7 +47,7 @@ public class RemoteShuffleResultPartitionFactory
}
@Override
- public ResultPartition createResmoteShuffleResultPartitionInternal(
+ ResultPartition createRemoteShuffleResultPartitionInternal(
String taskNameWithSubtaskAndId,
int partitionIndex,
ResultPartitionID id,
diff --git a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
similarity index 96%
copy from client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
copy to client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
index 74deffa40..7776572cc 100644
--- a/client-flink/flink-1.15/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
@@ -32,7 +32,7 @@ public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFac
@Override
public ShuffleMaster<RemoteShuffleDescriptor> createShuffleMaster(
ShuffleMasterContext shuffleMasterContext) {
- return new RemoteShuffleMaster(shuffleMasterContext);
+ return new RemoteShuffleMaster(shuffleMasterContext, new SimpleResultPartitionAdapter());
}
@Override
diff --git a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java
new file mode 100644
index 000000000..3476b8fff
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/SimpleResultPartitionAdapter.java
@@ -0,0 +1,27 @@
+/*
+ * 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.plugin.flink;
+
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+
+public class SimpleResultPartitionAdapter implements ResultPartitionAdapter {
+ @Override
+ public boolean isBlockingResultPartition(ResultPartitionType partitionType) {
+ return partitionType.isBlockingOrBlockingPersistentResultPartition();
+ }
+}
diff --git a/client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
similarity index 87%
copy from client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
copy to client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
index b634d5a5c..9a9a242d3 100644
--- a/client-flink/flink-1.15/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterTest.java
@@ -31,10 +31,13 @@ import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.runtime.clusterframework.types.ResourceID;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.executiongraph.ExecutionGraphID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.apache.flink.runtime.shuffle.JobShuffleContext;
import org.apache.flink.runtime.shuffle.PartitionDescriptor;
import org.apache.flink.runtime.shuffle.ProducerDescriptor;
@@ -170,6 +173,8 @@ public class RemoteShuffleMasterTest {
public void testShuffleMemoryAnnouncing() {
Map<IntermediateDataSetID, Integer> numberOfInputGateChannels = new HashMap<>();
Map<IntermediateDataSetID, Integer> numbersOfResultSubpartitions = new HashMap<>();
+ Map<IntermediateDataSetID, Integer> subPartitionNums = new HashMap<>();
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes = new HashMap<>();
Map<IntermediateDataSetID, ResultPartitionType> resultPartitionTypes = new HashMap<>();
IntermediateDataSetID inputDataSetID0 = new IntermediateDataSetID();
IntermediateDataSetID inputDataSetID1 = new IntermediateDataSetID();
@@ -179,16 +184,29 @@ public class RemoteShuffleMasterTest {
Random random = new Random();
numberOfInputGateChannels.put(inputDataSetID0, random.nextInt(1000));
numberOfInputGateChannels.put(inputDataSetID1, random.nextInt(1000));
- numbersOfResultSubpartitions.put(outputDataSetID0, random.nextInt(1000));
- numbersOfResultSubpartitions.put(outputDataSetID1, random.nextInt(1000));
- numbersOfResultSubpartitions.put(outputDataSetID2, random.nextInt(1000));
+ int subPartitionNum1 = random.nextInt(1000);
+ int subPartitionNum2 = random.nextInt(1000);
+ int subPartitionNum3 = random.nextInt(1000);
+ numbersOfResultSubpartitions.put(outputDataSetID0, subPartitionNum1);
+ numbersOfResultSubpartitions.put(outputDataSetID1, subPartitionNum2);
+ numbersOfResultSubpartitions.put(outputDataSetID2, subPartitionNum3);
+ subPartitionNums.put(outputDataSetID0, subPartitionNum1);
+ subPartitionNums.put(outputDataSetID1, subPartitionNum2);
+ subPartitionNums.put(outputDataSetID2, subPartitionNum3);
+ inputPartitionTypes.put(inputDataSetID0, ResultPartitionType.BLOCKING);
+ inputPartitionTypes.put(inputDataSetID1, ResultPartitionType.BLOCKING);
resultPartitionTypes.put(outputDataSetID0, ResultPartitionType.BLOCKING);
resultPartitionTypes.put(outputDataSetID1, ResultPartitionType.BLOCKING);
resultPartitionTypes.put(outputDataSetID2, ResultPartitionType.BLOCKING);
MemorySize calculated =
remoteShuffleMaster.computeShuffleMemorySizeForTask(
TaskInputsOutputsDescriptor.from(
- numberOfInputGateChannels, numbersOfResultSubpartitions, resultPartitionTypes));
+ 3,
+ numberOfInputGateChannels,
+ numbersOfResultSubpartitions,
+ subPartitionNums,
+ inputPartitionTypes,
+ resultPartitionTypes));
long numBytesPerGate =
FlinkUtils.byteStringValueAsBytes(configuration, PluginConf.MEMORY_PER_INPUT_GATE);
@@ -226,7 +244,8 @@ public class RemoteShuffleMasterTest {
public void onFatalError(Throwable throwable) {
System.exit(-1);
}
- });
+ },
+ new SimpleResultPartitionAdapter());
return remoteShuffleMaster;
}
@@ -256,11 +275,15 @@ public class RemoteShuffleMasterTest {
intermediateResultPartitionId,
ResultPartitionType.BLOCKING,
5,
- 1);
+ 1,
+ false,
+ false);
}
public ProducerDescriptor createProducerDescriptor() throws UnknownHostException {
- ExecutionAttemptID executionAttemptId = new ExecutionAttemptID();
+ ExecutionAttemptID executionAttemptId =
+ new ExecutionAttemptID(
+ new ExecutionGraphID(), new ExecutionVertexID(new JobVertexID(0, 0), 0), 0);
return new ProducerDescriptor(
ResourceID.generate(), executionAttemptId, InetAddress.getLocalHost(), 100);
}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
new file mode 100644
index 000000000..38ccbb56a
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
@@ -0,0 +1,618 @@
+/*
+ * 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.plugin.flink;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.anyBoolean;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.Random;
+import java.util.Set;
+import java.util.stream.IntStream;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.core.memory.MemorySegment;
+import org.apache.flink.core.memory.MemorySegmentFactory;
+import org.apache.flink.runtime.io.network.api.EndOfPartitionEvent;
+import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+import org.apache.flink.util.function.SupplierWithException;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.plugin.flink.buffer.BufferPacker;
+import org.apache.celeborn.plugin.flink.buffer.SortBuffer;
+import org.apache.celeborn.plugin.flink.readclient.FlinkShuffleClientImpl;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
+
+public class RemoteShuffleResultPartitionSuiteJ {
+ private final int networkBufferSize = 32 * 1024;
+ private BufferCompressor bufferCompressor = new BufferCompressor(networkBufferSize, "lz4");
+ private RemoteShuffleOutputGate remoteShuffleOutputGate = mock(RemoteShuffleOutputGate.class);
+ private final String compressCodec = "LZ4";
+ private final CelebornConf conf = new CelebornConf();
+ BufferDecompressor bufferDecompressor = new BufferDecompressor(networkBufferSize, "LZ4");
+
+ private static final int totalBuffers = 1000;
+
+ private static final int bufferSize = 1024;
+
+ private NetworkBufferPool globalBufferPool;
+
+ private BufferPool sortBufferPool;
+
+ private BufferPool nettyBufferPool;
+
+ private RemoteShuffleResultPartition partitionWriter;
+
+ private FakedRemoteShuffleOutputGate outputGate;
+
+ @Before
+ public void setup() {
+ globalBufferPool = new NetworkBufferPool(totalBuffers, bufferSize);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ if (outputGate != null) {
+ outputGate.release();
+ }
+
+ if (sortBufferPool != null) {
+ sortBufferPool.lazyDestroy();
+ }
+ if (nettyBufferPool != null) {
+ nettyBufferPool.lazyDestroy();
+ }
+ assertEquals(totalBuffers, globalBufferPool.getNumberOfAvailableMemorySegments());
+ globalBufferPool.destroy();
+ }
+
+ @Test
+ public void tesSimpleFlush() throws IOException, InterruptedException {
+ List<SupplierWithException<BufferPool, IOException>> bufferPool = createBufferPoolFactory();
+ RemoteShuffleResultPartition remoteShuffleResultPartition =
+ new RemoteShuffleResultPartition(
+ "test",
+ 0,
+ new ResultPartitionID(),
+ ResultPartitionType.BLOCKING,
+ 2,
+ 2,
+ 32 * 1024,
+ new ResultPartitionManager(),
+ bufferCompressor,
+ bufferPool.get(0),
+ remoteShuffleOutputGate);
+ remoteShuffleResultPartition.setup();
+ doNothing().when(remoteShuffleOutputGate).regionStart(anyBoolean());
+ doNothing().when(remoteShuffleOutputGate).regionFinish();
+ when(remoteShuffleOutputGate.getBufferPool()).thenReturn(bufferPool.get(1).get());
+ SortBuffer sortBuffer = remoteShuffleResultPartition.getDelegation().getUnicastSortBuffer();
+ ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[] {1, 2, 3});
+ sortBuffer.append(byteBuffer, 0, Buffer.DataType.DATA_BUFFER);
+ remoteShuffleResultPartition.getDelegation().flushSortBuffer(sortBuffer, true);
+ }
+
+ private List<SupplierWithException<BufferPool, IOException>> createBufferPoolFactory() {
+ NetworkBufferPool networkBufferPool =
+ new NetworkBufferPool(256 * 8, 32 * 1024, Duration.ofMillis(1000));
+
+ int numBuffersPerPartition = 64 * 1024 / 32;
+ int numForResultPartition = numBuffersPerPartition * 7 / 8;
+ int numForOutputGate = numBuffersPerPartition - numForResultPartition;
+
+ List<SupplierWithException<BufferPool, IOException>> factories = new ArrayList<>();
+ factories.add(
+ () -> networkBufferPool.createBufferPool(numForResultPartition, numForResultPartition));
+ factories.add(() -> networkBufferPool.createBufferPool(numForOutputGate, numForOutputGate));
+ return factories;
+ }
+
+ @Test
+ public void testWriteNormalRecordWithCompressionEnabled() throws Exception {
+ testWriteNormalRecord(true);
+ }
+
+ @Test
+ public void testWriteNormalRecordWithCompressionDisabled() throws Exception {
+ testWriteNormalRecord(false);
+ }
+
+ @Test
+ public void testWriteLargeRecord() throws Exception {
+ int numSubpartitions = 2;
+ int numBuffers = 100;
+ initResultPartitionWriter(numSubpartitions, 10, 200, false, conf, 10);
+
+ partitionWriter.setup();
+
+ byte[] dataWritten = new byte[bufferSize * numBuffers];
+ Random random = new Random();
+ random.nextBytes(dataWritten);
+ ByteBuffer recordWritten = ByteBuffer.wrap(dataWritten);
+ partitionWriter.emitRecord(recordWritten, 0);
+ assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.finish();
+ partitionWriter.close();
+
+ List<Buffer> receivedBuffers = outputGate.getReceivedBuffers()[0];
+
+ ByteBuffer recordRead = ByteBuffer.allocate(bufferSize * numBuffers);
+ for (Buffer buffer : receivedBuffers) {
+ if (buffer.isBuffer()) {
+ recordRead.put(
+ buffer.getNioBuffer(
+ BufferUtils.HEADER_LENGTH, buffer.readableBytes() - BufferUtils.HEADER_LENGTH));
+ }
+ }
+ recordWritten.rewind();
+ recordRead.flip();
+ assertEquals(recordWritten, recordRead);
+ }
+
+ @Test
+ public void testBroadcastLargeRecord() throws Exception {
+ int numSubpartitions = 2;
+ int numBuffers = 100;
+ initResultPartitionWriter(numSubpartitions, 10, 200, false, conf, 10);
+
+ partitionWriter.setup();
+
+ byte[] dataWritten = new byte[bufferSize * numBuffers];
+ Random random = new Random();
+ random.nextBytes(dataWritten);
+ ByteBuffer recordWritten = ByteBuffer.wrap(dataWritten);
+ partitionWriter.broadcastRecord(recordWritten);
+ assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.finish();
+ partitionWriter.close();
+
+ ByteBuffer recordRead0 = ByteBuffer.allocate(bufferSize * numBuffers);
+ for (Buffer buffer : outputGate.getReceivedBuffers()[0]) {
+ if (buffer.isBuffer()) {
+ recordRead0.put(
+ buffer.getNioBuffer(
+ BufferUtils.HEADER_LENGTH, buffer.readableBytes() - BufferUtils.HEADER_LENGTH));
+ }
+ }
+ recordWritten.rewind();
+ recordRead0.flip();
+ assertEquals(recordWritten, recordRead0);
+
+ ByteBuffer recordRead1 = ByteBuffer.allocate(bufferSize * numBuffers);
+ for (Buffer buffer : outputGate.getReceivedBuffers()[1]) {
+ if (buffer.isBuffer()) {
+ recordRead1.put(
+ buffer.getNioBuffer(
+ BufferUtils.HEADER_LENGTH, buffer.readableBytes() - BufferUtils.HEADER_LENGTH));
+ }
+ }
+ recordWritten.rewind();
+ recordRead1.flip();
+ assertEquals(recordWritten, recordRead0);
+ }
+
+ @Test
+ public void testFlush() throws Exception {
+ int numSubpartitions = 10;
+
+ initResultPartitionWriter(numSubpartitions, 10, 20, false, conf, 10);
+ partitionWriter.setup();
+
+ partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 0);
+ partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 1);
+ assertEquals(3, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.broadcastRecord(ByteBuffer.allocate(bufferSize));
+ assertEquals(2, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.flush(0);
+ assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 2);
+ partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 3);
+ assertEquals(3, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.flushAll();
+ assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
+
+ partitionWriter.finish();
+ partitionWriter.close();
+ }
+
+ private void testWriteNormalRecord(boolean compressionEnabled) throws Exception {
+ int numSubpartitions = 4;
+ int numRecords = 100;
+ Random random = new Random();
+
+ initResultPartitionWriter(numSubpartitions, 100, 500, compressionEnabled, conf, 10);
+ partitionWriter.setup();
+ assertTrue(outputGate.isSetup());
+
+ Queue<DataAndType>[] dataWritten = new Queue[numSubpartitions];
+ IntStream.range(0, numSubpartitions).forEach(i -> dataWritten[i] = new ArrayDeque<>());
+ int[] numBytesWritten = new int[numSubpartitions];
+ Arrays.fill(numBytesWritten, 0);
+
+ for (int i = 0; i < numRecords; i++) {
+ byte[] data = new byte[random.nextInt(2 * bufferSize) + 1];
+ if (compressionEnabled) {
+ byte randomByte = (byte) random.nextInt();
+ Arrays.fill(data, randomByte);
+ } else {
+ random.nextBytes(data);
+ }
+ ByteBuffer record = ByteBuffer.wrap(data);
+ boolean isBroadCast = random.nextBoolean();
+
+ if (isBroadCast) {
+ partitionWriter.broadcastRecord(record);
+ IntStream.range(0, numSubpartitions)
+ .forEach(
+ subpartition ->
+ recordDataWritten(
+ record,
+ Buffer.DataType.DATA_BUFFER,
+ subpartition,
+ dataWritten,
+ numBytesWritten));
+ } else {
+ int subpartition = random.nextInt(numSubpartitions);
+ partitionWriter.emitRecord(record, subpartition);
+ recordDataWritten(
+ record, Buffer.DataType.DATA_BUFFER, subpartition, dataWritten, numBytesWritten);
+ }
+ }
+
+ partitionWriter.finish();
+ assertTrue(outputGate.isFinished());
+ partitionWriter.close();
+ assertTrue(outputGate.isClosed());
+
+ for (int subpartition = 0; subpartition < numSubpartitions; ++subpartition) {
+ ByteBuffer record = EventSerializer.toSerializedEvent(EndOfPartitionEvent.INSTANCE);
+ recordDataWritten(
+ record, Buffer.DataType.EVENT_BUFFER, subpartition, dataWritten, numBytesWritten);
+ }
+
+ outputGate
+ .getFinishedRegions()
+ .forEach(
+ regionIndex -> assertTrue(outputGate.getNumBuffersByRegion().containsKey(regionIndex)));
+
+ int[] numBytesRead = new int[numSubpartitions];
+ List<Buffer>[] receivedBuffers = outputGate.getReceivedBuffers();
+ List<Buffer>[] validateTarget = new List[numSubpartitions];
+ Arrays.fill(numBytesRead, 0);
+ for (int i = 0; i < numSubpartitions; i++) {
+ validateTarget[i] = new ArrayList<>();
+ for (Buffer buffer : receivedBuffers[i]) {
+ for (Buffer unpackedBuffer : BufferPacker.unpack(buffer.asByteBuf())) {
+ if (compressionEnabled && unpackedBuffer.isCompressed()) {
+ Buffer decompressedBuffer =
+ bufferDecompressor.decompressToIntermediateBuffer(unpackedBuffer);
+ ByteBuffer decompressed = decompressedBuffer.getNioBufferReadable();
+ int numBytes = decompressed.remaining();
+ MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(numBytes);
+ segment.put(0, decompressed, numBytes);
+ decompressedBuffer.recycleBuffer();
+ validateTarget[i].add(
+ new NetworkBuffer(segment, buf -> {}, unpackedBuffer.getDataType(), numBytes));
+ numBytesRead[i] += numBytes;
+ } else {
+ numBytesRead[i] += buffer.readableBytes();
+ validateTarget[i].add(buffer);
+ }
+ }
+ }
+ }
+ IntStream.range(0, numSubpartitions).forEach(subpartitions -> {});
+ checkWriteReadResult(
+ numSubpartitions, numBytesWritten, numBytesWritten, dataWritten, validateTarget);
+ }
+
+ private void initResultPartitionWriter(
+ int numSubpartitions,
+ int sortBufferPoolSize,
+ int nettyBufferPoolSize,
+ boolean compressionEnabled,
+ CelebornConf conf,
+ int numMappers)
+ throws Exception {
+
+ sortBufferPool = globalBufferPool.createBufferPool(sortBufferPoolSize, sortBufferPoolSize);
+ nettyBufferPool = globalBufferPool.createBufferPool(nettyBufferPoolSize, nettyBufferPoolSize);
+
+ outputGate =
+ new FakedRemoteShuffleOutputGate(
+ getShuffleDescriptor(), numSubpartitions, () -> nettyBufferPool, conf, numMappers);
+ outputGate.setup();
+
+ if (compressionEnabled) {
+ partitionWriter =
+ new RemoteShuffleResultPartition(
+ "RemoteShuffleResultPartitionWriterTest",
+ 0,
+ new ResultPartitionID(),
+ ResultPartitionType.BLOCKING,
+ numSubpartitions,
+ numSubpartitions,
+ bufferSize,
+ new ResultPartitionManager(),
+ bufferCompressor,
+ () -> sortBufferPool,
+ outputGate);
+ } else {
+ partitionWriter =
+ new RemoteShuffleResultPartition(
+ "RemoteShuffleResultPartitionWriterTest",
+ 0,
+ new ResultPartitionID(),
+ ResultPartitionType.BLOCKING,
+ numSubpartitions,
+ numSubpartitions,
+ bufferSize,
+ new ResultPartitionManager(),
+ null,
+ () -> sortBufferPool,
+ outputGate);
+ }
+ }
+
+ private void recordDataWritten(
+ ByteBuffer record,
+ Buffer.DataType dataType,
+ int subpartition,
+ Queue<DataAndType>[] dataWritten,
+ int[] numBytesWritten) {
+
+ record.rewind();
+ dataWritten[subpartition].add(new DataAndType(record, dataType));
+ numBytesWritten[subpartition] += record.remaining();
+ }
+
+ private static class FakedRemoteShuffleOutputGate extends RemoteShuffleOutputGate {
+
+ private boolean isSetup;
+ private boolean isFinished;
+ private boolean isClosed;
+ private final List<Buffer>[] receivedBuffers;
+ private final Map<Integer, Integer> numBuffersByRegion;
+ private final Set<Integer> finishedRegions;
+ private int currentRegionIndex;
+ private boolean currentIsBroadcast;
+
+ FakedRemoteShuffleOutputGate(
+ RemoteShuffleDescriptor shuffleDescriptor,
+ int numSubpartitions,
+ SupplierWithException<BufferPool, IOException> bufferPoolFactory,
+ CelebornConf celebornConf,
+ int numMappers) {
+
+ super(
+ shuffleDescriptor,
+ numSubpartitions,
+ bufferSize,
+ bufferPoolFactory,
+ celebornConf,
+ numMappers);
+ isSetup = false;
+ isFinished = false;
+ isClosed = false;
+ numBuffersByRegion = new HashMap<>();
+ finishedRegions = new HashSet<>();
+ currentRegionIndex = -1;
+ receivedBuffers = new ArrayList[numSubpartitions];
+ IntStream.range(0, numSubpartitions).forEach(i -> receivedBuffers[i] = new ArrayList<>());
+ currentIsBroadcast = false;
+ }
+
+ @Override
+ FlinkShuffleClientImpl getShuffleClient() {
+ FlinkShuffleClientImpl client = mock(FlinkShuffleClientImpl.class);
+ doNothing().when(client).cleanup(anyString(), anyInt(), anyInt(), anyInt());
+ return client;
+ }
+
+ @Override
+ public void setup() throws IOException, InterruptedException {
+ bufferPool = bufferPoolFactory.get();
+ isSetup = true;
+ }
+
+ @Override
+ public void write(Buffer buffer, int subIdx) {
+ if (currentIsBroadcast) {
+ assertEquals(0, subIdx);
+ ByteBuffer byteBuffer = buffer.getNioBufferReadable();
+ for (int i = 0; i < numSubs; i++) {
+ int numBytes = buffer.readableBytes();
+ MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(numBytes);
+ byteBuffer.rewind();
+ segment.put(0, byteBuffer, numBytes);
+ receivedBuffers[i].add(
+ new NetworkBuffer(
+ segment, buf -> {}, buffer.getDataType(), buffer.isCompressed(), numBytes));
+ }
+ buffer.recycleBuffer();
+ } else {
+ receivedBuffers[subIdx].add(buffer);
+ }
+ if (numBuffersByRegion.containsKey(currentRegionIndex)) {
+ int prev = numBuffersByRegion.get(currentRegionIndex);
+ numBuffersByRegion.put(currentRegionIndex, prev + 1);
+ } else {
+ numBuffersByRegion.put(currentRegionIndex, 1);
+ }
+ }
+
+ @Override
+ public void regionStart(boolean isBroadcast) {
+ currentIsBroadcast = isBroadcast;
+ currentRegionIndex++;
+ }
+
+ @Override
+ public void regionFinish() {
+ if (finishedRegions.contains(currentRegionIndex)) {
+ throw new IllegalStateException("Unexpected region: " + currentRegionIndex);
+ }
+ finishedRegions.add(currentRegionIndex);
+ }
+
+ @Override
+ public void finish() throws InterruptedException {
+ isFinished = true;
+ }
+
+ @Override
+ public void close() {
+ isClosed = true;
+ }
+
+ public List<Buffer>[] getReceivedBuffers() {
+ return receivedBuffers;
+ }
+
+ public Map<Integer, Integer> getNumBuffersByRegion() {
+ return numBuffersByRegion;
+ }
+
+ public Set<Integer> getFinishedRegions() {
+ return finishedRegions;
+ }
+
+ public boolean isSetup() {
+ return isSetup;
+ }
+
+ public boolean isFinished() {
+ return isFinished;
+ }
+
+ public boolean isClosed() {
+ return isClosed;
+ }
+
+ public void release() throws Exception {
+ IntStream.range(0, numSubs)
+ .forEach(
+ subpartitionIndex -> {
+ receivedBuffers[subpartitionIndex].forEach(Buffer::recycleBuffer);
+ receivedBuffers[subpartitionIndex].clear();
+ });
+ numBuffersByRegion.clear();
+ finishedRegions.clear();
+ super.close();
+ }
+ }
+
+ private RemoteShuffleDescriptor getShuffleDescriptor() throws Exception {
+ Random random = new Random();
+ byte[] bytes = new byte[16];
+ random.nextBytes(bytes);
+ return new RemoteShuffleDescriptor(
+ new JobID(bytes).toString(),
+ new JobID(bytes),
+ new JobID(bytes).toString(),
+ new ResultPartitionID(),
+ new RemoteShuffleResource("1", 2, new ShuffleResourceDescriptor(1, 1, 1, 0)));
+ }
+
+ /** Data written and its {@link Buffer.DataType}. */
+ public static class DataAndType {
+ private final ByteBuffer data;
+ private final Buffer.DataType dataType;
+
+ DataAndType(ByteBuffer data, Buffer.DataType dataType) {
+ this.data = data;
+ this.dataType = dataType;
+ }
+ }
+
+ public static void checkWriteReadResult(
+ int numSubpartitions,
+ int[] numBytesWritten,
+ int[] numBytesRead,
+ Queue<DataAndType>[] dataWritten,
+ Collection<Buffer>[] buffersRead) {
+ for (int subpartitionIndex = 0; subpartitionIndex < numSubpartitions; ++subpartitionIndex) {
+ assertEquals(numBytesWritten[subpartitionIndex], numBytesRead[subpartitionIndex]);
+
+ List<DataAndType> eventsWritten = new ArrayList<>();
+ List<Buffer> eventsRead = new ArrayList<>();
+
+ ByteBuffer subpartitionDataWritten = ByteBuffer.allocate(numBytesWritten[subpartitionIndex]);
+ for (DataAndType dataAndType : dataWritten[subpartitionIndex]) {
+ subpartitionDataWritten.put(dataAndType.data);
+ dataAndType.data.rewind();
+ if (dataAndType.dataType.isEvent()) {
+ eventsWritten.add(dataAndType);
+ }
+ }
+
+ ByteBuffer subpartitionDataRead = ByteBuffer.allocate(numBytesRead[subpartitionIndex]);
+ for (Buffer buffer : buffersRead[subpartitionIndex]) {
+ subpartitionDataRead.put(buffer.getNioBufferReadable());
+ if (!buffer.isBuffer()) {
+ eventsRead.add(buffer);
+ }
+ }
+
+ subpartitionDataWritten.flip();
+ subpartitionDataRead.flip();
+ assertEquals(subpartitionDataWritten, subpartitionDataRead);
+
+ assertEquals(eventsWritten.size(), eventsRead.size());
+ for (int i = 0; i < eventsWritten.size(); ++i) {
+ assertEquals(eventsWritten.get(i).dataType, eventsRead.get(i).getDataType());
+ assertEquals(eventsWritten.get(i).data, eventsRead.get(i).getNioBufferReadable());
+ }
+ }
+ }
+}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuitJ.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuitJ.java
new file mode 100644
index 000000000..9a4f232f7
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuitJ.java
@@ -0,0 +1,58 @@
+/*
+ * 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.plugin.flink;
+
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.MemorySize;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
+import org.apache.flink.runtime.io.network.partition.consumer.IndexedInputGate;
+import org.apache.flink.runtime.shuffle.ShuffleEnvironment;
+import org.apache.flink.runtime.shuffle.ShuffleEnvironmentContext;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class RemoteShuffleServiceFactorySuitJ {
+ @Test
+ public void testCreateShuffleEnvironment() {
+ RemoteShuffleServiceFactory remoteShuffleServiceFactory = new RemoteShuffleServiceFactory();
+ ShuffleEnvironmentContext shuffleEnvironmentContext = mock(ShuffleEnvironmentContext.class);
+ when(shuffleEnvironmentContext.getConfiguration()).thenReturn(new Configuration());
+ when(shuffleEnvironmentContext.getNetworkMemorySize())
+ .thenReturn(new MemorySize(64 * 1024 * 1024));
+ MetricGroup parentMeric = mock(MetricGroup.class);
+ when(shuffleEnvironmentContext.getParentMetricGroup()).thenReturn(parentMeric);
+ MetricGroup childGroup = mock(MetricGroup.class);
+ MetricGroup childChildGroup = mock(MetricGroup.class);
+ when(parentMeric.addGroup(anyString())).thenReturn(childGroup);
+ when(childGroup.addGroup(any())).thenReturn(childChildGroup);
+ when(childChildGroup.gauge(any(), any())).thenReturn(null);
+ ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> shuffleEnvironment =
+ remoteShuffleServiceFactory.createShuffleEnvironment(shuffleEnvironmentContext);
+ Assert.assertEquals(
+ 32 * 1024,
+ ((RemoteShuffleEnvironment) shuffleEnvironment)
+ .getResultPartitionFactory()
+ .getNetworkBufferSize());
+ }
+}
diff --git a/dev/reformat b/dev/reformat
index 26d74fed1..e7ff64ab3 100755
--- a/dev/reformat
+++ b/dev/reformat
@@ -21,5 +21,6 @@ set -x
PROJECT_DIR="$(cd "`dirname "$0"`/.."; pwd)"
${PROJECT_DIR}/build/mvn spotless:apply -Pflink-1.14
${PROJECT_DIR}/build/mvn spotless:apply -Pflink-1.15
+${PROJECT_DIR}/build/mvn spotless:apply -Pflink-1.17
${PROJECT_DIR}/build/mvn spotless:apply -Pspark-2.4
${PROJECT_DIR}/build/mvn spotless:apply -Pspark-3.3
diff --git a/pom.xml b/pom.xml
index 4a264d3e4..4b1d274fc 100644
--- a/pom.xml
+++ b/pom.xml
@@ -1084,6 +1084,27 @@
</properties>
</profile>
+ <profile>
+ <id>flink-1.17</id>
+ <modules>
+ <module>client-flink/common</module>
+ <module>client-flink/flink-1.17</module>
+ <module>client-flink/flink-1.17-shaded</module>
+ <module>tests/flink-it</module>
+ </modules>
+ <properties>
+ <flink.version>1.17.0</flink.version>
+ <flink.binary.version>1.17</flink.binary.version>
+ <scala.version>2.12.7</scala.version>
+ <scala.binary.version>2.12</scala.binary.version>
+ <celeborn.flink.plugin.artifact>celeborn-client-flink-1.17_${scala.binary.version}</celeborn.flink.plugin.artifact>
+ <flink.streamig.artifact>flink-streaming-java</flink.streamig.artifact>
+ <flink.clients.artifact>flink-clients</flink.clients.artifact>
+ <flink.scala.artifact>flink-scala_${scala.binary.version}</flink.scala.artifact>
+ <flink.runtime.web.artifact>flink-runtime-web</flink.runtime.web.artifact>
+ </properties>
+ </profile>
+
<profile>
<id>google-mirror</id>
<properties>
@@ -1106,6 +1127,7 @@
<execution>
<id>source-release-assembly</id>
<phase>none</phase>
+
</execution>
</executions>
</plugin>
diff --git a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
index 14d247ba4..e18a47cd8 100644
--- a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
+++ b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.tests.flink;
+package org.apache.celeborn.tests.flink
import org.apache.flink.api.common.{ExecutionMode, InputDependencyConstraint, RuntimeExecutionMode}
import org.apache.flink.configuration.{ConfigConstants, Configuration, ExecutionOptions, RestOptions}
@@ -26,7 +26,7 @@ import org.scalatest.BeforeAndAfterAll
import org.scalatest.funsuite.AnyFunSuite
import org.apache.celeborn.common.internal.Logging
-import org.apache.celeborn.service.deploy.MiniClusterFeature;
+import org.apache.celeborn.service.deploy.MiniClusterFeature
class WordCountTest extends AnyFunSuite with Logging with MiniClusterFeature
with BeforeAndAfterAll {
@@ -58,6 +58,9 @@ class WordCountTest extends AnyFunSuite with Logging with MiniClusterFeature
configuration.set(ExecutionOptions.RUNTIME_MODE, RuntimeExecutionMode.BATCH)
configuration.setString("taskmanager.memory.network.min", "1024m")
configuration.setString(RestOptions.BIND_PORT, "8081-8089")
+ configuration.setString(
+ "execution.batch.adaptive.auto-parallelism.min-parallelism",
+ "" + parallelism)
val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(configuration)
env.getConfig.setExecutionMode(ExecutionMode.BATCH)
env.getConfig.setParallelism(parallelism)