[#2472] feat(spark): Add an rpc method to obtain the uniffleId and delete the Write Stage for retry at the same time.

Author: yl09099

client-mr/core/src/test/java/org/apache/hadoop/mapred/SortWriteBufferManagerTest.java renamed to client-mr/core/src/test/resources/SortWriteBufferManagerTest.java

@@ -54,6 +54,20 @@ public ShuffleIdMappingManager() {
     shuffleDeterminateMap = JavaUtils.newConcurrentMap();
   }
 
+  /**
+   * Record the Determinate status of each shuffle.
+   *
+   * @param shuffleId
+   * @param isDeterminate
+   */
+  public void recordShuffleIdDeterminate(int shuffleId, boolean isDeterminate) {
+    shuffleDeterminateMap.put(shuffleId, isDeterminate);
+  }
+
+  public boolean getShuffleIdDeterminate(int shuffleId) {
+    return shuffleDeterminateMap.get(shuffleId);
+  }
+
   /**
    * Create the shuffleId of uniffle based on the ShuffleID of Spark. When registerShuffle is being
    * performed, the default number of attempts by our stage is 0.
@@ -74,7 +88,48 @@ public int createUniffleShuffleId(int shuffleId) {
         .get(0);
   }
 
-  public void recordShuffleIdDeterminate(int shuffleId, boolean isDeterminate) {
-    shuffleDeterminateMap.put(shuffleId, isDeterminate);
+  /**
+   * A new shuffleId is generated based on the shuffleId of Spark and the stage Attempt number of
+   * Spark, and it is generally called only when the stage retry is triggered.
+   *
+   * @param shuffleId
+   * @param stageAttemptNumber
+   * @return
+   */
+  public int createUniffleShuffleId(int shuffleId, int stageAttemptNumber) {
+    int generatorShuffleId = shuffleIdGenerator.incrementAndGet();
+    shuffleIdMapping.get(shuffleId).put(stageAttemptNumber, generatorShuffleId);
+    return generatorShuffleId;
+  }
+
+  /**
+   * Determine whether there is a record based on the passed-in Spark shuffleId and the number of
+   * Spark Stage attempts.
+   *
+   * @param shuffleId
+   * @param stageAttemptNumber
+   * @return
+   */
+  public boolean hasUniffleShuffleIdByStageNumber(int shuffleId, int stageAttemptNumber) {
+    return shuffleIdMapping.get(shuffleId).get(stageAttemptNumber) != null;
+  }
+
+  public int getUniffleShuffleIdByStageNumber(int shuffleId, int stageAttemptNumber) {
+    return shuffleIdMapping.get(shuffleId).get(stageAttemptNumber);
+  }
+
+  /**
+   * Based on the passed-in Spark shuffleId, obtain all the stage attemptNumber. Sort in reverse
+   * according to the stage attemptNumber to obtain the current Uniffle shuffleId
+   *
+   * @param shuffleId
+   * @return
+   */
+  public int getUniffleShuffleIdByMaxStageNumber(int shuffleId) {
+    return shuffleIdMapping.get(shuffleId).entrySet().stream()
+        .sorted(Map.Entry.<Integer, Integer>comparingByKey().reversed())
+        .findFirst()
+        .get()
+        .getValue();
   }
 }

client-spark/common/src/main/java/org/apache/ShuffleIdMappingManager.java

@@ -46,22 +46,6 @@ public class RssSparkConfig {
           .withDeprecatedKeys(RssClientConfig.RSS_RESUBMIT_STAGE)
           .withDescription("Whether to enable the resubmit stage for fetch/write failure");
 
-  public static final ConfigOption<Boolean> RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED =
-      ConfigOptions.key("rss.stageRetry.fetchFailureEnabled")
-          .booleanType()
-          .defaultValue(false)
-          .withFallbackKeys(RSS_RESUBMIT_STAGE_ENABLED.key(), RssClientConfig.RSS_RESUBMIT_STAGE)
-          .withDescription(
-              "If set to true, the stage retry mechanism will be enabled when a fetch failure occurs.");
-
-  public static final ConfigOption<Boolean> RSS_RESUBMIT_STAGE_WITH_WRITE_FAILURE_ENABLED =
-      ConfigOptions.key("rss.stageRetry.writeFailureEnabled")
-          .booleanType()
-          .defaultValue(false)
-          .withFallbackKeys(RSS_RESUBMIT_STAGE_ENABLED.key(), RssClientConfig.RSS_RESUBMIT_STAGE)
-          .withDescription(
-              "If set to true, the stage retry mechanism will be enabled when a write failure occurs.");
-
   public static final ConfigOption<Boolean> RSS_BLOCK_ID_SELF_MANAGEMENT_ENABLED =
       ConfigOptions.key("rss.blockId.selfManagementEnabled")
           .booleanType()

client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java

@@ -54,7 +54,7 @@
 import org.apache.uniffle.common.exception.RssFetchFailedException;
 import org.apache.uniffle.common.util.Constants;
 
-import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_ENABLED;
 
 public class RssSparkShuffleUtils {
 
@@ -366,17 +366,21 @@ public static RssException reportRssFetchFailedException(
       SparkConf sparkConf,
       String appId,
       int shuffleId,
+      int uniffleShuffleId,
       int stageAttemptId,
+      int stageAttemptNumber,
       Set<Integer> failedPartitions) {
     RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
-    if (rssConf.getBoolean(RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED)
+    if (rssConf.getBoolean(RSS_RESUBMIT_STAGE_ENABLED)
         && RssSparkShuffleUtils.isStageResubmitSupported()) {
       for (int partitionId : failedPartitions) {
         RssReportShuffleFetchFailureRequest req =
             new RssReportShuffleFetchFailureRequest(
                 appId,
                 shuffleId,
+                uniffleShuffleId,
                 stageAttemptId,
+                stageAttemptNumber,
                 partitionId,
                 rssFetchFailedException.getMessage());
         RssReportShuffleFetchFailureResponse response =

client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java

@@ -21,26 +21,26 @@
 import java.io.IOException;
 import java.util.Map;
 
-import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.MutableShuffleHandleInfo;
 
 import org.apache.uniffle.common.util.JavaUtils;
 
 public class ShuffleHandleInfoManager implements Closeable {
-  private Map<Integer, ShuffleHandleInfo> shuffleIdToShuffleHandleInfo;
+  private Map<Integer, MutableShuffleHandleInfo> shuffleIdToShuffleHandleInfo;
 
   public ShuffleHandleInfoManager() {
     this.shuffleIdToShuffleHandleInfo = JavaUtils.newConcurrentMap();
   }
 
-  public ShuffleHandleInfo get(int shuffleId) {
+  public MutableShuffleHandleInfo get(int shuffleId) {
     return shuffleIdToShuffleHandleInfo.get(shuffleId);
   }
 
   public void remove(int shuffleId) {
     shuffleIdToShuffleHandleInfo.remove(shuffleId);
   }
 
-  public void register(int shuffleId, ShuffleHandleInfo handle) {
+  public void register(int shuffleId, MutableShuffleHandleInfo handle) {
     shuffleIdToShuffleHandleInfo.put(shuffleId, handle);
   }
 

client-spark/common/src/main/java/org/apache/spark/shuffle/ShuffleHandleInfoManager.java

@@ -48,8 +48,10 @@ private RssFetchFailedIterator(Builder builder, Iterator<Product2<K, C>> iterato
   public static class Builder {
     private String appId;
     private int shuffleId;
+    private int uniffleShuffleId;
     private int partitionId;
     private int stageAttemptId;
+    private int stageAttemptNumber;
     private Supplier<ShuffleManagerClient> managerClientSupplier;
 
     private Builder() {}
@@ -64,6 +66,11 @@ Builder shuffleId(int shuffleId) {
       return this;
     }
 
+    Builder uniffleShuffleId(int uniffleShuffleId) {
+      this.uniffleShuffleId = uniffleShuffleId;
+      return this;
+    }
+
     Builder partitionId(int partitionId) {
       this.partitionId = partitionId;
       return this;
@@ -74,6 +81,11 @@ Builder stageAttemptId(int stageAttemptId) {
       return this;
     }
 
+    Builder stageAttemptNumber(int stageAttemptNumber) {
+      this.stageAttemptNumber = stageAttemptNumber;
+      return this;
+    }
+
     Builder managerClientSupplier(Supplier<ShuffleManagerClient> managerClientSupplier) {
       this.managerClientSupplier = managerClientSupplier;
       return this;
@@ -95,7 +107,9 @@ private RssException generateFetchFailedIfNecessary(RssFetchFailedException e) {
         new RssReportShuffleFetchFailureRequest(
             builder.appId,
             builder.shuffleId,
+            builder.uniffleShuffleId,
             builder.stageAttemptId,
+            builder.stageAttemptNumber,
             builder.partitionId,
             e.getMessage());
     RssReportShuffleFetchFailureResponse response = client.reportShuffleFetchFailure(req);

client-spark/common/src/main/java/org/apache/spark/shuffle/handle/StageAttemptShuffleHandleInfo.java

@@ -25,22 +25,19 @@
 public class AddBlockEvent {
 
   private String taskId;
-  private int stageAttemptNumber;
   private List<ShuffleBlockInfo> shuffleDataInfoList;
   private List<Runnable> processedCallbackChain;
   private WriteBufferManager bufferManager;
 
   public AddBlockEvent(String taskId, List<ShuffleBlockInfo> shuffleDataInfoList) {
-    this(taskId, 0, shuffleDataInfoList, null);
+    this(taskId, shuffleDataInfoList, null);
   }
 
   public AddBlockEvent(
       String taskId,
-      int stageAttemptNumber,
       List<ShuffleBlockInfo> shuffleDataInfoList,
       WriteBufferManager writeBufferManager) {
     this.taskId = taskId;
-    this.stageAttemptNumber = stageAttemptNumber;
     this.shuffleDataInfoList = shuffleDataInfoList;
     this.processedCallbackChain = new ArrayList<>();
     this.bufferManager = writeBufferManager;
@@ -55,10 +52,6 @@ public String getTaskId() {
     return taskId;
   }
 
-  public int getStageAttemptNumber() {
-    return stageAttemptNumber;
-  }
-
   public List<ShuffleBlockInfo> getShuffleDataInfoList() {
     return shuffleDataInfoList;
   }

client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java

@@ -93,10 +93,7 @@ public CompletableFuture<Long> send(AddBlockEvent event) {
               try {
                 result =
                     shuffleWriteClient.sendShuffleData(
-                        rssAppId,
-                        event.getStageAttemptNumber(),
-                        shuffleBlockInfoList,
-                        () -> !isValidTask(taskId));
+                        rssAppId, shuffleBlockInfoList, () -> !isValidTask(taskId));
                 putBlockId(taskToSuccessBlockIds, taskId, result.getSuccessBlockIds());
                 putFailedBlockSendTracker(
                     taskToFailedBlockSendTracker, taskId, result.getFailedBlockSendTracker());