(celeborn) branch main updated: [CELEBORN-2426] Correct the calculation of the count of unreleased partitions at shutdown

[ni...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new 08a2da489 [CELEBORN-2426] Correct the calculation of the count of unreleased partitions at shutdown
08a2da489 is described below

commit 08a2da48972b2983590ddcddeea67a0948759921
Author: Rajeev Kumar <ra...@iitk.ac.in>
AuthorDate: Tue Aug 18 17:52:34 2026 +0800

    [CELEBORN-2426] Correct the calculation of the count of unreleased partitions at shutdown
    
    ### What changes were proposed in this pull request?
    This PR corrects the calculation of the number of unreleased partitions at shutdown.
    
    ### Why are the changes needed?
    
    #### Bug
    The current calculation of the number of unreleased partitions at shutdown is incorrect ([source](https://github.com/rjvkr2021/celeborn/blob/main/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala#L488-L489)).
    ```
    partitionLocationInfo.primaryPartitionLocations.size() +
            partitionLocationInfo.replicaPartitionLocations.size()
    ```
    
    #### Why so?
    
    _primaryPartitionLocations__ and _replicaPartitionLocations__ are maps shaped as:
    ```
    shuffleKey -> (uniquePartitionLocationId -> PartitionLocation)
    ```
    Therefore, the current calculation counts the number of shuffles containing unreleased partitions at shutdown, not the number of unreleased partitions itself.
    
    #### Fix
    
    It should sum the sizes of the inner maps.
    ```
    partitionLocationInfo.primaryPartitionLocations.values().asScala.map(_.size()).sum +
     partitionLocationInfo.replicaPartitionLocations.values().asScala.map(_.size()).sum
    ```
    
    ### Does this PR resolve a correctness bug?
    
    - [x] Yes
    
    ### Does this PR introduce _any_ user-facing change?
    
    No
    
    ### How was this patch tested?
    
    NA
    
    Closes #3807 from rjvkr2021/rajeevkumar/fix_unreleased_partition_count_calculation.
    
    Authored-by: Rajeev Kumar <ra...@iitk.ac.in>
    Signed-off-by: Nicholas Jiang <pr...@163.com>
---
 .../scala/org/apache/celeborn/service/deploy/worker/Worker.scala | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

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 91e145ba6..09999bad9 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
@@ -485,8 +485,8 @@ private[celeborn] class Worker(
   // Unreleased partition location count when worker is restarting
   workerSource.addGauge(WorkerSource.UNRELEASED_PARTITION_LOCATION_COUNT) { () =>
     if (shutdown.get()) {
-      partitionLocationInfo.primaryPartitionLocations.size() +
-        partitionLocationInfo.replicaPartitionLocations.size()
+      partitionLocationCount(partitionLocationInfo.primaryPartitionLocations) +
+        partitionLocationCount(partitionLocationInfo.replicaPartitionLocations)
     } else {
       0
     }
@@ -509,6 +509,11 @@ private[celeborn] class Worker(
     }
   }
 
+  private def partitionLocationCount(
+      partitionLocations: WorkerPartitionLocationInfo#PartitionInfo): Int = {
+    partitionLocations.values().asScala.map(_.size()).sum
+  }
+
   private def heartbeatToMaster(): Unit = {
     val activeShuffleKeys = new JHashSet[String]()
     activeShuffleKeys.addAll(partitionLocationInfo.shuffleKeySet)