[Bug] Long tail tasks in the Write Stage retry phase results in data …

…loss.
apache · Dec 19, 2024 · 0c3ccff · 0c3ccff
1 parent 0481f21
commit 0c3ccff
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/...-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java b/...-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java
@@ -45,6 +45,7 @@
 import org.apache.spark.MapOutputTracker;
 import org.apache.spark.MapOutputTrackerMaster;
 import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
 import org.apache.spark.SparkEnv;
 import org.apache.spark.SparkException;
 import org.apache.spark.shuffle.RssShuffleHandle;
@@ -1131,4 +1132,8 @@ public Map<String, String> sparkConfToMap(SparkConf sparkConf) {
   public ShuffleWriteClient getShuffleWriteClient() {
     return shuffleWriteClient;
   }
+
+  public void killAllTaskByStageId(int stageAttemptId, String cancelReason) {
+    SparkContext.getOrCreate().cancelStage(stageAttemptId, cancelReason);
+  }
 }
diff --git a/...k/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerInterface.java b/...k/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerInterface.java
@@ -96,4 +96,12 @@ MutableShuffleHandleInfo reassignOnBlockSendFailure(
    * @return
    */
   ShuffleWriteClient getShuffleWriteClient();
+
+  /**
+   * To cancel all tasks under current stage by stageId.
+   *
+   * @param stageAttemptId
+   * @param cancelReason
+   */
+  void killAllTaskByStageId(int stageAttemptId, String cancelReason);
 }
diff --git a/...rk/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerGrpcService.java b/...rk/common/src/main/java/org/apache/uniffle/shuffle/manager/ShuffleManagerGrpcService.java
@@ -115,9 +115,15 @@ public void reportShuffleWriteFailure(
                     shuffleManager.getMaxFetchFailures());
             if (!shuffleServerWriterFailureRecord.isClearedMapTrackerBlock()) {
               try {
+                // Cancel all tasks under the current StageId.
+                shuffleManager.killAllTaskByStageId(
+                    stageAttemptId, "Write failure triggers server reallocation.");
                 // Clear the metadata of the completed task, otherwise some of the stage's data will
                 // be lost.
                 shuffleManager.unregisterAllMapOutput(shuffleId);
+                // Need to clear the mapStatus twice to prevent partition data loss due to the
+                // long-tail task performed before the stage retry.
+                shuffleManager.unregisterAllMapOutput(shuffleId);
                 // Deregister the shuffleId corresponding to the Shuffle Server.
                 shuffleManager.getShuffleWriteClient().unregisterShuffle(appId, shuffleId);
                 shuffleServerWriterFailureRecord.setClearedMapTrackerBlock(true);