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

[nicholasjiang]

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

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


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

commit 99f021c63fc3a6f572f9775ee81ce302671827d1
Author: Rajeev Kumar <rajeevk...@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 <rajeevk...@iitk.ac.in>
    Signed-off-by: Nicholas Jiang <programg...@163.com>
    (cherry picked from commit 08a2da48972b2983590ddcddeea67a0948759921)
    Signed-off-by: Nicholas Jiang <programg...@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)