[#2644] feat(spark): Involve shuffle failure into the event logs (#2645)

### What changes were proposed in this pull request?

Involve shuffle failure into the event logs

### Why are the changes needed?

for #2644

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Internal job tests

Author: zuston

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

@@ -29,6 +29,7 @@ public class TaskShuffleReadInfoEvent extends UniffleEvent {
   private boolean isShuffleReadFailed;
   private String failureReason;
   private ShuffleReadTimes shuffleReadTimes;
+  private long taskAttemptNumber;
 
   public TaskShuffleReadInfoEvent(
       int stageId,
@@ -37,14 +38,16 @@ public TaskShuffleReadInfoEvent(
       Map<String, ShuffleReadMetric> metrics,
       boolean isShuffleReadFailed,
       String failureReason,
-      ShuffleReadTimes shuffleReadTimes) {
+      ShuffleReadTimes shuffleReadTimes,
+      long taskAttemptNumber) {
     this.stageId = stageId;
     this.shuffleId = shuffleId;
     this.taskId = taskId;
     this.metrics = metrics;
     this.isShuffleReadFailed = isShuffleReadFailed;
     this.failureReason = failureReason;
     this.shuffleReadTimes = shuffleReadTimes;
+    this.taskAttemptNumber = taskAttemptNumber;
   }
 
   public int getStageId() {
@@ -74,4 +77,8 @@ public String getFailureReason() {
   public ShuffleReadTimes getShuffleReadTimes() {
     return shuffleReadTimes;
   }
+
+  public long getTaskAttemptNumber() {
+    return taskAttemptNumber;
+  }
 }

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

@@ -28,6 +28,7 @@ public class TaskShuffleWriteInfoEvent extends UniffleEvent {
   private boolean isShuffleWriteFailed;
   private String failureReason;
   private long uncompressedByteSize;
+  private long taskAttemptNumber;
 
   public TaskShuffleWriteInfoEvent(
       int stageId,
@@ -37,7 +38,8 @@ public TaskShuffleWriteInfoEvent(
       ShuffleWriteTimes writeTimes,
       boolean isShuffleWriteFailed,
       String failureReason,
-      long uncompressedByteSize) {
+      long uncompressedByteSize,
+      long taskAttemptNumber) {
     this.stageId = stageId;
     this.shuffleId = shuffleId;
     this.taskId = taskId;
@@ -46,6 +48,7 @@ public TaskShuffleWriteInfoEvent(
     this.isShuffleWriteFailed = isShuffleWriteFailed;
     this.failureReason = failureReason;
     this.uncompressedByteSize = uncompressedByteSize;
+    this.taskAttemptNumber = taskAttemptNumber;
   }
 
   public int getStageId() {
@@ -79,4 +82,8 @@ public String getFailureReason() {
   public long getUncompressedByteSize() {
     return uncompressedByteSize;
   }
+
+  public long getTaskAttemptNumber() {
+    return taskAttemptNumber;
+  }
 }

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

@@ -735,7 +735,8 @@ public void reportShuffleWriteMetric(
             ShuffleWriteTimes.fromProto(request.getShuffleWriteTimes()),
             request.getIsTaskWriteFailed(),
             request.getShuffleWriteFailureReason(),
-            request.getUncompressedByteSize());
+            request.getUncompressedByteSize(),
+            request.getTaskAttemptNumber());
     RssSparkShuffleUtils.getActiveSparkContext().listenerBus().post(event);
     RssProtos.ReportShuffleWriteMetricResponse reply =
         RssProtos.ReportShuffleWriteMetricResponse.newBuilder()
@@ -770,7 +771,8 @@ public void reportShuffleReadMetric(
                                 x.getValue().getHadoopByteSize()))),
             request.getIsTaskReadFailed(),
             request.getShuffleReadFailureReason(),
-            ShuffleReadTimes.fromProto(request.getShuffleReadTimes()));
+            ShuffleReadTimes.fromProto(request.getShuffleReadTimes()),
+            request.getTaskAttemptNumber());
     RssSparkShuffleUtils.getActiveSparkContext().listenerBus().post(event);
     RssProtos.ReportShuffleReadMetricResponse reply =
         RssProtos.ReportShuffleReadMetricResponse.newBuilder()

client-spark/extension/src/main/scala/org/apache/spark/UniffleListener.scala

@@ -31,6 +31,9 @@ import scala.collection.JavaConverters.mapAsScalaMapConverter
 class UniffleListener(conf: SparkConf, kvstore: ElementTrackingStore)
   extends SparkListener with Logging {
 
+  private val writeTaskInfo = ShuffleTaskSummary(shuffleType = ShuffleType.WRITE)
+  private val readTaskInfo = ShuffleTaskSummary(shuffleType = ShuffleType.READ)
+
   private val aggregatedShuffleReadTimes = new ShuffleReadTimes()
   private val aggregatedShuffleWriteTimes = new ShuffleWriteTimes()
   private val aggregatedShuffleWriteMetric = new ConcurrentHashMap[String, AggregatedShuffleWriteMetric]
@@ -70,6 +73,8 @@ class UniffleListener(conf: SparkConf, kvstore: ElementTrackingStore)
       kvstore.write(
         AggregatedShuffleReadTimesUIData(aggregatedShuffleReadTimes)
       )
+      kvstore.write(writeTaskInfo)
+      kvstore.write(readTaskInfo)
     }
   }
 
@@ -122,6 +127,14 @@ class UniffleListener(conf: SparkConf, kvstore: ElementTrackingStore)
     }
     aggregatedShuffleWriteTimes.inc(event.getWriteTimes)
     totalUncompressedShuffleBytes.addAndGet(event.getUncompressedByteSize)
+
+    if (event.isShuffleWriteFailed) {
+      writeTaskInfo.failedTaskNumber += 1
+      if (event.getTaskAttemptNumber > writeTaskInfo.failedTaskMaxAttemptNumber) {
+        writeTaskInfo.failedTaskMaxAttemptNumber = event.getTaskAttemptNumber
+      }
+      writeTaskInfo.failureReasons.add(event.getFailureReason)
+    }
   }
 
   private def onTaskShuffleReadInfo(event: TaskShuffleReadInfoEvent): Unit = {
@@ -143,6 +156,13 @@ class UniffleListener(conf: SparkConf, kvstore: ElementTrackingStore)
       agg_metric.hadoopDurationMillis += rmetric.getHadoopDurationMillis
     }
     aggregatedShuffleReadTimes.merge(event.getShuffleReadTimes)
+
+    if (event.isShuffleReadFailed) {
+      readTaskInfo.failedTaskNumber += 1
+      if (event.getTaskAttemptNumber > readTaskInfo.failedTaskMaxAttemptNumber) {}
+      readTaskInfo.failedTaskMaxAttemptNumber = event.getTaskAttemptNumber
+    }
+    readTaskInfo.failureReasons.add(event.getFailureReason)
   }
 
   override def onOtherEvent(event: SparkListenerEvent): Unit = event match {

client-spark/extension/src/main/scala/org/apache/spark/UniffleStatusStore.scala

@@ -26,12 +26,22 @@ import org.apache.uniffle.common.ShuffleReadTimes
 
 import java.util.concurrent.ConcurrentHashMap
 import scala.collection.JavaConverters.asScalaIteratorConverter
+import scala.collection.mutable
 
 class UniffleStatusStore(store: KVStore) {
   private def viewToSeq[T](view: KVStoreView[T]): Seq[T] = {
     Utils.tryWithResource(view.closeableIterator())(iter => iter.asScala.toList)
   }
 
+  def shuffleTaskSummary(shuffleType: ShuffleType): ShuffleTaskSummary = {
+    val kClass = classOf[ShuffleTaskSummary]
+    try {
+      store.read(kClass, s"${kClass.getName}_$shuffleType")
+    } catch {
+      case _: NoSuchElementException => new ShuffleTaskSummary(shuffleType = shuffleType)
+    }
+  }
+
   def uniffleProperties(): UniffleProperties = {
     val kClass = classOf[UniffleProperties]
     try {
@@ -184,4 +194,19 @@ case class ReassignInfoUIData(event: TaskReassignInfoEvent) {
   @JsonIgnore
   @KVIndex
   def id: String = classOf[ReassignInfoUIData].getName()
+}
+
+sealed abstract class ShuffleType private ()
+object ShuffleType {
+  val READ: ShuffleType = new ShuffleType {}
+  val WRITE: ShuffleType = new ShuffleType {}
+}
+
+case class ShuffleTaskSummary(shuffleType: ShuffleType,
+                              var failureReasons: mutable.HashSet[String] = new mutable.HashSet[String](),
+                              var failedTaskNumber: Long = -1,
+                              var failedTaskMaxAttemptNumber: Long = -1) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = s"${classOf[ShuffleTaskSummary].getName}_${shuffleType.toString}"
 }

client-spark/extension/src/main/scala/org/apache/spark/ui/ShufflePage.scala

@@ -20,7 +20,7 @@ package org.apache.spark.ui
 import org.apache.spark.internal.Logging
 import org.apache.spark.shuffle.events.ShuffleWriteTimes
 import org.apache.spark.util.Utils
-import org.apache.spark.{AggregatedShuffleMetric, AggregatedShuffleReadMetric, AggregatedShuffleWriteMetric, AggregatedTaskInfoUIData}
+import org.apache.spark.{AggregatedShuffleMetric, AggregatedShuffleReadMetric, AggregatedShuffleWriteMetric, AggregatedTaskInfoUIData, ShuffleType}
 
 import java.util.concurrent.ConcurrentHashMap
 import javax.servlet.http.HttpServletRequest
@@ -139,6 +139,10 @@ class ShufflePage(parent: ShuffleTab) extends WebUIPage("") with Logging {
     // reassign info
     val reassignInfo = runtimeStatusStore.reassignInfo().event
 
+    // task failure summary
+    val writeSummary = runtimeStatusStore.shuffleTaskSummary(shuffleType = ShuffleType.WRITE)
+    val readSummary = runtimeStatusStore.shuffleTaskSummary(shuffleType = ShuffleType.READ)
+
     // render build info
     val buildInfo = runtimeStatusStore.buildInfo()
     val buildInfoTableUI = UIUtils.listingTable(
@@ -344,6 +348,14 @@ class ShufflePage(parent: ShuffleTab) extends WebUIPage("") with Logging {
               </a>
               partitionSplit={reassignInfo.isReassignTriggeredOnPartitionSplit}, blockSentFailure={reassignInfo.isReassignTriggeredOnBlockSendFailure}, stageRetry={reassignInfo.isReassignTriggeredOnStageRetry}
             </li>
+
+            <li>
+              <a>
+                <strong>Shuffle Task Failure Summary (failure write/read):</strong>
+              </a>
+              {writeSummary.failedTaskNumber > 0} / {readSummary.failedTaskNumber > 0}
+            </li>
+
           </ul>
         </div>
 

client-spark/spark3/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java

@@ -414,7 +414,8 @@ private void postShuffleReadMetricsToDriver() {
                                           x.getValue().getHadoopReadLocalFileBytes()))),
                       isShuffleReadFailed,
                       shuffleReadReason,
-                      shuffleReadTimes));
+                      shuffleReadTimes,
+                      context.attemptNumber()));
           if (response != null && response.getStatusCode() != StatusCode.SUCCESS) {
             LOG.error("Errors on reporting shuffle read metrics to driver");
           }

client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java

@@ -980,7 +980,8 @@ public Option<MapStatus> stop(boolean success) {
                         writeTimes,
                         isShuffleWriteFailed,
                         shuffleWriteFailureReason,
-                        bufferManager.getUncompressedDataLen()));
+                        bufferManager.getUncompressedDataLen(),
+                        taskContext.attemptNumber()));
             if (response.getStatusCode() != StatusCode.SUCCESS) {
               LOG.error(
                   "Errors on reporting shuffle write metrics to driver. status_code: {}",

internal-client/src/main/java/org/apache/uniffle/client/request/RssReportShuffleReadMetricRequest.java

@@ -32,6 +32,7 @@ public class RssReportShuffleReadMetricRequest {
   private boolean isShuffleReadFailed;
   private Optional<String> shuffleReadReason;
   private ShuffleReadTimes shuffleReadTimes;
+  private long taskAttemptNumber;
 
   public RssReportShuffleReadMetricRequest(
       int stageId,
@@ -40,14 +41,16 @@ public RssReportShuffleReadMetricRequest(
       Map<String, TaskShuffleReadMetric> metrics,
       boolean isShuffleReadFailed,
       Optional<String> shuffleReadReason,
-      ShuffleReadTimes shuffleReadTimes) {
+      ShuffleReadTimes shuffleReadTimes,
+      long taskAttemptNumber) {
     this.stageId = stageId;
     this.shuffleId = shuffleId;
     this.taskId = taskId;
     this.metrics = metrics;
     this.isShuffleReadFailed = isShuffleReadFailed;
     this.shuffleReadReason = shuffleReadReason;
     this.shuffleReadTimes = shuffleReadTimes;
+    this.taskAttemptNumber = taskAttemptNumber;
   }
 
   public RssProtos.ReportShuffleReadMetricRequest toProto() {
@@ -61,6 +64,7 @@ public RssProtos.ReportShuffleReadMetricRequest toProto() {
         .setIsTaskReadFailed(request.isShuffleReadFailed)
         .setShuffleReadFailureReason(request.shuffleReadReason.orElse(""))
         .setShuffleReadTimes(shuffleReadTimes.toProto())
+        .setTaskAttemptNumber(taskAttemptNumber)
         .putAllMetrics(
             request.metrics.entrySet().stream()
                 .collect(

internal-client/src/main/java/org/apache/uniffle/client/request/RssReportShuffleWriteMetricRequest.java

@@ -27,6 +27,7 @@ public class RssReportShuffleWriteMetricRequest {
   private int stageId;
   private int shuffleId;
   private long taskId;
+  private long taskAttemptNumber;
   private Map<String, TaskShuffleWriteMetric> metrics;
   private TaskShuffleWriteTimes writeTimes;
 
@@ -43,7 +44,8 @@ public RssReportShuffleWriteMetricRequest(
       TaskShuffleWriteTimes writeTimes,
       boolean isShuffleWriteFailed,
       Optional<String> shuffleWriteFailureReason,
-      long uncompressedByteSize) {
+      long uncompressedByteSize,
+      long taskAttemptNumber) {
     this.stageId = stageId;
     this.shuffleId = shuffleId;
     this.taskId = taskId;
@@ -52,6 +54,7 @@ public RssReportShuffleWriteMetricRequest(
     this.isShuffleWriteFailed = isShuffleWriteFailed;
     this.shuffleWriteFailureReason = shuffleWriteFailureReason;
     this.uncompressedByteSize = uncompressedByteSize;
+    this.taskAttemptNumber = taskAttemptNumber;
   }
 
   public RssProtos.ReportShuffleWriteMetricRequest toProto() {
@@ -66,6 +69,7 @@ public RssProtos.ReportShuffleWriteMetricRequest toProto() {
         .setIsTaskWriteFailed(isShuffleWriteFailed)
         .setShuffleWriteFailureReason(shuffleWriteFailureReason.orElse(""))
         .setUncompressedByteSize(request.uncompressedByteSize)
+        .setTaskAttemptNumber(taskAttemptNumber)
         .putAllMetrics(
             request.metrics.entrySet().stream()
                 .collect(