You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2022/11/22 07:23:25 UTC
[incubator-celeborn] branch main updated: [CELEBORN-14] [ISSUE-955] support register attempt map task (#984)
This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 16564587 [CELEBORN-14] [ISSUE-955] support register attempt map task (#984)
16564587 is described below
commit 16564587887daefc70434e3be85c748cb9619937
Author: Shuang <lv...@alibaba-inc.com>
AuthorDate: Tue Nov 22 15:23:20 2022 +0800
[CELEBORN-14] [ISSUE-955] support register attempt map task (#984)
---
.../apache/celeborn/client/ShuffleClientImpl.java | 11 --
.../apache/celeborn/client/LifecycleManager.scala | 153 ++++++++++++---------
.../client/RequestLocationCallContext.scala | 43 ++++++
.../celeborn/tests/client/ShuffleClientSuite.scala | 22 ++-
4 files changed, 152 insertions(+), 77 deletions(-)
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 751ac07f..332e8f80 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -279,17 +279,6 @@ public class ShuffleClientImpl extends ShuffleClient {
mapId,
attemptId,
partitionId);
- if (attemptId == 0) {
- return registerMapPartitionTaskWithFirstAttempt(
- appId, shuffleId, numMappers, mapId, attemptId, partitionId);
- }
-
- // TODO
- throw new UnsupportedOperationException("can not register shuffle task with attempt beyond 0");
- }
-
- private PartitionLocation registerMapPartitionTaskWithFirstAttempt(
- String appId, int shuffleId, int numMappers, int mapId, int attemptId, int partitionId) {
ConcurrentHashMap<Integer, PartitionLocation> partitionLocationMap =
registerShuffleInternal(
shuffleId,
diff --git a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 7a854f33..e03f9707 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -19,7 +19,7 @@ package org.apache.celeborn.client
import java.nio.ByteBuffer
import java.util
-import java.util.{function, List => JList, Map}
+import java.util.{function, List => JList}
import java.util.concurrent.{Callable, ConcurrentHashMap, ScheduledExecutorService, ScheduledFuture, TimeUnit}
import java.util.concurrent.atomic.LongAdder
@@ -104,7 +104,7 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
}
case class ChangePartitionRequest(
- context: RpcCallContext,
+ context: RequestLocationCallContext,
applicationId: String,
shuffleId: Int,
partitionId: Int,
@@ -112,6 +112,12 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
oldPartition: PartitionLocation,
causes: Option[StatusCode])
+ case class RegisterCallContext(context: RpcCallContext, partitionId: Int = -1) {
+ def reply(response: PbRegisterShuffleResponse) = {
+ context.reply(response)
+ }
+ }
+
// shuffleId -> (partitionId -> set of ChangePartition)
private val changePartitionRequests =
new ConcurrentHashMap[Int, ConcurrentHashMap[Integer, util.Set[ChangePartitionRequest]]]()
@@ -119,7 +125,8 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
private val inBatchPartitions = new ConcurrentHashMap[Int, util.Set[Integer]]()
// register shuffle request waiting for response
- private val registeringShuffleRequest = new ConcurrentHashMap[Int, util.Set[RpcCallContext]]()
+ private val registeringShuffleRequest =
+ new ConcurrentHashMap[Int, util.Set[RegisterCallContext]]()
// blacklist
private val blacklist = new ConcurrentHashMap[WorkerInfo, (StatusCode, Long)]()
@@ -218,7 +225,7 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
request.asScala.toArray.maxBy(_.epoch)
}.toArray
if (distinctPartitions.nonEmpty) {
- batchHandleChangePartitions(
+ batchHandleRequestPartitions(
distinctPartitions.head.applicationId,
shuffleId,
distinctPartitions)
@@ -324,7 +331,12 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
val numPartitions = pb.getNumPartitions
logDebug(s"Received RegisterShuffle request, " +
s"$applicationId, $shuffleId, $numMappers, $numPartitions.")
- handleRegisterShuffle(context, applicationId, shuffleId, numMappers, numPartitions)
+ offerAndReserveSlots(
+ RegisterCallContext(context),
+ applicationId,
+ shuffleId,
+ numMappers,
+ numPartitions)
case pb: PbRegisterMapPartitionTask =>
val applicationId = pb.getApplicationId
@@ -335,11 +347,13 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
val partitionId = pb.getPartitionId
logDebug(s"Received Register map partition task request, " +
s"$applicationId, $shuffleId, $numMappers, $mapId, $attemptId, $partitionId.")
- handleRegisterMapPartitionTask(
- context,
+ shufflePartitionType.putIfAbsent(shuffleId, PartitionType.MAP)
+ offerAndReserveSlots(
+ RegisterCallContext(context, partitionId),
applicationId,
shuffleId,
numMappers,
+ numMappers,
partitionId)
case pb: PbRevive =>
@@ -373,8 +387,8 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
val oldPartition = PbSerDeUtils.fromPbPartitionLocation(pb.getOldPartition)
logTrace(s"Received split request, " +
s"$applicationId, $shuffleId, $partitionId, $epoch, $oldPartition")
- handleChangePartitionLocation(
- context,
+ handleRequestPartitionLocation(
+ ChangeLocationCallContext(context),
applicationId,
shuffleId,
partitionId,
@@ -392,37 +406,8 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
handleGetReducerFileGroup(context, shuffleId)
}
- /* ========================================================== *
- | START OF EVENT HANDLER |
- * ========================================================== */
-
- private def handleRegisterShuffle(
- context: RpcCallContext,
- applicationId: String,
- shuffleId: Int,
- numMappers: Int,
- numReducers: Int): Unit = {
- handleOfferAndReserveSlots(context, applicationId, shuffleId, numMappers, numReducers)
- }
-
- private def handleRegisterMapPartitionTask(
- context: RpcCallContext,
- applicationId: String,
- shuffleId: Int,
- numMappers: Int,
- partitionId: Int): Unit = {
- shufflePartitionType.putIfAbsent(shuffleId, PartitionType.MAP)
- handleOfferAndReserveSlots(
- context,
- applicationId,
- shuffleId,
- numMappers,
- numMappers,
- partitionId)
- }
-
- private def handleOfferAndReserveSlots(
- context: RpcCallContext,
+ private def offerAndReserveSlots(
+ context: RegisterCallContext,
applicationId: String,
shuffleId: Int,
numMappers: Int,
@@ -447,12 +432,22 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
.filter(p =>
(partitionType == PartitionType.REDUCE && p.getEpoch == 0) || (partitionType == PartitionType.MAP && p.getId == partitionId))
.toArray
- context.reply(RegisterShuffleResponse(StatusCode.SUCCESS, initialLocs))
+ partitionType match {
+ case PartitionType.MAP => processMapTaskReply(
+ applicationId,
+ shuffleId,
+ context.context,
+ partitionId,
+ initialLocs)
+ case PartitionType.REDUCE =>
+ context.reply(RegisterShuffleResponse(StatusCode.SUCCESS, initialLocs))
+ }
return
}
+
logInfo(s"New shuffle request, shuffleId $shuffleId, partitionType: $partitionType " +
s"numMappers: $numMappers, numReducers: $numReducers.")
- val set = new util.HashSet[RpcCallContext]()
+ val set = new util.HashSet[RegisterCallContext]()
set.add(context)
registeringShuffleRequest.put(shuffleId, set)
}
@@ -463,7 +458,21 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
registeringShuffleRequest.synchronized {
registeringShuffleRequest.asScala
.get(shuffleId)
- .foreach(_.asScala.foreach(_.reply(response)))
+ .foreach(_.asScala.foreach(context => {
+ partitionType match {
+ case PartitionType.MAP =>
+ val partitionLocations =
+ response.getPartitionLocationsList.asScala.filter(_.getId == partitionId).map(r =>
+ PbSerDeUtils.fromPbPartitionLocation(r)).toArray
+ processMapTaskReply(
+ applicationId,
+ shuffleId,
+ context.context,
+ partitionId,
+ partitionLocations)
+ case PartitionType.REDUCE => context.reply(response)
+ }
+ }))
registeringShuffleRequest.remove(shuffleId)
}
}
@@ -564,13 +573,32 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
// Fifth, reply the allocated partition location to ShuffleClient.
logInfo(s"Handle RegisterShuffle Success for $shuffleId.")
- val allMasterPartitionLocations = slots.asScala.flatMap(_._2._1.asScala).filter(p =>
- partitionType ==
- PartitionType.REDUCE || (partitionType == PartitionType.MAP && p.getId == partitionId)).toArray
+ val allMasterPartitionLocations = slots.asScala.flatMap(_._2._1.asScala).toArray
reply(RegisterShuffleResponse(StatusCode.SUCCESS, allMasterPartitionLocations))
}
}
+ private def processMapTaskReply(
+ applicationId: String,
+ shuffleId: Int,
+ context: RpcCallContext,
+ partitionId: Int,
+ partitionLocations: Array[PartitionLocation]): Unit = {
+ // if any partition location resource exist just reply
+ if (partitionLocations.size > 0) {
+ context.reply(RegisterShuffleResponse(StatusCode.SUCCESS, partitionLocations))
+ } else {
+ // request new resource for this task
+ handleRequestPartitionLocation(
+ ApplyNewLocationCallContext(context),
+ applicationId,
+ shuffleId,
+ partitionId,
+ -1,
+ null)
+ }
+ }
+
private def blacklistPartition(
shuffleId: Int,
oldPartition: PartitionLocation,
@@ -619,8 +647,8 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
logWarning(s"Do Revive for shuffle ${Utils.makeShuffleKey(applicationId, shuffleId)}, " +
s"oldPartition: $oldPartition, cause: $cause")
- handleChangePartitionLocation(
- context,
+ handleRequestPartitionLocation(
+ ChangeLocationCallContext(context),
applicationId,
shuffleId,
partitionId,
@@ -641,8 +669,8 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
override def apply(s: Int): util.Set[Integer] = new util.HashSet[Integer]()
}
- private def handleChangePartitionLocation(
- context: RpcCallContext,
+ private def handleRequestPartitionLocation(
+ context: RequestLocationCallContext,
applicationId: String,
shuffleId: Int,
partitionId: Int,
@@ -664,14 +692,14 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
requests.synchronized {
if (requests.containsKey(partitionId)) {
requests.get(partitionId).add(changePartition)
- logTrace(s"[handleChangePartitionLocation] For $shuffleId, request for same partition" +
+ logTrace(s"[handleRequestPartitionLocation] For $shuffleId, request for same partition" +
s"$partitionId-$oldEpoch exists, register context.")
return
} else {
// If new slot for the partition has been allocated, reply and return.
// Else register and allocate for it.
getLatestPartition(shuffleId, partitionId, oldEpoch).foreach { latestLoc =>
- context.reply(ChangeLocationResponse(StatusCode.SUCCESS, Some(latestLoc)))
+ context.reply(StatusCode.SUCCESS, Some(latestLoc))
logDebug(s"New partition found, old partition $partitionId-$oldEpoch return it." +
s" shuffleId: $shuffleId $latestLoc")
return
@@ -682,11 +710,11 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
}
}
if (!batchHandleChangePartitionEnabled) {
- batchHandleChangePartitions(applicationId, shuffleId, Array(changePartition))
+ batchHandleRequestPartitions(applicationId, shuffleId, Array(changePartition))
}
}
- def batchHandleChangePartitions(
+ def batchHandleRequestPartitions(
applicationId: String,
shuffleId: Int,
changePartitions: Array[ChangePartitionRequest]): Unit = {
@@ -716,13 +744,14 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
location -> Option(requestsMap.remove(location.getId))
}
}.foreach { case (newLocation, requests) =>
- val response = ChangeLocationResponse(StatusCode.SUCCESS, Option(newLocation))
- requests.map(_.asScala.toList.foreach(_.context.reply(response)))
+ requests.map(_.asScala.toList.foreach(_.context.reply(
+ StatusCode.SUCCESS,
+ Option(newLocation))))
}
}
// remove together to reduce lock time
- def replyFailure(response: PbChangeLocationResponse): Unit = {
+ def replyFailure(status: StatusCode): Unit = {
requestsMap.synchronized {
changePartitions.map { changePartition =>
if (batchHandleChangePartitionEnabled) {
@@ -731,14 +760,14 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
Option(requestsMap.remove(changePartition.partitionId))
}
}.foreach { requests =>
- requests.map(_.asScala.toList.foreach(_.context.reply(response)))
+ requests.map(_.asScala.toList.foreach(_.context.reply(status, None)))
}
}
val candidates = workersNotBlacklisted(shuffleId)
if (candidates.size < 1 || (pushReplicateEnabled && candidates.size < 2)) {
logError("[Update partition] failed for not enough candidates for revive.")
- replyFailure(ChangeLocationResponse(StatusCode.SLOT_NOT_AVAILABLE, None))
+ replyFailure(StatusCode.SLOT_NOT_AVAILABLE)
return
}
@@ -748,13 +777,13 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
if (!registeredShuffle.contains(shuffleId)) {
logError(s"[handleChangePartition] shuffle $shuffleId not registered!")
- replyFailure(ChangeLocationResponse(StatusCode.SHUFFLE_NOT_REGISTERED, None))
+ replyFailure(StatusCode.SHUFFLE_NOT_REGISTERED)
return
}
if (stageEndShuffleSet.contains(shuffleId)) {
logError(s"[handleChangePartition] shuffle $shuffleId already ended!")
- replyFailure(ChangeLocationResponse(StatusCode.STAGE_ENDED, None))
+ replyFailure(StatusCode.STAGE_ENDED)
return
}
@@ -764,7 +793,7 @@ class LifecycleManager(appId: String, val conf: CelebornConf) extends RpcEndpoin
new util.HashSet(candidates.toSet.asJava),
newlyAllocatedLocations)) {
logError(s"[Update partition] failed for $shuffleId.")
- replyFailure(ChangeLocationResponse(StatusCode.RESERVE_SLOTS_FAILED, None))
+ replyFailure(StatusCode.RESERVE_SLOTS_FAILED)
return
}
diff --git a/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala b/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala
new file mode 100644
index 00000000..91cbf6df
--- /dev/null
+++ b/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala
@@ -0,0 +1,43 @@
+/*
+ * 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.client
+
+import org.apache.celeborn.common.protocol.PartitionLocation
+import org.apache.celeborn.common.protocol.message.ControlMessages.{ChangeLocationResponse, RegisterShuffleResponse}
+import org.apache.celeborn.common.protocol.message.StatusCode
+import org.apache.celeborn.common.rpc.RpcCallContext
+
+trait RequestLocationCallContext {
+ def reply(status: StatusCode, partitionLocationOpt: Option[PartitionLocation]): Unit
+}
+
+case class ChangeLocationCallContext(context: RpcCallContext) extends RequestLocationCallContext {
+ override def reply(status: StatusCode, partitionLocationOpt: Option[PartitionLocation]): Unit = {
+ context.reply(ChangeLocationResponse(status, partitionLocationOpt))
+ }
+}
+
+case class ApplyNewLocationCallContext(context: RpcCallContext) extends RequestLocationCallContext {
+ override def reply(status: StatusCode, partitionLocationOpt: Option[PartitionLocation]): Unit = {
+ partitionLocationOpt match {
+ case Some(partitionLocation) =>
+ context.reply(RegisterShuffleResponse(status, Array(partitionLocation)))
+ case None => context.reply(RegisterShuffleResponse(status, Array.empty))
+ }
+ }
+}
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ShuffleClientSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ShuffleClientSuite.scala
index e4b00d17..ea7e7fe2 100644
--- a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ShuffleClientSuite.scala
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/client/ShuffleClientSuite.scala
@@ -67,16 +67,30 @@ class ShuffleClientSuite extends AnyFunSuite with MiniClusterFeature
location = shuffleClient.registerMapPartitionTask(APP, shuffleId, numMappers, mapId, attemptId)
Assert.assertEquals(location.getId, PackedPartitionId.packedPartitionId(mapId, attemptId))
+ // check all allocated slots
+ var partitionLocationInfos = lifecycleManager.workerSnapshots(shuffleId).values().asScala
+ var count =
+ partitionLocationInfos.map(r => r.getAllMasterLocations(shuffleId.toString).size()).sum
+ Assert.assertEquals(count, numMappers)
+
// another mapId
location =
shuffleClient.registerMapPartitionTask(APP, shuffleId, numMappers, mapId + 1, attemptId)
Assert.assertEquals(location.getId, PackedPartitionId.packedPartitionId(mapId + 1, attemptId))
- // offer and reserve all slots
- val partitionLocationInfos = lifecycleManager.workerSnapshots(shuffleId).values().asScala
- val count =
+ // another mapId with another attemptId
+ location =
+ shuffleClient.registerMapPartitionTask(APP, shuffleId, numMappers, mapId + 1, attemptId + 1)
+ Assert.assertEquals(
+ location.getId,
+ PackedPartitionId.packedPartitionId(mapId + 1, attemptId + 1))
+
+ // check all allocated all slots
+ partitionLocationInfos = lifecycleManager.workerSnapshots(shuffleId).values().asScala
+ print(partitionLocationInfos)
+ count =
partitionLocationInfos.map(r => r.getAllMasterLocations(shuffleId.toString).size()).sum
- Assert.assertEquals(count, numMappers)
+ Assert.assertEquals(count, numMappers + 1)
}
override def afterAll(): Unit = {