[apache#1608][part-8] feat(spark3): add a limit to the number of retr…

…ies when block access is denied
dingshun3016 · May 15, 2024 · b683d1b · b683d1b
1 parent a3461fb
commit b683d1b
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 1 deletion.
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkConfig.java
@@ -90,6 +90,14 @@ public class RssSparkConfig {
           .defaultValue(1)
           .withDescription("The block retry max times when partition reassign is enabled.");
 
+  public static final ConfigOption<Integer>
+      RSS_PARTITION_REASSIGN_BLOCK_ACCESS_DENIED_RETRY_MAX_TIMES =
+          ConfigOptions.key("rss.client.reassign.blockAccessDeniedRetryMaxTimes")
+              .intType()
+              .defaultValue(1)
+              .withDescription(
+                  "The block access denied retry max times when partition reassign is enabled.");
+
   public static final String SPARK_RSS_CONFIG_PREFIX = "spark.";
 
   public static final ConfigEntry<Integer> RSS_PARTITION_NUM_PER_RANGE =

diff --git a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
@@ -90,6 +90,7 @@
 import org.apache.uniffle.common.rpc.StatusCode;
 import org.apache.uniffle.storage.util.StorageType;
 
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_PARTITION_REASSIGN_BLOCK_ACCESS_DENIED_RETRY_MAX_TIMES;
 import static org.apache.spark.shuffle.RssSparkConfig.RSS_PARTITION_REASSIGN_BLOCK_RETRY_MAX_TIMES;
 
 public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
@@ -122,6 +123,7 @@ public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
   private SparkConf sparkConf;
   private boolean blockFailSentRetryEnabled;
   private int blockFailSentRetryMaxTimes = 1;
+  private int blockAccessDeniedRetryMaxTimes = 1;
 
   /** used by columnar rss shuffle writer implementation */
   protected final long taskAttemptId;
@@ -212,6 +214,9 @@ private RssShuffleWriter(
             RssClientConf.RSS_CLIENT_BLOCK_SEND_FAILURE_RETRY_ENABLED.defaultValue());
     this.blockFailSentRetryMaxTimes =
         RssSparkConfig.toRssConf(sparkConf).get(RSS_PARTITION_REASSIGN_BLOCK_RETRY_MAX_TIMES);
+    this.blockAccessDeniedRetryMaxTimes =
+        RssSparkConfig.toRssConf(sparkConf)
+            .get(RSS_PARTITION_REASSIGN_BLOCK_ACCESS_DENIED_RETRY_MAX_TIMES);
   }
 
   public RssShuffleWriter(
@@ -502,6 +507,11 @@ private void collectFailedBlocksToResend() {
               .map(x -> x.getShuffleBlockInfo().getRetryCnt())
               .max(Comparator.comparing(Integer::valueOf))
               .get();
+      int accessDeniedRetryCnt =
+          failedBlockStatus.stream()
+              .map(x -> x.getShuffleBlockInfo().getAccessDeniedRetryCnt())
+              .max(Comparator.comparing(Integer::valueOf))
+              .get();
       if (retryIndex >= blockFailSentRetryMaxTimes) {
         LOG.error(
             "Partial blocks for taskId: [{}] retry exceeding the max retry times: [{}]. Fast fail! faulty server list: {}",
@@ -513,6 +523,17 @@ private void collectFailedBlocksToResend() {
         isFastFail = true;
         break;
       }
+      if (accessDeniedRetryCnt >= blockAccessDeniedRetryMaxTimes) {
+        LOG.error(
+            "Partial blocks for taskId: [{}] access denied retry exceeding the max retry times: [{}]. Fast fail! access denied server list: {}",
+            taskId,
+            blockAccessDeniedRetryMaxTimes,
+            failedBlockStatus.stream()
+                .map(x -> x.getShuffleServerInfo())
+                .collect(Collectors.toSet()));
+        isFastFail = true;
+        break;
+      }
 
       for (TrackingBlockStatus status : failedBlockStatus) {
         StatusCode code = status.getStatusCode();
@@ -626,7 +647,11 @@ private void reassignAndResendBlocks(Set<TrackingBlockStatus> blocks) {
       // clear the previous retry state of block
       clearFailedBlockState(block);
       final ShuffleBlockInfo newBlock = block;
-      newBlock.incrRetryCnt();
+      if (StatusCode.ACCESS_DENIED.equals(blockStatus.getStatusCode())) {
+        newBlock.incrAccessDeniedRetryCnt();
+      } else {
+        newBlock.incrRetryCnt();
+      }
       newBlock.reassignShuffleServers(Arrays.asList(replacement));
       resendCandidates.add(newBlock);
     }

diff --git a/common/src/main/java/org/apache/uniffle/common/ShuffleBlockInfo.java b/common/src/main/java/org/apache/uniffle/common/ShuffleBlockInfo.java
@@ -37,6 +37,7 @@ public class ShuffleBlockInfo {
   private int uncompressLength;
   private long freeMemory;
   private int retryCnt = 0;
+  private int accessDeniedRetryCnt = 0;
 
   private transient BlockCompletionCallback completionCallback;
 
@@ -164,6 +165,14 @@ public int getRetryCnt() {
     return retryCnt;
   }
 
+  public void incrAccessDeniedRetryCnt() {
+    this.accessDeniedRetryCnt += 1;
+  }
+
+  public int getAccessDeniedRetryCnt() {
+    return accessDeniedRetryCnt;
+  }
+
   public void reassignShuffleServers(List<ShuffleServerInfo> replacements) {
     this.shuffleServerInfos = replacements;
   }