You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by ni...@apache.org on 2026/01/22 02:35:55 UTC
(celeborn) branch main updated: [CELEBORN-2251] Introducing a shim layer and a common-tiered submodule for Flink clients
This is an automated email from the ASF dual-hosted git repository.
nicholasjiang pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new f04ddddc2 [CELEBORN-2251] Introducing a shim layer and a common-tiered submodule for Flink clients
f04ddddc2 is described below
commit f04ddddc2d4a7f5311e4ef802cb49d84cb0eea95
Author: luogen.lg <lu...@alibaba-inc.com>
AuthorDate: Thu Jan 22 10:35:44 2026 +0800
[CELEBORN-2251] Introducing a shim layer and a common-tiered submodule for Flink clients
### What changes were proposed in this pull request?
1. Introduce a client-flink-common-tiered submodule to host the shared tiered shuffle logic.
2. Introduce a shim layer to further unify the implementations and make version-specific changes more explicit.
3. Unify the tests as well, so that versioned clients can simply run the common test suite with their own shims.
### Why are the changes needed?
Though Celeborn now has a client-flink-common submodule, clients still have to copy a lot of code from version to version, with small but necessary changes buried in the duplicates. The tests don’t share a common implementation either. Even worse, with Flink introducing tiered shuffle from 1.20 onward, all tiered shuffle implementations must be placed inside the versioned clients to ensure that client-flink-common can be compiled against all Flink versions. This makes it much harder t [...]
For the maintenance of Flink clients in the future, we need better organization of the submodules, consolidating the implementations and tests.
### Does this PR resolve a correctness bug?
No.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The patch consolidates all current Flink client tests, and can be tested with them.
Closes #3589 from pltbkd/flink-shim-0.7.
Lead-authored-by: luogen.lg <lu...@alibaba-inc.com>
Co-authored-by: SteNicholas <pr...@163.com>
Signed-off-by: SteNicholas <pr...@163.com>
---
client-flink/{flink-2.0 => common-tiered}/pom.xml | 21 +-
.../flink/tiered/BaseCelebornTierFactory.java} | 70 +--
.../flink/tiered/CelebornChannelBufferManager.java | 0
.../flink/tiered/CelebornChannelBufferReader.java | 8 +-
.../plugin/flink/tiered/CelebornFlinkTierShim.java | 37 ++
.../flink/tiered/CelebornTierConsumerAgent.java | 0
.../flink/tiered/CelebornTierMasterAgent.java | 41 +-
.../flink/tiered/CelebornTierProducerAgent.java | 2 +-
.../flink/tiered/TierShuffleDescriptorImpl.java | 0
.../tiered/CelebornTierMasterAgentSuiteJ.java | 5 +-
.../celeborn/plugin/flink/CelebornFlinkShim.java | 172 ++++++
.../plugin/flink/RemoteShuffleEnvironment.java | 12 +-
...eInputGate.java => RemoteShuffleInputGate.java} | 26 +-
...ory.java => RemoteShuffleInputGateFactory.java} | 35 +-
...tion.java => RemoteShuffleResultPartition.java} | 12 +-
...va => RemoteShuffleResultPartitionFactory.java} | 39 +-
...ctory.java => RemoteShuffleServiceFactory.java} | 31 +-
.../plugin/flink/ResultPartitionBytesCounter.java | 28 +
.../plugin/flink/utils/ReflectionUtils.java | 84 +++
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 40 +-
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 7 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 21 +-
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 8 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 27 +-
.../apache/celeborn/plugin/flink/TestUtils.java | 46 ++
client-flink/flink-1.16/pom.xml | 7 +
.../plugin/flink/CelebornFlinkShimV116.java | 155 +++++
.../plugin/flink/RemoteShuffleInputGate.java | 181 ------
.../flink/RemoteShuffleInputGateFactory.java | 62 --
.../plugin/flink/RemoteShuffleResultPartition.java | 73 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 89 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 63 --
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 366 ------------
.../flink/RemoteShuffleMasterSuiteJV116.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...oteShuffleResultPartitionFactorySuiteJV116.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 638 --------------------
.../RemoteShuffleResultPartitionSuiteJV116.java | 20 +
.../flink/RemoteShuffleServiceFactorySuitJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV116.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 138 -----
.../flink/ShuffleResourceTrackerSuiteJV116.java | 20 +
client-flink/flink-1.17/pom.xml | 7 +
.../plugin/flink/CelebornFlinkShimV117.java | 164 ++++++
.../plugin/flink/RemoteShuffleInputGate.java | 181 ------
.../flink/RemoteShuffleInputGateFactory.java | 63 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 89 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 63 --
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 378 ------------
.../flink/RemoteShuffleMasterSuiteJV117.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...oteShuffleResultPartitionFactorySuiteJV117.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 638 --------------------
.../RemoteShuffleResultPartitionSuiteJV117.java | 20 +
.../RemoteShuffleServiceFactorySuiteJV117.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 138 -----
.../flink/ShuffleResourceTrackerSuiteJV117.java | 20 +
client-flink/flink-1.18/pom.xml | 7 +
.../plugin/flink/CelebornFlinkShimV118.java | 167 ++++++
.../plugin/flink/RemoteShuffleInputGate.java | 187 ------
.../flink/RemoteShuffleInputGateFactory.java | 63 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 89 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 63 --
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../flink/RemoteShuffleMasterSuiteJV118.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...oteShuffleResultPartitionFactorySuiteJV118.java | 22 +
.../RemoteShuffleResultPartitionSuiteJV118.java | 20 +
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV118.java | 20 +
.../flink/ShuffleResourceTrackerSuiteJV118.java | 20 +
client-flink/flink-1.19/pom.xml | 7 +
.../plugin/flink/CelebornFlinkShimV119.java | 166 ++++++
.../plugin/flink/RemoteShuffleInputGate.java | 195 -------
.../flink/RemoteShuffleInputGateFactory.java | 63 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 89 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 63 --
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 381 ------------
.../flink/RemoteShuffleMasterSuiteJV119.java | 20 +
...oteShuffleResultPartitionFactorySuiteJV119.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 640 --------------------
.../RemoteShuffleResultPartitionSuiteJV119.java | 20 +
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV119.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 138 -----
.../flink/ShuffleResourceTrackerSuiteJV119.java | 20 +
client-flink/flink-1.20/pom.xml | 19 +
.../plugin/flink/CelebornFlinkShimV120.java | 207 +++++++
.../plugin/flink/RemoteShuffleInputGate.java | 185 ------
.../flink/RemoteShuffleInputGateFactory.java | 64 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 93 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 61 --
.../plugin/flink/tiered/CelebornTierFactory.java | 56 +-
.../flink/tiered/CelebornTierMasterAgent.java | 219 -------
.../flink/tiered/CelebornTierProducerAgent.java | 503 ----------------
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 395 -------------
.../flink/RemoteShuffleMasterSuiteJV120.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...oteShuffleResultPartitionFactorySuiteJV120.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 643 ---------------------
.../RemoteShuffleResultPartitionSuiteJV120.java | 20 +
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV120.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 152 -----
.../flink/ShuffleResourceTrackerSuiteJV120.java | 20 +
.../tiered/CelebornTierMasterAgentSuiteJV120.java | 20 +
client-flink/flink-2.0/pom.xml | 19 +
.../plugin/flink/CelebornFlinkShimV20.java | 207 +++++++
.../plugin/flink/RemoteShuffleInputGate.java | 195 -------
.../flink/RemoteShuffleInputGateFactory.java | 64 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 93 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 61 --
.../flink/tiered/CelebornChannelBufferManager.java | 165 ------
.../flink/tiered/CelebornChannelBufferReader.java | 337 -----------
.../flink/tiered/CelebornTierConsumerAgent.java | 556 ------------------
.../plugin/flink/tiered/CelebornTierFactory.java | 61 +-
.../flink/tiered/CelebornTierMasterAgent.java | 219 -------
.../flink/tiered/CelebornTierProducerAgent.java | 503 ----------------
.../flink/tiered/TierShuffleDescriptorImpl.java | 72 ---
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 395 -------------
.../plugin/flink/RemoteShuffleMasterSuiteJV20.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...moteShuffleResultPartitionFactorySuiteJV20.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 643 ---------------------
.../RemoteShuffleResultPartitionSuiteJV20.java | 20 +
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV20.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 152 -----
.../flink/ShuffleResourceTrackerSuiteJV20.java | 20 +
.../tiered/CelebornTierMasterAgentSuiteJ.java | 201 -------
.../tiered/CelebornTierMasterAgentSuiteJV20.java | 20 +
client-flink/flink-2.1/pom.xml | 19 +
.../plugin/flink/CelebornFlinkShimV21.java | 207 +++++++
.../plugin/flink/RemoteShuffleInputGate.java | 195 -------
.../flink/RemoteShuffleInputGateFactory.java | 64 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 93 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 61 --
.../flink/tiered/CelebornChannelBufferManager.java | 165 ------
.../flink/tiered/CelebornChannelBufferReader.java | 337 -----------
.../flink/tiered/CelebornTierConsumerAgent.java | 556 ------------------
.../plugin/flink/tiered/CelebornTierFactory.java | 61 +-
.../flink/tiered/TierShuffleDescriptorImpl.java | 72 ---
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 395 -------------
.../plugin/flink/RemoteShuffleMasterSuiteJV21.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...moteShuffleResultPartitionFactorySuiteJV21.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 643 ---------------------
.../RemoteShuffleResultPartitionSuiteJV21.java | 20 +
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV21.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 152 -----
.../flink/ShuffleResourceTrackerSuiteJV21.java | 20 +
.../tiered/CelebornTierMasterAgentSuiteJ.java | 201 -------
.../tiered/CelebornTierMasterAgentSuiteJV21.java | 20 +
client-flink/flink-2.2/pom.xml | 19 +
.../plugin/flink/CelebornFlinkShimV22.java | 207 +++++++
.../plugin/flink/RemoteShuffleInputGate.java | 195 -------
.../flink/RemoteShuffleInputGateFactory.java | 64 --
.../plugin/flink/RemoteShuffleResultPartition.java | 77 ---
.../flink/RemoteShuffleResultPartitionFactory.java | 93 ---
.../plugin/flink/RemoteShuffleServiceFactory.java | 61 --
.../flink/tiered/CelebornChannelBufferManager.java | 165 ------
.../flink/tiered/CelebornChannelBufferReader.java | 337 -----------
.../flink/tiered/CelebornTierConsumerAgent.java | 556 ------------------
.../plugin/flink/tiered/CelebornTierFactory.java | 61 +-
.../flink/tiered/CelebornTierMasterAgent.java | 219 -------
.../flink/tiered/CelebornTierProducerAgent.java | 503 ----------------
.../flink/tiered/TierShuffleDescriptorImpl.java | 72 ---
....apache.celeborn.plugin.flink.CelebornFlinkShim | 16 +
.../plugin/flink/RemoteShuffleMasterSuiteJ.java | 395 -------------
.../plugin/flink/RemoteShuffleMasterSuiteJV22.java | 20 +
.../RemoteShuffleResultPartitionFactorySuiteJ.java | 51 --
...moteShuffleResultPartitionFactorySuiteJV22.java | 22 +
.../flink/RemoteShuffleResultPartitionSuiteJ.java | 643 ---------------------
.../RemoteShuffleResultPartitionSuiteJV22.java | 20 +
.../flink/RemoteShuffleServiceFactorySuiteJ.java | 58 --
.../RemoteShuffleServiceFactorySuiteJV22.java | 20 +
.../plugin/flink/ShuffleResourceTrackerSuiteJ.java | 152 -----
.../flink/ShuffleResourceTrackerSuiteJV22.java | 20 +
.../tiered/CelebornTierMasterAgentSuiteJ.java | 201 -------
.../tiered/CelebornTierMasterAgentSuiteJV22.java | 20 +
pom.xml | 4 +
project/CelebornBuild.scala | 40 +-
195 files changed, 3215 insertions(+), 19442 deletions(-)
diff --git a/client-flink/flink-2.0/pom.xml b/client-flink/common-tiered/pom.xml
similarity index 84%
copy from client-flink/flink-2.0/pom.xml
copy to client-flink/common-tiered/pom.xml
index c5ef0d677..40c0422ad 100644
--- a/client-flink/flink-2.0/pom.xml
+++ b/client-flink/common-tiered/pom.xml
@@ -24,24 +24,24 @@
<relativePath>../../pom.xml</relativePath>
</parent>
- <artifactId>celeborn-client-flink-2.0_${scala.binary.version}</artifactId>
+ <artifactId>celeborn-client-flink-common-tiered_${scala.binary.version}</artifactId>
<packaging>jar</packaging>
- <name>Celeborn Client for Flink 2.0</name>
+ <name>Celeborn Client for Flink Common Tiered</name>
<dependencies>
<dependency>
<groupId>org.apache.celeborn</groupId>
- <artifactId>celeborn-common_${scala.binary.version}</artifactId>
+ <artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.apache.celeborn</groupId>
- <artifactId>celeborn-client_${scala.binary.version}</artifactId>
+ <artifactId>celeborn-common_${scala.binary.version}</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.apache.celeborn</groupId>
- <artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
+ <artifactId>celeborn-client_${scala.binary.version}</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
@@ -55,21 +55,10 @@
</exclusion>
</exclusions>
</dependency>
- <!-- Test dependencies -->
<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.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/BaseCelebornTierFactory.java
similarity index 50%
copy from client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java
copy to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/BaseCelebornTierFactory.java
index 3c821363e..9bb3b1ca1 100644
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java
+++ b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/BaseCelebornTierFactory.java
@@ -18,22 +18,10 @@
package org.apache.celeborn.plugin.flink.tiered;
-import java.util.List;
-import java.util.concurrent.ScheduledExecutorService;
-
import org.apache.flink.configuration.Configuration;
-import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.TieredStorageNettyService;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageConsumerSpec;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageMemoryManager;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageMemorySpec;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageResourceRegistry;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierConsumerAgent;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierFactory;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierProducerAgent;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleDescriptor;
import org.apache.flink.runtime.util.ConfigurationParserUtils;
import org.apache.celeborn.common.CelebornConf;
@@ -41,88 +29,54 @@ import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
/**
* The factory class of the Celeborn client, used as a tier of flink hybrid shuffle tiered storage.
+ * The name 'CelebornTierFactory' is reserved for versioned factories so configurations can be same
+ * for each version.
*/
-public class CelebornTierFactory implements TierFactory {
+public class BaseCelebornTierFactory {
- private CelebornConf conf;
+ protected CelebornConf conf;
/**
* The bytes size of a single buffer, default value is 32KB, it will be set according to the flink
- * configuration in {@link CelebornTierFactory#setup}.
+ * configuration in {@link BaseCelebornTierFactory#setup}.
*/
- private int bufferSizeBytes = -1;
+ protected int bufferSizeBytes = -1;
/**
* The max bytes size of a single segment, it will determine how many buffer can save in a single
* segment.
*/
- private static int MAX_BYTES_PER_SEGMENT = 8 * 1024 * 1024;
+ protected static int MAX_BYTES_PER_SEGMENT = 8 * 1024 * 1024;
- private static final String CELEBORN_TIER_NAME = CelebornTierFactory.class.getSimpleName();
+ protected static final String CELEBORN_TIER_NAME = "CelebornTierFactory";
- @Override
public void setup(Configuration configuration) {
conf = FlinkUtils.toCelebornConf(configuration);
this.bufferSizeBytes = ConfigurationParserUtils.getPageSize(configuration);
}
- @Override
public TieredStorageMemorySpec getMasterAgentMemorySpec() {
return new TieredStorageMemorySpec(getCelebornTierName(), 0);
}
- @Override
public TieredStorageMemorySpec getProducerAgentMemorySpec() {
return new TieredStorageMemorySpec(getCelebornTierName(), 1);
}
- @Override
public TieredStorageMemorySpec getConsumerAgentMemorySpec() {
return new TieredStorageMemorySpec(getCelebornTierName(), 0);
}
- @Override
public TierMasterAgent createMasterAgent(
TieredStorageResourceRegistry tieredStorageResourceRegistry) {
return new CelebornTierMasterAgent(conf);
}
- @Override
- public TierProducerAgent createProducerAgent(
- int numPartitions,
- int numSubpartitions,
- TieredStoragePartitionId partitionId,
- String dataFileBasePath,
- boolean isBroadcastOnly,
- TieredStorageMemoryManager storageMemoryManager,
- TieredStorageNettyService nettyService,
- TieredStorageResourceRegistry resourceRegistry,
- BatchShuffleReadBufferPool bufferPool,
- ScheduledExecutorService ioExecutor,
- List<TierShuffleDescriptor> shuffleDescriptors,
- int maxRequestedBuffers) {
- return new CelebornTierProducerAgent(
- conf,
- partitionId,
- numPartitions,
- numSubpartitions,
- MAX_BYTES_PER_SEGMENT,
- bufferSizeBytes,
- storageMemoryManager,
- resourceRegistry,
- shuffleDescriptors);
- }
-
- @Override
- public TierConsumerAgent createConsumerAgent(
- List<TieredStorageConsumerSpec> tieredStorageConsumerSpecs,
- List<TierShuffleDescriptor> shuffleDescriptors,
- TieredStorageNettyService nettyService) {
- return new CelebornTierConsumerAgent(
- conf, tieredStorageConsumerSpecs, shuffleDescriptors, bufferSizeBytes);
- }
-
public static String getCelebornTierName() {
return CELEBORN_TIER_NAME;
}
+
+ public String identifier() {
+ return "celeborn";
+ }
}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferManager.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferManager.java
similarity index 100%
rename from client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferManager.java
rename to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferManager.java
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
similarity index 98%
rename from client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
rename to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
index 883d90c89..e033d1283 100644
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
+++ b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
@@ -25,6 +25,8 @@ import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
+import javax.annotation.Nullable;
+
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageInputChannelId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
@@ -58,7 +60,7 @@ public class CelebornChannelBufferReader {
private final int partitionId;
- private final TieredStorageInputChannelId inputChannelId;
+ @Nullable private final TieredStorageInputChannelId inputChannelId;
private final int subPartitionIndexStart;
@@ -84,7 +86,7 @@ public class CelebornChannelBufferReader {
public CelebornChannelBufferReader(
FlinkShuffleClientImpl client,
ShuffleResourceDescriptor shuffleDescriptor,
- TieredStorageInputChannelId inputChannelId,
+ @Nullable TieredStorageInputChannelId inputChannelId,
int startSubIdx,
int endSubIdx,
BiConsumer<ByteBuf, TieredStorageSubpartitionId> dataListener,
@@ -312,8 +314,6 @@ public class CelebornChannelBufferReader {
long streamId = streamEnd.getStreamId();
LOG.debug("Buffer stream reader get stream end for {}", streamId);
if (!closed && !CelebornBufferStream.isEmptyStream(bufferStream)) {
- // TOOD: Update the partition locations here if support reading and writing shuffle data
- // simultaneously
bufferStream.moveToNextPartitionIfPossible(streamId, this::sendRequireSegmentId, true);
}
}
diff --git a/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornFlinkTierShim.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornFlinkTierShim.java
new file mode 100644
index 000000000..fce98d248
--- /dev/null
+++ b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornFlinkTierShim.java
@@ -0,0 +1,37 @@
+/*
+ * 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.tiered;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleHandler;
+import org.apache.flink.runtime.shuffle.JobShuffleContext;
+
+import org.apache.celeborn.plugin.flink.CelebornFlinkShim;
+
+/** Shim for Celeborn and Flink with tiered shuffle. */
+public abstract class CelebornFlinkTierShim extends CelebornFlinkShim {
+
+ public static JobShuffleContext getJobShuffleContext(
+ JobID jobID, TierShuffleHandler tierShuffleHandler) {
+ return ((CelebornFlinkTierShim) getInstance())
+ .getJobShuffleContextImpl(jobID, tierShuffleHandler);
+ }
+
+ protected abstract JobShuffleContext getJobShuffleContextImpl(
+ JobID jobID, TierShuffleHandler tierShuffleHandler);
+}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
similarity index 100%
rename from client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
rename to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
diff --git a/client-flink/flink-2.1/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
similarity index 85%
rename from client-flink/flink-2.1/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
rename to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
index 36e343244..aab285f2f 100644
--- a/client-flink/flink-2.1/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
+++ b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
@@ -20,25 +20,18 @@ package org.apache.celeborn.plugin.flink.tiered;
import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
-import java.time.Duration;
-import java.util.Collection;
-import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
-import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageIdMappingUtils;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleDescriptor;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleHandler;
import org.apache.flink.runtime.shuffle.JobShuffleContext;
-import org.apache.flink.runtime.shuffle.PartitionWithMetrics;
import org.apache.flink.util.ExecutorUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -117,7 +110,11 @@ public class CelebornTierMasterAgent implements TierMasterAgent {
}
}
- @Override
+ public TierShuffleDescriptor addPartitionAndGetShuffleDescriptor(
+ JobID jobID, ResultPartitionID resultPartitionID) {
+ return addPartitionAndGetShuffleDescriptor(jobID, -1, resultPartitionID);
+ }
+
public TierShuffleDescriptor addPartitionAndGetShuffleDescriptor(
JobID jobID, int numSubpartitions, ResultPartitionID resultPartitionID) {
Set<Integer> shuffleIds = jobShuffleIds.get(jobID);
@@ -188,32 +185,6 @@ public class CelebornTierMasterAgent implements TierMasterAgent {
private JobShuffleContext getJobShuffleContext(
JobID jobID, TierShuffleHandler tierShuffleHandler) {
- return new JobShuffleContext() {
- @Override
- public JobID getJobId() {
- return jobID;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> resultPartitionIds) {
- return tierShuffleHandler.onReleasePartitions(
- resultPartitionIds.stream()
- .map(TieredStorageIdMappingUtils::convertId)
- .collect(Collectors.toList()));
- }
-
- @Override
- public CompletableFuture<Collection<PartitionWithMetrics>> getPartitionWithMetrics(
- Duration duration, Set<ResultPartitionID> set) {
- // TODO we should impl this when we support JM failover.
- return CompletableFuture.completedFuture(Collections.emptyList());
- }
-
- @Override
- public void notifyPartitionRecoveryStarted() {
- // TODO we should impl this when we support JM failover.
- }
- };
+ return CelebornFlinkTierShim.getJobShuffleContext(jobID, tierShuffleHandler);
}
}
diff --git a/client-flink/flink-2.1/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierProducerAgent.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierProducerAgent.java
similarity index 99%
rename from client-flink/flink-2.1/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierProducerAgent.java
rename to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierProducerAgent.java
index 3b10d93f5..11f338b46 100644
--- a/client-flink/flink-2.1/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierProducerAgent.java
+++ b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierProducerAgent.java
@@ -217,7 +217,7 @@ public class CelebornTierProducerAgent implements TierProducerAgent {
if (buffer.isBuffer()) {
memoryManager.transferBufferOwnership(
- bufferOwner, CelebornTierFactory.getCelebornTierName(), buffer);
+ bufferOwner, BaseCelebornTierFactory.getCelebornTierName(), buffer);
}
// write buffer to BufferPacker and record buffer count per subPartition per segment
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/TierShuffleDescriptorImpl.java b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/TierShuffleDescriptorImpl.java
similarity index 100%
rename from client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/TierShuffleDescriptorImpl.java
rename to client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/TierShuffleDescriptorImpl.java
diff --git a/client-flink/flink-1.20/src/test/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgentSuiteJ.java b/client-flink/common-tiered/src/test/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgentSuiteJ.java
similarity index 97%
rename from client-flink/flink-1.20/src/test/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgentSuiteJ.java
rename to client-flink/common-tiered/src/test/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgentSuiteJ.java
index 1ee5b8a19..e92cc7d9f 100644
--- a/client-flink/flink-1.20/src/test/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgentSuiteJ.java
+++ b/client-flink/common-tiered/src/test/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgentSuiteJ.java
@@ -35,6 +35,7 @@ import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.junit.After;
import org.junit.Assert;
+import org.junit.Assume;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
@@ -42,6 +43,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.util.Utils$;
+import org.apache.celeborn.plugin.flink.CelebornFlinkShim;
import org.apache.celeborn.plugin.flink.ShuffleResource;
import org.apache.celeborn.plugin.flink.ShuffleResourceDescriptor;
import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
@@ -52,9 +54,10 @@ public class CelebornTierMasterAgentSuiteJ {
@Before
public void setUp() {
+ Assume.assumeTrue(CelebornFlinkShim.isShimLoaded());
Configuration configuration = new Configuration();
int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
- configuration.setInteger("celeborn.master.port", startPort);
+ configuration.setString("celeborn.master.port", String.valueOf(startPort));
configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
masterAgent = createMasterAgent(configuration);
}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShim.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShim.java
new file mode 100644
index 000000000..bb0e75645
--- /dev/null
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShim.java
@@ -0,0 +1,172 @@
+/*
+ * 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.time.Duration;
+import java.util.Map;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+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.partition.ResultPartitionType;
+import org.apache.flink.runtime.io.network.partition.consumer.InputChannel;
+import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
+import org.apache.flink.runtime.shuffle.PartitionDescriptor;
+import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
+
+/** Shim for Celeborn and various versions of Flink. The shim is loaded via a ServiceLoader. */
+public abstract class CelebornFlinkShim {
+
+ private static volatile CelebornFlinkShim INSTANCE;
+
+ // Exposed production shims
+
+ public static BufferCompressor createBufferCompressor(
+ int networkBufferSize, String compressionCodec) {
+ return getInstance().createBufferCompressorImpl(networkBufferSize, compressionCodec);
+ }
+
+ public static BufferDecompressor createBufferDecompressor(
+ int networkBufferSize, String compressionCodec) {
+ return getInstance().createBufferDecompressorImpl(networkBufferSize, compressionCodec);
+ }
+
+ public static Duration getRequestSegmentsTimeout(Configuration configuration) {
+ return getInstance().getRequestSegmentsTimeoutImpl(configuration);
+ }
+
+ public static InputChannel createFakeChannel(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex) {
+ return getInstance().createFakeChannelImpl(inputGateDelegation, channelIndex);
+ }
+
+ public static Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ return getInstance().getConsumedSubpartitionIndexRangeImpl(gateDescriptor);
+ }
+
+ public static ResultPartitionBytesCounter createResultPartitionBytesCounter(
+ RemoteShuffleResultPartition remoteShuffleResultPartition) {
+ return getInstance().createResultPartitionBytesCounterImpl(remoteShuffleResultPartition);
+ }
+
+ // Exposed test shims
+
+ @VisibleForTesting
+ public static Buffer.DataType getEndOfPartitionDataType() {
+ return getInstance().getEndOfPartitionDataTypeImpl();
+ }
+
+ @VisibleForTesting
+ public static PartitionDescriptor createPartitionDescriptor(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum) {
+ return getInstance().createPartitionDescriptorImpl(intermediateDataSetId, partitionNum);
+ }
+
+ @VisibleForTesting
+ public static TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptor(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {
+ return getInstance()
+ .createTaskInputsOutputsDescriptorImpl(
+ inputGateNums,
+ inputChannelNums,
+ partitionReuseCount,
+ subpartitionNums,
+ inputPartitionTypes,
+ partitionTypes);
+ }
+
+ // Loading and management methods
+
+ @VisibleForTesting
+ public static boolean isShimLoaded() {
+ try {
+ getInstance();
+ } catch (Exception e) {
+ return false;
+ }
+ return true;
+ }
+
+ protected static CelebornFlinkShim getInstance() {
+ if (INSTANCE != null) {
+ return INSTANCE;
+ }
+ synchronized (CelebornFlinkShim.class) {
+ java.util.ServiceLoader<CelebornFlinkShim> loader =
+ java.util.ServiceLoader.load(CelebornFlinkShim.class);
+ for (CelebornFlinkShim shim : loader) {
+ if (INSTANCE == null) {
+ INSTANCE = shim;
+ } else {
+ throw new IllegalStateException(
+ "Multiple CelebornFlinkShim implementations found in classpath:"
+ + INSTANCE.getClass().getSimpleName()
+ + ", "
+ + shim.getClass().getSimpleName());
+ }
+ }
+ }
+ if (INSTANCE != null) {
+ return INSTANCE;
+ } else {
+ throw new IllegalStateException("No CelebornFlinkShim implementation found in classpath");
+ }
+ }
+
+ // Implementations
+
+ protected abstract BufferCompressor createBufferCompressorImpl(
+ int networkBufferSize, String compressionCodec);
+
+ protected abstract BufferDecompressor createBufferDecompressorImpl(
+ int networkBufferSize, String compressionCodec);
+
+ protected abstract Duration getRequestSegmentsTimeoutImpl(Configuration configuration);
+
+ protected abstract InputChannel createFakeChannelImpl(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex);
+
+ protected abstract Tuple2<Integer, Integer> getConsumedSubpartitionIndexRangeImpl(
+ InputGateDeploymentDescriptor gateDescriptor);
+
+ protected abstract ResultPartitionBytesCounter createResultPartitionBytesCounterImpl(
+ RemoteShuffleResultPartition remoteShuffleResultPartition);
+
+ protected abstract Buffer.DataType getEndOfPartitionDataTypeImpl();
+
+ protected abstract PartitionDescriptor createPartitionDescriptorImpl(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum);
+
+ protected abstract TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptorImpl(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes);
+}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
index e64b1202e..738350298 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleEnvironment.java
@@ -73,10 +73,10 @@ public class RemoteShuffleEnvironment
/** A trivial {@link ResultPartitionManager}. */
private final ResultPartitionManager resultPartitionManager;
- /** Factory class to create {@link AbstractRemoteShuffleResultPartition}. */
- private final AbstractRemoteShuffleResultPartitionFactory resultPartitionFactory;
+ /** Factory class to create {@link RemoteShuffleResultPartition}. */
+ private final RemoteShuffleResultPartitionFactory resultPartitionFactory;
- private final AbstractRemoteShuffleInputGateFactory inputGateFactory;
+ private final RemoteShuffleInputGateFactory inputGateFactory;
private final CelebornConf conf;
@@ -103,8 +103,8 @@ public class RemoteShuffleEnvironment
public RemoteShuffleEnvironment(
NetworkBufferPool networkBufferPool,
ResultPartitionManager resultPartitionManager,
- AbstractRemoteShuffleResultPartitionFactory resultPartitionFactory,
- AbstractRemoteShuffleInputGateFactory inputGateFactory,
+ RemoteShuffleResultPartitionFactory resultPartitionFactory,
+ RemoteShuffleInputGateFactory inputGateFactory,
CelebornConf conf,
NettyShuffleEnvironmentWrapper shuffleEnvironmentWrapper) {
this.networkBufferPool = networkBufferPool;
@@ -277,7 +277,7 @@ public class RemoteShuffleEnvironment
}
@VisibleForTesting
- AbstractRemoteShuffleResultPartitionFactory getResultPartitionFactory() {
+ RemoteShuffleResultPartitionFactory getResultPartitionFactory() {
return resultPartitionFactory;
}
}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleInputGate.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
similarity index 86%
rename from client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleInputGate.java
rename to client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
index 9f3b8b9a9..cad869eff 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleInputGate.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
@@ -31,6 +31,7 @@ import org.apache.flink.runtime.event.TaskEvent;
import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
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.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;
@@ -42,11 +43,12 @@ import org.apache.flink.util.function.SupplierWithException;
import org.apache.celeborn.common.CelebornConf;
/** An abstract {@link IndexedInputGate} which ingest data from remote shuffle workers. */
-public abstract class AbstractRemoteShuffleInputGate extends IndexedInputGate {
+public class RemoteShuffleInputGate extends IndexedInputGate {
public final RemoteShuffleInputGateDelegation inputGateDelegation;
+ private final ResultPartitionType consumedPartitionType;
- public AbstractRemoteShuffleInputGate(
+ public RemoteShuffleInputGate(
CelebornConf celebornConf,
ShuffleIOOwnerContext ownerContext,
int gateIndex,
@@ -69,6 +71,7 @@ public abstract class AbstractRemoteShuffleInputGate extends IndexedInputGate {
indexRange.f0,
indexRange.f1,
shuffleIOMetricGroups);
+ this.consumedPartitionType = gateDescriptor.getConsumedPartitionType();
}
/** Setup gate and build network connections. */
@@ -106,13 +109,13 @@ public abstract class AbstractRemoteShuffleInputGate extends IndexedInputGate {
return inputGateDelegation.pollNext();
}
- /** Close all reading channels inside this {@link AbstractRemoteShuffleInputGate}. */
+ /** Close all reading channels inside this {@link RemoteShuffleInputGate}. */
@Override
public void close() throws Exception {
inputGateDelegation.close();
}
- /** Get {@link InputChannelInfo}s of this {@link AbstractRemoteShuffleInputGate}. */
+ /** Get {@link InputChannelInfo}s of this {@link RemoteShuffleInputGate}. */
@Override
public List<InputChannelInfo> getChannelInfos() {
return inputGateDelegation.getChannelsInfo();
@@ -159,7 +162,9 @@ public abstract class AbstractRemoteShuffleInputGate extends IndexedInputGate {
}
@Override
- public abstract InputChannel getChannel(int channelIndex);
+ public InputChannel getChannel(int channelIndex) {
+ return CelebornFlinkShim.createFakeChannel(inputGateDelegation, channelIndex);
+ }
@Override
public void sendTaskEvent(TaskEvent event) {
@@ -192,6 +197,13 @@ public abstract class AbstractRemoteShuffleInputGate extends IndexedInputGate {
inputGateDelegation.getGateDescriptor().toString());
}
- public abstract Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
- InputGateDeploymentDescriptor gateDescriptor);
+ public Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ return CelebornFlinkShim.getConsumedSubpartitionIndexRange(gateDescriptor);
+ }
+
+ public ResultPartitionType getConsumedPartitionType() {
+ // TODO only used in 1.19
+ return consumedPartitionType;
+ }
}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleInputGateFactory.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
similarity index 81%
rename from client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleInputGateFactory.java
rename to client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
index 22e39ffdf..71b2a41b8 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleInputGateFactory.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
@@ -22,6 +22,7 @@ import java.io.IOException;
import java.util.Map;
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.BufferPoolFactory;
import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
@@ -35,13 +36,12 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.plugin.flink.utils.Utils;
-/** Factory class to create {@link AbstractRemoteShuffleInputGate}. */
-public abstract class AbstractRemoteShuffleInputGateFactory {
+/** Factory class to create {@link RemoteShuffleInputGate}. */
+public class RemoteShuffleInputGateFactory {
public static final int MIN_BUFFERS_PER_GATE = 16;
- private static final Logger LOG =
- LoggerFactory.getLogger(AbstractRemoteShuffleInputGateFactory.class);
+ private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleInputGateFactory.class);
/** Number of max concurrent reading channels. */
protected final int numConcurrentReading;
@@ -62,7 +62,7 @@ public abstract class AbstractRemoteShuffleInputGateFactory {
protected CelebornConf celebornConf;
- public AbstractRemoteShuffleInputGateFactory(
+ public RemoteShuffleInputGateFactory(
CelebornConf conf, NetworkBufferPool networkBufferPool, int networkBufferSize) {
this.celebornConf = conf;
this.numBuffersPerGate =
@@ -96,22 +96,39 @@ public abstract class AbstractRemoteShuffleInputGateFactory {
SupplierWithException<BufferPool, IOException> bufferPoolFactory =
createBufferPoolFactory(networkBufferPool, numBuffersPerGate, supportFloatingBuffers);
- return createInputGate(
+ BufferDecompressor bufferDecompressor =
+ CelebornFlinkShim.createBufferDecompressor(
+ networkBufferSize, celebornConf.shuffleCompressionCodec().name());
+ return new RemoteShuffleInputGate(
+ celebornConf,
ownerContext,
gateIndex,
igdd,
bufferPoolFactory,
- celebornConf.shuffleCompressionCodec().name(),
+ bufferDecompressor,
+ numConcurrentReading,
shuffleIOMetricGroups);
}
- protected abstract IndexedInputGate createInputGate(
+ protected IndexedInputGate createInputGate(
ShuffleIOOwnerContext ownerContext,
int gateIndex,
InputGateDeploymentDescriptor igdd,
SupplierWithException<BufferPool, IOException> bufferPoolFactory,
String compressionCodec,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroupMap);
+ Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroupMap) {
+ BufferDecompressor bufferDecompressor =
+ CelebornFlinkShim.createBufferDecompressor(networkBufferSize, compressionCodec);
+ return new RemoteShuffleInputGate(
+ celebornConf,
+ ownerContext,
+ gateIndex,
+ igdd,
+ bufferPoolFactory,
+ bufferDecompressor,
+ numConcurrentReading,
+ shuffleIOMetricGroupMap);
+ }
private SupplierWithException<BufferPool, IOException> createBufferPoolFactory(
BufferPoolFactory bufferPoolFactory, int numBuffers, boolean supportFloatingBuffers) {
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartition.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
similarity index 95%
rename from client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartition.java
rename to client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
index bafc4ffcc..91b28b979 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartition.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
@@ -56,13 +56,15 @@ import org.apache.celeborn.plugin.flink.utils.Utils;
* spilled to the remote shuffle service in subpartition index order sequentially. Large records
* that can not be appended to an empty {@link DataBuffer} will be spilled directly.
*/
-public abstract class AbstractRemoteShuffleResultPartition extends ResultPartition {
+public class RemoteShuffleResultPartition extends ResultPartition {
private final RemoteShuffleResultPartitionDelegation delegation;
private final SupplierWithException<BufferPool, IOException> bufferPoolFactory;
- public AbstractRemoteShuffleResultPartition(
+ private final ResultPartitionBytesCounter delegateResultPartitionBytes;
+
+ public RemoteShuffleResultPartition(
String owningTaskName,
int partitionIndex,
ResultPartitionID partitionId,
@@ -88,6 +90,7 @@ public abstract class AbstractRemoteShuffleResultPartition extends ResultPartiti
new RemoteShuffleResultPartitionDelegation(
networkBufferSize, outputGate, this::updateStatistics, numSubpartitions);
this.bufferPoolFactory = bufferPoolFactory;
+ this.delegateResultPartitionBytes = CelebornFlinkShim.createResultPartitionBytesCounter(this);
}
@Override
@@ -231,10 +234,7 @@ public abstract class AbstractRemoteShuffleResultPartition extends ResultPartiti
numBuffersOut.inc(isBroadcast ? numSubpartitions : 1);
long readableBytes =
(long) bufferWithSubpartition.getBuffer().readableBytes() - BufferUtils.HEADER_LENGTH;
- updateReadableBytes(bufferWithSubpartition, isBroadcast, readableBytes);
+ delegateResultPartitionBytes.update(bufferWithSubpartition, isBroadcast, readableBytes);
numBytesOut.inc(isBroadcast ? readableBytes * numSubpartitions : readableBytes);
}
-
- public abstract void updateReadableBytes(
- BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes);
}
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/RemoteShuffleResultPartitionFactory.java
similarity index 85%
rename from client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleResultPartitionFactory.java
rename to client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
index 55a7d09e1..ff05dfd6d 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/RemoteShuffleResultPartitionFactory.java
@@ -36,13 +36,14 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.CompressionCodec;
import org.apache.celeborn.plugin.flink.utils.Utils;
-/** Factory class to create {@link AbstractRemoteShuffleResultPartition}. */
-public abstract class AbstractRemoteShuffleResultPartitionFactory {
+/** Factory class to create {@link RemoteShuffleResultPartition}. */
+public class RemoteShuffleResultPartitionFactory {
private static final Logger LOG =
- LoggerFactory.getLogger(AbstractRemoteShuffleResultPartitionFactory.class);
+ LoggerFactory.getLogger(RemoteShuffleResultPartitionFactory.class);
public static final int MIN_BUFFERS_PER_PARTITION = 16;
@@ -65,7 +66,7 @@ public abstract class AbstractRemoteShuffleResultPartitionFactory {
protected String compressionCodec;
- public AbstractRemoteShuffleResultPartitionFactory(
+ public RemoteShuffleResultPartitionFactory(
CelebornConf celebornConf,
ResultPartitionManager partitionManager,
BufferPoolFactory bufferPoolFactory,
@@ -144,7 +145,7 @@ public abstract class AbstractRemoteShuffleResultPartitionFactory {
return partition;
}
- abstract ResultPartition createRemoteShuffleResultPartitionInternal(
+ ResultPartition createRemoteShuffleResultPartitionInternal(
String taskNameWithSubtaskAndId,
int partitionIndex,
ResultPartitionID id,
@@ -156,7 +157,27 @@ public abstract class AbstractRemoteShuffleResultPartitionFactory {
int numMappers,
BufferCompressor bufferCompressor,
RemoteShuffleDescriptor rsd,
- ShuffleIOMetricGroup shuffleIOMetricGroup);
+ ShuffleIOMetricGroup shuffleIOMetricGroup) {
+ return new RemoteShuffleResultPartition(
+ taskNameWithSubtaskAndId,
+ partitionIndex,
+ id,
+ type,
+ numSubpartitions,
+ maxParallelism,
+ networkBufferSize,
+ partitionManager,
+ bufferCompressor,
+ bufferPoolFactories.get(0),
+ new RemoteShuffleOutputGate(
+ rsd,
+ numSubpartitions,
+ networkBufferSize,
+ bufferPoolFactories.get(1),
+ celebornConf,
+ numMappers,
+ shuffleIOMetricGroup));
+ }
/**
* Used to create 2 buffer pools -- sorting buffer pool (7/8), transportation buffer pool (1/8).
@@ -183,5 +204,9 @@ public abstract class AbstractRemoteShuffleResultPartitionFactory {
}
@VisibleForTesting
- protected abstract BufferCompressor getBufferCompressor();
+ protected BufferCompressor getBufferCompressor() {
+ return CompressionCodec.NONE.name().equals(compressionCodec)
+ ? null
+ : CelebornFlinkShim.createBufferCompressor(networkBufferSize, compressionCodec);
+ }
}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleServiceFactory.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
similarity index 78%
rename from client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleServiceFactory.java
rename to client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
index b2f8c10a0..6559fa129 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleServiceFactory.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
@@ -39,9 +39,10 @@ import org.apache.flink.runtime.shuffle.ShuffleServiceFactory;
import org.apache.flink.runtime.util.ConfigurationParserUtils;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;
import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
-public abstract class AbstractRemoteShuffleServiceFactory
+public class RemoteShuffleServiceFactory
implements ShuffleServiceFactory<ShuffleDescriptor, ResultPartitionWriter, IndexedInputGate> {
public final NettyShuffleServiceFactory nettyShuffleServiceFactory =
@@ -54,8 +55,28 @@ public abstract class AbstractRemoteShuffleServiceFactory
}
@Override
- public abstract ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate>
- createShuffleEnvironment(ShuffleEnvironmentContext shuffleEnvironmentContext);
+ public ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> createShuffleEnvironment(
+ ShuffleEnvironmentContext shuffleEnvironmentContext) {
+ AbstractRemoteShuffleServiceParameters parameters =
+ initializePreCreateShuffleEnvironment(shuffleEnvironmentContext);
+ RemoteShuffleResultPartitionFactory resultPartitionFactory =
+ new RemoteShuffleResultPartitionFactory(
+ parameters.celebornConf,
+ parameters.resultPartitionManager,
+ parameters.networkBufferPool,
+ parameters.bufferSize);
+ RemoteShuffleInputGateFactory inputGateFactory =
+ new RemoteShuffleInputGateFactory(
+ parameters.celebornConf, parameters.networkBufferPool, parameters.bufferSize);
+
+ return new RemoteShuffleEnvironment(
+ parameters.networkBufferPool,
+ parameters.resultPartitionManager,
+ resultPartitionFactory,
+ inputGateFactory,
+ parameters.celebornConf,
+ new NettyShuffleEnvironmentWrapper(nettyShuffleServiceFactory, shuffleEnvironmentContext));
+ }
private int calculateNumberOfNetworkBuffers(MemorySize memorySize, int bufferSize) {
long numBuffersLong = memorySize.getBytes() / bufferSize;
@@ -91,7 +112,9 @@ public abstract class AbstractRemoteShuffleServiceFactory
return result;
}
- abstract Duration getRequestSegmentsTimeout(Configuration configuration);
+ Duration getRequestSegmentsTimeout(Configuration configuration) {
+ return CelebornFlinkShim.getRequestSegmentsTimeout(configuration);
+ }
protected static class AbstractRemoteShuffleServiceParameters {
public final Configuration configuration;
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ResultPartitionBytesCounter.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ResultPartitionBytesCounter.java
new file mode 100644
index 000000000..4b938f2cb
--- /dev/null
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ResultPartitionBytesCounter.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.celeborn.plugin.flink.buffer.BufferWithSubpartition;
+
+/** This counter will count the data size of a partition. */
+public interface ResultPartitionBytesCounter {
+
+ void update(
+ BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes);
+}
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/utils/ReflectionUtils.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/utils/ReflectionUtils.java
new file mode 100644
index 000000000..0a9b75282
--- /dev/null
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/utils/ReflectionUtils.java
@@ -0,0 +1,84 @@
+/*
+ * 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.utils;
+
+import java.lang.reflect.Field;
+
+/** Utility methods for reflection. */
+public class ReflectionUtils {
+
+ @SuppressWarnings("unchecked")
+ public static <T> T readDeclaredFieldRecursiveNoThrow(Object object, String fieldName) {
+ try {
+ Field field = findFieldRecursive(object.getClass(), fieldName);
+ field.setAccessible(true);
+ return (T) field.get(object);
+ } catch (NoSuchFieldException | IllegalAccessException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private static Field findFieldRecursive(Class<?> clazz, String fieldName)
+ throws NoSuchFieldException {
+ Class<?> cur = clazz;
+ while (cur != null) {
+ try {
+ return cur.getDeclaredField(fieldName);
+ } catch (NoSuchFieldException e) {
+ cur = cur.getSuperclass();
+ }
+ }
+ throw new NoSuchFieldException(
+ "Field '" + fieldName + "' not found in class hierarchy of " + clazz.getSimpleName());
+ }
+
+ @SuppressWarnings("unchecked")
+ public static <T> T createEmptyImplementation(Class<T> interfaceClass) {
+ if (!interfaceClass.isInterface()) {
+ throw new IllegalArgumentException("Class must be an interface: " + interfaceClass.getName());
+ }
+
+ return (T)
+ java.lang.reflect.Proxy.newProxyInstance(
+ interfaceClass.getClassLoader(),
+ new Class<?>[] {interfaceClass},
+ (proxy, method, args) -> {
+ // Return default value based on return type
+ if (method.getReturnType() == boolean.class) {
+ return false;
+ } else if (method.getReturnType() == byte.class) {
+ return (byte) 0;
+ } else if (method.getReturnType() == char.class) {
+ return (char) 0;
+ } else if (method.getReturnType() == short.class) {
+ return (short) 0;
+ } else if (method.getReturnType() == int.class) {
+ return 0;
+ } else if (method.getReturnType() == long.class) {
+ return 0L;
+ } else if (method.getReturnType() == float.class) {
+ return 0.0f;
+ } else if (method.getReturnType() == double.class) {
+ return 0.0d;
+ } else {
+ // For object types, return null
+ return null;
+ }
+ });
+ }
+}
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
similarity index 92%
rename from client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
rename to client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
index 046af4938..018bd4b4d 100644
--- a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
@@ -17,13 +17,13 @@
package org.apache.celeborn.plugin.flink;
+import static org.apache.celeborn.plugin.flink.TestUtils.createJobShuffleContext;
+
import java.net.InetAddress;
import java.net.UnknownHostException;
-import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
-import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import org.apache.flink.api.common.BatchShuffleMode;
@@ -35,10 +35,8 @@ 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.NettyShuffleServiceFactory;
-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;
@@ -47,9 +45,9 @@ import org.apache.flink.runtime.shuffle.PartitionDescriptor;
import org.apache.flink.runtime.shuffle.ProducerDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleMasterContext;
-import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
import org.junit.After;
import org.junit.Assert;
+import org.junit.Assume;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
@@ -69,9 +67,10 @@ public class RemoteShuffleMasterSuiteJ {
@Before
public void setUp() {
+ Assume.assumeTrue(CelebornFlinkShim.isShimLoaded());
configuration = new Configuration();
int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
- configuration.setInteger("celeborn.master.port", startPort);
+ configuration.setString("celeborn.master.port", String.valueOf(startPort));
configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
configuration.setString("celeborn.client.application.heartbeatInterval", "30s");
remoteShuffleMaster = createShuffleMaster(configuration);
@@ -259,7 +258,7 @@ public class RemoteShuffleMasterSuiteJ {
resultPartitionTypes.put(outputDataSetID2, ResultPartitionType.BLOCKING);
MemorySize calculated =
remoteShuffleMaster.computeShuffleMemorySizeForTask(
- TaskInputsOutputsDescriptor.from(
+ CelebornFlinkShim.createTaskInputsOutputsDescriptor(
3,
numberOfInputGateChannels,
numbersOfResultSubpartitions,
@@ -338,34 +337,9 @@ public class RemoteShuffleMasterSuiteJ {
return remoteShuffleMaster;
}
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public org.apache.flink.api.common.JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
public PartitionDescriptor createPartitionDescriptor(
IntermediateDataSetID intermediateDataSetId, int partitionNum) {
- IntermediateResultPartitionID intermediateResultPartitionId =
- new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
- return new PartitionDescriptor(
- intermediateDataSetId,
- 10,
- intermediateResultPartitionId,
- ResultPartitionType.BLOCKING,
- 5,
- 1,
- false,
- false);
+ return CelebornFlinkShim.createPartitionDescriptor(intermediateDataSetId, partitionNum);
}
public ProducerDescriptor createProducerDescriptor() throws UnknownHostException {
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
similarity index 93%
rename from client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
rename to client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
index 548336db6..f80eab0b7 100644
--- a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
@@ -22,6 +22,8 @@ import static org.mockito.Mockito.mock;
import org.apache.flink.runtime.io.network.buffer.BufferPoolFactory;
import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
import org.junit.Test;
import org.apache.celeborn.common.CelebornConf;
@@ -30,6 +32,11 @@ import org.apache.celeborn.common.protocol.CompressionCodec;
/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
public class RemoteShuffleResultPartitionFactorySuiteJ {
+ @Before
+ public void setup() {
+ Assume.assumeTrue(CelebornFlinkShim.isShimLoaded());
+ }
+
@Test
public void testGetBufferCompressor() {
CelebornConf celebornConf = new CelebornConf();
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
similarity index 97%
rename from client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
rename to client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
index 7e1946e76..470a18aca 100644
--- a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
@@ -60,6 +60,7 @@ import org.apache.flink.runtime.metrics.groups.ShuffleMetricGroup;
import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
import org.apache.flink.util.function.SupplierWithException;
import org.junit.After;
+import org.junit.Assume;
import org.junit.Before;
import org.junit.Test;
@@ -70,12 +71,13 @@ import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
import org.apache.celeborn.plugin.flink.utils.BufferUtils;
public class RemoteShuffleResultPartitionSuiteJ {
+
private final int networkBufferSize = 32 * 1024;
- private final BufferCompressor bufferCompressor = new BufferCompressor(networkBufferSize, "lz4");
+ private BufferCompressor bufferCompressor;
private final RemoteShuffleOutputGate remoteShuffleOutputGate =
mock(RemoteShuffleOutputGate.class);
private final CelebornConf conf = new CelebornConf();
- BufferDecompressor bufferDecompressor = new BufferDecompressor(networkBufferSize, "LZ4");
+ private BufferDecompressor bufferDecompressor;
private static final int totalBuffers = 1000;
@@ -93,7 +95,10 @@ public class RemoteShuffleResultPartitionSuiteJ {
@Before
public void setup() {
+ Assume.assumeTrue(CelebornFlinkShim.isShimLoaded());
globalBufferPool = new NetworkBufferPool(totalBuffers, bufferSize);
+ bufferCompressor = CelebornFlinkShim.createBufferCompressor(networkBufferSize, "LZ4");
+ bufferDecompressor = CelebornFlinkShim.createBufferDecompressor(networkBufferSize, "LZ4");
}
@After
@@ -108,8 +113,10 @@ public class RemoteShuffleResultPartitionSuiteJ {
if (nettyBufferPool != null) {
nettyBufferPool.lazyDestroy();
}
- assertEquals(totalBuffers, globalBufferPool.getNumberOfAvailableMemorySegments());
- globalBufferPool.destroy();
+ if (globalBufferPool != null) {
+ assertEquals(totalBuffers, globalBufferPool.getNumberOfAvailableMemorySegments());
+ globalBufferPool.destroy();
+ }
}
@Test
@@ -331,7 +338,11 @@ public class RemoteShuffleResultPartitionSuiteJ {
for (int subpartition = 0; subpartition < numSubpartitions; ++subpartition) {
ByteBuffer record = EventSerializer.toSerializedEvent(EndOfPartitionEvent.INSTANCE);
recordDataWritten(
- record, Buffer.DataType.EVENT_BUFFER, subpartition, dataWritten, numBytesWritten);
+ record,
+ CelebornFlinkShim.getEndOfPartitionDataType(),
+ subpartition,
+ dataWritten,
+ numBytesWritten);
}
outputGate
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
similarity index 95%
rename from client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
rename to client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
index 49affdee4..43ffc814c 100644
--- a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
@@ -30,9 +30,17 @@ 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.Assume;
+import org.junit.Before;
import org.junit.Test;
public class RemoteShuffleServiceFactorySuiteJ {
+
+ @Before
+ public void setup() {
+ Assume.assumeTrue(CelebornFlinkShim.isShimLoaded());
+ }
+
@Test
public void testCreateShuffleEnvironment() {
RemoteShuffleServiceFactory remoteShuffleServiceFactory = new RemoteShuffleServiceFactory();
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
similarity index 90%
rename from client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
rename to client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
index b93ae3ea1..f1d20af72 100644
--- a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
@@ -17,18 +17,19 @@
package org.apache.celeborn.plugin.flink;
+import static org.apache.celeborn.plugin.flink.TestUtils.createJobShuffleContext;
+
import java.util.ArrayList;
-import java.util.Collection;
import java.util.List;
-import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import com.google.common.collect.Sets;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
-import org.apache.flink.runtime.shuffle.JobShuffleContext;
import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;
@@ -41,6 +42,11 @@ import org.apache.celeborn.common.util.JavaUtils;
public class ShuffleResourceTrackerSuiteJ {
+ @Before
+ public void setup() {
+ Assume.assumeTrue(CelebornFlinkShim.isShimLoaded());
+ }
+
@Test
public void testNotifyUnknownWorkers() {
LifecycleManager lifecycleManager = Mockito.mock(LifecycleManager.class);
@@ -116,21 +122,6 @@ public class ShuffleResourceTrackerSuiteJ {
return shufflePartitionLocationInfo;
}
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
private PartitionLocation mockShufflePartitionLocationInfo(int partitionId) {
return new PartitionLocation(
partitionId, -1, "mock", -1, -1, -1, -1, PartitionLocation.Mode.PRIMARY);
diff --git a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/TestUtils.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/TestUtils.java
new file mode 100644
index 000000000..24f32d35b
--- /dev/null
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/TestUtils.java
@@ -0,0 +1,46 @@
+/*
+ * 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.util.concurrent.CompletableFuture;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.shuffle.JobShuffleContext;
+
+public class TestUtils {
+
+ public static JobShuffleContext createJobShuffleContext(JobID jobId) {
+ return (JobShuffleContext)
+ java.lang.reflect.Proxy.newProxyInstance(
+ JobShuffleContext.class.getClassLoader(),
+ new Class<?>[] {JobShuffleContext.class},
+ (proxy, method, args) -> {
+ if (method.getName().equals("getJobId")) {
+ return jobId;
+ } else if (method.getName().equals("stopTrackingAndReleasePartitions")) {
+ return CompletableFuture.completedFuture(null);
+ } else if (method.getName().equals("getPartitionWithMetrics")) {
+ return CompletableFuture.completedFuture(null);
+ } else if (method.getName().equals("notifyPartitionRecoveryStarted")) {
+ return null;
+ } else {
+ throw new UnsupportedOperationException();
+ }
+ });
+ }
+}
diff --git a/client-flink/flink-1.16/pom.xml b/client-flink/flink-1.16/pom.xml
index 1caac0b4b..6835fc6bf 100644
--- a/client-flink/flink-1.16/pom.xml
+++ b/client-flink/flink-1.16/pom.xml
@@ -56,6 +56,13 @@
</exclusions>
</dependency>
<!-- Test dependencies -->
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
diff --git a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV116.java b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV116.java
new file mode 100644
index 000000000..4271a6e6e
--- /dev/null
+++ b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV116.java
@@ -0,0 +1,155 @@
+/*
+ * 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.net.InetAddress;
+import java.time.Duration;
+import java.util.Map;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+import org.apache.flink.runtime.deployment.SubpartitionIndexRange;
+import org.apache.flink.runtime.io.network.ConnectionID;
+import org.apache.flink.runtime.io.network.LocalConnectionManager;
+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.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+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.jobgraph.IntermediateResultPartitionID;
+import org.apache.flink.runtime.shuffle.PartitionDescriptor;
+import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.runtime.throughput.ThroughputCalculator;
+import org.apache.flink.util.clock.SystemClock;
+
+import org.apache.celeborn.plugin.flink.utils.ReflectionUtils;
+
+public class CelebornFlinkShimV116 extends CelebornFlinkShim {
+ @Override
+ protected BufferCompressor createBufferCompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferCompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected BufferDecompressor createBufferDecompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferDecompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected Duration getRequestSegmentsTimeoutImpl(Configuration configuration) {
+ return Duration.ofMillis(
+ configuration.get(
+ NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
+ }
+
+ @Override
+ protected InputChannel createFakeChannelImpl(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex) {
+ return new RemoteInputChannel(
+ new SingleInputGate(
+ inputGateDelegation.getTaskName(),
+ inputGateDelegation.getGateIndex(),
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ new SubpartitionIndexRange(0, 1),
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ ReflectionUtils.createEmptyImplementation(MemorySegmentProvider.class),
+ 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(),
+ ReflectionUtils.createEmptyImplementation(ChannelStateWriter.class));
+ }
+
+ @Override
+ protected Tuple2<Integer, Integer> getConsumedSubpartitionIndexRangeImpl(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ SubpartitionIndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
+ return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
+ }
+
+ @Override
+ protected ResultPartitionBytesCounter createResultPartitionBytesCounterImpl(
+ RemoteShuffleResultPartition remoteShuffleResultPartition) {
+ Counter numBytesProduced =
+ ReflectionUtils.readDeclaredFieldRecursiveNoThrow(
+ remoteShuffleResultPartition, "numBytesProduced");
+ return (bufferWithSubpartition, isBroadcast, readableBytes) ->
+ numBytesProduced.inc(readableBytes);
+ }
+
+ @Override
+ protected Buffer.DataType getEndOfPartitionDataTypeImpl() {
+ return Buffer.DataType.EVENT_BUFFER;
+ }
+
+ @Override
+ protected PartitionDescriptor createPartitionDescriptorImpl(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum) {
+ IntermediateResultPartitionID intermediateResultPartitionId =
+ new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
+ return new PartitionDescriptor(
+ intermediateDataSetId,
+ 10,
+ intermediateResultPartitionId,
+ ResultPartitionType.BLOCKING,
+ 5,
+ 1,
+ false,
+ false);
+ }
+
+ @Override
+ protected TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptorImpl(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {
+ return TaskInputsOutputsDescriptor.from(
+ inputGateNums, inputChannelNums, partitionReuseCount, partitionTypes);
+ }
+}
diff --git a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
deleted file mode 100644
index 6e663fd68..000000000
--- a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
+++ /dev/null
@@ -1,181 +0,0 @@
-/*
- * 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.List;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-
-import org.apache.flink.api.java.tuple.Tuple2;
-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.deployment.SubpartitionIndexRange;
-import org.apache.flink.runtime.io.network.ConnectionID;
-import org.apache.flink.runtime.io.network.LocalConnectionManager;
-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.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.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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.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 AbstractRemoteShuffleInputGate {
-
- public RemoteShuffleInputGate(
- CelebornConf celebornConf,
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor gateDescriptor,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- BufferDecompressor bufferDecompressor,
- int numConcurrentReading,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- super(
- celebornConf,
- ownerContext,
- gateIndex,
- gateDescriptor,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-
- @Override
- public InputChannel getChannel(int channelIndex) {
- return new FakedRemoteInputChannel(channelIndex);
- }
-
- @Override
- public Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
- InputGateDeploymentDescriptor gateDescriptor) {
- SubpartitionIndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
- return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
- }
-
- /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
- private class FakedRemoteInputChannel extends RemoteInputChannel {
- FakedRemoteInputChannel(int channelIndex) {
- super(
- new SingleInputGate(
- inputGateDelegation.getTaskName(),
- inputGateDelegation.getGateIndex(),
- new IntermediateDataSetID(),
- ResultPartitionType.BLOCKING,
- new SubpartitionIndexRange(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 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() {}
-
- @Override
- public void addOutputDataFuture(
- long l,
- ResultSubpartitionInfo resultSubpartitionInfo,
- int i,
- CompletableFuture<List<Buffer>> completableFuture)
- throws IllegalArgumentException {}
- }
-}
diff --git a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
deleted file mode 100644
index b19c77238..000000000
--- a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * 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.util.Map;
-
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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(
- CelebornConf conf, NetworkBufferPool networkBufferPool, int networkBufferSize) {
- super(conf, networkBufferPool, networkBufferSize);
- }
-
- @Override
- protected RemoteShuffleInputGate createInputGate(
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor igdd,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- String compressionCodec,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- BufferDecompressor bufferDecompressor =
- new BufferDecompressor(networkBufferSize, compressionCodec);
- return new RemoteShuffleInputGate(
- celebornConf,
- ownerContext,
- gateIndex,
- igdd,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-}
diff --git a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
deleted file mode 100644
index cc33a6464..000000000
--- a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
+++ /dev/null
@@ -1,73 +0,0 @@
-/*
- * 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 javax.annotation.Nullable;
-
-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.ResultPartition;
-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.apache.celeborn.plugin.flink.buffer.BufferWithSubpartition;
-import org.apache.celeborn.plugin.flink.buffer.DataBuffer;
-
-/**
- * A {@link ResultPartition} which appends records and events to {@link DataBuffer} and after the
- * {@link DataBuffer} is full, all data in the {@link DataBuffer} 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 DataBuffer} will be spilled directly.
- */
-public class RemoteShuffleResultPartition extends AbstractRemoteShuffleResultPartition {
-
- 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,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactory,
- outputGate);
- }
-
- public void updateReadableBytes(
- BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes) {
- numBytesProduced.inc(readableBytes);
- }
-}
diff --git a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
deleted file mode 100644
index d49827ed9..000000000
--- a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ /dev/null
@@ -1,89 +0,0 @@
-/*
- * 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.util.List;
-
-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.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartition;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.util.function.SupplierWithException;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Factory class to create {@link RemoteShuffleResultPartition}. */
-public class RemoteShuffleResultPartitionFactory
- extends AbstractRemoteShuffleResultPartitionFactory {
-
- public RemoteShuffleResultPartitionFactory(
- CelebornConf celebornConf,
- ResultPartitionManager partitionManager,
- BufferPoolFactory bufferPoolFactory,
- int networkBufferSize) {
-
- super(celebornConf, partitionManager, bufferPoolFactory, networkBufferSize);
- }
-
- @Override
- ResultPartition createRemoteShuffleResultPartitionInternal(
- String taskNameWithSubtaskAndId,
- int partitionIndex,
- ResultPartitionID id,
- ResultPartitionType type,
- int numSubpartitions,
- int maxParallelism,
- List<SupplierWithException<BufferPool, IOException>> bufferPoolFactories,
- CelebornConf celebornConf,
- int numMappers,
- BufferCompressor bufferCompressor,
- RemoteShuffleDescriptor rsd,
- ShuffleIOMetricGroup shuffleIOMetricGroup) {
- return new RemoteShuffleResultPartition(
- taskNameWithSubtaskAndId,
- partitionIndex,
- id,
- type,
- numSubpartitions,
- maxParallelism,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactories.get(0),
- new RemoteShuffleOutputGate(
- rsd,
- numSubpartitions,
- networkBufferSize,
- bufferPoolFactories.get(1),
- celebornConf,
- numMappers,
- shuffleIOMetricGroup));
- }
-
- protected BufferCompressor getBufferCompressor() {
- return CompressionCodec.NONE.name().equals(compressionCodec)
- ? null
- : new BufferCompressor(networkBufferSize, compressionCodec);
- }
-}
diff --git a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
deleted file mode 100644
index 4e0652472..000000000
--- a/client-flink/flink-1.16/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.time.Duration;
-
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
-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.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;
-
-public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFactory {
-
- @Override
- public ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> createShuffleEnvironment(
- ShuffleEnvironmentContext shuffleEnvironmentContext) {
- AbstractRemoteShuffleServiceParameters parameters =
- initializePreCreateShuffleEnvironment(shuffleEnvironmentContext);
- RemoteShuffleResultPartitionFactory resultPartitionFactory =
- new RemoteShuffleResultPartitionFactory(
- parameters.celebornConf,
- parameters.resultPartitionManager,
- parameters.networkBufferPool,
- parameters.bufferSize);
- RemoteShuffleInputGateFactory inputGateFactory =
- new RemoteShuffleInputGateFactory(
- parameters.celebornConf, parameters.networkBufferPool, parameters.bufferSize);
-
- return new RemoteShuffleEnvironment(
- parameters.networkBufferPool,
- parameters.resultPartitionManager,
- resultPartitionFactory,
- inputGateFactory,
- parameters.celebornConf,
- new NettyShuffleEnvironmentWrapper(nettyShuffleServiceFactory, shuffleEnvironmentContext));
- }
-
- @Override
- public Duration getRequestSegmentsTimeout(Configuration configuration) {
- return Duration.ofMillis(
- configuration.get(
- NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
- }
-}
diff --git a/client-flink/flink-1.16/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim b/client-flink/flink-1.16/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
new file mode 100644
index 000000000..58a85712d
--- /dev/null
+++ b/client-flink/flink-1.16/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
@@ -0,0 +1,16 @@
+# 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.
+
+org.apache.celeborn.plugin.flink.CelebornFlinkShimV116
\ No newline at end of file
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
deleted file mode 100644
index 1567b4bb7..000000000
--- a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
+++ /dev/null
@@ -1,366 +0,0 @@
-/*
- * 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.net.InetAddress;
-import java.net.UnknownHostException;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Random;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-
-import org.apache.flink.api.common.BatchShuffleMode;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.ExecutionOptions;
-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.NettyShuffleServiceFactory;
-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.NettyShuffleDescriptor;
-import org.apache.flink.runtime.shuffle.PartitionDescriptor;
-import org.apache.flink.runtime.shuffle.ProducerDescriptor;
-import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
-import org.apache.flink.runtime.shuffle.ShuffleMasterContext;
-import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
-import org.junit.After;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.FallbackPolicy;
-import org.apache.celeborn.common.util.Utils$;
-import org.apache.celeborn.plugin.flink.fallback.ForceFallbackPolicy;
-import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
-
-public class RemoteShuffleMasterSuiteJ {
-
- private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleMasterSuiteJ.class);
- private RemoteShuffleMaster remoteShuffleMaster;
- private Configuration configuration;
-
- @Before
- public void setUp() {
- configuration = new Configuration();
- int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
- configuration.setInteger("celeborn.master.port", startPort);
- configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
- configuration.setString("celeborn.client.application.heartbeatInterval", "30s");
- remoteShuffleMaster = createShuffleMaster(configuration);
- }
-
- @Test
- public void testRegisterJob() {
- JobShuffleContext jobShuffleContext = createJobShuffleContext(JobID.generate());
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- // reRunRegister job
- try {
- remoteShuffleMaster.registerJob(jobShuffleContext);
- } catch (Exception e) {
- Assert.assertTrue(true);
- }
-
- // unRegister job
- remoteShuffleMaster.unregisterJob(jobShuffleContext.getJobId());
- remoteShuffleMaster.registerJob(jobShuffleContext);
- }
-
- @Test
- public void testRegisterJobWithForceFallbackPolicy() {
- configuration.setString(
- CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
- remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
- Assert.assertTrue(remoteShuffleMaster.jobFallbackPolicies().containsKey(jobID));
- remoteShuffleMaster.unregisterJob(jobShuffleContext.getJobId());
- Assert.assertTrue(remoteShuffleMaster.jobFallbackPolicies().isEmpty());
- }
-
- @Test
- public void testRegisterPartitionWithProducer()
- throws UnknownHostException, ExecutionException, InterruptedException {
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- RemoteShuffleDescriptor remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(1, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- ShuffleResource shuffleResource = remoteShuffleDescriptor.getShuffleResource();
- ShuffleResourceDescriptor mapPartitionShuffleDescriptor =
- shuffleResource.getMapPartitionShuffleDescriptor();
-
- LOG.info("remoteShuffleDescriptor:{}", remoteShuffleDescriptor);
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getPartitionId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getAttemptId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getMapId());
-
- // use same dataset id
- partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 1);
- remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(2, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- mapPartitionShuffleDescriptor =
- remoteShuffleDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
-
- // use another attemptId
- producerDescriptor = createProducerDescriptor();
- remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(3, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- mapPartitionShuffleDescriptor =
- remoteShuffleDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getAttemptId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
- }
-
- @Test
- public void testRegisterPartitionWithProducerForForceFallbackPolicy()
- throws UnknownHostException, ExecutionException, InterruptedException {
- configuration.setString(
- CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
- remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- ShuffleDescriptor shuffleDescriptor =
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertTrue(shuffleDescriptor instanceof NettyShuffleDescriptor);
- Assert.assertEquals(1, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- Map<String, Long> shuffleFallbackCounts =
- remoteShuffleMaster.lifecycleManager().shuffleFallbackCounts();
- Assert.assertEquals(1, shuffleFallbackCounts.size());
- Assert.assertEquals(
- 1L, shuffleFallbackCounts.get(ForceFallbackPolicy.class.getName()).longValue());
- }
-
- @Test
- public void testRegisterMultipleJobs()
- throws UnknownHostException, ExecutionException, InterruptedException {
- JobID jobID1 = JobID.generate();
- JobShuffleContext jobShuffleContext1 = createJobShuffleContext(jobID1);
- remoteShuffleMaster.registerJob(jobShuffleContext1);
-
- JobID jobID2 = JobID.generate();
- JobShuffleContext jobShuffleContext2 = createJobShuffleContext(jobID2);
- remoteShuffleMaster.registerJob(jobShuffleContext2);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- RemoteShuffleDescriptor remoteShuffleDescriptor1 =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID1, partitionDescriptor, producerDescriptor)
- .get();
-
- // use same datasetId but different jobId
- RemoteShuffleDescriptor remoteShuffleDescriptor2 =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID2, partitionDescriptor, producerDescriptor)
- .get();
-
- Assert.assertEquals(
- remoteShuffleDescriptor1
- .getShuffleResource()
- .getMapPartitionShuffleDescriptor()
- .getShuffleId(),
- 0);
- Assert.assertEquals(
- remoteShuffleDescriptor2
- .getShuffleResource()
- .getMapPartitionShuffleDescriptor()
- .getShuffleId(),
- 1);
- }
-
- @Test
- public void testShuffleMemoryAnnouncing() {
- Map<IntermediateDataSetID, Integer> numberOfInputGateChannels = new HashMap<>();
- Map<IntermediateDataSetID, Integer> numbersOfResultSubpartitions = new HashMap<>();
- Map<IntermediateDataSetID, ResultPartitionType> resultPartitionTypes = new HashMap<>();
- IntermediateDataSetID inputDataSetID0 = new IntermediateDataSetID();
- IntermediateDataSetID inputDataSetID1 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID0 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID1 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID2 = new IntermediateDataSetID();
- Random random = new Random();
- numberOfInputGateChannels.put(inputDataSetID0, random.nextInt(1000));
- numberOfInputGateChannels.put(inputDataSetID1, 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);
- resultPartitionTypes.put(outputDataSetID0, ResultPartitionType.BLOCKING);
- resultPartitionTypes.put(outputDataSetID1, ResultPartitionType.BLOCKING);
- resultPartitionTypes.put(outputDataSetID2, ResultPartitionType.BLOCKING);
- MemorySize calculated =
- remoteShuffleMaster.computeShuffleMemorySizeForTask(
- TaskInputsOutputsDescriptor.from(
- 3, numberOfInputGateChannels, numbersOfResultSubpartitions, resultPartitionTypes));
-
- CelebornConf celebornConf = FlinkUtils.toCelebornConf(configuration);
-
- long numBytesPerGate = celebornConf.clientFlinkMemoryPerInputGate();
- long expectedInput = 2 * numBytesPerGate;
-
- long numBytesPerResultPartition = celebornConf.clientFlinkMemoryPerResultPartition();
- long expectedOutput = 3 * numBytesPerResultPartition;
- MemorySize expected = new MemorySize(expectedInput + expectedOutput);
-
- Assert.assertEquals(expected, calculated);
- }
-
- @Test
- public void testInvalidShuffleConfig() {
- Assert.assertThrows(
- String.format(
- "The config option %s should configure as %s",
- ExecutionOptions.BATCH_SHUFFLE_MODE.key(),
- BatchShuffleMode.ALL_EXCHANGES_BLOCKING.name()),
- IllegalArgumentException.class,
- () ->
- createShuffleMaster(
- new Configuration()
- .set(
- ExecutionOptions.BATCH_SHUFFLE_MODE,
- BatchShuffleMode.ALL_EXCHANGES_PIPELINED)));
- Configuration configuration = new Configuration();
- configuration.setString(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED().key(), "true");
- Assert.assertThrows(
- String.format(
- "Flink does not support replicate shuffle data. Please check the config %s.",
- CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED().key()),
- IllegalArgumentException.class,
- () -> createShuffleMaster(configuration));
- }
-
- @After
- public void tearDown() {
- if (remoteShuffleMaster != null) {
- try {
- remoteShuffleMaster.close();
- } catch (Exception e) {
- LOG.warn(e.getMessage(), e);
- }
- }
- }
-
- public RemoteShuffleMaster createShuffleMaster(Configuration configuration) {
- return createShuffleMaster(configuration, null);
- }
-
- public RemoteShuffleMaster createShuffleMaster(
- Configuration configuration, NettyShuffleServiceFactory nettyShuffleServiceFactory) {
- remoteShuffleMaster =
- new RemoteShuffleMaster(
- new ShuffleMasterContext() {
- @Override
- public Configuration getConfiguration() {
- return configuration;
- }
-
- @Override
- public void onFatalError(Throwable throwable) {
- System.exit(-1);
- }
- },
- nettyShuffleServiceFactory);
-
- return remoteShuffleMaster;
- }
-
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public org.apache.flink.api.common.JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
- public PartitionDescriptor createPartitionDescriptor(
- IntermediateDataSetID intermediateDataSetId, int partitionNum) {
- IntermediateResultPartitionID intermediateResultPartitionId =
- new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
- return new PartitionDescriptor(
- intermediateDataSetId,
- 10,
- intermediateResultPartitionId,
- ResultPartitionType.BLOCKING,
- 5,
- 1,
- false,
- false);
- }
-
- public ProducerDescriptor createProducerDescriptor() throws UnknownHostException {
- 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.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV116.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV116.java
new file mode 100644
index 000000000..2a9a73da8
--- /dev/null
+++ b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV116.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleMasterSuiteJV116 extends RemoteShuffleMasterSuiteJ {}
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
deleted file mode 100644
index 548336db6..000000000
--- a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * 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.Mockito.mock;
-
-import org.apache.flink.runtime.io.network.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
-import org.junit.Assert;
-import org.junit.Test;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
-public class RemoteShuffleResultPartitionFactorySuiteJ {
-
- @Test
- public void testGetBufferCompressor() {
- CelebornConf celebornConf = new CelebornConf();
- for (CompressionCodec compressionCodec : CompressionCodec.values()) {
- RemoteShuffleResultPartitionFactory partitionFactory =
- new RemoteShuffleResultPartitionFactory(
- celebornConf.set(
- CelebornConf.SHUFFLE_COMPRESSION_CODEC().key(), compressionCodec.name()),
- mock(ResultPartitionManager.class),
- mock(BufferPoolFactory.class),
- 1);
- if (CompressionCodec.NONE.equals(compressionCodec)) {
- Assert.assertNull(partitionFactory.getBufferCompressor());
- } else {
- Assert.assertNotNull(partitionFactory.getBufferCompressor());
- }
- }
- }
-}
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV116.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV116.java
new file mode 100644
index 000000000..7a09c935e
--- /dev/null
+++ b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV116.java
@@ -0,0 +1,22 @@
+/*
+ * 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;
+
+/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
+public class RemoteShuffleResultPartitionFactorySuiteJV116
+ extends RemoteShuffleResultPartitionFactorySuiteJ {}
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
deleted file mode 100644
index 7e1946e76..000000000
--- a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
+++ /dev/null
@@ -1,638 +0,0 @@
-/*
- * 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.ArgumentMatchers.anyBoolean;
-import static org.mockito.ArgumentMatchers.anyInt;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.metrics.groups.ShuffleMetricGroup;
-import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
-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.DataBuffer;
-import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
-import org.apache.celeborn.plugin.flink.utils.BufferUtils;
-
-public class RemoteShuffleResultPartitionSuiteJ {
- private final int networkBufferSize = 32 * 1024;
- private final BufferCompressor bufferCompressor = new BufferCompressor(networkBufferSize, "lz4");
- private final RemoteShuffleOutputGate remoteShuffleOutputGate =
- mock(RemoteShuffleOutputGate.class);
- 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());
- DataBuffer dataBuffer = remoteShuffleResultPartition.getDelegation().getUnicastDataBuffer();
- ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[] {1, 2, 3});
- dataBuffer.append(byteBuffer, 0, Buffer.DataType.DATA_BUFFER);
- remoteShuffleResultPartition.getDelegation().flushDataBuffer(dataBuffer, true);
- }
-
- private List<SupplierWithException<BufferPool, IOException>> createBufferPoolFactory() {
- NetworkBufferPool networkBufferPool =
- new NetworkBufferPool(256 * 8, 32 * 1024, Duration.ofSeconds(1));
-
- 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());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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());
- assertEquals(2, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(2, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.broadcastRecord(ByteBuffer.allocate(bufferSize));
- assertEquals(2, sortBufferPool.bestEffortGetNumOfUsedBuffers());
- assertEquals(3, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(3, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.flush(0);
- assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
-
- partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 2);
- partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 3);
- assertEquals(3, sortBufferPool.bestEffortGetNumOfUsedBuffers());
- assertEquals(5, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(5, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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);
- }
- }
- assertEquals(numRecords, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(numRecords, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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,
- new ShuffleIOMetricGroup(
- new ShuffleMetricGroup(
- UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
- 1,
- CelebornConf.CLIENT_METRICS_SCOPE_NAMING_SHUFFLE().defaultValueString())));
- 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(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() {
- 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, System.currentTimeMillis(), 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.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV116.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV116.java
new file mode 100644
index 000000000..dbae0bb86
--- /dev/null
+++ b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV116.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleResultPartitionSuiteJV116 extends RemoteShuffleResultPartitionSuiteJ {}
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuitJ.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuitJ.java
deleted file mode 100644
index 73dbb6a9d..000000000
--- a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuitJ.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * 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.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.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/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV116.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV116.java
new file mode 100644
index 000000000..090e976e8
--- /dev/null
+++ b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV116.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleServiceFactorySuiteJV116 extends RemoteShuffleServiceFactorySuiteJ {}
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
deleted file mode 100644
index b93ae3ea1..000000000
--- a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
+++ /dev/null
@@ -1,138 +0,0 @@
-/*
- * 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.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-
-import com.google.common.collect.Sets;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
-import org.apache.flink.runtime.shuffle.JobShuffleContext;
-import org.junit.Assert;
-import org.junit.Test;
-import org.mockito.Mockito;
-
-import org.apache.celeborn.client.LifecycleManager;
-import org.apache.celeborn.client.listener.WorkersStatus;
-import org.apache.celeborn.common.meta.ShufflePartitionLocationInfo;
-import org.apache.celeborn.common.meta.WorkerInfo;
-import org.apache.celeborn.common.protocol.PartitionLocation;
-import org.apache.celeborn.common.util.JavaUtils;
-
-public class ShuffleResourceTrackerSuiteJ {
-
- @Test
- public void testNotifyUnknownWorkers() {
- LifecycleManager lifecycleManager = Mockito.mock(LifecycleManager.class);
- ScheduledThreadPoolExecutor executor = Mockito.mock(ScheduledThreadPoolExecutor.class);
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map = JavaUtils.newConcurrentHashMap();
- WorkerInfo workerInfo = new WorkerInfo("mock", -1, -1, -1, -1);
- map.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map2 = JavaUtils.newConcurrentHashMap();
- map2.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map3 = JavaUtils.newConcurrentHashMap();
- map3.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- Mockito.when(lifecycleManager.workerSnapshots(Mockito.anyInt())).thenReturn(map, map2, map3);
-
- ShuffleResourceTracker shuffleResourceTracker =
- new ShuffleResourceTracker(executor, lifecycleManager);
-
- JobID jobID1 = new JobID();
- shuffleResourceTracker.registerJob(createJobShuffleContext(jobID1));
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 1, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 2, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 3, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 2, 3, new ResultPartitionID());
-
- JobID jobID2 = new JobID();
- shuffleResourceTracker.registerJob(createJobShuffleContext(jobID2));
- shuffleResourceTracker.addPartitionResource(jobID2, 3, 1, new ResultPartitionID());
-
- List<WorkerInfo> workerInfoList = new ArrayList<>();
- workerInfoList.add(workerInfo);
- shuffleResourceTracker.notifyChangedWorkersStatus(new WorkersStatus(workerInfoList, null));
-
- Assert.assertEquals(
- Sets.newHashSet(3),
- shuffleResourceTracker
- .getJobResourceListener(jobID1)
- .getResultPartitionMap()
- .get(2)
- .keySet());
- Assert.assertEquals(
- Sets.newHashSet(3),
- shuffleResourceTracker
- .getJobResourceListener(jobID1)
- .getResultPartitionMap()
- .get(2)
- .keySet());
-
- Assert.assertTrue(
- shuffleResourceTracker
- .getJobResourceListener(jobID2)
- .getResultPartitionMap()
- .get(3)
- .isEmpty());
- }
-
- public ShufflePartitionLocationInfo mockShufflePartitionLocationInfo(WorkerInfo workerInfo) {
- ShufflePartitionLocationInfo shufflePartitionLocationInfo =
- new ShufflePartitionLocationInfo(workerInfo);
-
- List<PartitionLocation> primaryLocations = new ArrayList<>();
- primaryLocations.add(mockShufflePartitionLocationInfo(1));
- primaryLocations.add(mockShufflePartitionLocationInfo(2));
-
- List<PartitionLocation> replicaLocations = new ArrayList<>();
- replicaLocations.add(mockShufflePartitionLocationInfo(3));
- replicaLocations.add(mockShufflePartitionLocationInfo(4));
-
- shufflePartitionLocationInfo.addPrimaryPartitions(primaryLocations);
- shufflePartitionLocationInfo.addReplicaPartitions(replicaLocations);
- return shufflePartitionLocationInfo;
- }
-
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
- private PartitionLocation mockShufflePartitionLocationInfo(int partitionId) {
- return new PartitionLocation(
- partitionId, -1, "mock", -1, -1, -1, -1, PartitionLocation.Mode.PRIMARY);
- }
-}
diff --git a/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV116.java b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV116.java
new file mode 100644
index 000000000..a31ade8af
--- /dev/null
+++ b/client-flink/flink-1.16/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV116.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class ShuffleResourceTrackerSuiteJV116 extends ShuffleResourceTrackerSuiteJ {}
diff --git a/client-flink/flink-1.17/pom.xml b/client-flink/flink-1.17/pom.xml
index ae4ed501b..501f420a8 100644
--- a/client-flink/flink-1.17/pom.xml
+++ b/client-flink/flink-1.17/pom.xml
@@ -44,6 +44,13 @@
<artifactId>celeborn-client-flink-common_${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>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-runtime</artifactId>
diff --git a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV117.java b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV117.java
new file mode 100644
index 000000000..41e8c49a2
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV117.java
@@ -0,0 +1,164 @@
+/*
+ * 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.net.InetAddress;
+import java.time.Duration;
+import java.util.Map;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+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.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.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+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.jobgraph.IntermediateResultPartitionID;
+import org.apache.flink.runtime.shuffle.PartitionDescriptor;
+import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.runtime.throughput.ThroughputCalculator;
+import org.apache.flink.util.clock.SystemClock;
+
+import org.apache.celeborn.plugin.flink.utils.ReflectionUtils;
+
+public class CelebornFlinkShimV117 extends CelebornFlinkShim {
+ @Override
+ protected BufferCompressor createBufferCompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferCompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected BufferDecompressor createBufferDecompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferDecompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected Duration getRequestSegmentsTimeoutImpl(Configuration configuration) {
+ return Duration.ofMillis(
+ configuration.get(
+ NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
+ }
+
+ @Override
+ protected InputChannel createFakeChannelImpl(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex) {
+ return new RemoteInputChannel(
+ new SingleInputGate(
+ inputGateDelegation.getTaskName(),
+ inputGateDelegation.getGateIndex(),
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ new IndexRange(0, 1),
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ ReflectionUtils.createEmptyImplementation(MemorySegmentProvider.class),
+ 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(),
+ ReflectionUtils.createEmptyImplementation(ChannelStateWriter.class));
+ }
+
+ @Override
+ protected Tuple2<Integer, Integer> getConsumedSubpartitionIndexRangeImpl(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
+ return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
+ }
+
+ @Override
+ protected ResultPartitionBytesCounter createResultPartitionBytesCounterImpl(
+ RemoteShuffleResultPartition remoteShuffleResultPartition) {
+ org.apache.flink.runtime.io.network.metrics.ResultPartitionBytesCounter resultPartitionBytes =
+ ReflectionUtils.readDeclaredFieldRecursiveNoThrow(
+ remoteShuffleResultPartition, "resultPartitionBytes");
+ return (bufferWithSubpartition, isBroadcast, readableBytes) -> {
+ if (isBroadcast) {
+ resultPartitionBytes.incAll(readableBytes);
+ } else {
+ resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
+ }
+ };
+ }
+
+ @Override
+ protected Buffer.DataType getEndOfPartitionDataTypeImpl() {
+ return Buffer.DataType.EVENT_BUFFER;
+ }
+
+ @Override
+ protected PartitionDescriptor createPartitionDescriptorImpl(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum) {
+ IntermediateResultPartitionID intermediateResultPartitionId =
+ new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
+ return new PartitionDescriptor(
+ intermediateDataSetId,
+ 10,
+ intermediateResultPartitionId,
+ ResultPartitionType.BLOCKING,
+ 5,
+ 1,
+ false,
+ false);
+ }
+
+ @Override
+ protected TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptorImpl(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {
+ return TaskInputsOutputsDescriptor.from(
+ inputGateNums,
+ inputChannelNums,
+ partitionReuseCount,
+ subpartitionNums,
+ inputPartitionTypes,
+ partitionTypes);
+ }
+}
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
deleted file mode 100644
index d340a74e6..000000000
--- a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
+++ /dev/null
@@ -1,181 +0,0 @@
-/*
- * 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.List;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-
-import org.apache.flink.api.java.tuple.Tuple2;
-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.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.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.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.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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.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 AbstractRemoteShuffleInputGate {
-
- public RemoteShuffleInputGate(
- CelebornConf celebornConf,
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor gateDescriptor,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- BufferDecompressor bufferDecompressor,
- int numConcurrentReading,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- super(
- celebornConf,
- ownerContext,
- gateIndex,
- gateDescriptor,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-
- @Override
- public InputChannel getChannel(int channelIndex) {
- return new FakedRemoteInputChannel(channelIndex);
- }
-
- @Override
- public Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
- InputGateDeploymentDescriptor gateDescriptor) {
- IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
- return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
- }
-
- /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
- private class FakedRemoteInputChannel extends RemoteInputChannel {
- FakedRemoteInputChannel(int channelIndex) {
- super(
- new SingleInputGate(
- inputGateDelegation.getTaskName(),
- 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
deleted file mode 100644
index 38c5292bf..000000000
--- a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.util.Map;
-
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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(
- CelebornConf conf, NetworkBufferPool networkBufferPool, int networkBufferSize) {
- super(conf, networkBufferPool, networkBufferSize);
- }
-
- // For testing.
- @Override
- protected RemoteShuffleInputGate createInputGate(
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor igdd,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- String compressionCodec,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- BufferDecompressor bufferDecompressor =
- new BufferDecompressor(networkBufferSize, compressionCodec);
- return new RemoteShuffleInputGate(
- celebornConf,
- ownerContext,
- gateIndex,
- igdd,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-}
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
deleted file mode 100644
index 256acb828..000000000
--- a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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 javax.annotation.Nullable;
-
-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.ResultPartition;
-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.apache.celeborn.plugin.flink.buffer.BufferWithSubpartition;
-import org.apache.celeborn.plugin.flink.buffer.DataBuffer;
-
-/**
- * A {@link ResultPartition} which appends records and events to {@link DataBuffer} and after the
- * {@link DataBuffer} is full, all data in the {@link DataBuffer} 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 DataBuffer} will be spilled directly.
- */
-public class RemoteShuffleResultPartition extends AbstractRemoteShuffleResultPartition {
-
- 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,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactory,
- outputGate);
- }
-
- public void updateReadableBytes(
- BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes) {
- if (isBroadcast) {
- resultPartitionBytes.incAll(readableBytes);
- } else {
- resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
- }
- }
-}
diff --git a/client-flink/flink-1.17/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
deleted file mode 100644
index d49827ed9..000000000
--- a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ /dev/null
@@ -1,89 +0,0 @@
-/*
- * 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.util.List;
-
-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.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartition;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.util.function.SupplierWithException;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Factory class to create {@link RemoteShuffleResultPartition}. */
-public class RemoteShuffleResultPartitionFactory
- extends AbstractRemoteShuffleResultPartitionFactory {
-
- public RemoteShuffleResultPartitionFactory(
- CelebornConf celebornConf,
- ResultPartitionManager partitionManager,
- BufferPoolFactory bufferPoolFactory,
- int networkBufferSize) {
-
- super(celebornConf, partitionManager, bufferPoolFactory, networkBufferSize);
- }
-
- @Override
- ResultPartition createRemoteShuffleResultPartitionInternal(
- String taskNameWithSubtaskAndId,
- int partitionIndex,
- ResultPartitionID id,
- ResultPartitionType type,
- int numSubpartitions,
- int maxParallelism,
- List<SupplierWithException<BufferPool, IOException>> bufferPoolFactories,
- CelebornConf celebornConf,
- int numMappers,
- BufferCompressor bufferCompressor,
- RemoteShuffleDescriptor rsd,
- ShuffleIOMetricGroup shuffleIOMetricGroup) {
- return new RemoteShuffleResultPartition(
- taskNameWithSubtaskAndId,
- partitionIndex,
- id,
- type,
- numSubpartitions,
- maxParallelism,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactories.get(0),
- new RemoteShuffleOutputGate(
- rsd,
- numSubpartitions,
- networkBufferSize,
- bufferPoolFactories.get(1),
- celebornConf,
- numMappers,
- shuffleIOMetricGroup));
- }
-
- protected BufferCompressor getBufferCompressor() {
- return CompressionCodec.NONE.name().equals(compressionCodec)
- ? null
- : new BufferCompressor(networkBufferSize, compressionCodec);
- }
-}
diff --git a/client-flink/flink-1.17/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
deleted file mode 100644
index 4e0652472..000000000
--- a/client-flink/flink-1.17/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.time.Duration;
-
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
-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.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;
-
-public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFactory {
-
- @Override
- public ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> createShuffleEnvironment(
- ShuffleEnvironmentContext shuffleEnvironmentContext) {
- AbstractRemoteShuffleServiceParameters parameters =
- initializePreCreateShuffleEnvironment(shuffleEnvironmentContext);
- RemoteShuffleResultPartitionFactory resultPartitionFactory =
- new RemoteShuffleResultPartitionFactory(
- parameters.celebornConf,
- parameters.resultPartitionManager,
- parameters.networkBufferPool,
- parameters.bufferSize);
- RemoteShuffleInputGateFactory inputGateFactory =
- new RemoteShuffleInputGateFactory(
- parameters.celebornConf, parameters.networkBufferPool, parameters.bufferSize);
-
- return new RemoteShuffleEnvironment(
- parameters.networkBufferPool,
- parameters.resultPartitionManager,
- resultPartitionFactory,
- inputGateFactory,
- parameters.celebornConf,
- new NettyShuffleEnvironmentWrapper(nettyShuffleServiceFactory, shuffleEnvironmentContext));
- }
-
- @Override
- public Duration getRequestSegmentsTimeout(Configuration configuration) {
- return Duration.ofMillis(
- configuration.get(
- NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
- }
-}
diff --git a/client-flink/flink-1.17/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim b/client-flink/flink-1.17/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
new file mode 100644
index 000000000..ec0a625df
--- /dev/null
+++ b/client-flink/flink-1.17/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
@@ -0,0 +1,16 @@
+# 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.
+
+org.apache.celeborn.plugin.flink.CelebornFlinkShimV117
\ No newline at end of file
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
deleted file mode 100644
index 046af4938..000000000
--- a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
+++ /dev/null
@@ -1,378 +0,0 @@
-/*
- * 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.net.InetAddress;
-import java.net.UnknownHostException;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Random;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-
-import org.apache.flink.api.common.BatchShuffleMode;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.ExecutionOptions;
-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.NettyShuffleServiceFactory;
-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.NettyShuffleDescriptor;
-import org.apache.flink.runtime.shuffle.PartitionDescriptor;
-import org.apache.flink.runtime.shuffle.ProducerDescriptor;
-import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
-import org.apache.flink.runtime.shuffle.ShuffleMasterContext;
-import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
-import org.junit.After;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.FallbackPolicy;
-import org.apache.celeborn.common.util.Utils$;
-import org.apache.celeborn.plugin.flink.fallback.ForceFallbackPolicy;
-import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
-
-public class RemoteShuffleMasterSuiteJ {
-
- private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleMasterSuiteJ.class);
- private RemoteShuffleMaster remoteShuffleMaster;
- private Configuration configuration;
-
- @Before
- public void setUp() {
- configuration = new Configuration();
- int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
- configuration.setInteger("celeborn.master.port", startPort);
- configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
- configuration.setString("celeborn.client.application.heartbeatInterval", "30s");
- remoteShuffleMaster = createShuffleMaster(configuration);
- }
-
- @Test
- public void testRegisterJob() {
- JobShuffleContext jobShuffleContext = createJobShuffleContext(JobID.generate());
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- // reRunRegister job
- try {
- remoteShuffleMaster.registerJob(jobShuffleContext);
- } catch (Exception e) {
- Assert.assertTrue(true);
- }
-
- // unRegister job
- remoteShuffleMaster.unregisterJob(jobShuffleContext.getJobId());
- remoteShuffleMaster.registerJob(jobShuffleContext);
- }
-
- @Test
- public void testRegisterJobWithForceFallbackPolicy() {
- configuration.setString(
- CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
- remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
- Assert.assertTrue(remoteShuffleMaster.jobFallbackPolicies().containsKey(jobID));
- remoteShuffleMaster.unregisterJob(jobShuffleContext.getJobId());
- Assert.assertTrue(remoteShuffleMaster.jobFallbackPolicies().isEmpty());
- }
-
- @Test
- public void testRegisterPartitionWithProducer()
- throws UnknownHostException, ExecutionException, InterruptedException {
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- RemoteShuffleDescriptor remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(1, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- ShuffleResource shuffleResource = remoteShuffleDescriptor.getShuffleResource();
- ShuffleResourceDescriptor mapPartitionShuffleDescriptor =
- shuffleResource.getMapPartitionShuffleDescriptor();
-
- LOG.info("remoteShuffleDescriptor:{}", remoteShuffleDescriptor);
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getPartitionId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getAttemptId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getMapId());
-
- // use same dataset id
- partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 1);
- remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(2, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- mapPartitionShuffleDescriptor =
- remoteShuffleDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
-
- // use another attemptId
- producerDescriptor = createProducerDescriptor();
- remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(3, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- mapPartitionShuffleDescriptor =
- remoteShuffleDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getAttemptId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
- }
-
- @Test
- public void testRegisterPartitionWithProducerForForceFallbackPolicy()
- throws UnknownHostException, ExecutionException, InterruptedException {
- configuration.setString(
- CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
- remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- ShuffleDescriptor shuffleDescriptor =
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertTrue(shuffleDescriptor instanceof NettyShuffleDescriptor);
- Assert.assertEquals(1, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- Map<String, Long> shuffleFallbackCounts =
- remoteShuffleMaster.lifecycleManager().shuffleFallbackCounts();
- Assert.assertEquals(1, shuffleFallbackCounts.size());
- Assert.assertEquals(
- 1L, shuffleFallbackCounts.get(ForceFallbackPolicy.class.getName()).longValue());
- }
-
- @Test
- public void testRegisterMultipleJobs()
- throws UnknownHostException, ExecutionException, InterruptedException {
- JobID jobID1 = JobID.generate();
- JobShuffleContext jobShuffleContext1 = createJobShuffleContext(jobID1);
- remoteShuffleMaster.registerJob(jobShuffleContext1);
-
- JobID jobID2 = JobID.generate();
- JobShuffleContext jobShuffleContext2 = createJobShuffleContext(jobID2);
- remoteShuffleMaster.registerJob(jobShuffleContext2);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- RemoteShuffleDescriptor remoteShuffleDescriptor1 =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID1, partitionDescriptor, producerDescriptor)
- .get();
-
- // use same datasetId but different jobId
- RemoteShuffleDescriptor remoteShuffleDescriptor2 =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID2, partitionDescriptor, producerDescriptor)
- .get();
-
- Assert.assertEquals(
- remoteShuffleDescriptor1
- .getShuffleResource()
- .getMapPartitionShuffleDescriptor()
- .getShuffleId(),
- 0);
- Assert.assertEquals(
- remoteShuffleDescriptor2
- .getShuffleResource()
- .getMapPartitionShuffleDescriptor()
- .getShuffleId(),
- 1);
- }
-
- @Test
- 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();
- IntermediateDataSetID outputDataSetID0 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID1 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID2 = new IntermediateDataSetID();
- Random random = new Random();
- numberOfInputGateChannels.put(inputDataSetID0, random.nextInt(1000));
- numberOfInputGateChannels.put(inputDataSetID1, 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(
- 3,
- numberOfInputGateChannels,
- numbersOfResultSubpartitions,
- subPartitionNums,
- inputPartitionTypes,
- resultPartitionTypes));
-
- CelebornConf celebornConf = FlinkUtils.toCelebornConf(configuration);
-
- long numBytesPerGate = celebornConf.clientFlinkMemoryPerInputGate();
- long expectedInput = 2 * numBytesPerGate;
-
- long numBytesPerResultPartition = celebornConf.clientFlinkMemoryPerResultPartition();
- long expectedOutput = 3 * numBytesPerResultPartition;
- MemorySize expected = new MemorySize(expectedInput + expectedOutput);
-
- Assert.assertEquals(expected, calculated);
- }
-
- @Test
- public void testInvalidShuffleConfig() {
- Assert.assertThrows(
- String.format(
- "The config option %s should configure as %s",
- ExecutionOptions.BATCH_SHUFFLE_MODE.key(),
- BatchShuffleMode.ALL_EXCHANGES_BLOCKING.name()),
- IllegalArgumentException.class,
- () ->
- createShuffleMaster(
- new Configuration()
- .set(
- ExecutionOptions.BATCH_SHUFFLE_MODE,
- BatchShuffleMode.ALL_EXCHANGES_PIPELINED)));
- Configuration configuration = new Configuration();
- configuration.setString(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED().key(), "true");
- Assert.assertThrows(
- String.format(
- "Flink does not support replicate shuffle data. Please check the config %s.",
- CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED().key()),
- IllegalArgumentException.class,
- () -> createShuffleMaster(configuration));
- }
-
- @After
- public void tearDown() {
- if (remoteShuffleMaster != null) {
- try {
- remoteShuffleMaster.close();
- } catch (Exception e) {
- LOG.warn(e.getMessage(), e);
- }
- }
- }
-
- public RemoteShuffleMaster createShuffleMaster(Configuration configuration) {
- return createShuffleMaster(configuration, null);
- }
-
- public RemoteShuffleMaster createShuffleMaster(
- Configuration configuration, NettyShuffleServiceFactory nettyShuffleServiceFactory) {
- remoteShuffleMaster =
- new RemoteShuffleMaster(
- new ShuffleMasterContext() {
- @Override
- public Configuration getConfiguration() {
- return configuration;
- }
-
- @Override
- public void onFatalError(Throwable throwable) {
- System.exit(-1);
- }
- },
- nettyShuffleServiceFactory);
-
- return remoteShuffleMaster;
- }
-
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public org.apache.flink.api.common.JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
- public PartitionDescriptor createPartitionDescriptor(
- IntermediateDataSetID intermediateDataSetId, int partitionNum) {
- IntermediateResultPartitionID intermediateResultPartitionId =
- new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
- return new PartitionDescriptor(
- intermediateDataSetId,
- 10,
- intermediateResultPartitionId,
- ResultPartitionType.BLOCKING,
- 5,
- 1,
- false,
- false);
- }
-
- public ProducerDescriptor createProducerDescriptor() throws UnknownHostException {
- 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/RemoteShuffleMasterSuiteJV117.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV117.java
new file mode 100644
index 000000000..1eeef35e3
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV117.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleMasterSuiteJV117 extends RemoteShuffleMasterSuiteJ {}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
deleted file mode 100644
index 548336db6..000000000
--- a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * 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.Mockito.mock;
-
-import org.apache.flink.runtime.io.network.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
-import org.junit.Assert;
-import org.junit.Test;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
-public class RemoteShuffleResultPartitionFactorySuiteJ {
-
- @Test
- public void testGetBufferCompressor() {
- CelebornConf celebornConf = new CelebornConf();
- for (CompressionCodec compressionCodec : CompressionCodec.values()) {
- RemoteShuffleResultPartitionFactory partitionFactory =
- new RemoteShuffleResultPartitionFactory(
- celebornConf.set(
- CelebornConf.SHUFFLE_COMPRESSION_CODEC().key(), compressionCodec.name()),
- mock(ResultPartitionManager.class),
- mock(BufferPoolFactory.class),
- 1);
- if (CompressionCodec.NONE.equals(compressionCodec)) {
- Assert.assertNull(partitionFactory.getBufferCompressor());
- } else {
- Assert.assertNotNull(partitionFactory.getBufferCompressor());
- }
- }
- }
-}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV117.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV117.java
new file mode 100644
index 000000000..e4d3bfc5f
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV117.java
@@ -0,0 +1,22 @@
+/*
+ * 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;
+
+/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
+public class RemoteShuffleResultPartitionFactorySuiteJV117
+ extends RemoteShuffleResultPartitionFactorySuiteJ {}
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
deleted file mode 100644
index 7e1946e76..000000000
--- a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
+++ /dev/null
@@ -1,638 +0,0 @@
-/*
- * 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.ArgumentMatchers.anyBoolean;
-import static org.mockito.ArgumentMatchers.anyInt;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.metrics.groups.ShuffleMetricGroup;
-import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
-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.DataBuffer;
-import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
-import org.apache.celeborn.plugin.flink.utils.BufferUtils;
-
-public class RemoteShuffleResultPartitionSuiteJ {
- private final int networkBufferSize = 32 * 1024;
- private final BufferCompressor bufferCompressor = new BufferCompressor(networkBufferSize, "lz4");
- private final RemoteShuffleOutputGate remoteShuffleOutputGate =
- mock(RemoteShuffleOutputGate.class);
- 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());
- DataBuffer dataBuffer = remoteShuffleResultPartition.getDelegation().getUnicastDataBuffer();
- ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[] {1, 2, 3});
- dataBuffer.append(byteBuffer, 0, Buffer.DataType.DATA_BUFFER);
- remoteShuffleResultPartition.getDelegation().flushDataBuffer(dataBuffer, true);
- }
-
- private List<SupplierWithException<BufferPool, IOException>> createBufferPoolFactory() {
- NetworkBufferPool networkBufferPool =
- new NetworkBufferPool(256 * 8, 32 * 1024, Duration.ofSeconds(1));
-
- 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());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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());
- assertEquals(2, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(2, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.broadcastRecord(ByteBuffer.allocate(bufferSize));
- assertEquals(2, sortBufferPool.bestEffortGetNumOfUsedBuffers());
- assertEquals(3, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(3, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.flush(0);
- assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
-
- partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 2);
- partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 3);
- assertEquals(3, sortBufferPool.bestEffortGetNumOfUsedBuffers());
- assertEquals(5, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(5, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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);
- }
- }
- assertEquals(numRecords, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(numRecords, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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,
- new ShuffleIOMetricGroup(
- new ShuffleMetricGroup(
- UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
- 1,
- CelebornConf.CLIENT_METRICS_SCOPE_NAMING_SHUFFLE().defaultValueString())));
- 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(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() {
- 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, System.currentTimeMillis(), 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/RemoteShuffleResultPartitionSuiteJV117.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV117.java
new file mode 100644
index 000000000..7dec6a97f
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV117.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleResultPartitionSuiteJV117 extends RemoteShuffleResultPartitionSuiteJ {}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV117.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV117.java
new file mode 100644
index 000000000..7f48b55b4
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV117.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleServiceFactorySuiteJV117 extends RemoteShuffleServiceFactorySuiteJ {}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
deleted file mode 100644
index b93ae3ea1..000000000
--- a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
+++ /dev/null
@@ -1,138 +0,0 @@
-/*
- * 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.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-
-import com.google.common.collect.Sets;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
-import org.apache.flink.runtime.shuffle.JobShuffleContext;
-import org.junit.Assert;
-import org.junit.Test;
-import org.mockito.Mockito;
-
-import org.apache.celeborn.client.LifecycleManager;
-import org.apache.celeborn.client.listener.WorkersStatus;
-import org.apache.celeborn.common.meta.ShufflePartitionLocationInfo;
-import org.apache.celeborn.common.meta.WorkerInfo;
-import org.apache.celeborn.common.protocol.PartitionLocation;
-import org.apache.celeborn.common.util.JavaUtils;
-
-public class ShuffleResourceTrackerSuiteJ {
-
- @Test
- public void testNotifyUnknownWorkers() {
- LifecycleManager lifecycleManager = Mockito.mock(LifecycleManager.class);
- ScheduledThreadPoolExecutor executor = Mockito.mock(ScheduledThreadPoolExecutor.class);
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map = JavaUtils.newConcurrentHashMap();
- WorkerInfo workerInfo = new WorkerInfo("mock", -1, -1, -1, -1);
- map.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map2 = JavaUtils.newConcurrentHashMap();
- map2.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map3 = JavaUtils.newConcurrentHashMap();
- map3.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- Mockito.when(lifecycleManager.workerSnapshots(Mockito.anyInt())).thenReturn(map, map2, map3);
-
- ShuffleResourceTracker shuffleResourceTracker =
- new ShuffleResourceTracker(executor, lifecycleManager);
-
- JobID jobID1 = new JobID();
- shuffleResourceTracker.registerJob(createJobShuffleContext(jobID1));
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 1, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 2, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 3, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 2, 3, new ResultPartitionID());
-
- JobID jobID2 = new JobID();
- shuffleResourceTracker.registerJob(createJobShuffleContext(jobID2));
- shuffleResourceTracker.addPartitionResource(jobID2, 3, 1, new ResultPartitionID());
-
- List<WorkerInfo> workerInfoList = new ArrayList<>();
- workerInfoList.add(workerInfo);
- shuffleResourceTracker.notifyChangedWorkersStatus(new WorkersStatus(workerInfoList, null));
-
- Assert.assertEquals(
- Sets.newHashSet(3),
- shuffleResourceTracker
- .getJobResourceListener(jobID1)
- .getResultPartitionMap()
- .get(2)
- .keySet());
- Assert.assertEquals(
- Sets.newHashSet(3),
- shuffleResourceTracker
- .getJobResourceListener(jobID1)
- .getResultPartitionMap()
- .get(2)
- .keySet());
-
- Assert.assertTrue(
- shuffleResourceTracker
- .getJobResourceListener(jobID2)
- .getResultPartitionMap()
- .get(3)
- .isEmpty());
- }
-
- public ShufflePartitionLocationInfo mockShufflePartitionLocationInfo(WorkerInfo workerInfo) {
- ShufflePartitionLocationInfo shufflePartitionLocationInfo =
- new ShufflePartitionLocationInfo(workerInfo);
-
- List<PartitionLocation> primaryLocations = new ArrayList<>();
- primaryLocations.add(mockShufflePartitionLocationInfo(1));
- primaryLocations.add(mockShufflePartitionLocationInfo(2));
-
- List<PartitionLocation> replicaLocations = new ArrayList<>();
- replicaLocations.add(mockShufflePartitionLocationInfo(3));
- replicaLocations.add(mockShufflePartitionLocationInfo(4));
-
- shufflePartitionLocationInfo.addPrimaryPartitions(primaryLocations);
- shufflePartitionLocationInfo.addReplicaPartitions(replicaLocations);
- return shufflePartitionLocationInfo;
- }
-
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
- private PartitionLocation mockShufflePartitionLocationInfo(int partitionId) {
- return new PartitionLocation(
- partitionId, -1, "mock", -1, -1, -1, -1, PartitionLocation.Mode.PRIMARY);
- }
-}
diff --git a/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV117.java b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV117.java
new file mode 100644
index 000000000..a3ba4e3c0
--- /dev/null
+++ b/client-flink/flink-1.17/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV117.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class ShuffleResourceTrackerSuiteJV117 extends ShuffleResourceTrackerSuiteJ {}
diff --git a/client-flink/flink-1.18/pom.xml b/client-flink/flink-1.18/pom.xml
index d0e527efa..f50dc9976 100644
--- a/client-flink/flink-1.18/pom.xml
+++ b/client-flink/flink-1.18/pom.xml
@@ -56,6 +56,13 @@
</exclusions>
</dependency>
<!-- Test dependencies -->
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
diff --git a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV118.java b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV118.java
new file mode 100644
index 000000000..71a068fc1
--- /dev/null
+++ b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV118.java
@@ -0,0 +1,167 @@
+/*
+ * 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.net.InetAddress;
+import java.time.Duration;
+import java.util.Map;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+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.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.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+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.jobgraph.IntermediateResultPartitionID;
+import org.apache.flink.runtime.shuffle.PartitionDescriptor;
+import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.runtime.throughput.ThroughputCalculator;
+import org.apache.flink.util.clock.SystemClock;
+
+import org.apache.celeborn.plugin.flink.utils.ReflectionUtils;
+
+public class CelebornFlinkShimV118 extends CelebornFlinkShim {
+ @Override
+ protected BufferCompressor createBufferCompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferCompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected BufferDecompressor createBufferDecompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferDecompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected Duration getRequestSegmentsTimeoutImpl(Configuration configuration) {
+ return Duration.ofMillis(
+ configuration.get(
+ NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
+ }
+
+ @Override
+ protected InputChannel createFakeChannelImpl(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex) {
+ return new RemoteInputChannel(
+ new SingleInputGate(
+ inputGateDelegation.getTaskName(),
+ inputGateDelegation.getGateIndex(),
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ new IndexRange(0, 1),
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ ReflectionUtils.createEmptyImplementation(MemorySegmentProvider.class),
+ 0,
+ new ThroughputCalculator(SystemClock.getInstance()),
+ null,
+ null,
+ null,
+ null),
+ channelIndex,
+ new ResultPartitionID(),
+ 0,
+ new ConnectionID(
+ new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1), 0),
+ new LocalConnectionManager(),
+ 0,
+ 0,
+ 0,
+ new SimpleCounter(),
+ new SimpleCounter(),
+ ReflectionUtils.createEmptyImplementation(ChannelStateWriter.class));
+ }
+
+ @Override
+ protected Tuple2<Integer, Integer> getConsumedSubpartitionIndexRangeImpl(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
+ return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
+ }
+
+ @Override
+ protected ResultPartitionBytesCounter createResultPartitionBytesCounterImpl(
+ RemoteShuffleResultPartition remoteShuffleResultPartition) {
+ org.apache.flink.runtime.io.network.metrics.ResultPartitionBytesCounter resultPartitionBytes =
+ ReflectionUtils.readDeclaredFieldRecursiveNoThrow(
+ remoteShuffleResultPartition, "resultPartitionBytes");
+ return (bufferWithSubpartition, isBroadcast, readableBytes) -> {
+ if (isBroadcast) {
+ resultPartitionBytes.incAll(readableBytes);
+ } else {
+ resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
+ }
+ };
+ }
+
+ @Override
+ protected Buffer.DataType getEndOfPartitionDataTypeImpl() {
+ return Buffer.DataType.EVENT_BUFFER;
+ }
+
+ @Override
+ protected PartitionDescriptor createPartitionDescriptorImpl(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum) {
+ IntermediateResultPartitionID intermediateResultPartitionId =
+ new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
+ return new PartitionDescriptor(
+ intermediateDataSetId,
+ 10,
+ intermediateResultPartitionId,
+ ResultPartitionType.BLOCKING,
+ 5,
+ 1,
+ false,
+ false);
+ }
+
+ @Override
+ protected TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptorImpl(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {
+ return TaskInputsOutputsDescriptor.from(
+ inputGateNums,
+ inputChannelNums,
+ partitionReuseCount,
+ subpartitionNums,
+ inputPartitionTypes,
+ partitionTypes);
+ }
+}
diff --git a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
deleted file mode 100644
index bfba41530..000000000
--- a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
+++ /dev/null
@@ -1,187 +0,0 @@
-/*
- * 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.List;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-
-import org.apache.flink.api.java.tuple.Tuple2;
-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.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.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.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.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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.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 AbstractRemoteShuffleInputGate {
-
- public RemoteShuffleInputGate(
- CelebornConf celebornConf,
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor gateDescriptor,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- BufferDecompressor bufferDecompressor,
- int numConcurrentReading,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- super(
- celebornConf,
- ownerContext,
- gateIndex,
- gateDescriptor,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-
- @Override
- public InputChannel getChannel(int channelIndex) {
- return new FakedRemoteInputChannel(channelIndex);
- }
-
- @Override
- public Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
- InputGateDeploymentDescriptor gateDescriptor) {
- IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
- return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
- }
-
- /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
- private class FakedRemoteInputChannel extends RemoteInputChannel {
- FakedRemoteInputChannel(int channelIndex) {
- // Flink 1.18.0
- // [FLINK-31638][network] Introduce the TieredStorageConsumerClient to SingleInputGate
- // [FLINK-31642][network] Introduce the MemoryTierConsumerAgent to TieredStorageConsumerClient
- super(
- new SingleInputGate(
- inputGateDelegation.getTaskName(),
- inputGateDelegation.getGateIndex(),
- new IntermediateDataSetID(),
- ResultPartitionType.BLOCKING,
- new IndexRange(0, 0),
- 1,
- (a, b, c) -> {},
- () -> null,
- null,
- new FakedMemorySegmentProvider(),
- 0,
- new ThroughputCalculator(SystemClock.getInstance()),
- null,
- null,
- null,
- 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.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
deleted file mode 100644
index 38c5292bf..000000000
--- a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.util.Map;
-
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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(
- CelebornConf conf, NetworkBufferPool networkBufferPool, int networkBufferSize) {
- super(conf, networkBufferPool, networkBufferSize);
- }
-
- // For testing.
- @Override
- protected RemoteShuffleInputGate createInputGate(
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor igdd,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- String compressionCodec,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- BufferDecompressor bufferDecompressor =
- new BufferDecompressor(networkBufferSize, compressionCodec);
- return new RemoteShuffleInputGate(
- celebornConf,
- ownerContext,
- gateIndex,
- igdd,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-}
diff --git a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
deleted file mode 100644
index 256acb828..000000000
--- a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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 javax.annotation.Nullable;
-
-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.ResultPartition;
-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.apache.celeborn.plugin.flink.buffer.BufferWithSubpartition;
-import org.apache.celeborn.plugin.flink.buffer.DataBuffer;
-
-/**
- * A {@link ResultPartition} which appends records and events to {@link DataBuffer} and after the
- * {@link DataBuffer} is full, all data in the {@link DataBuffer} 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 DataBuffer} will be spilled directly.
- */
-public class RemoteShuffleResultPartition extends AbstractRemoteShuffleResultPartition {
-
- 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,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactory,
- outputGate);
- }
-
- public void updateReadableBytes(
- BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes) {
- if (isBroadcast) {
- resultPartitionBytes.incAll(readableBytes);
- } else {
- resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
- }
- }
-}
diff --git a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
deleted file mode 100644
index d49827ed9..000000000
--- a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ /dev/null
@@ -1,89 +0,0 @@
-/*
- * 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.util.List;
-
-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.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartition;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.util.function.SupplierWithException;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Factory class to create {@link RemoteShuffleResultPartition}. */
-public class RemoteShuffleResultPartitionFactory
- extends AbstractRemoteShuffleResultPartitionFactory {
-
- public RemoteShuffleResultPartitionFactory(
- CelebornConf celebornConf,
- ResultPartitionManager partitionManager,
- BufferPoolFactory bufferPoolFactory,
- int networkBufferSize) {
-
- super(celebornConf, partitionManager, bufferPoolFactory, networkBufferSize);
- }
-
- @Override
- ResultPartition createRemoteShuffleResultPartitionInternal(
- String taskNameWithSubtaskAndId,
- int partitionIndex,
- ResultPartitionID id,
- ResultPartitionType type,
- int numSubpartitions,
- int maxParallelism,
- List<SupplierWithException<BufferPool, IOException>> bufferPoolFactories,
- CelebornConf celebornConf,
- int numMappers,
- BufferCompressor bufferCompressor,
- RemoteShuffleDescriptor rsd,
- ShuffleIOMetricGroup shuffleIOMetricGroup) {
- return new RemoteShuffleResultPartition(
- taskNameWithSubtaskAndId,
- partitionIndex,
- id,
- type,
- numSubpartitions,
- maxParallelism,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactories.get(0),
- new RemoteShuffleOutputGate(
- rsd,
- numSubpartitions,
- networkBufferSize,
- bufferPoolFactories.get(1),
- celebornConf,
- numMappers,
- shuffleIOMetricGroup));
- }
-
- protected BufferCompressor getBufferCompressor() {
- return CompressionCodec.NONE.name().equals(compressionCodec)
- ? null
- : new BufferCompressor(networkBufferSize, compressionCodec);
- }
-}
diff --git a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
deleted file mode 100644
index 4e0652472..000000000
--- a/client-flink/flink-1.18/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.time.Duration;
-
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
-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.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;
-
-public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFactory {
-
- @Override
- public ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> createShuffleEnvironment(
- ShuffleEnvironmentContext shuffleEnvironmentContext) {
- AbstractRemoteShuffleServiceParameters parameters =
- initializePreCreateShuffleEnvironment(shuffleEnvironmentContext);
- RemoteShuffleResultPartitionFactory resultPartitionFactory =
- new RemoteShuffleResultPartitionFactory(
- parameters.celebornConf,
- parameters.resultPartitionManager,
- parameters.networkBufferPool,
- parameters.bufferSize);
- RemoteShuffleInputGateFactory inputGateFactory =
- new RemoteShuffleInputGateFactory(
- parameters.celebornConf, parameters.networkBufferPool, parameters.bufferSize);
-
- return new RemoteShuffleEnvironment(
- parameters.networkBufferPool,
- parameters.resultPartitionManager,
- resultPartitionFactory,
- inputGateFactory,
- parameters.celebornConf,
- new NettyShuffleEnvironmentWrapper(nettyShuffleServiceFactory, shuffleEnvironmentContext));
- }
-
- @Override
- public Duration getRequestSegmentsTimeout(Configuration configuration) {
- return Duration.ofMillis(
- configuration.get(
- NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
- }
-}
diff --git a/client-flink/flink-1.18/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim b/client-flink/flink-1.18/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
new file mode 100644
index 000000000..f3b3ab048
--- /dev/null
+++ b/client-flink/flink-1.18/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
@@ -0,0 +1,16 @@
+# 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.
+
+org.apache.celeborn.plugin.flink.CelebornFlinkShimV118
\ No newline at end of file
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV118.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV118.java
new file mode 100644
index 000000000..2f9e95f35
--- /dev/null
+++ b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV118.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleMasterSuiteJV118 extends RemoteShuffleMasterSuiteJ {}
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
deleted file mode 100644
index 548336db6..000000000
--- a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJ.java
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * 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.Mockito.mock;
-
-import org.apache.flink.runtime.io.network.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionManager;
-import org.junit.Assert;
-import org.junit.Test;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
-public class RemoteShuffleResultPartitionFactorySuiteJ {
-
- @Test
- public void testGetBufferCompressor() {
- CelebornConf celebornConf = new CelebornConf();
- for (CompressionCodec compressionCodec : CompressionCodec.values()) {
- RemoteShuffleResultPartitionFactory partitionFactory =
- new RemoteShuffleResultPartitionFactory(
- celebornConf.set(
- CelebornConf.SHUFFLE_COMPRESSION_CODEC().key(), compressionCodec.name()),
- mock(ResultPartitionManager.class),
- mock(BufferPoolFactory.class),
- 1);
- if (CompressionCodec.NONE.equals(compressionCodec)) {
- Assert.assertNull(partitionFactory.getBufferCompressor());
- } else {
- Assert.assertNotNull(partitionFactory.getBufferCompressor());
- }
- }
- }
-}
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV118.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV118.java
new file mode 100644
index 000000000..dccef51e3
--- /dev/null
+++ b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV118.java
@@ -0,0 +1,22 @@
+/*
+ * 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;
+
+/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
+public class RemoteShuffleResultPartitionFactorySuiteJV118
+ extends RemoteShuffleResultPartitionFactorySuiteJ {}
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV118.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV118.java
new file mode 100644
index 000000000..d003a66bf
--- /dev/null
+++ b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV118.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleResultPartitionSuiteJV118 extends RemoteShuffleResultPartitionSuiteJ {}
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
deleted file mode 100644
index 49affdee4..000000000
--- a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * 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.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.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 RemoteShuffleServiceFactorySuiteJ {
- @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 parentMetric = mock(MetricGroup.class);
- when(shuffleEnvironmentContext.getParentMetricGroup()).thenReturn(parentMetric);
- MetricGroup childGroup = mock(MetricGroup.class);
- MetricGroup childChildGroup = mock(MetricGroup.class);
- when(parentMetric.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/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV118.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV118.java
new file mode 100644
index 000000000..638557e3f
--- /dev/null
+++ b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV118.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleServiceFactorySuiteJV118 extends RemoteShuffleServiceFactorySuiteJ {}
diff --git a/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV118.java b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV118.java
new file mode 100644
index 000000000..766f1c48d
--- /dev/null
+++ b/client-flink/flink-1.18/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV118.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class ShuffleResourceTrackerSuiteJV118 extends ShuffleResourceTrackerSuiteJ {}
diff --git a/client-flink/flink-1.19/pom.xml b/client-flink/flink-1.19/pom.xml
index 72e197692..de8f16718 100644
--- a/client-flink/flink-1.19/pom.xml
+++ b/client-flink/flink-1.19/pom.xml
@@ -56,6 +56,13 @@
</exclusions>
</dependency>
<!-- Test dependencies -->
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
diff --git a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV119.java b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV119.java
new file mode 100644
index 000000000..640c67f24
--- /dev/null
+++ b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV119.java
@@ -0,0 +1,166 @@
+/*
+ * 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.net.InetAddress;
+import java.time.Duration;
+import java.util.Map;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+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.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.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+import org.apache.flink.runtime.io.network.partition.ResultSubpartitionIndexSet;
+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.jobgraph.IntermediateResultPartitionID;
+import org.apache.flink.runtime.shuffle.PartitionDescriptor;
+import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.runtime.throughput.ThroughputCalculator;
+import org.apache.flink.util.clock.SystemClock;
+
+import org.apache.celeborn.plugin.flink.utils.ReflectionUtils;
+
+public class CelebornFlinkShimV119 extends CelebornFlinkShim {
+ @Override
+ protected BufferCompressor createBufferCompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferCompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected BufferDecompressor createBufferDecompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferDecompressor(networkBufferSize, compressionCodec);
+ }
+
+ @Override
+ protected Duration getRequestSegmentsTimeoutImpl(Configuration configuration) {
+ return Duration.ofMillis(
+ configuration.get(
+ NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
+ }
+
+ @Override
+ protected InputChannel createFakeChannelImpl(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex) {
+ return new RemoteInputChannel(
+ new SingleInputGate(
+ inputGateDelegation.getTaskName(),
+ inputGateDelegation.getGateIndex(),
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ ReflectionUtils.createEmptyImplementation(MemorySegmentProvider.class),
+ 0,
+ new ThroughputCalculator(SystemClock.getInstance()),
+ null),
+ channelIndex,
+ new ResultPartitionID(),
+ new ResultSubpartitionIndexSet(new IndexRange(0, 0)),
+ new ConnectionID(
+ new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1), 0),
+ new LocalConnectionManager(),
+ 0,
+ 0,
+ 0,
+ 0,
+ new SimpleCounter(),
+ new SimpleCounter(),
+ ReflectionUtils.createEmptyImplementation(ChannelStateWriter.class));
+ }
+
+ @Override
+ protected Tuple2<Integer, Integer> getConsumedSubpartitionIndexRangeImpl(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
+ return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
+ }
+
+ @Override
+ protected ResultPartitionBytesCounter createResultPartitionBytesCounterImpl(
+ RemoteShuffleResultPartition remoteShuffleResultPartition) {
+ org.apache.flink.runtime.io.network.metrics.ResultPartitionBytesCounter resultPartitionBytes =
+ ReflectionUtils.readDeclaredFieldRecursiveNoThrow(
+ remoteShuffleResultPartition, "resultPartitionBytes");
+ return (bufferWithSubpartition, isBroadcast, readableBytes) -> {
+ if (isBroadcast) {
+ resultPartitionBytes.incAll(readableBytes);
+ } else {
+ resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
+ }
+ };
+ }
+
+ @Override
+ protected Buffer.DataType getEndOfPartitionDataTypeImpl() {
+ return Buffer.DataType.END_OF_PARTITION;
+ }
+
+ @Override
+ protected PartitionDescriptor createPartitionDescriptorImpl(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum) {
+ IntermediateResultPartitionID intermediateResultPartitionId =
+ new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
+ return new PartitionDescriptor(
+ intermediateDataSetId,
+ 10,
+ intermediateResultPartitionId,
+ ResultPartitionType.BLOCKING,
+ 5,
+ 1,
+ false,
+ false,
+ false);
+ }
+
+ @Override
+ protected TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptorImpl(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {
+ return TaskInputsOutputsDescriptor.from(
+ inputGateNums,
+ inputChannelNums,
+ partitionReuseCount,
+ subpartitionNums,
+ inputPartitionTypes,
+ partitionTypes);
+ }
+}
diff --git a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
deleted file mode 100644
index c6f909b4a..000000000
--- a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
+++ /dev/null
@@ -1,195 +0,0 @@
-/*
- * 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.List;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-
-import org.apache.flink.api.java.tuple.Tuple2;
-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.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.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.ResultSubpartitionIndexSet;
-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.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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.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 AbstractRemoteShuffleInputGate {
-
- /** The type of the partition the input gate is consuming. */
- private final ResultPartitionType consumedPartitionType;
-
- public RemoteShuffleInputGate(
- CelebornConf celebornConf,
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor gateDescriptor,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- BufferDecompressor bufferDecompressor,
- int numConcurrentReading,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- super(
- celebornConf,
- ownerContext,
- gateIndex,
- gateDescriptor,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- this.consumedPartitionType = gateDescriptor.getConsumedPartitionType();
- }
-
- @Override
- public InputChannel getChannel(int channelIndex) {
- return new FakedRemoteInputChannel(channelIndex);
- }
-
- @Override
- public Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
- InputGateDeploymentDescriptor gateDescriptor) {
- IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
- return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
- }
-
- public ResultPartitionType getConsumedPartitionType() {
- // Flink 1.19.3
- // [FLINK-37783] Auto-disable buffer debloating for tiered shuffle
- return consumedPartitionType;
- }
-
- /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
- private class FakedRemoteInputChannel extends RemoteInputChannel {
- FakedRemoteInputChannel(int channelIndex) {
- // Flink 1.19.0
- // [FLINK-25055][network] Support listen and notify mechanism for partition request
- // [FLINK-33743][runtime] Support consuming multiple subpartitions on a single channel
- super(
- new SingleInputGate(
- inputGateDelegation.getTaskName(),
- inputGateDelegation.getGateIndex(),
- new IntermediateDataSetID(),
- ResultPartitionType.BLOCKING,
- 1,
- (a, b, c) -> {},
- () -> null,
- null,
- new FakedMemorySegmentProvider(),
- 0,
- new ThroughputCalculator(SystemClock.getInstance()),
- null),
- channelIndex,
- new ResultPartitionID(),
- new ResultSubpartitionIndexSet(new IndexRange(0, 0)),
- new ConnectionID(
- new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1),
- 0),
- new LocalConnectionManager(),
- 0,
- 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.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
deleted file mode 100644
index 38c5292bf..000000000
--- a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.util.Map;
-
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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(
- CelebornConf conf, NetworkBufferPool networkBufferPool, int networkBufferSize) {
- super(conf, networkBufferPool, networkBufferSize);
- }
-
- // For testing.
- @Override
- protected RemoteShuffleInputGate createInputGate(
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor igdd,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- String compressionCodec,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- BufferDecompressor bufferDecompressor =
- new BufferDecompressor(networkBufferSize, compressionCodec);
- return new RemoteShuffleInputGate(
- celebornConf,
- ownerContext,
- gateIndex,
- igdd,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-}
diff --git a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
deleted file mode 100644
index 256acb828..000000000
--- a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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 javax.annotation.Nullable;
-
-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.ResultPartition;
-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.apache.celeborn.plugin.flink.buffer.BufferWithSubpartition;
-import org.apache.celeborn.plugin.flink.buffer.DataBuffer;
-
-/**
- * A {@link ResultPartition} which appends records and events to {@link DataBuffer} and after the
- * {@link DataBuffer} is full, all data in the {@link DataBuffer} 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 DataBuffer} will be spilled directly.
- */
-public class RemoteShuffleResultPartition extends AbstractRemoteShuffleResultPartition {
-
- 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,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactory,
- outputGate);
- }
-
- public void updateReadableBytes(
- BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes) {
- if (isBroadcast) {
- resultPartitionBytes.incAll(readableBytes);
- } else {
- resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
- }
- }
-}
diff --git a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
deleted file mode 100644
index d49827ed9..000000000
--- a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ /dev/null
@@ -1,89 +0,0 @@
-/*
- * 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.util.List;
-
-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.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartition;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.util.function.SupplierWithException;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Factory class to create {@link RemoteShuffleResultPartition}. */
-public class RemoteShuffleResultPartitionFactory
- extends AbstractRemoteShuffleResultPartitionFactory {
-
- public RemoteShuffleResultPartitionFactory(
- CelebornConf celebornConf,
- ResultPartitionManager partitionManager,
- BufferPoolFactory bufferPoolFactory,
- int networkBufferSize) {
-
- super(celebornConf, partitionManager, bufferPoolFactory, networkBufferSize);
- }
-
- @Override
- ResultPartition createRemoteShuffleResultPartitionInternal(
- String taskNameWithSubtaskAndId,
- int partitionIndex,
- ResultPartitionID id,
- ResultPartitionType type,
- int numSubpartitions,
- int maxParallelism,
- List<SupplierWithException<BufferPool, IOException>> bufferPoolFactories,
- CelebornConf celebornConf,
- int numMappers,
- BufferCompressor bufferCompressor,
- RemoteShuffleDescriptor rsd,
- ShuffleIOMetricGroup shuffleIOMetricGroup) {
- return new RemoteShuffleResultPartition(
- taskNameWithSubtaskAndId,
- partitionIndex,
- id,
- type,
- numSubpartitions,
- maxParallelism,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactories.get(0),
- new RemoteShuffleOutputGate(
- rsd,
- numSubpartitions,
- networkBufferSize,
- bufferPoolFactories.get(1),
- celebornConf,
- numMappers,
- shuffleIOMetricGroup));
- }
-
- protected BufferCompressor getBufferCompressor() {
- return CompressionCodec.NONE.name().equals(compressionCodec)
- ? null
- : new BufferCompressor(networkBufferSize, compressionCodec);
- }
-}
diff --git a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
deleted file mode 100644
index 4e0652472..000000000
--- a/client-flink/flink-1.19/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * 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.time.Duration;
-
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
-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.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;
-
-public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFactory {
-
- @Override
- public ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> createShuffleEnvironment(
- ShuffleEnvironmentContext shuffleEnvironmentContext) {
- AbstractRemoteShuffleServiceParameters parameters =
- initializePreCreateShuffleEnvironment(shuffleEnvironmentContext);
- RemoteShuffleResultPartitionFactory resultPartitionFactory =
- new RemoteShuffleResultPartitionFactory(
- parameters.celebornConf,
- parameters.resultPartitionManager,
- parameters.networkBufferPool,
- parameters.bufferSize);
- RemoteShuffleInputGateFactory inputGateFactory =
- new RemoteShuffleInputGateFactory(
- parameters.celebornConf, parameters.networkBufferPool, parameters.bufferSize);
-
- return new RemoteShuffleEnvironment(
- parameters.networkBufferPool,
- parameters.resultPartitionManager,
- resultPartitionFactory,
- inputGateFactory,
- parameters.celebornConf,
- new NettyShuffleEnvironmentWrapper(nettyShuffleServiceFactory, shuffleEnvironmentContext));
- }
-
- @Override
- public Duration getRequestSegmentsTimeout(Configuration configuration) {
- return Duration.ofMillis(
- configuration.get(
- NettyShuffleEnvironmentOptions.NETWORK_EXCLUSIVE_BUFFERS_REQUEST_TIMEOUT_MILLISECONDS));
- }
-}
diff --git a/client-flink/flink-1.19/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim b/client-flink/flink-1.19/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
new file mode 100644
index 000000000..b9790bd26
--- /dev/null
+++ b/client-flink/flink-1.19/src/main/resources/META-INF/services/org.apache.celeborn.plugin.flink.CelebornFlinkShim
@@ -0,0 +1,16 @@
+# 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.
+
+org.apache.celeborn.plugin.flink.CelebornFlinkShimV119
\ No newline at end of file
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
deleted file mode 100644
index b702f6063..000000000
--- a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJ.java
+++ /dev/null
@@ -1,381 +0,0 @@
-/*
- * 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.net.InetAddress;
-import java.net.UnknownHostException;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Random;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-
-import org.apache.flink.api.common.BatchShuffleMode;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.ExecutionOptions;
-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.NettyShuffleServiceFactory;
-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.NettyShuffleDescriptor;
-import org.apache.flink.runtime.shuffle.PartitionDescriptor;
-import org.apache.flink.runtime.shuffle.ProducerDescriptor;
-import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
-import org.apache.flink.runtime.shuffle.ShuffleMasterContext;
-import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
-import org.junit.After;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.FallbackPolicy;
-import org.apache.celeborn.common.util.Utils$;
-import org.apache.celeborn.plugin.flink.fallback.ForceFallbackPolicy;
-import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
-
-public class RemoteShuffleMasterSuiteJ {
-
- private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleMasterSuiteJ.class);
- private RemoteShuffleMaster remoteShuffleMaster;
- private Configuration configuration;
-
- @Before
- public void setUp() {
- configuration = new Configuration();
- int startPort = Utils$.MODULE$.selectRandomInt(1024, 65535);
- configuration.setInteger("celeborn.master.port", startPort);
- configuration.setString("celeborn.master.endpoints", "localhost:" + startPort);
- configuration.setString("celeborn.client.application.heartbeatInterval", "30s");
- remoteShuffleMaster = createShuffleMaster(configuration);
- }
-
- @Test
- public void testRegisterJob() {
- JobShuffleContext jobShuffleContext = createJobShuffleContext(JobID.generate());
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- // reRunRegister job
- try {
- remoteShuffleMaster.registerJob(jobShuffleContext);
- } catch (Exception e) {
- Assert.assertTrue(true);
- }
-
- // unRegister job
- remoteShuffleMaster.unregisterJob(jobShuffleContext.getJobId());
- remoteShuffleMaster.registerJob(jobShuffleContext);
- }
-
- @Test
- public void testRegisterJobWithForceFallbackPolicy() {
- configuration.setString(
- CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
- remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
- Assert.assertTrue(remoteShuffleMaster.jobFallbackPolicies().containsKey(jobID));
- remoteShuffleMaster.unregisterJob(jobShuffleContext.getJobId());
- Assert.assertTrue(remoteShuffleMaster.jobFallbackPolicies().isEmpty());
- }
-
- @Test
- public void testRegisterPartitionWithProducer()
- throws UnknownHostException, ExecutionException, InterruptedException {
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- RemoteShuffleDescriptor remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(1, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- ShuffleResource shuffleResource = remoteShuffleDescriptor.getShuffleResource();
- ShuffleResourceDescriptor mapPartitionShuffleDescriptor =
- shuffleResource.getMapPartitionShuffleDescriptor();
-
- LOG.info("remoteShuffleDescriptor:{}", remoteShuffleDescriptor);
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getPartitionId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getAttemptId());
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getMapId());
-
- // use same dataset id
- partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 1);
- remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(2, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- mapPartitionShuffleDescriptor =
- remoteShuffleDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
-
- // use another attemptId
- producerDescriptor = createProducerDescriptor();
- remoteShuffleDescriptor =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertEquals(3, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- mapPartitionShuffleDescriptor =
- remoteShuffleDescriptor.getShuffleResource().getMapPartitionShuffleDescriptor();
- Assert.assertEquals(0, mapPartitionShuffleDescriptor.getShuffleId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getAttemptId());
- Assert.assertEquals(1, mapPartitionShuffleDescriptor.getMapId());
- }
-
- @Test
- public void testRegisterPartitionWithProducerForForceFallbackPolicy()
- throws UnknownHostException, ExecutionException, InterruptedException {
- configuration.setString(
- CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(), FallbackPolicy.ALWAYS.name());
- remoteShuffleMaster = createShuffleMaster(configuration, new NettyShuffleServiceFactory());
- JobID jobID = JobID.generate();
- JobShuffleContext jobShuffleContext = createJobShuffleContext(jobID);
- remoteShuffleMaster.registerJob(jobShuffleContext);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- ShuffleDescriptor shuffleDescriptor =
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
- .get();
- Assert.assertTrue(shuffleDescriptor instanceof NettyShuffleDescriptor);
- Assert.assertEquals(1, remoteShuffleMaster.lifecycleManager().shuffleCount().sum());
- Map<String, Long> shuffleFallbackCounts =
- remoteShuffleMaster.lifecycleManager().shuffleFallbackCounts();
- Assert.assertEquals(1, shuffleFallbackCounts.size());
- Assert.assertEquals(
- 1L, shuffleFallbackCounts.get(ForceFallbackPolicy.class.getName()).longValue());
- }
-
- @Test
- public void testRegisterMultipleJobs()
- throws UnknownHostException, ExecutionException, InterruptedException {
- JobID jobID1 = JobID.generate();
- JobShuffleContext jobShuffleContext1 = createJobShuffleContext(jobID1);
- remoteShuffleMaster.registerJob(jobShuffleContext1);
-
- JobID jobID2 = JobID.generate();
- JobShuffleContext jobShuffleContext2 = createJobShuffleContext(jobID2);
- remoteShuffleMaster.registerJob(jobShuffleContext2);
-
- IntermediateDataSetID intermediateDataSetID = new IntermediateDataSetID();
- PartitionDescriptor partitionDescriptor = createPartitionDescriptor(intermediateDataSetID, 0);
- ProducerDescriptor producerDescriptor = createProducerDescriptor();
- RemoteShuffleDescriptor remoteShuffleDescriptor1 =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID1, partitionDescriptor, producerDescriptor)
- .get();
-
- // use same datasetId but different jobId
- RemoteShuffleDescriptor remoteShuffleDescriptor2 =
- (RemoteShuffleDescriptor)
- remoteShuffleMaster
- .registerPartitionWithProducer(jobID2, partitionDescriptor, producerDescriptor)
- .get();
-
- Assert.assertEquals(
- remoteShuffleDescriptor1
- .getShuffleResource()
- .getMapPartitionShuffleDescriptor()
- .getShuffleId(),
- 0);
- Assert.assertEquals(
- remoteShuffleDescriptor2
- .getShuffleResource()
- .getMapPartitionShuffleDescriptor()
- .getShuffleId(),
- 1);
- }
-
- @Test
- 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();
- IntermediateDataSetID outputDataSetID0 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID1 = new IntermediateDataSetID();
- IntermediateDataSetID outputDataSetID2 = new IntermediateDataSetID();
- Random random = new Random();
- numberOfInputGateChannels.put(inputDataSetID0, random.nextInt(1000));
- numberOfInputGateChannels.put(inputDataSetID1, 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(
- 3,
- numberOfInputGateChannels,
- numbersOfResultSubpartitions,
- subPartitionNums,
- inputPartitionTypes,
- resultPartitionTypes));
-
- CelebornConf celebornConf = FlinkUtils.toCelebornConf(configuration);
-
- long numBytesPerGate = celebornConf.clientFlinkMemoryPerInputGate();
- long expectedInput = 2 * numBytesPerGate;
-
- long numBytesPerResultPartition = celebornConf.clientFlinkMemoryPerResultPartition();
- long expectedOutput = 3 * numBytesPerResultPartition;
- MemorySize expected = new MemorySize(expectedInput + expectedOutput);
-
- Assert.assertEquals(expected, calculated);
- }
-
- @Test
- public void testInvalidShuffleConfig() {
- Assert.assertThrows(
- String.format(
- "The config option %s should configure as %s",
- ExecutionOptions.BATCH_SHUFFLE_MODE.key(),
- BatchShuffleMode.ALL_EXCHANGES_BLOCKING.name()),
- IllegalArgumentException.class,
- () ->
- createShuffleMaster(
- new Configuration()
- .set(
- ExecutionOptions.BATCH_SHUFFLE_MODE,
- BatchShuffleMode.ALL_EXCHANGES_PIPELINED)));
- Configuration configuration = new Configuration();
- configuration.setString(CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED().key(), "true");
- Assert.assertThrows(
- String.format(
- "Flink does not support replicate shuffle data. Please check the config %s.",
- CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED().key()),
- IllegalArgumentException.class,
- () -> createShuffleMaster(configuration));
- }
-
- @After
- public void tearDown() {
- if (remoteShuffleMaster != null) {
- try {
- remoteShuffleMaster.close();
- } catch (Exception e) {
- LOG.warn(e.getMessage(), e);
- }
- }
- }
-
- public RemoteShuffleMaster createShuffleMaster(Configuration configuration) {
- return createShuffleMaster(configuration, null);
- }
-
- public RemoteShuffleMaster createShuffleMaster(
- Configuration configuration, NettyShuffleServiceFactory nettyShuffleServiceFactory) {
- remoteShuffleMaster =
- new RemoteShuffleMaster(
- new ShuffleMasterContext() {
- @Override
- public Configuration getConfiguration() {
- return configuration;
- }
-
- @Override
- public void onFatalError(Throwable throwable) {
- System.exit(-1);
- }
- },
- nettyShuffleServiceFactory);
-
- return remoteShuffleMaster;
- }
-
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public org.apache.flink.api.common.JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
- public PartitionDescriptor createPartitionDescriptor(
- IntermediateDataSetID intermediateDataSetId, int partitionNum) {
- IntermediateResultPartitionID intermediateResultPartitionId =
- new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
- // Flink 1.19.0
- // [FLINK-33743][runtime] Align watermark at subpartition granularity
- return new PartitionDescriptor(
- intermediateDataSetId,
- 10,
- intermediateResultPartitionId,
- ResultPartitionType.BLOCKING,
- 5,
- 1,
- false,
- false,
- false);
- }
-
- public ProducerDescriptor createProducerDescriptor() throws UnknownHostException {
- 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.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV119.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV119.java
new file mode 100644
index 000000000..40dd01f4d
--- /dev/null
+++ b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleMasterSuiteJV119.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleMasterSuiteJV119 extends RemoteShuffleMasterSuiteJ {}
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV119.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV119.java
new file mode 100644
index 000000000..f1636bf9c
--- /dev/null
+++ b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactorySuiteJV119.java
@@ -0,0 +1,22 @@
+/*
+ * 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;
+
+/** Tests for {@link RemoteShuffleResultPartitionFactory}. */
+public class RemoteShuffleResultPartitionFactorySuiteJV119
+ extends RemoteShuffleResultPartitionFactorySuiteJ {}
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
deleted file mode 100644
index 40a703db9..000000000
--- a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJ.java
+++ /dev/null
@@ -1,640 +0,0 @@
-/*
- * 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.ArgumentMatchers.anyBoolean;
-import static org.mockito.ArgumentMatchers.anyInt;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.metrics.groups.ShuffleMetricGroup;
-import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
-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.DataBuffer;
-import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
-import org.apache.celeborn.plugin.flink.utils.BufferUtils;
-
-public class RemoteShuffleResultPartitionSuiteJ {
- private final int networkBufferSize = 32 * 1024;
- private final BufferCompressor bufferCompressor = new BufferCompressor(networkBufferSize, "lz4");
- private final RemoteShuffleOutputGate remoteShuffleOutputGate =
- mock(RemoteShuffleOutputGate.class);
- 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());
- DataBuffer dataBuffer = remoteShuffleResultPartition.getDelegation().getUnicastDataBuffer();
- ByteBuffer byteBuffer = ByteBuffer.wrap(new byte[] {1, 2, 3});
- dataBuffer.append(byteBuffer, 0, Buffer.DataType.DATA_BUFFER);
- remoteShuffleResultPartition.getDelegation().flushDataBuffer(dataBuffer, true);
- }
-
- private List<SupplierWithException<BufferPool, IOException>> createBufferPoolFactory() {
- NetworkBufferPool networkBufferPool =
- new NetworkBufferPool(256 * 8, 32 * 1024, Duration.ofSeconds(1));
-
- 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());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(1, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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());
- assertEquals(2, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(2, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.broadcastRecord(ByteBuffer.allocate(bufferSize));
- assertEquals(2, sortBufferPool.bestEffortGetNumOfUsedBuffers());
- assertEquals(3, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(3, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.flush(0);
- assertEquals(0, sortBufferPool.bestEffortGetNumOfUsedBuffers());
-
- partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 2);
- partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 3);
- assertEquals(3, sortBufferPool.bestEffortGetNumOfUsedBuffers());
- assertEquals(5, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(5, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- 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);
- }
- }
- assertEquals(numRecords, outputGate.shuffleIOMetricGroup.getNumRecordsOut().getCount());
- assertEquals(numRecords, outputGate.shuffleIOMetricGroup.getNumRecordsOutRate().getCount());
-
- partitionWriter.finish();
- assertTrue(outputGate.isFinished());
- partitionWriter.close();
- assertTrue(outputGate.isClosed());
-
- for (int subpartition = 0; subpartition < numSubpartitions; ++subpartition) {
- ByteBuffer record = EventSerializer.toSerializedEvent(EndOfPartitionEvent.INSTANCE);
- // Flink 1.19.0
- // [FLINK-33743][runtime] Support consuming multiple subpartitions on a single channel
- recordDataWritten(
- record, Buffer.DataType.END_OF_PARTITION, 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,
- new ShuffleIOMetricGroup(
- new ShuffleMetricGroup(
- UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
- 1,
- CelebornConf.CLIENT_METRICS_SCOPE_NAMING_SHUFFLE().defaultValueString())));
- 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(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() {
- 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, System.currentTimeMillis(), 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.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV119.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV119.java
new file mode 100644
index 000000000..8b8f6bb45
--- /dev/null
+++ b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionSuiteJV119.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleResultPartitionSuiteJV119 extends RemoteShuffleResultPartitionSuiteJ {}
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
deleted file mode 100644
index 49affdee4..000000000
--- a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJ.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * 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.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.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 RemoteShuffleServiceFactorySuiteJ {
- @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 parentMetric = mock(MetricGroup.class);
- when(shuffleEnvironmentContext.getParentMetricGroup()).thenReturn(parentMetric);
- MetricGroup childGroup = mock(MetricGroup.class);
- MetricGroup childChildGroup = mock(MetricGroup.class);
- when(parentMetric.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/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV119.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV119.java
new file mode 100644
index 000000000..03a096883
--- /dev/null
+++ b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactorySuiteJV119.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class RemoteShuffleServiceFactorySuiteJV119 extends RemoteShuffleServiceFactorySuiteJ {}
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
deleted file mode 100644
index b93ae3ea1..000000000
--- a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJ.java
+++ /dev/null
@@ -1,138 +0,0 @@
-/*
- * 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.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-
-import com.google.common.collect.Sets;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
-import org.apache.flink.runtime.shuffle.JobShuffleContext;
-import org.junit.Assert;
-import org.junit.Test;
-import org.mockito.Mockito;
-
-import org.apache.celeborn.client.LifecycleManager;
-import org.apache.celeborn.client.listener.WorkersStatus;
-import org.apache.celeborn.common.meta.ShufflePartitionLocationInfo;
-import org.apache.celeborn.common.meta.WorkerInfo;
-import org.apache.celeborn.common.protocol.PartitionLocation;
-import org.apache.celeborn.common.util.JavaUtils;
-
-public class ShuffleResourceTrackerSuiteJ {
-
- @Test
- public void testNotifyUnknownWorkers() {
- LifecycleManager lifecycleManager = Mockito.mock(LifecycleManager.class);
- ScheduledThreadPoolExecutor executor = Mockito.mock(ScheduledThreadPoolExecutor.class);
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map = JavaUtils.newConcurrentHashMap();
- WorkerInfo workerInfo = new WorkerInfo("mock", -1, -1, -1, -1);
- map.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map2 = JavaUtils.newConcurrentHashMap();
- map2.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- ConcurrentHashMap<String, ShufflePartitionLocationInfo> map3 = JavaUtils.newConcurrentHashMap();
- map3.put(workerInfo.toUniqueId(), mockShufflePartitionLocationInfo(workerInfo));
-
- Mockito.when(lifecycleManager.workerSnapshots(Mockito.anyInt())).thenReturn(map, map2, map3);
-
- ShuffleResourceTracker shuffleResourceTracker =
- new ShuffleResourceTracker(executor, lifecycleManager);
-
- JobID jobID1 = new JobID();
- shuffleResourceTracker.registerJob(createJobShuffleContext(jobID1));
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 1, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 2, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 1, 3, new ResultPartitionID());
- shuffleResourceTracker.addPartitionResource(jobID1, 2, 3, new ResultPartitionID());
-
- JobID jobID2 = new JobID();
- shuffleResourceTracker.registerJob(createJobShuffleContext(jobID2));
- shuffleResourceTracker.addPartitionResource(jobID2, 3, 1, new ResultPartitionID());
-
- List<WorkerInfo> workerInfoList = new ArrayList<>();
- workerInfoList.add(workerInfo);
- shuffleResourceTracker.notifyChangedWorkersStatus(new WorkersStatus(workerInfoList, null));
-
- Assert.assertEquals(
- Sets.newHashSet(3),
- shuffleResourceTracker
- .getJobResourceListener(jobID1)
- .getResultPartitionMap()
- .get(2)
- .keySet());
- Assert.assertEquals(
- Sets.newHashSet(3),
- shuffleResourceTracker
- .getJobResourceListener(jobID1)
- .getResultPartitionMap()
- .get(2)
- .keySet());
-
- Assert.assertTrue(
- shuffleResourceTracker
- .getJobResourceListener(jobID2)
- .getResultPartitionMap()
- .get(3)
- .isEmpty());
- }
-
- public ShufflePartitionLocationInfo mockShufflePartitionLocationInfo(WorkerInfo workerInfo) {
- ShufflePartitionLocationInfo shufflePartitionLocationInfo =
- new ShufflePartitionLocationInfo(workerInfo);
-
- List<PartitionLocation> primaryLocations = new ArrayList<>();
- primaryLocations.add(mockShufflePartitionLocationInfo(1));
- primaryLocations.add(mockShufflePartitionLocationInfo(2));
-
- List<PartitionLocation> replicaLocations = new ArrayList<>();
- replicaLocations.add(mockShufflePartitionLocationInfo(3));
- replicaLocations.add(mockShufflePartitionLocationInfo(4));
-
- shufflePartitionLocationInfo.addPrimaryPartitions(primaryLocations);
- shufflePartitionLocationInfo.addReplicaPartitions(replicaLocations);
- return shufflePartitionLocationInfo;
- }
-
- public JobShuffleContext createJobShuffleContext(JobID jobId) {
- return new JobShuffleContext() {
- @Override
- public JobID getJobId() {
- return jobId;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> collection) {
- return CompletableFuture.completedFuture(null);
- }
- };
- }
-
- private PartitionLocation mockShufflePartitionLocationInfo(int partitionId) {
- return new PartitionLocation(
- partitionId, -1, "mock", -1, -1, -1, -1, PartitionLocation.Mode.PRIMARY);
- }
-}
diff --git a/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV119.java b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV119.java
new file mode 100644
index 000000000..d09468c63
--- /dev/null
+++ b/client-flink/flink-1.19/src/test/java/org/apache/celeborn/plugin/flink/ShuffleResourceTrackerSuiteJV119.java
@@ -0,0 +1,20 @@
+/*
+ * 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;
+
+public class ShuffleResourceTrackerSuiteJV119 extends ShuffleResourceTrackerSuiteJ {}
diff --git a/client-flink/flink-1.20/pom.xml b/client-flink/flink-1.20/pom.xml
index ffac3efd1..2736b7eae 100644
--- a/client-flink/flink-1.20/pom.xml
+++ b/client-flink/flink-1.20/pom.xml
@@ -44,6 +44,25 @@
<artifactId>celeborn-client-flink-common_${scala.binary.version}</artifactId>
<version>${project.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-common-tiered_${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>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.celeborn</groupId>
+ <artifactId>celeborn-client-flink-common-tiered_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-runtime</artifactId>
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV120.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV120.java
new file mode 100644
index 000000000..c31e21908
--- /dev/null
+++ b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/CelebornFlinkShimV120.java
@@ -0,0 +1,207 @@
+/*
+ * 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.net.InetAddress;
+import java.time.Duration;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
+import org.apache.flink.core.memory.MemorySegmentProvider;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
+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.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.partition.ResultPartitionID;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
+import org.apache.flink.runtime.io.network.partition.ResultSubpartitionIndexSet;
+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.io.network.partition.hybrid.tiered.common.TieredStorageIdMappingUtils;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleHandler;
+import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
+import org.apache.flink.runtime.jobgraph.IntermediateResultPartitionID;
+import org.apache.flink.runtime.shuffle.JobShuffleContext;
+import org.apache.flink.runtime.shuffle.PartitionDescriptor;
+import org.apache.flink.runtime.shuffle.PartitionWithMetrics;
+import org.apache.flink.runtime.shuffle.TaskInputsOutputsDescriptor;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.runtime.throughput.ThroughputCalculator;
+import org.apache.flink.util.clock.SystemClock;
+
+import org.apache.celeborn.plugin.flink.tiered.CelebornFlinkTierShim;
+import org.apache.celeborn.plugin.flink.utils.ReflectionUtils;
+
+public class CelebornFlinkShimV120 extends CelebornFlinkTierShim {
+ @Override
+ protected BufferCompressor createBufferCompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferCompressor(
+ networkBufferSize,
+ NettyShuffleEnvironmentOptions.CompressionCodec.valueOf(compressionCodec));
+ }
+
+ @Override
+ protected BufferDecompressor createBufferDecompressorImpl(
+ int networkBufferSize, String compressionCodec) {
+ return new BufferDecompressor(
+ networkBufferSize,
+ NettyShuffleEnvironmentOptions.CompressionCodec.valueOf(compressionCodec));
+ }
+
+ @Override
+ protected Duration getRequestSegmentsTimeoutImpl(Configuration configuration) {
+ return configuration.get(NettyShuffleEnvironmentOptions.NETWORK_BUFFERS_REQUEST_TIMEOUT);
+ }
+
+ @Override
+ protected InputChannel createFakeChannelImpl(
+ RemoteShuffleInputGateDelegation inputGateDelegation, int channelIndex) {
+ return new RemoteInputChannel(
+ new SingleInputGate(
+ inputGateDelegation.getTaskName(),
+ inputGateDelegation.getGateIndex(),
+ new IntermediateDataSetID(),
+ ResultPartitionType.BLOCKING,
+ 1,
+ (a, b, c) -> {},
+ () -> null,
+ null,
+ ReflectionUtils.createEmptyImplementation(MemorySegmentProvider.class),
+ 0,
+ new ThroughputCalculator(SystemClock.getInstance()),
+ null),
+ channelIndex,
+ new ResultPartitionID(),
+ new ResultSubpartitionIndexSet(new IndexRange(0, 0)),
+ new ConnectionID(
+ new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1), 0),
+ new LocalConnectionManager(),
+ 0,
+ 0,
+ 0,
+ 0,
+ new SimpleCounter(),
+ new SimpleCounter(),
+ ReflectionUtils.createEmptyImplementation(ChannelStateWriter.class));
+ }
+
+ @Override
+ protected Tuple2<Integer, Integer> getConsumedSubpartitionIndexRangeImpl(
+ InputGateDeploymentDescriptor gateDescriptor) {
+ IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
+ return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
+ }
+
+ @Override
+ protected ResultPartitionBytesCounter createResultPartitionBytesCounterImpl(
+ RemoteShuffleResultPartition remoteShuffleResultPartition) {
+ org.apache.flink.runtime.io.network.metrics.ResultPartitionBytesCounter resultPartitionBytes =
+ ReflectionUtils.readDeclaredFieldRecursiveNoThrow(
+ remoteShuffleResultPartition, "resultPartitionBytes");
+ return (bufferWithSubpartition, isBroadcast, readableBytes) -> {
+ if (isBroadcast) {
+ resultPartitionBytes.incAll(readableBytes);
+ } else {
+ resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
+ }
+ };
+ }
+
+ @Override
+ protected Buffer.DataType getEndOfPartitionDataTypeImpl() {
+ return Buffer.DataType.END_OF_PARTITION;
+ }
+
+ @Override
+ protected PartitionDescriptor createPartitionDescriptorImpl(
+ IntermediateDataSetID intermediateDataSetId, int partitionNum) {
+ IntermediateResultPartitionID intermediateResultPartitionId =
+ new IntermediateResultPartitionID(intermediateDataSetId, partitionNum);
+ return new PartitionDescriptor(
+ intermediateDataSetId,
+ 10,
+ intermediateResultPartitionId,
+ ResultPartitionType.BLOCKING,
+ 5,
+ 1,
+ false,
+ false,
+ false);
+ }
+
+ @Override
+ protected TaskInputsOutputsDescriptor createTaskInputsOutputsDescriptorImpl(
+ int inputGateNums,
+ Map<IntermediateDataSetID, Integer> inputChannelNums,
+ Map<IntermediateDataSetID, Integer> partitionReuseCount,
+ Map<IntermediateDataSetID, Integer> subpartitionNums,
+ Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
+ Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {
+ return TaskInputsOutputsDescriptor.from(
+ inputGateNums,
+ inputChannelNums,
+ partitionReuseCount,
+ subpartitionNums,
+ inputPartitionTypes,
+ partitionTypes);
+ }
+
+ @Override
+ protected JobShuffleContext getJobShuffleContextImpl(
+ JobID jobID, TierShuffleHandler tierShuffleHandler) {
+ return new JobShuffleContext() {
+ @Override
+ public JobID getJobId() {
+ return jobID;
+ }
+
+ @Override
+ public CompletableFuture<?> stopTrackingAndReleasePartitions(
+ Collection<ResultPartitionID> resultPartitionIds) {
+ return tierShuffleHandler.onReleasePartitions(
+ resultPartitionIds.stream()
+ .map(TieredStorageIdMappingUtils::convertId)
+ .collect(Collectors.toList()));
+ }
+
+ @Override
+ public CompletableFuture<Collection<PartitionWithMetrics>> getPartitionWithMetrics(
+ Duration duration, Set<ResultPartitionID> set) {
+ return CompletableFuture.completedFuture(null);
+ }
+
+ @Override
+ public void notifyPartitionRecoveryStarted() {}
+ };
+ }
+}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
deleted file mode 100644
index c89f54489..000000000
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java
+++ /dev/null
@@ -1,185 +0,0 @@
-/*
- * 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.List;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-
-import org.apache.flink.api.java.tuple.Tuple2;
-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.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.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.ResultSubpartitionIndexSet;
-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.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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.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 AbstractRemoteShuffleInputGate {
-
- public RemoteShuffleInputGate(
- CelebornConf celebornConf,
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor gateDescriptor,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- BufferDecompressor bufferDecompressor,
- int numConcurrentReading,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- super(
- celebornConf,
- ownerContext,
- gateIndex,
- gateDescriptor,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-
- @Override
- public InputChannel getChannel(int channelIndex) {
- return new FakedRemoteInputChannel(channelIndex);
- }
-
- @Override
- public Tuple2<Integer, Integer> getConsumedSubpartitionIndexRange(
- InputGateDeploymentDescriptor gateDescriptor) {
- IndexRange indexRange = gateDescriptor.getConsumedSubpartitionIndexRange();
- return Tuple2.of(indexRange.getStartIndex(), indexRange.getEndIndex());
- }
-
- /** Accommodation for the incompleteness of Flink pluggable shuffle service. */
- private class FakedRemoteInputChannel extends RemoteInputChannel {
- FakedRemoteInputChannel(int channelIndex) {
- // Flink 1.19.0
- // [FLINK-25055][network] Support listen and notify mechanism for partition request
- // [FLINK-33743][runtime] Support consuming multiple subpartitions on a single channel
- super(
- new SingleInputGate(
- inputGateDelegation.getTaskName(),
- inputGateDelegation.getGateIndex(),
- new IntermediateDataSetID(),
- ResultPartitionType.BLOCKING,
- 1,
- (a, b, c) -> {},
- () -> null,
- null,
- new FakedMemorySegmentProvider(),
- 0,
- new ThroughputCalculator(SystemClock.getInstance()),
- null),
- channelIndex,
- new ResultPartitionID(),
- new ResultSubpartitionIndexSet(new IndexRange(0, 0)),
- new ConnectionID(
- new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1),
- 0),
- new LocalConnectionManager(),
- 0,
- 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.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
deleted file mode 100644
index a5f3781d7..000000000
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateFactory.java
+++ /dev/null
@@ -1,64 +0,0 @@
-/*
- * 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.util.Map;
-
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions.CompressionCodec;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.runtime.shuffle.ShuffleIOOwnerContext;
-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(
- CelebornConf conf, NetworkBufferPool networkBufferPool, int networkBufferSize) {
- super(conf, networkBufferPool, networkBufferSize);
- }
-
- // For testing.
- @Override
- protected RemoteShuffleInputGate createInputGate(
- ShuffleIOOwnerContext ownerContext,
- int gateIndex,
- InputGateDeploymentDescriptor igdd,
- SupplierWithException<BufferPool, IOException> bufferPoolFactory,
- String compressionCodec,
- Map<Integer, ShuffleIOMetricGroup> shuffleIOMetricGroups) {
- BufferDecompressor bufferDecompressor =
- new BufferDecompressor(networkBufferSize, CompressionCodec.valueOf(compressionCodec));
- return new RemoteShuffleInputGate(
- celebornConf,
- ownerContext,
- gateIndex,
- igdd,
- bufferPoolFactory,
- bufferDecompressor,
- numConcurrentReading,
- shuffleIOMetricGroups);
- }
-}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
deleted file mode 100644
index 256acb828..000000000
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartition.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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 javax.annotation.Nullable;
-
-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.ResultPartition;
-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.apache.celeborn.plugin.flink.buffer.BufferWithSubpartition;
-import org.apache.celeborn.plugin.flink.buffer.DataBuffer;
-
-/**
- * A {@link ResultPartition} which appends records and events to {@link DataBuffer} and after the
- * {@link DataBuffer} is full, all data in the {@link DataBuffer} 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 DataBuffer} will be spilled directly.
- */
-public class RemoteShuffleResultPartition extends AbstractRemoteShuffleResultPartition {
-
- 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,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactory,
- outputGate);
- }
-
- public void updateReadableBytes(
- BufferWithSubpartition bufferWithSubpartition, boolean isBroadcast, long readableBytes) {
- if (isBroadcast) {
- resultPartitionBytes.incAll(readableBytes);
- } else {
- resultPartitionBytes.inc(bufferWithSubpartition.getSubpartitionIndex(), readableBytes);
- }
- }
-}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
deleted file mode 100644
index 691c70655..000000000
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleResultPartitionFactory.java
+++ /dev/null
@@ -1,93 +0,0 @@
-/*
- * 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.util.List;
-
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
-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.buffer.BufferPoolFactory;
-import org.apache.flink.runtime.io.network.partition.ResultPartition;
-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.runtime.metrics.groups.ShuffleIOMetricGroup;
-import org.apache.flink.util.function.SupplierWithException;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.protocol.CompressionCodec;
-
-/** Factory class to create {@link RemoteShuffleResultPartition}. */
-public class RemoteShuffleResultPartitionFactory
- extends AbstractRemoteShuffleResultPartitionFactory {
-
- public RemoteShuffleResultPartitionFactory(
- CelebornConf celebornConf,
- ResultPartitionManager partitionManager,
- BufferPoolFactory bufferPoolFactory,
- int networkBufferSize) {
-
- super(celebornConf, partitionManager, bufferPoolFactory, networkBufferSize);
- }
-
- @Override
- ResultPartition createRemoteShuffleResultPartitionInternal(
- String taskNameWithSubtaskAndId,
- int partitionIndex,
- ResultPartitionID id,
- ResultPartitionType type,
- int numSubpartitions,
- int maxParallelism,
- List<SupplierWithException<BufferPool, IOException>> bufferPoolFactories,
- CelebornConf celebornConf,
- int numMappers,
- BufferCompressor bufferCompressor,
- RemoteShuffleDescriptor rsd,
- ShuffleIOMetricGroup shuffleIOMetricGroup) {
- return new RemoteShuffleResultPartition(
- taskNameWithSubtaskAndId,
- partitionIndex,
- id,
- type,
- numSubpartitions,
- maxParallelism,
- networkBufferSize,
- partitionManager,
- bufferCompressor,
- bufferPoolFactories.get(0),
- new RemoteShuffleOutputGate(
- rsd,
- numSubpartitions,
- networkBufferSize,
- bufferPoolFactories.get(1),
- celebornConf,
- numMappers,
- shuffleIOMetricGroup));
- }
-
- @Override
- protected BufferCompressor getBufferCompressor() {
- return CompressionCodec.NONE.name().equals(compressionCodec)
- ? null
- : new BufferCompressor(
- networkBufferSize,
- NettyShuffleEnvironmentOptions.CompressionCodec.valueOf(compressionCodec));
- }
-}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
deleted file mode 100644
index 4c24fe5f9..000000000
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleServiceFactory.java
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * 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.time.Duration;
-
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.NettyShuffleEnvironmentOptions;
-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.apache.celeborn.plugin.flink.netty.NettyShuffleEnvironmentWrapper;
-
-public class RemoteShuffleServiceFactory extends AbstractRemoteShuffleServiceFactory {
-
- @Override
- public ShuffleEnvironment<ResultPartitionWriter, IndexedInputGate> createShuffleEnvironment(
- ShuffleEnvironmentContext shuffleEnvironmentContext) {
- AbstractRemoteShuffleServiceParameters parameters =
- initializePreCreateShuffleEnvironment(shuffleEnvironmentContext);
- RemoteShuffleResultPartitionFactory resultPartitionFactory =
- new RemoteShuffleResultPartitionFactory(
- parameters.celebornConf,
- parameters.resultPartitionManager,
- parameters.networkBufferPool,
- parameters.bufferSize);
- RemoteShuffleInputGateFactory inputGateFactory =
- new RemoteShuffleInputGateFactory(
- parameters.celebornConf, parameters.networkBufferPool, parameters.bufferSize);
-
- return new RemoteShuffleEnvironment(
- parameters.networkBufferPool,
- parameters.resultPartitionManager,
- resultPartitionFactory,
- inputGateFactory,
- parameters.celebornConf,
- new NettyShuffleEnvironmentWrapper(nettyShuffleServiceFactory, shuffleEnvironmentContext));
- }
-
- @Override
- public Duration getRequestSegmentsTimeout(Configuration configuration) {
- return configuration.get(NettyShuffleEnvironmentOptions.NETWORK_BUFFERS_REQUEST_TIMEOUT);
- }
-}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java
index 3c821363e..bf5f1261a 100644
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java
+++ b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierFactory.java
@@ -21,71 +21,21 @@ package org.apache.celeborn.plugin.flink.tiered;
import java.util.List;
import java.util.concurrent.ScheduledExecutorService;
-import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.io.disk.BatchShuffleReadBufferPool;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStoragePartitionId;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.netty.TieredStorageNettyService;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageConsumerSpec;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageMemoryManager;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageMemorySpec;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage.TieredStorageResourceRegistry;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierConsumerAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierFactory;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierProducerAgent;
import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleDescriptor;
-import org.apache.flink.runtime.util.ConfigurationParserUtils;
-
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
/**
* The factory class of the Celeborn client, used as a tier of flink hybrid shuffle tiered storage.
*/
-public class CelebornTierFactory implements TierFactory {
-
- private CelebornConf conf;
-
- /**
- * The bytes size of a single buffer, default value is 32KB, it will be set according to the flink
- * configuration in {@link CelebornTierFactory#setup}.
- */
- private int bufferSizeBytes = -1;
-
- /**
- * The max bytes size of a single segment, it will determine how many buffer can save in a single
- * segment.
- */
- private static int MAX_BYTES_PER_SEGMENT = 8 * 1024 * 1024;
-
- private static final String CELEBORN_TIER_NAME = CelebornTierFactory.class.getSimpleName();
-
- @Override
- public void setup(Configuration configuration) {
- conf = FlinkUtils.toCelebornConf(configuration);
- this.bufferSizeBytes = ConfigurationParserUtils.getPageSize(configuration);
- }
-
- @Override
- public TieredStorageMemorySpec getMasterAgentMemorySpec() {
- return new TieredStorageMemorySpec(getCelebornTierName(), 0);
- }
-
- @Override
- public TieredStorageMemorySpec getProducerAgentMemorySpec() {
- return new TieredStorageMemorySpec(getCelebornTierName(), 1);
- }
-
- @Override
- public TieredStorageMemorySpec getConsumerAgentMemorySpec() {
- return new TieredStorageMemorySpec(getCelebornTierName(), 0);
- }
-
- @Override
- public TierMasterAgent createMasterAgent(
- TieredStorageResourceRegistry tieredStorageResourceRegistry) {
- return new CelebornTierMasterAgent(conf);
- }
+public class CelebornTierFactory extends BaseCelebornTierFactory implements TierFactory {
@Override
public TierProducerAgent createProducerAgent(
@@ -121,8 +71,4 @@ public class CelebornTierFactory implements TierFactory {
return new CelebornTierConsumerAgent(
conf, tieredStorageConsumerSpecs, shuffleDescriptors, bufferSizeBytes);
}
-
- public static String getCelebornTierName() {
- return CELEBORN_TIER_NAME;
- }
}
diff --git a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java b/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
deleted file mode 100644
index e0de46ae3..000000000
--- a/client-flink/flink-1.20/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierMasterAgent.java
+++ /dev/null
@@ -1,219 +0,0 @@
-/*
- * 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.tiered;
-
-import static org.apache.celeborn.plugin.flink.utils.Utils.checkState;
-
-import java.time.Duration;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
-
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageIdMappingUtils;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierMasterAgent;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleDescriptor;
-import org.apache.flink.runtime.io.network.partition.hybrid.tiered.tier.TierShuffleHandler;
-import org.apache.flink.runtime.shuffle.JobShuffleContext;
-import org.apache.flink.runtime.shuffle.PartitionWithMetrics;
-import org.apache.flink.util.ExecutorUtils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.celeborn.client.LifecycleManager;
-import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.util.JavaUtils;
-import org.apache.celeborn.common.util.ThreadUtils;
-import org.apache.celeborn.plugin.flink.FlinkResultPartitionInfo;
-import org.apache.celeborn.plugin.flink.RemoteShuffleResource;
-import org.apache.celeborn.plugin.flink.ShuffleResourceDescriptor;
-import org.apache.celeborn.plugin.flink.ShuffleResourceTracker;
-import org.apache.celeborn.plugin.flink.ShuffleTaskInfo;
-import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
-
-public class CelebornTierMasterAgent implements TierMasterAgent {
-
- private static final Logger LOG = LoggerFactory.getLogger(CelebornTierMasterAgent.class);
-
- // Flink JobId -> Celeborn register shuffleIds
- private final Map<JobID, Set<Integer>> jobShuffleIds = JavaUtils.newConcurrentHashMap();
-
- private final ShuffleTaskInfo shuffleTaskInfo = new ShuffleTaskInfo();
-
- private final ScheduledExecutorService executor =
- ThreadUtils.newDaemonSingleThreadScheduledExecutor("celeborn-client-tier-master-executor");
- private final long lifecycleManagerTimestamp;
-
- private final CelebornConf conf;
-
- private ShuffleResourceTracker shuffleResourceTracker;
-
- private String celebornAppId;
-
- private volatile LifecycleManager lifecycleManager;
-
- public CelebornTierMasterAgent(CelebornConf conf) {
- this.conf = conf;
- this.lifecycleManagerTimestamp = System.currentTimeMillis();
- }
-
- @Override
- public void registerJob(JobID jobID, TierShuffleHandler tierShuffleHandler) {
- if (lifecycleManager == null) {
- synchronized (CelebornTierMasterAgent.class) {
- if (lifecycleManager == null) {
- celebornAppId = FlinkUtils.toCelebornAppId(lifecycleManagerTimestamp, jobID);
- LOG.info("CelebornAppId: {}", celebornAppId);
- lifecycleManager = new LifecycleManager(celebornAppId, conf);
- this.shuffleResourceTracker = new ShuffleResourceTracker(executor, lifecycleManager);
- }
- }
- }
-
- Set<Integer> previousShuffleIds = jobShuffleIds.putIfAbsent(jobID, new HashSet<>());
- LOG.info("Register job: {}.", jobID);
- if (previousShuffleIds != null) {
- throw new RuntimeException("Duplicated registration job: " + jobID);
- }
- shuffleResourceTracker.registerJob(getJobShuffleContext(jobID, tierShuffleHandler));
- }
-
- @Override
- public void unregisterJob(JobID jobID) {
- LOG.info("Unregister job: {}.", jobID);
- Set<Integer> shuffleIds = jobShuffleIds.remove(jobID);
- if (shuffleIds != null) {
- executor.execute(
- () -> {
- for (Integer shuffleId : shuffleIds) {
- lifecycleManager.unregisterShuffle(shuffleId);
- shuffleTaskInfo.removeExpiredShuffle(shuffleId);
- }
- shuffleResourceTracker.unRegisterJob(jobID);
- });
- }
- }
-
- @Override
- public TierShuffleDescriptor addPartitionAndGetShuffleDescriptor(
- JobID jobID, ResultPartitionID resultPartitionID) {
- Set<Integer> shuffleIds = jobShuffleIds.get(jobID);
- if (shuffleIds == null) {
- throw new RuntimeException("Can not find job in master agent, job: " + jobID);
- }
- FlinkResultPartitionInfo resultPartitionInfo =
- new FlinkResultPartitionInfo(jobID, resultPartitionID);
- ShuffleResourceDescriptor shuffleResourceDescriptor =
- shuffleTaskInfo.genShuffleResourceDescriptor(
- resultPartitionInfo.getShuffleId(),
- resultPartitionInfo.getTaskId(),
- resultPartitionInfo.getAttemptId());
- shuffleIds.add(shuffleResourceDescriptor.getShuffleId());
- shuffleResourceTracker.addPartitionResource(
- jobID,
- shuffleResourceDescriptor.getShuffleId(),
- shuffleResourceDescriptor.getPartitionId(),
- resultPartitionID);
-
- RemoteShuffleResource remoteShuffleResource =
- new RemoteShuffleResource(
- lifecycleManager.getHost(),
- lifecycleManager.getPort(),
- lifecycleManagerTimestamp,
- shuffleResourceDescriptor);
- return new TierShuffleDescriptorImpl(
- celebornAppId,
- jobID,
- resultPartitionInfo.getShuffleId(),
- resultPartitionID,
- remoteShuffleResource);
- }
-
- @Override
- public void releasePartition(TierShuffleDescriptor shuffleDescriptor) {
- checkState(shuffleDescriptor instanceof TierShuffleDescriptorImpl, "Wrong descriptor type.");
- try {
- TierShuffleDescriptorImpl descriptor = (TierShuffleDescriptorImpl) shuffleDescriptor;
- RemoteShuffleResource shuffleResource = descriptor.getShuffleResource();
- ShuffleResourceDescriptor resourceDescriptor =
- shuffleResource.getMapPartitionShuffleDescriptor();
- LOG.debug("release partition resource: {}.", resourceDescriptor);
- lifecycleManager.releasePartition(
- resourceDescriptor.getShuffleId(), resourceDescriptor.getPartitionId());
- shuffleResourceTracker.removePartitionResource(
- descriptor.getJobId(),
- resourceDescriptor.getShuffleId(),
- resourceDescriptor.getPartitionId());
- } catch (Throwable throwable) {
- LOG.debug("Failed to release data partition {}.", shuffleDescriptor, throwable);
- }
- }
-
- @Override
- public void close() {
- try {
- jobShuffleIds.clear();
- if (null != lifecycleManager) {
- lifecycleManager.stop();
- }
- } catch (Exception e) {
- LOG.warn("Encounter exception when shutdown: {}", e.getMessage(), e);
- }
-
- ExecutorUtils.gracefulShutdown(10, TimeUnit.SECONDS, executor);
- }
-
- private JobShuffleContext getJobShuffleContext(
- JobID jobID, TierShuffleHandler tierShuffleHandler) {
- return new JobShuffleContext() {
- @Override
- public JobID getJobId() {
- return jobID;
- }
-
- @Override
- public CompletableFuture<?> stopTrackingAndReleasePartitions(
- Collection<ResultPartitionID> resultPartitionIds) {
... 14815 lines suppressed ...