[apache#1796] fix(spark): Implicitly unregister map output on fetch f…

…ailure
zuston · Jun 17, 2024 · c9777e7 · c9777e7
1 parent 8eff38e
commit c9777e7
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/RssSparkShuffleUtils.java
@@ -33,6 +33,7 @@
 import org.apache.hadoop.conf.Configuration;
 import org.apache.spark.SparkConf;
 import org.apache.spark.SparkContext;
+import org.apache.spark.TaskContext;
 import org.apache.spark.broadcast.Broadcast;
 import org.apache.spark.deploy.SparkHadoopUtil;
 import org.apache.spark.shuffle.handle.SimpleShuffleHandleInfo;
@@ -44,7 +45,9 @@
 import org.apache.uniffle.client.api.ShuffleManagerClient;
 import org.apache.uniffle.client.factory.CoordinatorClientFactory;
 import org.apache.uniffle.client.factory.ShuffleManagerClientFactory;
+import org.apache.uniffle.client.request.RssReassignServersRequest;
 import org.apache.uniffle.client.request.RssReportShuffleFetchFailureRequest;
+import org.apache.uniffle.client.response.RssReassignServersResponse;
 import org.apache.uniffle.client.response.RssReportShuffleFetchFailureResponse;
 import org.apache.uniffle.client.util.ClientUtils;
 import org.apache.uniffle.client.util.RssClientConfig;
@@ -371,6 +374,19 @@ public static RssException reportRssFetchFailedException(
                   rssFetchFailedException.getMessage());
           RssReportShuffleFetchFailureResponse response = client.reportShuffleFetchFailure(req);
           if (response.getReSubmitWholeStage()) {
+            TaskContext taskContext = TaskContext.get();
+            RssReassignServersRequest rssReassignServersRequest =
+                new RssReassignServersRequest(
+                    taskContext.stageId(),
+                    taskContext.stageAttemptNumber(),
+                    shuffleId,
+                    taskContext.numPartitions());
+            RssReassignServersResponse reassignServersResponse =
+                client.reassignShuffleServers(rssReassignServersRequest);
+            LOG.info(
+                "Reassign servers for stage retry due to the fetch failure, result: {}",
+                reassignServersResponse.isNeedReassign());
+
             // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1
             // is provided.
             FetchFailedException ffe =

diff --git a/...nt-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java b/...nt-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java
@@ -24,14 +24,17 @@
 import scala.collection.AbstractIterator;
 import scala.collection.Iterator;
 
+import org.apache.spark.TaskContext;
 import org.apache.spark.shuffle.FetchFailedException;
 import org.apache.spark.shuffle.RssSparkShuffleUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.apache.uniffle.client.api.ShuffleManagerClient;
 import org.apache.uniffle.client.factory.ShuffleManagerClientFactory;
+import org.apache.uniffle.client.request.RssReassignServersRequest;
 import org.apache.uniffle.client.request.RssReportShuffleFetchFailureRequest;
+import org.apache.uniffle.client.response.RssReassignServersResponse;
 import org.apache.uniffle.client.response.RssReportShuffleFetchFailureResponse;
 import org.apache.uniffle.common.ClientType;
 import org.apache.uniffle.common.exception.RssException;
@@ -120,6 +123,19 @@ private RssException generateFetchFailedIfNecessary(RssFetchFailedException e) {
               e.getMessage());
       RssReportShuffleFetchFailureResponse response = client.reportShuffleFetchFailure(req);
       if (response.getReSubmitWholeStage()) {
+        TaskContext taskContext = TaskContext.get();
+        RssReassignServersRequest rssReassignServersRequest =
+            new RssReassignServersRequest(
+                taskContext.stageId(),
+                taskContext.stageAttemptNumber(),
+                builder.shuffleId,
+                taskContext.numPartitions());
+        RssReassignServersResponse reassignServersResponse =
+            client.reassignShuffleServers(rssReassignServersRequest);
+        LOG.info(
+            "Reassign servers for stage retry due to the fetch failure, result: {}",
+            reassignServersResponse.isNeedReassign());
+
         // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1 is
         // provided.
         FetchFailedException ffe =