This is an automated email from the ASF dual-hosted git repository.
ethanfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 935806f03 [CELEBORN-341][Flink] cache file group for map partition in
Flink plugin (#1277)
935806f03 is described below
commit 935806f03618cdfe79314c83733809e9b7b5e834
Author: Shuang <[email protected]>
AuthorDate: Sun Feb 26 20:31:20 2023 +0800
[CELEBORN-341][Flink] cache file group for map partition in Flink plugin
(#1277)
---
.../flink/readclient/FlinkShuffleClientImpl.java | 27 ++++++++++++++++++----
.../apache/celeborn/client/ShuffleClientImpl.java | 25 +++++++++++++++-----
.../client/commit/MapPartitionCommitHandler.scala | 19 ++++++++++++---
common/src/main/proto/TransportMessages.proto | 3 +++
.../common/protocol/message/ControlMessages.scala | 10 +++++---
5 files changed, 68 insertions(+), 16 deletions(-)
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
index d11b0dcb2..3de2bb1f7 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
@@ -109,10 +109,29 @@ public class FlinkShuffleClientImpl extends
ShuffleClientImpl {
@Override
protected ReduceFileGroups updateFileGroup(
- String applicationId, String shuffleKey, int shuffleId) {
- ReduceFileGroups newGroup = loadFileGroupInternal(applicationId,
shuffleKey, shuffleId);
- reduceFileGroupsMap.put(shuffleId, newGroup);
- return newGroup;
+ String applicationId, String shuffleKey, int shuffleId, int partitionId)
throws IOException {
+ ReduceFileGroups reduceFileGroups =
+ reduceFileGroupsMap.computeIfAbsent(shuffleId, (id) -> new
ReduceFileGroups());
+ synchronized (reduceFileGroups) {
+ if (reduceFileGroups.partitionIds != null
+ && reduceFileGroups.partitionIds.contains(partitionId)) {
+ logger.debug(
+ "use cached file groups for partition: {}",
+ Utils.makeReducerKey(applicationId, shuffleId, partitionId));
+ return reduceFileGroups;
+ } else {
+ // refresh file groups
+ ReduceFileGroups newGroups = loadFileGroupInternal(applicationId,
shuffleKey, shuffleId);
+ if (newGroups == null ||
!newGroups.partitionIds.contains(partitionId)) {
+ throw new IOException(
+ "shuffle data lost for partition: "
+ + Utils.makeReducerKey(applicationId, shuffleId,
partitionId));
+ }
+
+ reduceFileGroupsMap.put(shuffleId, newGroups);
+ return newGroups;
+ }
+ }
}
public ReadClientHandler getReadClientHandler() {
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 6c5a1ad01..185f76cd6 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -120,10 +120,21 @@ public class ShuffleClientImpl extends ShuffleClient {
protected static class ReduceFileGroups {
public final Map<Integer, Set<PartitionLocation>> partitionGroups;
public final int[] mapAttempts;
+ public final Set<Integer> partitionIds;
- ReduceFileGroups(Map<Integer, Set<PartitionLocation>> partitionGroups,
int[] mapAttempts) {
+ ReduceFileGroups(
+ Map<Integer, Set<PartitionLocation>> partitionGroups,
+ int[] mapAttempts,
+ Set<Integer> partitionIds) {
this.partitionGroups = partitionGroups;
this.mapAttempts = mapAttempts;
+ this.partitionIds = partitionIds;
+ }
+
+ public ReduceFileGroups() {
+ this.partitionGroups = null;
+ this.mapAttempts = null;
+ this.partitionIds = null;
}
}
@@ -1365,11 +1376,12 @@ public class ShuffleClientImpl extends ShuffleClient {
if (response.status() == StatusCode.SUCCESS) {
logger.info(
- "Shuffle {} request reducer file group success using {} ms,
result partition ids {}.",
+ "Shuffle {} request reducer file group success using {} ms,
result partition size {}.",
shuffleId,
(System.nanoTime() - getReducerFileGroupStartTime) / 1000_000,
- response.fileGroup().keySet());
- return new ReduceFileGroups(response.fileGroup(),
response.attempts());
+ response.fileGroup().size());
+ return new ReduceFileGroups(
+ response.fileGroup(), response.attempts(),
response.partitionIds());
} else if (response.status() == StatusCode.STAGE_END_TIME_OUT) {
logger.warn(
"Request {} return {} for {}.",
@@ -1391,14 +1403,15 @@ public class ShuffleClientImpl extends ShuffleClient {
}
protected ReduceFileGroups updateFileGroup(
- String applicationId, String shuffleKey, int shuffleId) {
+ String applicationId, String shuffleKey, int shuffleId, int partitionId)
throws IOException {
return reduceFileGroupsMap.computeIfAbsent(
shuffleId, (id) -> loadFileGroupInternal(applicationId, shuffleKey,
shuffleId));
}
protected ReduceFileGroups loadFileGroup(
String applicationId, String shuffleKey, int shuffleId, int partitionId)
throws IOException {
- ReduceFileGroups reduceFileGroups = updateFileGroup(applicationId,
shuffleKey, shuffleId);
+ ReduceFileGroups reduceFileGroups =
+ updateFileGroup(applicationId, shuffleKey, shuffleId, partitionId);
if (reduceFileGroups == null) {
String msg =
"Shuffle data lost for shuffle " + shuffleId + " partitionId " +
partitionId + "!";
diff --git
a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
index 1a4f1e128..04c8f0649 100644
---
a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
@@ -54,8 +54,10 @@ class MapPartitionCommitHandler(
extends CommitHandler(appId, conf, allocatedWorkers, committedPartitionInfo)
with Logging {
+ private val shuffleSuccessPartitionIds = new ConcurrentHashMap[Int,
util.Set[Integer]]()
+
// shuffleId -> in processing partitionId set
- private val inProcessMapPartitionEndIds = new ConcurrentHashMap[Int,
util.Set[Int]]()
+ private val inProcessMapPartitionEndIds = new ConcurrentHashMap[Int,
util.Set[Integer]]()
override def getPartitionType(): PartitionType = {
PartitionType.MAP
@@ -189,7 +191,9 @@ class MapPartitionCommitHandler(
partitionId: Int,
recordWorkerFailure: ShuffleFailedWorkers => Unit): (Boolean, Boolean) =
{
val inProcessingPartitionIds =
- inProcessMapPartitionEndIds.computeIfAbsent(shuffleId, (k: Int) => new
util.HashSet[Int]())
+ inProcessMapPartitionEndIds.computeIfAbsent(
+ shuffleId,
+ (k: Int) => ConcurrentHashMap.newKeySet[Integer]())
inProcessingPartitionIds.add(partitionId)
val partitionAllocatedWorkers =
allocatedWorkers.get(shuffleId).asScala.filter(p =>
@@ -212,6 +216,14 @@ class MapPartitionCommitHandler(
}
inProcessingPartitionIds.remove(partitionId)
+ if (dataCommitSuccess) {
+ val resultPartitions =
+ shuffleSuccessPartitionIds.computeIfAbsent(
+ shuffleId,
+ (k: Int) => ConcurrentHashMap.newKeySet[Integer]())
+ resultPartitions.add(partitionId)
+ }
+
(dataCommitSuccess, false)
}
@@ -219,6 +231,7 @@ class MapPartitionCommitHandler(
context.reply(GetReducerFileGroupResponse(
StatusCode.SUCCESS,
reducerFileGroupsMap.getOrDefault(shuffleId, new ConcurrentHashMap()),
- getMapperAttempts(shuffleId)))
+ getMapperAttempts(shuffleId),
+ shuffleSuccessPartitionIds.get(shuffleId)))
}
}
diff --git a/common/src/main/proto/TransportMessages.proto
b/common/src/main/proto/TransportMessages.proto
index a9785590c..c27b2499a 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -252,6 +252,9 @@ message PbGetReducerFileGroupResponse {
// only reduce partition mode need know valid attempts
repeated int32 attempts = 3;
+
+ // only map partition mode has succeed partitionIds
+ repeated int32 partitionIds = 4;
}
message PbUnregisterShuffle {
diff --git
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
index 664c4253b..0cb3a3a8f 100644
---
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -255,7 +255,8 @@ object ControlMessages extends Logging {
case class GetReducerFileGroupResponse(
status: StatusCode,
fileGroup: util.Map[Integer, util.Set[PartitionLocation]],
- attempts: Array[Int])
+ attempts: Array[Int],
+ partitionIds: util.Set[Integer] = new util.HashSet[Integer]())
extends MasterMessage
object WorkerLost {
@@ -549,7 +550,7 @@ object ControlMessages extends Logging {
.build().toByteArray
new TransportMessage(MessageType.GET_REDUCER_FILE_GROUP, payload)
- case GetReducerFileGroupResponse(status, fileGroup, attempts) =>
+ case GetReducerFileGroupResponse(status, fileGroup, attempts,
partitionIds) =>
val builder = PbGetReducerFileGroupResponse
.newBuilder()
.setStatus(status.getValue)
@@ -561,6 +562,7 @@ object ControlMessages extends Logging {
.toPbPartitionLocation).toList.asJava).build())
}.asJava)
builder.addAllAttempts(attempts.map(new Integer(_)).toIterable.asJava)
+ builder.addAllPartitionIds(partitionIds)
val payload = builder.build().toByteArray
new TransportMessage(MessageType.GET_REDUCER_FILE_GROUP_RESPONSE,
payload)
@@ -912,10 +914,12 @@ object ControlMessages extends Logging {
}.asJava
val attempts =
pbGetReducerFileGroupResponse.getAttemptsList.asScala.map(_.toInt).toArray
+ val partitionIds = new
util.HashSet(pbGetReducerFileGroupResponse.getPartitionIdsList)
GetReducerFileGroupResponse(
Utils.toStatusCode(pbGetReducerFileGroupResponse.getStatus),
fileGroup,
- attempts)
+ attempts,
+ partitionIds)
case UNREGISTER_SHUFFLE =>
PbUnregisterShuffle.parseFrom(message.getPayload)
