Triggering Stage retry requires reassigning the shuffle server in the retry Stage

Author: yl09099

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java

@@ -34,6 +34,8 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
+import scala.Tuple2;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -684,8 +686,6 @@ public boolean reassignOnStageResubmit(
     int requiredShuffleServerNumber =
         RssSparkShuffleUtils.getRequiredShuffleServerNumber(sparkConf);
     int estimateTaskConcurrency = RssSparkShuffleUtils.estimateTaskConcurrency(sparkConf);
-    // Deregister the shuffleId corresponding to the Shuffle Server.
-    shuffleWriteClient.unregisterShuffle(appId, shuffleId);
     Map<Integer, List<ShuffleServerInfo>> partitionToServers =
         requestShuffleAssignment(
             shuffleId,
@@ -1042,7 +1042,7 @@ protected void registerShuffleServers(
     }
     LOG.info("Start to register shuffleId {}", shuffleId);
     long start = System.currentTimeMillis();
-    Map<String, String> sparkConfMap = RssSparkConfig.sparkConfToMap(getSparkConf());
+    Map<String, String> sparkConfMap = sparkConfToMap(getSparkConf());
     serverToPartitionRanges.entrySet().stream()
         .forEach(
             entry -> {
@@ -1073,7 +1073,7 @@ protected void registerShuffleServers(
     }
     LOG.info("Start to register shuffleId[{}]", shuffleId);
     long start = System.currentTimeMillis();
-    Map<String, String> sparkConfMap = RssSparkConfig.sparkConfToMap(getSparkConf());
+    Map<String, String> sparkConfMap = sparkConfToMap(getSparkConf());
     Set<Map.Entry<ShuffleServerInfo, List<PartitionRange>>> entries =
         serverToPartitionRanges.entrySet();
     entries.stream()
@@ -1119,4 +1119,19 @@ public boolean isRssStageRetryForFetchFailureEnabled() {
   public SparkConf getSparkConf() {
     return sparkConf;
   }
+
+  public Map<String, String> sparkConfToMap(SparkConf sparkConf) {
+    Map<String, String> map = new HashMap<>();
+
+    for (Tuple2<String, String> tuple : sparkConf.getAll()) {
+      String key = tuple._1;
+      map.put(key, tuple._2);
+    }
+
+    return map;
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
 }

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerInterface.java

@@ -24,6 +24,7 @@
 import org.apache.spark.shuffle.handle.MutableShuffleHandleInfo;
 import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
 
+import org.apache.uniffle.client.api.ShuffleWriteClient;
 import org.apache.uniffle.common.ReceivingFailureServer;
 import org.apache.uniffle.shuffle.BlockIdManager;
 
@@ -88,4 +89,11 @@ MutableShuffleHandleInfo reassignOnBlockSendFailure(
       int shuffleId,
       Map<Integer, List<ReceivingFailureServer>> partitionToFailureServers,
       boolean partitionSplit);
+
+  /**
+   * Driver Obtains ShuffleWriteClient.
+   *
+   * @return
+   */
+  ShuffleWriteClient getShuffleWriteClient();
 }

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerGrpcService.java

@@ -118,6 +118,8 @@ public void reportShuffleWriteFailure(
                 // Clear the metadata of the completed task, otherwise some of the stage's data will
                 // be lost.
                 shuffleManager.unregisterAllMapOutput(shuffleId);
+                // Deregister the shuffleId corresponding to the Shuffle Server.
+                shuffleManager.getShuffleWriteClient().unregisterShuffle(appId, shuffleId);
                 shuffleServerWriterFailureRecord.setClearedMapTrackerBlock(true);
                 LOG.info(
                     "Clear shuffle result in shuffleId:{}, stageId:{}, stageAttemptNumber:{}.",

client-spark/common/src/test/java/org/apache/uniffle/shuffle/manager/DummyRssShuffleManager.java

@@ -25,6 +25,7 @@
 import org.apache.spark.shuffle.handle.MutableShuffleHandleInfo;
 import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
 
+import org.apache.uniffle.client.api.ShuffleWriteClient;
 import org.apache.uniffle.common.ReceivingFailureServer;
 import org.apache.uniffle.shuffle.BlockIdManager;
 
@@ -84,4 +85,9 @@ public MutableShuffleHandleInfo reassignOnBlockSendFailure(
       boolean partitionSplit) {
     return null;
   }
+
+  @Override
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return null;
+  }
 }