This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch branch-0.2
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/branch-0.2 by this push:
new 730815e8a [CELEBORN-301] Refactor PartitionLocationInfo to use
ConcurrentHashMap (#1236)
730815e8a is described below
commit 730815e8a33c385e11358ff2d60ec9c37f5756f1
Author: Keyong Zhou <[email protected]>
AuthorDate: Sun Feb 26 00:08:20 2023 +0800
[CELEBORN-301] Refactor PartitionLocationInfo to use ConcurrentHashMap
(#1236)
---
.../spark/shuffle/celeborn/SortBasedPusher.java | 3 +-
.../common/meta/PartitionLocationInfo.scala | 228 ++++++++-------------
.../celeborn/service/deploy/worker/Worker.scala | 1 +
3 files changed, 83 insertions(+), 149 deletions(-)
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
index 82dd6c9ab..a23ac2646 100644
---
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
@@ -52,6 +52,7 @@ public class SortBasedPusher extends MemoryConsumer {
private final int pushBufferMaxSize;
private final long pushSortMemoryThreshold;
final int uaoSize = UnsafeAlignedOffset.getUaoSize();
+ static final long size8k = Utils.byteStringAsBytes("8k");
String appId;
int shuffleId;
@@ -182,7 +183,7 @@ public class SortBasedPusher extends MemoryConsumer {
throws IOException {
if (getUsed() > pushSortMemoryThreshold
- && pageCursor + Utils.byteStringAsBytes("8k")
+ && pageCursor + size8k
> currentPage.getBaseOffset() + currentPage.size()) {
logger.info(
"Memory Used across threshold, trigger push. Memory: "
diff --git
a/common/src/main/scala/org/apache/celeborn/common/meta/PartitionLocationInfo.scala
b/common/src/main/scala/org/apache/celeborn/common/meta/PartitionLocationInfo.scala
index 4f29c63a7..277b86add 100644
---
a/common/src/main/scala/org/apache/celeborn/common/meta/PartitionLocationInfo.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/meta/PartitionLocationInfo.scala
@@ -18,6 +18,7 @@
package org.apache.celeborn.common.meta
import java.util
+import java.util.concurrent.ConcurrentHashMap
import java.util.function.BiFunction
import scala.collection.JavaConverters._
@@ -27,33 +28,34 @@ import org.apache.celeborn.common.protocol.PartitionLocation
class PartitionLocationInfo extends Logging {
- // key: ShuffleKey, values: (partitionId -> partition locations)
- type PartitionInfo = util.HashMap[String, util.Map[Int,
util.List[PartitionLocation]]]
+ // key: ShuffleKey, values: (encodeUniqueId -> partition location)
+ type PartitionInfo = ConcurrentHashMap[String, ConcurrentHashMap[Long,
PartitionLocation]]
private val masterPartitionLocations = new PartitionInfo
private val slavePartitionLocations = new PartitionInfo
+ def encode(partitionId: Int, epoch: Int): Long = {
+ partitionId.toLong << 32 | epoch
+ }
+
+ def encodeUniqueId(uniqueId: String): Long = {
+ val tokens = uniqueId.split("-", 2)
+ val partitionId = tokens(0).toInt
+ val epoch = tokens(1).toInt
+ encode(partitionId, epoch)
+ }
+
def shuffleKeySet: util.HashSet[String] = {
val shuffleKeySet = new util.HashSet[String]()
- this.synchronized {
- shuffleKeySet.addAll(masterPartitionLocations.keySet())
- shuffleKeySet.addAll(slavePartitionLocations.keySet())
- }
+ shuffleKeySet.addAll(masterPartitionLocations.keySet())
+ shuffleKeySet.addAll(slavePartitionLocations.keySet())
shuffleKeySet
}
- def containsShuffle(shuffleKey: String): Boolean = this.synchronized {
+ def containsShuffle(shuffleKey: String): Boolean = {
masterPartitionLocations.containsKey(shuffleKey) ||
slavePartitionLocations.containsKey(shuffleKey)
}
- def addMasterPartition(shuffleKey: String, location: PartitionLocation): Int
= {
- addPartition(shuffleKey, location, masterPartitionLocations)
- }
-
- def addSlavePartition(shuffleKey: String, location: PartitionLocation): Int
= {
- addPartition(shuffleKey, location, slavePartitionLocations)
- }
-
def addMasterPartitions(
shuffleKey: String,
locations: util.List[PartitionLocation]): Unit = {
@@ -74,27 +76,20 @@ class PartitionLocationInfo extends Logging {
getLocation(shuffleKey, uniqueId, PartitionLocation.Mode.SLAVE)
}
- def getAllMasterLocations(shuffleKey: String): util.List[PartitionLocation]
= this.synchronized {
- if (masterPartitionLocations.containsKey(shuffleKey)) {
- masterPartitionLocations.get(shuffleKey)
- .values()
- .asScala
- .flatMap(_.asScala)
- .toList
- .asJava
- } else {
- new util.ArrayList[PartitionLocation]()
- }
+ def getAllMasterLocations(shuffleKey: String): util.List[PartitionLocation]
= {
+ getAllLocations(shuffleKey, masterPartitionLocations)
}
- def getAllSlaveLocations(shuffleKey: String): util.List[PartitionLocation] =
this.synchronized {
- if (slavePartitionLocations.containsKey(shuffleKey)) {
- slavePartitionLocations.get(shuffleKey)
- .values()
- .asScala
- .flatMap(_.asScala)
- .toList
- .asJava
+ def getAllSlaveLocations(shuffleKey: String): util.List[PartitionLocation] =
{
+ getAllLocations(shuffleKey, slavePartitionLocations)
+ }
+
+ def getAllLocations(
+ shuffleKey: String,
+ partitionInfo: PartitionInfo): util.List[PartitionLocation] = {
+ val partitionMap = partitionInfo.get(shuffleKey)
+ if (partitionMap != null) {
+ new util.ArrayList(partitionMap.values())
} else {
new util.ArrayList[PartitionLocation]()
}
@@ -122,81 +117,47 @@ class PartitionLocationInfo extends Logging {
removePartitions(shuffleKey, uniqueIds, slavePartitionLocations)
}
- def getAllMasterLocationsWithMinEpoch(shuffleKey: String):
util.List[PartitionLocation] =
- this.synchronized {
- getAllMasterLocationsWithExtremeEpoch(shuffleKey, (a, b) => a < b)
- }
-
- def getAllMasterLocationsWithExtremeEpoch(
- shuffleKey: String,
- order: (Int, Int) => Boolean): util.List[PartitionLocation] =
this.synchronized {
- if (masterPartitionLocations.containsKey(shuffleKey)) {
- masterPartitionLocations.get(shuffleKey)
- .values()
- .asScala
- .map { list =>
- var loc = list.get(0)
- 1 until list.size() foreach (ind => {
- if (order(list.get(ind).getEpoch, loc.getEpoch)) {
- loc = list.get(ind)
- }
- })
- loc
- }.toList.asJava
- } else {
- new util.ArrayList[PartitionLocation]()
- }
+ def removeShuffle(shuffleKey: String): Unit = {
+ masterPartitionLocations.remove(shuffleKey)
+ slavePartitionLocations.remove(shuffleKey)
}
- def getLocationWithMaxEpoch(
- shuffleKey: String,
- partitionId: Int): Option[PartitionLocation] = this.synchronized {
- if (!masterPartitionLocations.containsKey(shuffleKey) ||
- !masterPartitionLocations.get(shuffleKey).containsKey(partitionId)) {
- return None
- }
- val locations = masterPartitionLocations.get(shuffleKey).get(partitionId)
- if (locations == null || locations.size() == 0) {
- return None
- }
- var currentEpoch = -1
- var currentPartition: PartitionLocation = null
- locations.asScala.foreach(loc => {
- if (loc.getEpoch > currentEpoch) {
- currentEpoch = loc.getEpoch
- currentPartition = loc
- }
- })
- Some(currentPartition)
+ def getAllMasterLocationsWithMinEpoch(shuffleKey: String):
util.List[PartitionLocation] = {
+ getAllMasterLocationsWithExtremeEpoch(shuffleKey, (a, b) => a < b)
}
- private def addPartition(
+ // TODO: UT
+ def getAllMasterLocationsWithExtremeEpoch(
shuffleKey: String,
- location: PartitionLocation,
- partitionInfo: PartitionInfo): Int = this.synchronized {
- if (location != null) {
- partitionInfo.putIfAbsent(shuffleKey, new util.HashMap[Int,
util.List[PartitionLocation]]())
- val reduceLocMap = partitionInfo.get(shuffleKey)
- reduceLocMap.putIfAbsent(location.getId, new
util.ArrayList[PartitionLocation]())
- val locations = reduceLocMap.get(location.getId)
- locations.add(location)
- 1
+ order: (Int, Int) => Boolean): util.List[PartitionLocation] = {
+ val partitionMap = masterPartitionLocations.get(shuffleKey)
+ if (partitionMap != null) {
+ val resMap = new util.HashMap[Int, PartitionLocation]()
+ partitionMap.values().asScala.foreach(loc => {
+ val old = resMap.get(loc.getId)
+ if (old == null) {
+ resMap.put(loc.getId, loc)
+ } else {
+ if (order(loc.getEpoch, old.getEpoch)) {
+ resMap.put(loc.getId, loc)
+ }
+ }
+ })
+ new util.ArrayList(resMap.values())
} else {
- 0
+ new util.ArrayList[PartitionLocation]()
}
}
private def addPartitions(
shuffleKey: String,
locations: util.List[PartitionLocation],
- partitionInfo: PartitionInfo): Unit = this.synchronized {
+ partitionInfo: PartitionInfo): Unit = {
if (locations != null && locations.size() > 0) {
- partitionInfo.putIfAbsent(shuffleKey, new util.HashMap[Int,
util.List[PartitionLocation]]())
- val reduceLocMap = partitionInfo.get(shuffleKey)
+ partitionInfo.putIfAbsent(shuffleKey, new ConcurrentHashMap[Long,
PartitionLocation]())
+ val partitionMap = partitionInfo.get(shuffleKey)
locations.asScala.foreach { loc =>
- reduceLocMap.putIfAbsent(loc.getId, new
util.ArrayList[PartitionLocation]())
- val locations = reduceLocMap.get(loc.getId)
- locations.add(loc)
+ partitionMap.put(encode(loc.getId, loc.getEpoch), loc)
}
}
}
@@ -210,40 +171,27 @@ class PartitionLocationInfo extends Logging {
private def removePartitions(
shuffleKey: String,
uniqueIds: util.Collection[String],
- partitionInfo: PartitionInfo): (util.Map[String, Integer], Integer) =
this.synchronized {
- if (!partitionInfo.containsKey(shuffleKey)) {
+ partitionInfo: PartitionInfo): (util.Map[String, Integer], Integer) = {
+ val partitionMap = partitionInfo.get(shuffleKey)
+ if (partitionMap == null) {
return (Map.empty[String, Integer].asJava, 0)
}
val locMap = new util.HashMap[String, Integer]()
var numSlotsReleased: Int = 0
- val reduceLocMap = partitionInfo.get(shuffleKey)
uniqueIds.asScala.foreach { id =>
- val tokens = id.split("-", 2)
- val partitionId = tokens(0).toInt
- val epoch = tokens(1).toInt
- val locations = reduceLocMap.get(partitionId)
- if (locations != null) {
- val targetLocation = locations.asScala.find(_.getEpoch == epoch).orNull
- if (targetLocation != null) {
- locations.remove(targetLocation)
- numSlotsReleased += 1
- locMap.compute(
- targetLocation.getStorageInfo.getMountPoint,
- new BiFunction[String, Integer, Integer] {
- override def apply(t: String, u: Integer): Integer = {
- if (u == null) 1 else u + 1
- }
- })
- }
- }
- if (locations == null || locations.size() == 0) {
- reduceLocMap.remove(partitionId)
+ val loc = partitionMap.remove(encodeUniqueId(id))
+ if (loc != null) {
+ numSlotsReleased += 1
+ locMap.compute(
+ loc.getStorageInfo.getMountPoint,
+ new BiFunction[String, Integer, Integer] {
+ override def apply(t: String, u: Integer): Integer = {
+ if (u == null) 1 else u + 1
+ }
+ })
}
}
- if (reduceLocMap.size() == 0) {
- partitionInfo.remove(shuffleKey)
- }
// some locations might have no disk hint
(locMap, numSlotsReleased)
}
@@ -252,9 +200,6 @@ class PartitionLocationInfo extends Logging {
shuffleKey: String,
uniqueId: String,
mode: PartitionLocation.Mode): PartitionLocation = {
- val tokens = uniqueId.split("-", 2)
- val partitionId = tokens(0).toInt
- val epoch = tokens(1).toInt
val partitionInfo =
if (mode == PartitionLocation.Mode.MASTER) {
masterPartitionLocations
@@ -262,32 +207,19 @@ class PartitionLocationInfo extends Logging {
slavePartitionLocations
}
- this.synchronized {
- if (!partitionInfo.containsKey(shuffleKey)
- || !partitionInfo.get(shuffleKey).containsKey(partitionId)) {
- return null
- }
- partitionInfo.get(shuffleKey)
- .get(partitionId)
- .asScala
- .find(loc => loc.getEpoch == epoch)
- .orNull
- }
+ val partitionMap = partitionInfo.get(shuffleKey)
+ if (partitionMap != null) {
+ partitionMap.get(encodeUniqueId(uniqueId))
+ } else null
}
private def getAllIds(
shuffleKey: String,
- partitionInfo: PartitionInfo): util.List[String] = this.synchronized {
- if (!partitionInfo.containsKey(shuffleKey)) {
- return null
- }
- partitionInfo.get(shuffleKey)
- .values()
- .asScala
- .flatMap(_.asScala)
- .map(_.getUniqueId)
- .toList
- .asJava
+ partitionInfo: PartitionInfo): util.List[String] = {
+ val partitionMap = partitionInfo.get(shuffleKey)
+ if (partitionMap != null) {
+ partitionMap.values().asScala.map(_.getUniqueId).toList.asJava
+ } else null
}
private def getAllMasterIds(shuffleKey: String): util.List[String] = {
@@ -298,11 +230,11 @@ class PartitionLocationInfo extends Logging {
getAllIds(shuffleKey, slavePartitionLocations)
}
- def isEmpty: Boolean = this.synchronized {
+ def isEmpty: Boolean = {
masterPartitionLocations.isEmpty && slavePartitionLocations.isEmpty
}
- override def toString: String = this.synchronized {
+ override def toString: String = {
s"""
| Partition Location Info:
| master: ${masterPartitionLocations.asScala}
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index aed555878..920947d63 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -413,6 +413,7 @@ private[celeborn] class Worker(
expiredShuffleKeys.asScala.foreach { shuffleKey =>
partitionLocationInfo.removeMasterPartitions(shuffleKey)
partitionLocationInfo.removeSlavePartitions(shuffleKey)
+ partitionLocationInfo.removeShuffle(shuffleKey)
shufflePartitionType.remove(shuffleKey)
shuffleMapperAttempts.remove(shuffleKey)
shuffleCommitInfos.remove(shuffleKey)
