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 3c8c58e09 [CELEBORN-301] Refactor PartitionLocationInfo to use
ConcurrentHashMap (#1278)
3c8c58e09 is described below
commit 3c8c58e09d01d815d68e769970f9de03959d14ac
Author: Keyong Zhou <[email protected]>
AuthorDate: Sun Feb 26 16:46:30 2023 +0800
[CELEBORN-301] Refactor PartitionLocationInfo to use ConcurrentHashMap
(#1278)
---
.../common/meta/WorkerPartitionLocationInfo.scala | 142 +++++++++------------
.../celeborn/service/deploy/worker/Worker.scala | 1 +
2 files changed, 61 insertions(+), 82 deletions(-)
diff --git
a/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
b/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
index b716d65a7..5f172b4c9 100644
---
a/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.scala
+++
b/common/src/main/scala/org/apache/celeborn/common/meta/WorkerPartitionLocationInfo.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,21 +28,30 @@ import org.apache.celeborn.common.protocol.PartitionLocation
class WorkerPartitionLocationInfo extends Logging {
- // key: ShuffleKey, values: (partitionId -> partition locations)
- type PartitionInfo = util.HashMap[String, util.Map[Int,
util.List[PartitionLocation]]]
+ // key: ShuffleKey, values: (partitionId -> (encodedPartitionId ->
PartitionLocation))
+ 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)
}
@@ -66,11 +76,16 @@ class WorkerPartitionLocationInfo extends Logging {
getLocation(shuffleKey, uniqueId, PartitionLocation.Mode.SLAVE)
}
- def getAllMasterLocations(shuffleKey: String): util.List[PartitionLocation]
= this.synchronized {
+ def removeShuffle(shuffleKey: String): Unit = {
+ masterPartitionLocations.remove(shuffleKey)
+ slavePartitionLocations.remove(shuffleKey)
+ }
+
+ def getAllMasterLocations(shuffleKey: String): util.List[PartitionLocation]
= {
getMasterLocations(shuffleKey)
}
- def getAllSlaveLocations(shuffleKey: String): util.List[PartitionLocation] =
this.synchronized {
+ def getAllSlaveLocations(shuffleKey: String): util.List[PartitionLocation] =
{
getSlaveLocations(shuffleKey)
}
@@ -86,24 +101,18 @@ class WorkerPartitionLocationInfo extends Logging {
getLocations(shuffleKey, slavePartitionLocations, partitionIdOpt)
}
- private def getLocations(
+ def getLocations(
shuffleKey: String,
partitionInfo: PartitionInfo,
- partitionIdOpt: Option[Int] = None): util.List[PartitionLocation] =
this.synchronized {
- if (partitionInfo.containsKey(shuffleKey)) {
+ partitionIdOpt: Option[Int] = None): util.List[PartitionLocation] = {
+ val partitionMap = partitionInfo.get(shuffleKey)
+ if (partitionMap != null) {
partitionIdOpt match {
- case Some(partitionId) => partitionInfo.get(shuffleKey)
- .values()
- .asScala
- .flatMap(_.asScala)
- .filter(_.getId == partitionId)
+ case Some(partitionId) =>
+ partitionMap.values().asScala.filter(_.getId == partitionId)
.toList.asJava
case None =>
- partitionInfo.get(shuffleKey)
- .values()
- .asScala
- .flatMap(_.asScala)
- .toList.asJava
+ new util.ArrayList(partitionMap.values())
}
} else {
new util.ArrayList[PartitionLocation]()
@@ -132,17 +141,15 @@ class WorkerPartitionLocationInfo extends Logging {
removePartitions(shuffleKey, uniqueIds, slavePartitionLocations)
}
- private def addPartitions(
+ 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.putIfAbsent(encode(loc.getId, loc.getEpoch), loc)
}
}
}
@@ -156,40 +163,27 @@ class WorkerPartitionLocationInfo 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)
}
@@ -198,9 +192,6 @@ class WorkerPartitionLocationInfo 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
@@ -208,32 +199,19 @@ class WorkerPartitionLocationInfo 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] = {
@@ -244,11 +222,11 @@ class WorkerPartitionLocationInfo 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 7617d1c6d..68825aa4b 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
@@ -432,6 +432,7 @@ private[celeborn] class Worker(
expiredShuffleKeys.asScala.foreach { shuffleKey =>
partitionLocationInfo.removeMasterPartitions(shuffleKey)
partitionLocationInfo.removeSlavePartitions(shuffleKey)
+ partitionLocationInfo.removeShuffle(shuffleKey)
shufflePartitionType.remove(shuffleKey)
shufflePushDataTimeout.remove(shuffleKey)
shuffleMapperAttempts.remove(shuffleKey)
