[fix] distinguish sorting-file from sort-tasks waiting to be submitted.

apache · Dec 30, 2024 · 149dc32 · 149dc32
1 parent a572380
commit 149dc32
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
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
@@ -177,10 +177,14 @@ public PartitionFilesSorter(
         });
   }
 
-  public int getSortingCount() {
+  public int getPendingSortTaskCount() {
     return shuffleSortTaskDeque.size();
   }
 
+  public int getSortingCount() {
+    return sortingShuffleFiles.values().stream().map(Set::size).reduce(Integer::sum).orElse(0);
+  }
+
   public int getSortedCount() {
     return sortedFileCount.get();
   }

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
@@ -370,6 +370,9 @@ private[celeborn] class Worker(
   workerSource.addGauge(WorkerSource.SORT_MEMORY) { () =>
     memoryManager.getSortMemoryCounter.get()
   }
+  workerSource.addGauge(WorkerSource.PENDING_SORT_TASKS) { () =>
+    partitionsSorter.getPendingSortTaskCount
+  }
   workerSource.addGauge(WorkerSource.SORTING_FILES) { () =>
     partitionsSorter.getSortingCount
   }

diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
@@ -193,6 +193,7 @@ object WorkerSource {
   val SORT_TIME = "SortTime"
   val SORT_MEMORY = "SortMemory"
   val SORTING_FILES = "SortingFiles"
+  val PENDING_SORT_TASKS = "PendingSortTasks"
   val SORTED_FILES = "SortedFiles"
   val SORTED_FILE_SIZE = "SortedFileSize"
   val DISK_BUFFER = "DiskBuffer"