diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 19e03c63fd6..202eb62d3ad 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -729,7 +729,7 @@ Map<Integer, Integer> reviveBatch(
             pushExcludedWorkers.remove(loc.getPeer().hostAndPushPort());
           }
         } else if (StatusCode.STAGE_ENDED.getValue() == statusCode) {
-          stageEnded(shuffleId);
+          stageEndShuffleSet.add(shuffleId);
           return results;
         } else if (StatusCode.SHUFFLE_NOT_REGISTERED.getValue() == statusCode) {
           logger.error("SHUFFLE_NOT_REGISTERED!");
@@ -1650,10 +1650,10 @@ public void setupLifecycleManagerRef(RpcEndpointRef endpointRef) {
 
   boolean mapperEnded(int shuffleId, int mapId) {
     return (mapperEndMap.containsKey(shuffleId) && mapperEndMap.get(shuffleId).contains(mapId))
-        || stageEnded(shuffleId);
+        || isStageEnded(shuffleId);
   }
 
-  protected boolean stageEnded(int shuffleId) {
+  protected boolean isStageEnded(int shuffleId) {
     return stageEndShuffleSet.contains(shuffleId);
   }