[CELEBORN-2477] Enhance Celeborn Spark UI extension with detailed shuffle metrics

client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala

@@ -18,11 +18,16 @@
 package org.apache.spark.shuffle.celeborn
 
 import java.util
+import java.util.{Collections, LinkedHashMap => JLinkedHashMap}
+
+import scala.collection.JavaConverters._
 
 import org.apache.spark.SparkContext
 import org.apache.spark.api.plugin.{DriverPlugin, PluginContext, SparkPlugin}
 import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.celeborn.events.CelebornBuildInfoEvent
 import org.apache.spark.shuffle.celeborn.ui.CelebornUITab
+import org.apache.spark.shuffle.celeborn.ui.CelebornUIUtils
 
 /**
  * SparkPlugin entry point for the Celeborn UI extension.
@@ -31,6 +36,12 @@ import org.apache.spark.shuffle.celeborn.ui.CelebornUITab
  * {{{
  *   spark.plugins=org.apache.spark.shuffle.celeborn.CelebornPlugin
  * }}}
+ *
+ * [[init]] registers the Celeborn listener to the status queue and [[registerMetrics]] posts a
+ * build-info event and attaches the UI tab. These split so the tab's page servlets are added
+ * once the Jetty server is fully started: `registerMetrics` runs AFTER
+ * SparkContext.attachAllHandlers, while `init` runs before (attaching during init would race
+ * attachAllHandlers and throw IllegalStateException: STARTED on some Spark builds).
  */
 class CelebornPlugin extends SparkPlugin {
 
@@ -48,18 +59,53 @@ private class CelebornDriverPlugin extends DriverPlugin with Logging {
       ctx: PluginContext): util.Map[String, String] = {
     logInfo("Initializing CelebornDriverPlugin...")
     this.sc = sc
-    val kvStore = sc.statusStore.store
-    new CelebornListener(kvStore, sc.conf).register(sc)
-    java.util.Collections.emptyMap[String, String]()
+    if (CelebornUIUtils.isUIEnabled(sc.conf)) {
+      CelebornListener.register(sc)
+    } else {
+      logInfo("Celeborn Spark UI extension is disabled, skipping.")
+    }
+    Collections.emptyMap[String, String]()
   }
 
   override def registerMetrics(
       appId: String,
       ctx: PluginContext): Unit = {
-    sc.ui.foreach { ui =>
-      new CelebornUITab(new CelebornStatusStore(ui.store.store), ui)
+    if (sc == null || !CelebornUIUtils.isUIEnabled(sc.conf)) {
+      return
     }
+    // registerMetrics runs after SparkContext is fully initialized, so applicationId is
+    // available here (it's null during init). Post the build-info event now and attach UI.
+    sc.listenerBus.post(CelebornBuildInfoEvent(buildInfoMap(sc, appId)))
+    CelebornUIUtils.attachUI(sc)
   }
 
   override def shutdown(): Unit = {}
+
+  private def buildInfoMap(sc: SparkContext, appId: String): Map[String, String] = {
+    val info = new JLinkedHashMap[String, String]()
+    info.put("Spark Version", sc.version)
+    // Prefer the appId passed to registerMetrics (reliable); fall back to sc.applicationId.
+    info.put("Application Id", Option(appId).filter(_.nonEmpty).getOrElse(sc.applicationId))
+    info.put("Celeborn UI Enabled", "true")
+    // Derive key Celeborn runtime config from CelebornConf (rebuild from sc.conf via
+    // CelebornUIUtils.celebornConf, which does not touch SparkUtils: SparkUtils' class init
+    // drags in spark-sql's UnsafeRowSerializer and would fail in RDD-only applications).
+    // Celeborn has no ProjectConstants version class, so runtime config stands in for version —
+    // it is more useful for diagnosing shuffle behavior anyway.
+    try {
+      val celebornConf = CelebornUIUtils.celebornConf(sc.conf)
+      info.put("Celeborn Compression Codec", celebornConf.shuffleCompressionCodec.toString)
+      info.put("Celeborn Shuffle Writer Mode", celebornConf.shuffleWriterMode.toString)
+      info.put("Celeborn Push Replicate Enabled", celebornConf.clientPushReplicateEnabled.toString)
+      info.put("Celeborn Partition Split Mode", celebornConf.shufflePartitionSplitMode.toString)
+      info.put(
+        "Celeborn Partition Split Threshold",
+        org.apache.spark.util.Utils.bytesToString(celebornConf.shufflePartitionSplitThreshold))
+      info.put("Celeborn Fallback Policy", celebornConf.sparkShuffleFallbackPolicy.toString)
+    } catch {
+      case e: Throwable =>
+        logWarning("Failed to derive Celeborn conf for build info", e)
+    }
+    info.asScala.toMap
+  }
 }

client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.shuffle.celeborn
 
 import com.fasterxml.jackson.annotation.JsonIgnore
-import org.apache.spark.util.kvstore.{KVIndex, KVStore}
+import org.apache.spark.util.kvstore.{KVIndex, KVStore, KVStoreView}
 
 private[celeborn] case class AggregatedTaskInfoUIData(
     shuffleWriteBytes: Long,
@@ -83,4 +83,104 @@ private[celeborn] class CelebornStatusStore(store: KVStore) {
       case _: NoSuchElementException => false
     }
   }
+
+  /** Build info summary, or an empty entity if not yet written. */
+  def buildInfo(): CelebornBuildInfoUIData = {
+    val kClass = classOf[CelebornBuildInfoUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException => new CelebornBuildInfoUIData(Seq.empty)
+    }
+  }
+
+  /** All recorded shuffle assignments (shuffle -> worker topology), newest last. */
+  def assignmentInfos(): Seq[CelebornShuffleAssignmentUIData] = {
+    viewToSeq(store.view(classOf[CelebornShuffleAssignmentUIData]))
+  }
+
+  /** Per-policy fallback counts snapshot, or empty if no fallback recorded. */
+  def fallbackStats(): CelebornFallbackStatsUIData = {
+    val kClass = classOf[CelebornFallbackStatsUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException =>
+        new CelebornFallbackStatsUIData(new java.util.HashMap[String, java.lang.Long]())
+    }
+  }
+
+  /** Reassign status snapshot, or all-false if no reassign recorded. */
+  def reassignStats(): CelebornReassignStatsUIData = {
+    val kClass = classOf[CelebornReassignStatsUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException =>
+        new CelebornReassignStatsUIData(false, false, false, 0L)
+    }
+  }
+
+  /** Aggregated write-path timing breakdown, or all-zero if none recorded. */
+  def writeTimes(): CelebornWriteTimesUIData = {
+    val kClass = classOf[CelebornWriteTimesUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException => new CelebornWriteTimesUIData(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+    }
+  }
+
+  /** Per-worker push stats, newest last. */
+  def perWorkerWriteStats(): Seq[CelebornPerWorkerWriteStatsUIData] = {
+    viewToSeq(store.view(classOf[CelebornPerWorkerWriteStatsUIData]))
+  }
+
+  /** Aggregated read-path timing breakdown, or all-zero if none recorded. */
+  def readTimes(): CelebornReadTimesUIData = {
+    val kClass = classOf[CelebornReadTimesUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException =>
+        new CelebornReadTimesUIData(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+    }
+  }
+
+  /** Per-worker read stats, newest last. */
+  def perWorkerReadStats(): Seq[CelebornPerWorkerReadStatsUIData] = {
+    viewToSeq(store.view(classOf[CelebornPerWorkerReadStatsUIData]))
+  }
+
+  /** Per-shuffle write metrics snapshot, or empty if none recorded. */
+  def aggregatedWriteMetrics(): CelebornAggregatedWriteMetricsUIData = {
+    val kClass = classOf[CelebornAggregatedWriteMetricsUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException =>
+        new CelebornAggregatedWriteMetricsUIData(new java.util.HashMap[
+          Int,
+          AggregatedShuffleWriteMetric]())
+    }
+  }
+
+  /** Per-shuffle read metrics snapshot, or empty if none recorded. */
+  def aggregatedReadMetrics(): CelebornAggregatedReadMetricsUIData = {
+    val kClass = classOf[CelebornAggregatedReadMetricsUIData]
+    try {
+      store.read(kClass, kClass.getName)
+    } catch {
+      case _: NoSuchElementException =>
+        new CelebornAggregatedReadMetricsUIData(new java.util.HashMap[
+          Int,
+          AggregatedShuffleReadMetric]())
+    }
+  }
+
+  private def viewToSeq[T](view: KVStoreView[T]): Seq[T] = {
+    import scala.collection.JavaConverters._
+    org.apache.spark.util.Utils.tryWithResource(view.closeableIterator())(iter =>
+      iter.asScala.toList)
+  }
 }

client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornUIData.scala

@@ -0,0 +1,167 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.celeborn
+
+import com.fasterxml.jackson.annotation.JsonIgnore
+import org.apache.spark.status.KVUtils.KVIndexParam
+import org.apache.spark.util.kvstore.KVIndex
+
+/**
+ * KVStore entity holding Celeborn build/version info for the UI summary.
+ *  Singleton: keyed by the class name.
+ */
+private[celeborn] class CelebornBuildInfoUIData(val info: Seq[(String, String)]) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornBuildInfoUIData].getName
+}
+
+/**
+ * Per-shuffle assignment record. `appShuffleId` is the natural key (multiple rows).
+ *  @KVIndexParam makes appShuffleId the primary index; no separate @KVIndex id needed
+ *  (having both causes "Duplicate index __main__").
+ */
+private[celeborn] class CelebornShuffleAssignmentUIData(
+    @KVIndexParam val appShuffleId: Int,
+    val celebornShuffleId: Int,
+    val workers: java.util.List[String],
+    val numPartitions: Int,
+    val timestamp: Long)
+
+/** Snapshot of per-policy fallback counts. Singleton, overwritten on each fallback. */
+private[celeborn] class CelebornFallbackStatsUIData(
+    val counts: java.util.Map[String, java.lang.Long]) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornFallbackStatsUIData].getName
+}
+
+/**
+ * Reassign status snapshot: whether partition split / block send failure / stage retry
+ *  have been triggered. Singleton, overwritten on each reassign event. Mirrors Uniffle's
+ *  ReassignInfoUIData.
+ */
+private[celeborn] class CelebornReassignStatsUIData(
+    val partitionSplit: Boolean,
+    val reviveTriggered: Boolean,
+    val stageRetry: Boolean,
+    val timestamp: Long) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornReassignStatsUIData].getName
+}
+
+/** Per-shuffle write metrics aggregated from onTaskEnd TaskMetrics. Keyed by shuffleDepId. */
+private[celeborn] class CelebornAggregatedWriteMetricsUIData(
+    val metrics: java.util.Map[Int, AggregatedShuffleWriteMetric]) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornAggregatedWriteMetricsUIData].getName
+}
+
+/**
+ * Per-shuffle write metrics (sourced from Spark native ShuffleWriteMetrics).
+ *  Mutable fields updated by the single status-queue thread; serialized as a snapshot.
+ */
+private[celeborn] class AggregatedShuffleWriteMetric(
+    var bytesWritten: Long,
+    var recordsWritten: Long,
+    var writeTimeMs: Long) {
+  def this() = this(0L, 0L, 0L)
+}
+
+/** Per-shuffle read metrics aggregated from CelebornReadMetricsEvent. Keyed by shuffleId. */
+private[celeborn] class CelebornAggregatedReadMetricsUIData(
+    val metrics: java.util.Map[Int, AggregatedShuffleReadMetric]) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornAggregatedReadMetricsUIData].getName
+}
+
+/**
+ * Per-shuffle read metrics: bytes/chunks come from the event's WorkerReadCost entries,
+ *  fetchTimeMs sums the read-path timing breakdown (chunk wait + decompress +
+ *  deserialize + copy). Mutable fields updated by the single status-queue thread;
+ *  serialized as a snapshot. Mirrors [[AggregatedShuffleWriteMetric]].
+ */
+private[celeborn] class AggregatedShuffleReadMetric(
+    var bytesRead: Long,
+    var chunkCount: Long,
+    var fetchTimeMs: Long) {
+  def this() = this(0L, 0L, 0L)
+}
+
+/**
+ * Aggregated write-path timing breakdown (ms) across all tasks, mirroring Uniffle's
+ *  ShuffleWriteTimes. Singleton, overwritten on each CelebornWriteMetricsEvent.
+ */
+private[celeborn] class CelebornWriteTimesUIData(
+    val copyTimeMs: Long,
+    val serializeTimeMs: Long,
+    val compressTimeMs: Long,
+    val queueWaitTimeMs: Long,
+    val queueStallTimeMs: Long,
+    val inflightWaitTimeMs: Long,
+    val drainWaitTimeMs: Long,
+    val slowPushCount: Long,
+    val maxPushRttMs: Long,
+    val uncompressedBytes: Long) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornWriteTimesUIData].getName
+}
+
+/**
+ * Per-worker push stats aggregated from CelebornWriteMetricsEvent's PushWorkerStats.
+ *  `workerId` is the natural key (multiple rows).
+ */
+private[celeborn] class CelebornPerWorkerWriteStatsUIData(
+    @KVIndexParam val workerId: String,
+    val pushCount: Long,
+    val pushBytes: Long,
+    val totalPushRttNanos: Long,
+    val softSplitCount: Long,
+    val hardSplitCount: Long,
+    val primaryCongestedCount: Long,
+    val replicaCongestedCount: Long,
+    val lastPushFailureReason: String)
+
+/** Aggregated read-path timing breakdown (ms) from CelebornReadMetricsEvent. Singleton. */
+private[celeborn] class CelebornReadTimesUIData(
+    val decompressTimeMs: Long,
+    val chunkWaitTimeMs: Long,
+    val deserializeTimeMs: Long,
+    val copyTimeMs: Long,
+    val retryCount: Long,
+    val retryWaitTimeMs: Long,
+    val peerSwitchCount: Long,
+    val excludeCount: Long,
+    val slowChunkCount: Long,
+    val maxChunkRttMs: Long) {
+  @JsonIgnore
+  @KVIndex
+  def id: String = classOf[CelebornReadTimesUIData].getName
+}
+
+/** Per-worker read cost aggregated from CelebornReadMetricsEvent. `workerId` is the key. */
+private[celeborn] class CelebornPerWorkerReadStatsUIData(
+    @KVIndexParam val workerId: String,
+    val chunkCount: Long,
+    val bytes: Long,
+    val totalRttNanos: Long,
+    val maxRttNanos: Long)