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 2024/01/04 03:34:56 UTC
(incubator-celeborn) branch main updated: [CELEBORN-1201] Optimize memory usage of cache in partition sorter
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 8b6aae04d [CELEBORN-1201] Optimize memory usage of cache in partition sorter
8b6aae04d is described below
commit 8b6aae04d1032f7fba1ddad2cbe5c6deb5563064
Author: mingji <fe...@alibaba-inc.com>
AuthorDate: Thu Jan 4 11:34:48 2024 +0800
[CELEBORN-1201] Optimize memory usage of cache in partition sorter
### What changes were proposed in this pull request?
Add a cache in partition sorted and limit its max size.
### Why are the changes needed?
To reduce memory consumption in partition sort by tweak the index cache.
### Does this PR introduce _any_ user-facing change?
NO.
### How was this patch tested?
GA and cluster.
Closes #2194 from FMX/B1201.
Lead-authored-by: mingji <fe...@alibaba-inc.com>
Co-authored-by: Keyong Zhou <zk...@alibaba-inc.com>
Signed-off-by: zky.zhoukeyong <zk...@alibaba-inc.com>
---
.../org/apache/celeborn/common/CelebornConf.scala | 18 ++++
docs/configuration/worker.md | 2 +
.../worker/storage/PartitionFilesSorter.java | 116 +++++++++++++--------
3 files changed, 92 insertions(+), 44 deletions(-)
diff --git a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 589a8ab66..0d13a0d1d 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -689,6 +689,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
get(WORKER_PARTITION_SORTER_PER_PARTITION_RESERVED_MEMORY)
def partitionSorterThreads: Int =
get(PARTITION_SORTER_THREADS).getOrElse(Runtime.getRuntime.availableProcessors)
+ def partitionSorterIndexCacheMaxWeight: Long = get(PARTITION_SORTER_INDEX_CACHE_MAX_WEIGHT)
+ def partitionSorterIndexExpire: Long = get(PARTITION_SORTER_INDEX_CACHE_EXPIRE)
def workerPushHeartbeatEnabled: Boolean = get(WORKER_PUSH_HEARTBEAT_ENABLED)
def workerPushMaxComponents: Int = get(WORKER_PUSH_COMPOSITEBUFFER_MAXCOMPONENTS)
def workerFetchHeartbeatEnabled: Boolean = get(WORKER_FETCH_HEARTBEAT_ENABLED)
@@ -2418,6 +2420,22 @@ object CelebornConf extends Logging {
.intConf
.createOptional
+ val PARTITION_SORTER_INDEX_CACHE_MAX_WEIGHT: ConfigEntry[Long] =
+ buildConf("celeborn.worker.sortPartition.indexCache.maxWeight")
+ .categories("worker")
+ .doc("PartitionSorter's cache max weight for index buffer.")
+ .version("0.4.0")
+ .longConf
+ .createWithDefault(100000)
+
+ val PARTITION_SORTER_INDEX_CACHE_EXPIRE: ConfigEntry[Long] =
+ buildConf("celeborn.worker.sortPartition.indexCache.expire")
+ .categories("worker")
+ .doc("PartitionSorter's cache item expire time.")
+ .version("0.4.0")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("180s")
+
val WORKER_PARTITION_SORTER_PER_PARTITION_RESERVED_MEMORY: ConfigEntry[Long] =
buildConf("celeborn.worker.sortPartition.reservedMemoryPerPartition")
.withAlternative("celeborn.worker.partitionSorter.reservedMemoryPerPartition")
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index 8628d3d13..c29c44305 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -110,6 +110,8 @@ license: |
| celeborn.worker.shuffle.partitionSplit.enabled | true | enable the partition split on worker side | 0.3.0 |
| celeborn.worker.shuffle.partitionSplit.max | 2g | Specify the maximum partition size for splitting, and ensure that individual partition files are always smaller than this limit. | 0.3.0 |
| celeborn.worker.shuffle.partitionSplit.min | 1m | Min size for a partition to split | 0.3.0 |
+| celeborn.worker.sortPartition.indexCache.expire | 180s | PartitionSorter's cache item expire time. | 0.4.0 |
+| celeborn.worker.sortPartition.indexCache.maxWeight | 100000 | PartitionSorter's cache max weight for index buffer. | 0.4.0 |
| celeborn.worker.sortPartition.reservedMemoryPerPartition | 1mb | Reserved memory when sorting a shuffle file off-heap. | 0.3.0 |
| celeborn.worker.sortPartition.threads | <undefined> | PartitionSorter's thread counts. It's recommended to set at least `64` when `HDFS` is enabled in `celeborn.storage.activeTypes`. | 0.3.0 |
| celeborn.worker.sortPartition.timeout | 220s | Timeout for a shuffle file to sort. | 0.3.0 |
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
index f2b6e5f96..97d0d8498 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -30,10 +30,7 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Condition;
@@ -41,6 +38,8 @@ import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import com.google.common.annotations.VisibleForTesting;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
import org.apache.commons.io.IOUtils;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FSDataOutputStream;
@@ -77,8 +76,9 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
JavaUtils.newConcurrentHashMap();
private final ConcurrentHashMap<String, Set<String>> sortingShuffleFiles =
JavaUtils.newConcurrentHashMap();
- private final ConcurrentHashMap<String, Map<String, Map<Integer, List<ShuffleBlockInfo>>>>
- cachedIndexMaps = JavaUtils.newConcurrentHashMap();
+ private final Cache<String, Map<Integer, List<ShuffleBlockInfo>>> indexCache;
+ private final Map<String, Set<String>> indexCacheNames = JavaUtils.newConcurrentHashMap();
+
private final LinkedBlockingQueue<FileSorter> shuffleSortTaskDeque = new LinkedBlockingQueue<>();
private final PartitionFilesCleaner cleaner;
@@ -95,6 +95,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
private final ExecutorService fileSorterExecutors;
private final Thread fileSorterSchedulerThread;
+ private final long indexCacheMaxWeight;
public PartitionFilesSorter(
MemoryManager memoryManager, CelebornConf conf, AbstractSource source) {
@@ -103,6 +104,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
this.reservedMemoryPerPartition = conf.partitionSorterReservedMemoryPerPartition();
this.partitionSorterShutdownAwaitTime =
conf.workerGracefulShutdownPartitionSorterCloseAwaitTimeMs();
+ this.indexCacheMaxWeight = conf.partitionSorterIndexCacheMaxWeight();
this.source = source;
this.cleaner = new PartitionFilesCleaner(this);
this.gracefulShutdown = conf.workerGracefulShutdown();
@@ -129,6 +131,17 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
ThreadUtils.newDaemonCachedThreadPool(
"worker-file-sorter-execute", conf.partitionSorterThreads(), 120);
+ indexCache =
+ CacheBuilder.newBuilder()
+ .concurrencyLevel(conf.partitionSorterThreads())
+ .expireAfterAccess(conf.partitionSorterIndexExpire(), TimeUnit.MILLISECONDS)
+ .maximumWeight(indexCacheMaxWeight)
+ .weigher(
+ (key, cache) ->
+ ((Map<Integer, List<ShuffleBlockInfo>>) cache)
+ .values().stream().mapToInt(List::size).sum())
+ .build();
+
fileSorterSchedulerThread =
new Thread(
() -> {
@@ -253,7 +266,12 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
for (String expiredShuffleKey : expiredShuffleKeys) {
sortingShuffleFiles.remove(expiredShuffleKey);
deleteSortedShuffleFiles(expiredShuffleKey);
- cachedIndexMaps.remove(expiredShuffleKey);
+ Set<String> expiredIndexCacheItems = indexCacheNames.remove(expiredShuffleKey);
+ if (expiredIndexCacheItems != null) {
+ for (String expiredIndexCacheItem : expiredIndexCacheItems) {
+ indexCache.invalidate(expiredIndexCacheItem);
+ }
+ }
}
}
@@ -295,7 +313,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
}
}
}
- cachedIndexMaps.clear();
+ indexCache.invalidateAll();
}
private void reloadAndCleanSortedShuffleFiles(DB db) {
@@ -466,42 +484,52 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
int endMapIndex)
throws IOException {
Map<Integer, List<ShuffleBlockInfo>> indexMap;
- if (cachedIndexMaps.containsKey(shuffleKey)
- && cachedIndexMaps.get(shuffleKey).containsKey(fileId)) {
- indexMap = cachedIndexMaps.get(shuffleKey).get(fileId);
- } else {
- FileChannel indexChannel = null;
- FSDataInputStream hdfsIndexStream = null;
- boolean isHdfs = Utils.isHdfsPath(indexFilePath);
- int indexSize = 0;
- try {
- if (isHdfs) {
- hdfsIndexStream = StorageManager.hadoopFs().open(new Path(indexFilePath));
- indexSize =
- (int) StorageManager.hadoopFs().getFileStatus(new Path(indexFilePath)).getLen();
- } else {
- indexChannel = FileChannelUtils.openReadableFileChannel(indexFilePath);
- File indexFile = new File(indexFilePath);
- indexSize = (int) indexFile.length();
- }
- ByteBuffer indexBuf = ByteBuffer.allocate(indexSize);
- if (isHdfs) {
- readStreamFully(hdfsIndexStream, indexBuf, indexFilePath);
- } else {
- readChannelFully(indexChannel, indexBuf, indexFilePath);
- }
- indexBuf.rewind();
- indexMap = ShuffleBlockInfoUtils.parseShuffleBlockInfosFromByteBuffer(indexBuf);
- Map<String, Map<Integer, List<ShuffleBlockInfo>>> cacheMap =
- cachedIndexMaps.computeIfAbsent(shuffleKey, v -> JavaUtils.newConcurrentHashMap());
- cacheMap.put(fileId, indexMap);
- } catch (Exception e) {
- logger.error("Read sorted shuffle file index " + indexFilePath + " error, detail: ", e);
- throw new IOException("Read sorted shuffle file index failed.", e);
- } finally {
- IOUtils.closeQuietly(indexChannel, null);
- IOUtils.closeQuietly(hdfsIndexStream, null);
- }
+ try {
+ indexMap =
+ indexCache.get(
+ fileId,
+ () -> {
+ FileChannel indexChannel = null;
+ FSDataInputStream hdfsIndexStream = null;
+ boolean isHdfs = Utils.isHdfsPath(indexFilePath);
+ int indexSize = 0;
+ try {
+ if (isHdfs) {
+ hdfsIndexStream = StorageManager.hadoopFs().open(new Path(indexFilePath));
+ indexSize =
+ (int)
+ StorageManager.hadoopFs()
+ .getFileStatus(new Path(indexFilePath))
+ .getLen();
+ } else {
+ indexChannel = FileChannelUtils.openReadableFileChannel(indexFilePath);
+ File indexFile = new File(indexFilePath);
+ indexSize = (int) indexFile.length();
+ }
+ ByteBuffer indexBuf = ByteBuffer.allocate(indexSize);
+ if (isHdfs) {
+ readStreamFully(hdfsIndexStream, indexBuf, indexFilePath);
+ } else {
+ readChannelFully(indexChannel, indexBuf, indexFilePath);
+ }
+ indexBuf.rewind();
+ Map<Integer, List<ShuffleBlockInfo>> tIndexMap =
+ ShuffleBlockInfoUtils.parseShuffleBlockInfosFromByteBuffer(indexBuf);
+ Set<String> indexCacheItemsSet =
+ indexCacheNames.computeIfAbsent(shuffleKey, v -> new HashSet<>());
+ indexCacheItemsSet.add(fileId);
+ return tIndexMap;
+ } catch (Exception e) {
+ logger.error(
+ "Read sorted shuffle file index " + indexFilePath + " error, detail: ", e);
+ throw new IOException("Read sorted shuffle file index failed.", e);
+ } finally {
+ IOUtils.closeQuietly(indexChannel, null);
+ IOUtils.closeQuietly(hdfsIndexStream, null);
+ }
+ });
+ } catch (ExecutionException e) {
+ throw new IOException("Read sorted shuffle file index failed.", e);
}
ReduceFileMeta reduceFileMeta =
new ReduceFileMeta(