[CELEBORN-845][BUG] Sort memory counter won't decrease after sort fai…

…led.
apache · Jul 27, 2023 · ac9f8be · ac9f8be
1 parent e16b267
commit ac9f8be
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 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
@@ -2173,6 +2173,7 @@ object CelebornConf extends Logging {
       .doc("Reserved memory when sorting a shuffle file off-heap.")
       .version("0.3.0")
       .bytesConf(ByteUnit.BYTE)
+      .checkValue(v => v < Int.MaxValue, "Reserved memory per partition must be less than 2GB.")
       .createWithDefaultString("1mb")
 
   val WORKER_FLUSHER_BUFFER_SIZE: ConfigEntry[Long] =

diff --git a/...src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java b/...src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -133,6 +133,7 @@ public PartitionFilesSorter(
                       () -> {
                         try {
                           task.sort();
+                          memoryManager.releaseSortMemory(reservedMemoryPerPartition);
                         } catch (InterruptedException e) {
                           logger.warn(
                               "File sorter thread was interrupted when expanding padding buffer.");
@@ -586,8 +587,6 @@ public void sort() throws InterruptedException {
           sortedBlockInfoMap.put(mapId, sortedShuffleBlocks);
         }
 
-        memoryManager.releaseSortMemory(reserveMemory);
-
         writeIndex(sortedBlockInfoMap, indexFilePath, isHdfs);
         updateSortedShuffleFiles(shuffleKey, fileId, originFileLen);
         deleteOriginFiles();