diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
index 1348ee2f8f1..504e2a8ef50 100644
--- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
+++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala
@@ -22,6 +22,9 @@ import java.util.concurrent.atomic.AtomicLong
import org.apache.spark.SparkConf
import org.apache.spark.internal.Logging
import org.apache.spark.scheduler._
+import org.apache.spark.shuffle.celeborn.events._
+import org.apache.spark.shuffle.celeborn.ui.CelebornUIUtils
+import org.apache.spark.status.{ElementTrackingStore, KVUtils}
import org.apache.spark.util.Utils
import org.apache.spark.util.kvstore.KVStore
@@ -33,6 +36,16 @@ import org.apache.spark.util.kvstore.KVStore
* disabled until the application's recorded `spark.plugins` or
* `spark.plugins.defaultList` contains [[CelebornPlugin]], and an enable marker
* is persisted so `setupUI` can decide whether to attach the Celeborn tab.
+ *
+ * Data sources:
+ * - onTaskEnd: Spark native TaskMetrics (already written by Celeborn writer/reader),
+ * aggregated globally and per shuffle dependency.
+ * - onOtherEvent: CelebornBuildInfoEvent / CelebornShuffleAssignmentEvent /
+ * CelebornFallbackEvent / CelebornReassignEvent / CelebornWriteMetricsEvent /
+ * CelebornReadMetricsEvent posted from the driver-side plugin / SparkShuffleManager /
+ * FallbackPolicyRunner and forwarded through the MapperEnd / read-metrics RPCs.
+ * These custom events enter the event log via the status queue, so HistoryServer
+ * replay rebuilds them here as well.
*/
private[celeborn] class CelebornListener(
val kvstore: KVStore,
@@ -48,12 +61,95 @@ private[celeborn] class CelebornListener(
private val totalFetchWaitTimeMs = new AtomicLong(0L)
private val totalTaskDurationMs = new AtomicLong(0L)
+ // Per-shuffle write aggregation (key = shuffleDepId, only ShuffleMapStage has one).
+ private val perShuffleWrite =
+ new java.util.concurrent.ConcurrentHashMap[Int, AggregatedShuffleWriteMetric]
+ // Per-shuffle read aggregation (key = shuffleId from CelebornReadMetricsEvent), so the
+ // read side can be joined with perShuffleWrite on the per-shuffle UI table.
+ private val perShuffleRead =
+ new java.util.concurrent.ConcurrentHashMap[Int, AggregatedShuffleReadMetric]
+ // Aggregated write-path timing breakdown (ms) from CelebornWriteMetricsEvent. Uses AtomicLong
+ // so concurrent events on the status queue accumulate safely.
+ private val aggCopyTimeMs = new AtomicLong(0L)
+ private val aggSerializeTimeMs = new AtomicLong(0L)
+ private val aggCompressTimeMs = new AtomicLong(0L)
+ private val aggQueueWaitTimeMs = new AtomicLong(0L)
+ private val aggQueueStallTimeMs = new AtomicLong(0L)
+ private val aggInflightWaitTimeMs = new AtomicLong(0L)
+ private val aggDrainWaitTimeMs = new AtomicLong(0L)
+ private val aggSlowPushCount = new AtomicLong(0L)
+ private val aggMaxPushRttMs = new AtomicLong(0L)
+ private val aggUncompressedBytes = new AtomicLong(0L)
+ // Per-worker push stats (workerId -> mutable accumulator), merged from each event.
+ private val perWorkerWriteStats =
+ new java.util.concurrent.ConcurrentHashMap[String, PerWorkerWriteAccumulator]
+ // Aggregated read-path timing breakdown (ms) from CelebornReadMetricsEvent.
+ private val aggDecompressTimeMs = new AtomicLong(0L)
+ private val aggChunkWaitTimeMs = new AtomicLong(0L)
+ private val aggDeserializeTimeMs = new AtomicLong(0L)
+ private val aggCopyTimeMsRead = new AtomicLong(0L)
+ private val aggRetryCount = new AtomicLong(0L)
+ private val aggRetryWaitTimeMs = new AtomicLong(0L)
+ private val aggPeerSwitchCount = new AtomicLong(0L)
+ private val aggExcludeCount = new AtomicLong(0L)
+ private val aggSlowChunkCount = new AtomicLong(0L)
+ private val aggMaxChunkRttMs = new AtomicLong(0L)
+ private val perWorkerReadStats =
+ new java.util.concurrent.ConcurrentHashMap[String, PerWorkerReadAccumulator]
+ // stageId -> shuffleDepId (only ShuffleMapStage has a shuffleDepId). Instance-level (not
+ // companion-object static) so that different applications replayed on a shared HistoryServer
+ // do not cross-contaminate each other's stageId namespace.
+ private val stageToShuffleMappings =
+ new java.util.concurrent.ConcurrentHashMap[Int, Int]()
+
private val lastUpdateTimestamp = new AtomicLong(-1L)
private val updateIntervalMillis = 5000L
- def register(sc: org.apache.spark.SparkContext): Unit = {
- sc.addSparkListener(this)
- logInfo("CelebornListener registered successfully")
+ // Cap the per-shuffle assignment rows to bound KVStore / event-log growth on long-running
+ // jobs with many shuffles (mirrors Gluten's UI_RETAINED_EXECUTIONS trigger and Spark's
+ // retainedStages). When the count exceeds the threshold, evict the oldest rows. Read via
+ // the CelebornConf entry so the documented `spark.celeborn.client.spark.ui.retainedShuffles`
+ // key (and the entry's default and `retainedShuffle` alias handling) applies; a bare-key
+ // SparkConf lookup never matches the user-set key because SparkConf does exact matching.
+ private val retainedShuffles: Int =
+ CelebornUIUtils.celebornConf(conf).clientSparkUIRetainedShuffles
+
+ kvstore match {
+ case tracking: ElementTrackingStore =>
+ // Register a final flush so the last <=5s of aggregations (within the throttle window)
+ // are persisted on store close / replay end, mirroring AppStatusListener's onFlush hook.
+ // Without this, the trailing metrics would be lost from both the live store rebuild and
+ // the HistoryServer replay.
+ tracking.onFlush {
+ mayUpdate(force = true)
+ }
+ tracking.addTrigger(classOf[CelebornShuffleAssignmentUIData], retainedShuffles.toLong) {
+ count => cleanupAssignments(tracking, count)
+ }
+ case _ => // InMemoryStore (tests) or other stores: no lifecycle hooks needed.
+ }
+
+ private def cleanupAssignments(tracking: ElementTrackingStore, count: Long): Unit = {
+ val toDelete = count - retainedShuffles
+ if (toDelete <= 0) {
+ return
+ }
+ val view = tracking.view(classOf[CelebornShuffleAssignmentUIData])
+ KVUtils.viewToSeq(view, toDelete.toInt)(_ => true).foreach { e =>
+ tracking.delete(classOf[CelebornShuffleAssignmentUIData], e.appShuffleId)
+ }
+ }
+
+ /**
+ * Writes an entity through the trigger-checking path when the store supports it: the plain
+ * KVStore.write bypasses ElementTrackingStore's eviction triggers, so the
+ * retainedShuffles cap only takes effect via write(value, checkTriggers = true).
+ */
+ private def writeEntity(value: Any): Unit = kvstore match {
+ case tracking: ElementTrackingStore =>
+ tracking.write(value, checkTriggers = true)
+ case _ =>
+ kvstore.write(value)
}
override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
@@ -66,11 +162,33 @@ private[celeborn] class CelebornListener(
totalWriteTimeMs.addAndGet(metrics.shuffleWriteMetrics.writeTime / 1000000L)
totalReadBytes.addAndGet(metrics.shuffleReadMetrics.totalBytesRead)
totalFetchWaitTimeMs.addAndGet(metrics.shuffleReadMetrics.fetchWaitTime)
- totalTaskDurationMs.addAndGet(taskEnd.taskInfo.duration)
+ if (taskEnd.taskInfo != null) {
+ totalTaskDurationMs.addAndGet(taskEnd.taskInfo.duration)
+ }
+ // Per-shuffle write attribution: ShuffleMapTask carries shuffle write and its stage
+ // has a shuffleDepId mapped in onStageSubmitted.
+ if (stageToShuffleMappings.containsKey(taskEnd.stageId)
+ && taskEnd.taskType == "ShuffleMapTask") {
+ val shuffleId = stageToShuffleMappings.get(taskEnd.stageId)
+ val metric = perShuffleWrite.computeIfAbsent(
+ shuffleId,
+ _ => new AggregatedShuffleWriteMetric(0L, 0L, 0L))
+ metric.bytesWritten += metrics.shuffleWriteMetrics.bytesWritten
+ metric.recordsWritten += metrics.shuffleWriteMetrics.recordsWritten
+ metric.writeTimeMs += metrics.shuffleWriteMetrics.writeTime / 1000000L
+ }
}
mayUpdate()
}
+ override def onStageSubmitted(stageSubmitted: SparkListenerStageSubmitted): Unit = {
+ // ShuffleMapStage carries shuffleDepId = the Spark shuffle dependency id; result stages
+ // have None. Used in onTaskEnd to attribute write metrics to a shuffle.
+ stageSubmitted.stageInfo.shuffleDepId.foreach { sid =>
+ stageToShuffleMappings.put(stageSubmitted.stageInfo.stageId, sid.asInstanceOf[Int])
+ }
+ }
+
override def onJobEnd(jobEnd: SparkListenerJobEnd): Unit = {
// Flush per-job so that tasks finishing within the throttle interval of the
// previous flush are not lost when the application stays alive but idle.
@@ -104,6 +222,81 @@ private[celeborn] class CelebornListener(
}
}
+ override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
+ case _: CelebornEvent if !pluginEnabled => // gated, same semantics as onTaskEnd
+ case e: CelebornBuildInfoEvent =>
+ kvstore.write(new CelebornBuildInfoUIData(e.info.toSeq.sortBy(_._1)))
+ mayUpdate()
+ case e: CelebornShuffleAssignmentEvent =>
+ writeEntity(new CelebornShuffleAssignmentUIData(
+ e.appShuffleId,
+ e.celebornShuffleId,
+ e.workers,
+ e.numPartitions,
+ e.timestamp))
+ mayUpdate()
+ case e: CelebornFallbackEvent =>
+ kvstore.write(new CelebornFallbackStatsUIData(e.fallbackCounts))
+ mayUpdate()
+ case e: CelebornReassignEvent =>
+ kvstore.write(new CelebornReassignStatsUIData(
+ e.partitionSplit,
+ e.reviveTriggered,
+ e.stageRetry,
+ e.timestamp))
+ mayUpdate()
+ case e: CelebornWriteMetricsEvent =>
+ import scala.collection.JavaConverters._
+ val w = e.writeMetrics
+ aggCopyTimeMs.addAndGet(w.copyTimeMs)
+ aggSerializeTimeMs.addAndGet(w.serializeTimeMs)
+ aggCompressTimeMs.addAndGet(w.compressTimeMs)
+ aggQueueWaitTimeMs.addAndGet(w.queueWaitTimeMs)
+ aggQueueStallTimeMs.addAndGet(w.queueStallTimeMs)
+ aggInflightWaitTimeMs.addAndGet(w.inflightWaitTimeMs)
+ aggDrainWaitTimeMs.addAndGet(w.drainWaitTimeMs)
+ aggSlowPushCount.addAndGet(w.slowPushCount)
+ aggMaxPushRttMs.accumulateAndGet(w.maxPushRttMs, math.max)
+ aggUncompressedBytes.addAndGet(w.uncompressedBytes)
+ // Merge per-worker push stats.
+ e.pushWorkerStats.asScala.foreach { s =>
+ val acc =
+ perWorkerWriteStats.computeIfAbsent(s.workerId, _ => new PerWorkerWriteAccumulator)
+ acc.merge(s)
+ }
+ mayUpdate()
+ case e: CelebornReadMetricsEvent =>
+ import scala.collection.JavaConverters._
+ val rm = e.readMetrics
+ aggDecompressTimeMs.addAndGet(rm.decompressTimeMs)
+ aggChunkWaitTimeMs.addAndGet(rm.chunkWaitTimeMs)
+ aggDeserializeTimeMs.addAndGet(rm.deserializeTimeMs)
+ aggCopyTimeMsRead.addAndGet(rm.copyTimeMs)
+ aggRetryCount.addAndGet(rm.retryCount)
+ aggRetryWaitTimeMs.addAndGet(rm.retryWaitTimeMs)
+ aggPeerSwitchCount.addAndGet(rm.peerSwitchCount)
+ aggExcludeCount.addAndGet(rm.excludeCount)
+ aggSlowChunkCount.addAndGet(rm.slowChunkCount)
+ aggMaxChunkRttMs.accumulateAndGet(rm.maxChunkRttMs, math.max)
+ e.workerReadCosts.asScala.foreach { w =>
+ val acc = perWorkerReadStats.computeIfAbsent(w.workerId, _ => new PerWorkerReadAccumulator)
+ acc.merge(w)
+ }
+ // Per-shuffle read aggregation: bytes/chunks from WorkerReadCost, fetch time from
+ // the read-path timing breakdown. Same accumulation pattern as perShuffleWrite.
+ val shuffleRead = perShuffleRead.computeIfAbsent(
+ e.shuffleId,
+ _ => new AggregatedShuffleReadMetric(0L, 0L, 0L))
+ shuffleRead.fetchTimeMs +=
+ rm.chunkWaitTimeMs + rm.decompressTimeMs + rm.deserializeTimeMs + rm.copyTimeMs
+ e.workerReadCosts.asScala.foreach { w =>
+ shuffleRead.bytesRead += w.bytes
+ shuffleRead.chunkCount += w.chunkCount
+ }
+ mayUpdate()
+ case _ => // ignore unknown events
+ }
+
override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
mayUpdate(force = true)
logInfo("CelebornListener: application ended, final flush completed")
@@ -130,16 +323,124 @@ private[celeborn] class CelebornListener(
if (!pluginEnabled) {
return
}
+ import scala.collection.JavaConverters._
try {
+ val writeSnapshot = new java.util.HashMap[Int, AggregatedShuffleWriteMetric]()
+ perShuffleWrite.asScala.foreach { case (k, v) =>
+ writeSnapshot.put(
+ k,
+ new AggregatedShuffleWriteMetric(v.bytesWritten, v.recordsWritten, v.writeTimeMs))
+ }
+ kvstore.write(new CelebornAggregatedWriteMetricsUIData(writeSnapshot))
+ // Read-side counterpart of writeSnapshot, keyed by shuffleId.
+ val readSnapshot = new java.util.HashMap[Int, AggregatedShuffleReadMetric]()
+ perShuffleRead.asScala.foreach { case (k, v) =>
+ readSnapshot.put(
+ k,
+ new AggregatedShuffleReadMetric(v.bytesRead, v.chunkCount, v.fetchTimeMs))
+ }
+ kvstore.write(new CelebornAggregatedReadMetricsUIData(readSnapshot))
kvstore.write(AggregatedTaskInfoUIData(
totalWriteBytes.get(),
totalWriteTimeMs.get(),
totalReadBytes.get(),
totalFetchWaitTimeMs.get(),
totalTaskDurationMs.get()))
+ // Write-path timing breakdown (ms) + per-worker push stats, from CelebornWriteMetricsEvent.
+ kvstore.write(new CelebornWriteTimesUIData(
+ aggCopyTimeMs.get(),
+ aggSerializeTimeMs.get(),
+ aggCompressTimeMs.get(),
+ aggQueueWaitTimeMs.get(),
+ aggQueueStallTimeMs.get(),
+ aggInflightWaitTimeMs.get(),
+ aggDrainWaitTimeMs.get(),
+ aggSlowPushCount.get(),
+ aggMaxPushRttMs.get(),
+ aggUncompressedBytes.get()))
+ perWorkerWriteStats.asScala.foreach { case (workerId, acc) =>
+ kvstore.write(new CelebornPerWorkerWriteStatsUIData(
+ workerId,
+ acc.pushCount.get(),
+ acc.pushBytes.get(),
+ acc.totalPushRttNanos.get(),
+ acc.softSplitCount.get(),
+ acc.hardSplitCount.get(),
+ acc.primaryCongestedCount.get(),
+ acc.replicaCongestedCount.get(),
+ acc.lastPushFailureReason))
+ }
+ // Read-path timing breakdown (ms) + per-worker read stats, from CelebornReadMetricsEvent.
+ kvstore.write(new CelebornReadTimesUIData(
+ aggDecompressTimeMs.get(),
+ aggChunkWaitTimeMs.get(),
+ aggDeserializeTimeMs.get(),
+ aggCopyTimeMsRead.get(),
+ aggRetryCount.get(),
+ aggRetryWaitTimeMs.get(),
+ aggPeerSwitchCount.get(),
+ aggExcludeCount.get(),
+ aggSlowChunkCount.get(),
+ aggMaxChunkRttMs.get()))
+ perWorkerReadStats.asScala.foreach { case (workerId, acc) =>
+ kvstore.write(new CelebornPerWorkerReadStatsUIData(
+ workerId,
+ acc.chunkCount.get(),
+ acc.bytes.get(),
+ acc.totalRttNanos.get(),
+ acc.maxRttNanos.get()))
+ }
} catch {
case e: Exception =>
logWarning("Failed to flush CelebornListener aggregations", e)
}
}
}
+
+private[celeborn] object CelebornListener extends Logging {
+
+ /** Register a listener to the status queue so its events enter the event log. */
+ def register(sc: org.apache.spark.SparkContext): Unit = {
+ val listener = new CelebornListener(sc.statusStore.store, sc.conf)
+ sc.listenerBus.addToStatusQueue(listener)
+ logInfo("CelebornListener registered to the status queue")
+ }
+}
+
+private[celeborn] class PerWorkerWriteAccumulator {
+ private[celeborn] val pushCount = new AtomicLong(0)
+ private[celeborn] val pushBytes = new AtomicLong(0)
+ private[celeborn] val totalPushRttNanos = new AtomicLong(0)
+ private[celeborn] val softSplitCount = new AtomicLong(0)
+ private[celeborn] val hardSplitCount = new AtomicLong(0)
+ private[celeborn] val primaryCongestedCount = new AtomicLong(0)
+ private[celeborn] val replicaCongestedCount = new AtomicLong(0)
+ @volatile private[celeborn] var lastPushFailureReason: String = ""
+
+ def merge(s: org.apache.celeborn.common.protocol.message.PushWorkerStats): Unit = {
+ pushCount.addAndGet(s.pushCount)
+ pushBytes.addAndGet(s.pushBytes)
+ totalPushRttNanos.addAndGet(s.totalPushRttNanos)
+ softSplitCount.addAndGet(s.softSplitCount)
+ hardSplitCount.addAndGet(s.hardSplitCount)
+ primaryCongestedCount.addAndGet(s.primaryCongestedCount)
+ replicaCongestedCount.addAndGet(s.replicaCongestedCount)
+ if (s.lastPushFailureReason != null && s.lastPushFailureReason.nonEmpty) {
+ lastPushFailureReason = s.lastPushFailureReason
+ }
+ }
+}
+
+private[celeborn] class PerWorkerReadAccumulator {
+ private[celeborn] val chunkCount = new AtomicLong(0)
+ private[celeborn] val bytes = new AtomicLong(0)
+ private[celeborn] val totalRttNanos = new AtomicLong(0)
+ private[celeborn] val maxRttNanos = new AtomicLong(0)
+
+ def merge(w: org.apache.celeborn.common.protocol.message.WorkerReadCost): Unit = {
+ chunkCount.addAndGet(w.chunkCount)
+ bytes.addAndGet(w.bytes)
+ totalRttNanos.addAndGet(w.totalRttNanos)
+ maxRttNanos.accumulateAndGet(w.maxRttNanos, math.max)
+ }
+}
diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala
index d0e0caae507..2e008ee3a01 100644
--- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala
+++ b/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
+ }
}
diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala
index 9a8ffe5cf1f..269255c3455 100644
--- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala
+++ b/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)
+ }
}
diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornUIData.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornUIData.scala
new file mode 100644
index 00000000000..76929e28117
--- /dev/null
+++ b/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)
diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
index 4a57652fcf2..1d054c65590 100644
--- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
+++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala
@@ -17,21 +17,31 @@
package org.apache.spark.shuffle.celeborn.ui
+import scala.collection.JavaConverters._
import scala.xml.Node
import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.celeborn.CelebornStatusStore
import org.apache.spark.shuffle.celeborn.ui.SparkServletBridge.HttpServletRequest
import org.apache.spark.ui.{UIUtils, WebUIPage}
import org.apache.spark.util.Utils
+/**
+ * Page rendered under the Celeborn tab. Uses [[SparkServletBridge]] to bridge the
+ * javax.servlet (Spark 3.x) / jakarta.servlet (Spark 4.x) parameter type. The content
+ * sections are built by the companion object's buildContent so they can be unit-tested
+ * against an in-memory store without a live SparkUI or servlet request.
+ */
private[celeborn] class CelebornShufflePage(parent: CelebornUITab)
extends WebUIPage("") with Logging {
- private val store = parent.store
-
override def render(request: HttpServletRequest): Seq[Node] = {
try {
- renderBody(request)
+ UIUtils.headerSparkPage(
+ request,
+ "Celeborn Shuffle Service",
+ CelebornShufflePage.buildContent(parent.store),
+ parent)
} catch {
case e: Throwable =>
logError("Failed to render Celeborn Shuffle page", e)
@@ -39,88 +49,459 @@ private[celeborn] class CelebornShufflePage(parent: CelebornUITab)
- Celeborn Shuffle
+ Celeborn Shuffle Service
Failed to render the Celeborn page: {e.getMessage}
- UIUtils.headerSparkPage(request, "Celeborn Shuffle", errorContent, parent)
+ UIUtils.headerSparkPage(request, "Celeborn Shuffle Service", errorContent, parent)
}
}
+}
+
+private[celeborn] object CelebornShufflePage {
- private def renderBody(request: HttpServletRequest): Seq[Node] = {
- val taskInfo = store.aggregatedTaskInfo()
- val properties = store.celebornProperties()
+ /** Renders everything below the Spark UI header: stat cards, summary, sections. */
+ def buildContent(statusStore: CelebornStatusStore): Seq[Node] = {
+ val assignments = statusStore.assignmentInfos()
+ val properties = statusStore.celebornProperties()
+ val reassign = statusStore.reassignStats()
+ val taskInfo = statusStore.aggregatedTaskInfo()
+ val writeTimes = statusStore.writeTimes()
+ val perWorkerStats = statusStore.perWorkerWriteStats()
+ val readTimes = statusStore.readTimes()
+ val perWorkerReadStats = statusStore.perWorkerReadStats()
+ // Entities persisted by the listener but rendered only by the enhanced page:
+ // build/version info, per-policy fallback counts and per-shuffle write metrics.
+ val buildInfo = statusStore.buildInfo()
+ val fallback = statusStore.fallbackStats()
+ val perShuffleWrite = statusStore.aggregatedWriteMetrics()
+ val perShuffleRead = statusStore.aggregatedReadMetrics()
+ // --- Summary derived from onTaskEnd TaskMetrics (plan A: no executor->driver metric RPC) ---
val writeBytes = taskInfo.shuffleWriteBytes
val readBytes = taskInfo.shuffleReadBytes
val writeMs = taskInfo.shuffleWriteTimeMs
val readMs = taskInfo.shuffleFetchWaitTimeMs
val durationMs = taskInfo.taskDurationMs
-
def mbps(bytes: Long, ms: Long): String =
if (ms <= 0) "N/A" else f"${bytes.toDouble / 1000.0 / 1000.0 / (ms.toDouble / 1000.0)}%.2f"
def pct(part: Long, total: Long): String =
if (total <= 0) "N/A" else f"${part.toDouble * 100.0 / total.toDouble}%.1f%%"
+ val compressionRatio: String =
+ if (writeBytes > 0) f"${writeTimes.uncompressedBytes.toDouble / writeBytes.toDouble}%.2f"
+ else "N/A"
+
+ // Per-shuffle join: assignments (primary side) outer-joined with the write and read
+ // metric maps on the shuffle id. Assignments can be evicted by retainedShuffles while
+ // the metric aggregations (keyed by shuffleId, not row-capped) are still present, so
+ // metric-only ids are kept as rows with default columns instead of being dropped.
+ val writeMetrics = perShuffleWrite.metrics.asScala
+ val readMetrics = perShuffleRead.metrics.asScala
+ val assignmentIds = assignments.map(_.appShuffleId).toSet
+ val perShuffleIds =
+ (assignmentIds ++ (writeMetrics.keySet ++ readMetrics.keySet)).toSeq.sorted
- val summary =
+ val summary: Seq[Node] =
+ - Total Shuffle Write Bytes: {Utils.bytesToString(writeBytes)}
+ - Total Shuffle Read Bytes: {Utils.bytesToString(readBytes)}
+ - Compression Ratio: {compressionRatio}
+ (uncompressed {Utils.bytesToString(writeTimes.uncompressedBytes)}
+ / compressed {Utils.bytesToString(writeBytes)})
+ - Client Observed Write Speed: {mbps(writeBytes, writeMs)} MB/s
+ - Client Observed Read Speed: {mbps(readBytes, readMs)} MB/s
+ - Shuffle Write Time / Task Duration: {pct(writeMs, durationMs)}
+ - Shuffle Read Time / Task Duration: {pct(readMs, durationMs)}
+ - Shuffle Duration (write+read) / Task Duration: {pct(writeMs + readMs, durationMs)}
+ - Shuffle Adjustments: partitionSplit={reassign.partitionSplit},
+ reviveTriggered={reassign.reviveTriggered}, stageRetry={reassign.stageRetry}
-
- Shuffle Write:
- {
- s"${Utils.bytesToString(writeBytes)} | Time: ${UIUtils.formatDuration(
- writeMs)} | Speed: ${mbps(writeBytes, writeMs)} MB/s"
- }
+ Celeborn Build Information
-
- Shuffle Read:
- {
- s"${Utils.bytesToString(readBytes)} | Time: ${UIUtils.formatDuration(
- readMs)} | Speed: ${mbps(readBytes, readMs)} MB/s"
- }
+ Celeborn Properties ({properties.info.length} entries)
-
- Shuffle Duration (write+read) / Task Duration:
- {
- s"${pct(writeMs + readMs, durationMs)} (Write ${pct(
- writeMs,
- durationMs)}, Read ${pct(readMs, durationMs)})"
- }
+ Per-Shuffle Details ({perShuffleIds.size} shuffles)
+
+ -
+ Shuffle Throughput
+
+ -
+ Shuffle Write Times
+
+ -
+ Shuffle Read Times
+
+ -
+ Shuffle Write Servers
+
+ -
+ Shuffle Read Servers
+
+ -
+ Shuffle Assignments ({assignments.length} shuffles)
+
+ -
+ Shuffle Fallback ({fallback.counts.size} policies)
+ // Uniffle-style summary cards: headline numbers rendered as tiles above the detail
+ // sections. Only Spark UI's own grid classes are used (row-fluid/span2, valid through
+ // Spark 3.5); on Spark 4.x (Bootstrap 5) the unknown classes degrade to a vertical
+ // stack, which stays readable without shipping any extra CSS/JS.
+ val workerCount =
+ (perWorkerStats.map(_.workerId) ++ perWorkerReadStats.map(_.workerId)).toSet.size
+ val fallbackShuffleTotal = fallback.counts.asScala.values.map(_.toLong).sum
+ val summaryCards: Seq[Node] =
+
+ {Seq(
+ ("Shuffles", perShuffleIds.size.toString),
+ ("Workers", workerCount.toString),
+ ("Total Write Bytes", Utils.bytesToString(writeBytes)),
+ ("Total Read Bytes", Utils.bytesToString(readBytes)),
+ ("Compression Ratio", compressionRatio),
+ ("Fallback Shuffles", fallbackShuffleTotal.toString)).map { case (label, value) =>
+ statTile(label, value)
+ }}
+
+
+ // Build info uses the same two-column listing as the properties section; both are
+ // Seq[(String, String)] key/value summaries written once at startup.
+ val buildInfoTable = UIUtils.listingTable(
+ propertyHeader,
+ propertyRow,
+ buildInfo.info,
+ fixedWidth = true)
+
+ // Fallback counts are a policy -> count snapshot; sorted by policy name for a
+ // stable display. An empty map still renders the section (with a placeholder row)
+ // so the anchor stays valid.
+ val fallbackRows = fallback.counts.asScala.toSeq.sortBy(_._1).map { case (policy, count) =>
+
+ | {policy} |
+ {count.toString} |
+
+ }
+ val fallbackTable =
+
+
+
+ | Fallback Policy |
+ Count |
+
+
+
+ {
+ if (fallbackRows.isEmpty) {
+
+ | No fallback recorded. |
+ |
+
+ } else {
+ fallbackRows
+ }
+ }
+
+
+
+ // Per-Shuffle Details: assignments outer-joined with per-shuffle write and read metrics,
+ // with client-side derived columns (% of total write, write/read speed) in the style of
+ // Uniffle's per-application table. Rows are sorted by shuffle id (default ordering).
+ val perShuffleRows = perShuffleIds.map { shuffleId =>
+ val assignment = assignments.find(_.appShuffleId == shuffleId)
+ val w = writeMetrics.get(shuffleId)
+ val r = readMetrics.get(shuffleId)
+
+ | {shuffleId} |
+ {assignment.fold("N/A")(_.celebornShuffleId.toString)} |
+ {assignment.fold("N/A")(_.numPartitions.toString)} |
+ {assignment.fold("N/A")(_.workers.asScala.size.toString)} |
+ {w.fold("N/A")(m => Utils.bytesToString(m.bytesWritten))} |
+ {w.fold("N/A")(m => pct(m.bytesWritten, writeBytes))} |
+ {w.fold("N/A")(m => mbps(m.bytesWritten, m.writeTimeMs))} |
+ {w.fold("N/A")(_.recordsWritten.toString)} |
+ {w.fold("N/A")(m => Utils.msDurationToString(m.writeTimeMs))} |
+ {r.fold("N/A")(m => Utils.bytesToString(m.bytesRead))} |
+ {r.fold("N/A")(m => mbps(m.bytesRead, m.fetchTimeMs))} |
+
+ }
+ val perShuffleTable =
+
+
+
+ | App Shuffle Id |
+ Celeborn Shuffle Id |
+ Partitions |
+ Workers Count |
+ Write Bytes |
+ % of Total Write |
+ Write Speed (MB/s) |
+ Records |
+ Write Time |
+ Read Bytes |
+ Read Speed (MB/s) |
+
+
+
+ {perShuffleRows}
+
+
+
val propertiesTable = UIUtils.listingTable(
propertyHeader,
propertyRow,
properties.info,
- fixedWidth = true,
- headerClasses = headerClasses)
-
- val content =
-
- {summary}
-
-
-
-
- {propertiesTable}
-
-
+ fixedWidth = true)
+
+ val assignmentRows = assignments.map { a =>
+
+ | {a.appShuffleId} |
+ {a.celebornShuffleId} |
+ {a.workers.asScala.mkString(", ")} |
+ {a.numPartitions} |
+ {UIUtils.formatDate(new java.util.Date(a.timestamp))} |
+
+ }
+ val assignmentTable =
+
+
+
+ | App Shuffle Id |
+ Celeborn Shuffle Id |
+ Assigned Workers |
+ Num Partitions |
+ Timestamp |
+
+
+
+ {assignmentRows}
+
+
+
+ val throughputRows: Seq[(String, String)] = Seq(
+ ("Total Shuffle Write Bytes", Utils.bytesToString(taskInfo.shuffleWriteBytes)),
+ ("Total Shuffle Write Time", Utils.msDurationToString(taskInfo.shuffleWriteTimeMs)),
+ ("Total Shuffle Read Bytes", Utils.bytesToString(taskInfo.shuffleReadBytes)),
+ ("Total Fetch Wait Time", Utils.msDurationToString(taskInfo.shuffleFetchWaitTimeMs)),
+ ("Total Task Duration", Utils.msDurationToString(taskInfo.taskDurationMs)))
+ val throughputRowsXml = throughputRows.map { case (label, value) =>
+
+ | {label} |
+ {value} |
+
+ }
+ val throughputTable =
+
+
+
+ | Metric |
+ Value |
+
+
+
+ {throughputRowsXml}
+
+
+
+ // Vertical Metric | Duration | Count layout so timing breakdowns and occurrence
+ // counts stay in separate columns instead of a single mixed "Duration" row
+ // (mirrors Uniffle's per-metric rows).
+ val writeTimesRows = Seq(
+ ("Serialize", Utils.msDurationToString(writeTimes.serializeTimeMs), "–"),
+ ("Copy", Utils.msDurationToString(writeTimes.copyTimeMs), "–"),
+ ("Queue Wait", Utils.msDurationToString(writeTimes.queueWaitTimeMs), "–"),
+ ("Background Compress", Utils.msDurationToString(writeTimes.compressTimeMs), "–"),
+ ("Background Queue Stall", Utils.msDurationToString(writeTimes.queueStallTimeMs), "–"),
+ ("Background Inflight Wait", Utils.msDurationToString(writeTimes.inflightWaitTimeMs), "–"),
+ ("Drain Wait", Utils.msDurationToString(writeTimes.drainWaitTimeMs), "–"),
+ ("Max Push RTT", Utils.msDurationToString(writeTimes.maxPushRttMs), "–"),
+ ("Slow Push", "–", writeTimes.slowPushCount.toString))
+ val writeTimesTable = timesTable(writeTimesRows)
+
+ // Sorted by push bytes (descending) so the busiest workers surface first, mirroring
+ // Uniffle's sortable node list with a server-side default ordering.
+ val perWorkerRows = perWorkerStats.sortBy(-_.pushBytes).map { s =>
+ val totalPushRttMs = java.util.concurrent.TimeUnit.NANOSECONDS.toMillis(s.totalPushRttNanos)
+ val avgPushRttMs =
+ if (s.pushCount > 0) f"${totalPushRttMs.toDouble / s.pushCount}%.2f" else "N/A"
+
+ | {s.workerId} |
+ {Utils.bytesToString(s.pushBytes)} |
+ {s.pushCount} |
+ {Utils.msDurationToString(totalPushRttMs)} |
+ {avgPushRttMs} |
+ {s.softSplitCount} |
+ {s.hardSplitCount} |
+ {s.primaryCongestedCount} |
+ {s.replicaCongestedCount} |
+ {Option(s.lastPushFailureReason).getOrElse("")} |
+
+ }
+ val perWorkerTable =
+
+
+
+ | Worker Id |
+ Push Bytes |
+ Push Count |
+ Total Push RTT |
+ Avg Push RTT (ms) |
+ Soft Split |
+ Hard Split |
+ Primary Congested |
+ Replica Congested |
+ Last Push Failure Reason |
+
+
+
+ {perWorkerRows}
+
+
+
+ // All 10 read-side fields, including retry/peer-switch/exclude counts that the
+ // previous single-row layout dropped.
+ val readTimesRows = Seq(
+ ("Chunk Wait", Utils.msDurationToString(readTimes.chunkWaitTimeMs), "–"),
+ ("Decompress", Utils.msDurationToString(readTimes.decompressTimeMs), "–"),
+ ("Deserialize", Utils.msDurationToString(readTimes.deserializeTimeMs), "–"),
+ ("Copy", Utils.msDurationToString(readTimes.copyTimeMs), "–"),
+ ("Retry Wait", Utils.msDurationToString(readTimes.retryWaitTimeMs),
+ s"retries: ${readTimes.retryCount}"),
+ ("Peer Switch", "–", readTimes.peerSwitchCount.toString),
+ ("Exclude", "–", readTimes.excludeCount.toString),
+ ("Max Chunk RTT", Utils.msDurationToString(readTimes.maxChunkRttMs), "–"),
+ ("Slow Chunk", "–", readTimes.slowChunkCount.toString))
+ val readTimesTable = timesTable(readTimesRows)
- UIUtils.headerSparkPage(request, "Celeborn Shuffle", content, parent)
+ // Sorted by read bytes (descending), same default-ordering rationale as the write table.
+ val perWorkerReadRows = perWorkerReadStats.sortBy(-_.bytes).map { s =>
+ val totalRttMs = java.util.concurrent.TimeUnit.NANOSECONDS.toMillis(s.totalRttNanos)
+ val maxRttMs = java.util.concurrent.TimeUnit.NANOSECONDS.toMillis(s.maxRttNanos)
+ val avgRttMs =
+ if (s.chunkCount > 0) f"${totalRttMs.toDouble / s.chunkCount}%.2f" else "N/A"
+
+ | {s.workerId} |
+ {s.chunkCount} |
+ {Utils.bytesToString(s.bytes)} |
+ {Utils.msDurationToString(totalRttMs)} |
+ {avgRttMs} |
+ {Utils.msDurationToString(maxRttMs)} |
+
+ }
+ val perWorkerReadTable =
+
+
+
+ | Worker Id |
+ Chunk Count |
+ Read Bytes |
+ Total RTT |
+ Avg RTT (ms) |
+ Max RTT |
+
+
+
+ {perWorkerReadRows}
+
+
+
+ val content: Seq[Node] =
+
+ {summaryCards}
+
{summary}
+
+
+ {collapsible("build-info", "Celeborn Build Information", buildInfoTable)}
+
+ {collapsible("celeborn-properties", "Celeborn Properties", propertiesTable)}
+
+ {collapsible("per-shuffle", "Per-Shuffle Details", perShuffleTable)}
+
+ {collapsible("throughput", "Shuffle Throughput", throughputTable)}
+
+ {collapsible("write-times", "Shuffle Write Times", writeTimesTable)}
+
+ {collapsible("read-times", "Shuffle Read Times", readTimesTable)}
+
+ {collapsible("shuffle-write-servers", "Shuffle Write Servers", perWorkerTable)}
+
+ {collapsible("shuffle-read-servers", "Shuffle Read Servers", perWorkerReadTable)}
+
+ {collapsible("assignments", "Shuffle Assignments", assignmentTable)}
+
+ {collapsible("fallback", "Shuffle Fallback", fallbackTable)}
+
+ | {name} |
+ {duration} |
+ {count} |
+
+ }
+
+
+
+ | Metric |
+ Duration |
+ Count |
+
+
+
+ {rowXml}
+
+
+ }
+
+ /** A collapsible section (default collapsed), mirroring Uniffle's collapse-table pattern. */
+ private def collapsible(id: String, title: String, body: Seq[Node]): Seq[Node] = {
+ val spanClass = s"collapse-$id collapse-table"
+ val bodyClass = s"$id-table collapsible-table collapsed"
+
}
- private def propertyHeader: Seq[String] = Seq("Name", "Value")
- private def headerClasses: Seq[String] = Seq("sorttable_alpha", "sorttable_alpha")
+ private def propertyHeader: Seq[String] = Seq("Property", "Value")
private def propertyRow(kv: (String, String)): Seq[Node] =
diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUIUtils.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUIUtils.scala
new file mode 100644
index 00000000000..da0716cc1ac
--- /dev/null
+++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUIUtils.scala
@@ -0,0 +1,62 @@
+/*
+ * 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.ui
+
+import org.apache.spark.{SparkConf, SparkContext}
+import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.celeborn.CelebornStatusStore
+
+import org.apache.celeborn.common.CelebornConf
+
+/** Helpers for the Celeborn UI extension: enablement check and tab attachment. */
+private[celeborn] object CelebornUIUtils extends Logging {
+
+ /**
+ * Build a [[CelebornConf]] from the `spark.celeborn.*` entries of a SparkConf.
+ * SparkUtils.fromSparkConf does the same but its class initialization drags in
+ * spark-sql's UnsafeRowSerializer (via the AQE dataSize DynField), which is not on
+ * this module's classpath (spark-core only) and absent in RDD-only applications.
+ */
+ private[celeborn] def celebornConf(conf: SparkConf): CelebornConf = {
+ val celebornConf = new CelebornConf()
+ conf.getAll.foreach { case (key, value) =>
+ if (key.startsWith("spark.celeborn.")) {
+ celebornConf.set(key.substring("spark.".length), value)
+ }
+ }
+ celebornConf
+ }
+
+ /**
+ * The UI extension is on when both celeborn's flag and Spark's own UI flag are on.
+ * The celeborn flag is read through the CelebornConf entry so the documented
+ * `spark.celeborn.client.spark.ui.enabled` key (and the entry's default handling)
+ * applies; a bare-key SparkConf lookup never matches the user-set key because
+ * SparkConf does exact matching.
+ */
+ def isUIEnabled(conf: SparkConf): Boolean = {
+ celebornConf(conf).clientSparkUIEnabled &&
+ conf.getBoolean("spark.ui.enabled", true)
+ }
+
+ /** Attach the Celeborn tab to the live Spark WebUI. */
+ def attachUI(sc: SparkContext): Unit = {
+ val store = new CelebornStatusStore(sc.statusStore.store)
+ sc.ui.foreach(ui => new CelebornUITab(store, ui))
+ }
+}
diff --git a/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
index ed9ff78ad49..9d7c4d3090b 100644
--- a/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
+++ b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala
@@ -20,7 +20,8 @@ package org.apache.spark.shuffle.celeborn
import org.apache.spark.{SparkConf, Success, TaskState}
import org.apache.spark.executor.TaskMetrics
import org.apache.spark.internal.config.Status.ASYNC_TRACKING_ENABLED
-import org.apache.spark.scheduler.{JobSucceeded, SparkListenerEnvironmentUpdate, SparkListenerJobEnd, SparkListenerTaskEnd, TaskInfo, TaskLocality}
+import org.apache.spark.scheduler._
+import org.apache.spark.shuffle.celeborn.events._
import org.apache.spark.status.ElementTrackingStore
import org.apache.spark.util.Utils
import org.apache.spark.util.kvstore.InMemoryStore
@@ -29,6 +30,9 @@ import org.junit.Test
import org.junit.runner.RunWith
import org.junit.runners.JUnit4
+import org.apache.celeborn.common.protocol.message.{PushWorkerStats, ReadMetrics, WorkerReadCost, WriteMetrics}
+import org.apache.spark.shuffle.celeborn.ui.CelebornUIUtils
+
@RunWith(classOf[JUnit4])
class CelebornListenerSuite {
@@ -167,4 +171,259 @@ class CelebornListenerSuite {
// Non-Celeborn properties are filtered out.
assertFalse(props.contains("spark.executor.memory"))
}
+
+ @Test
+ def customEventsGatedOnPluginOptIn(): Unit = {
+ // Before opt-in: custom events are ignored, nothing is written to the store.
+ val store1 = new InMemoryStore()
+ val statusStore1 = new CelebornStatusStore(store1)
+ val listener1 = new CelebornListener(store1, new SparkConf(), requirePluginOptIn = true)
+ listener1.onOtherEvent(CelebornWriteMetricsEvent(
+ 0,
+ new WriteMetrics(1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L),
+ java.util.Collections.emptyList[PushWorkerStats](),
+ 1L))
+ listener1.onOtherEvent(CelebornShuffleAssignmentEvent(
+ 0,
+ 100,
+ java.util.Arrays.asList("host1:9097"),
+ 8,
+ 1L))
+ listener1.onOtherEvent(CelebornBuildInfoEvent(Map("Spark Version" -> "3.5")))
+ assertEquals(0L, statusStore1.writeTimes().copyTimeMs)
+ assertTrue(statusStore1.assignmentInfos().isEmpty)
+ assertTrue(statusStore1.buildInfo().info.isEmpty)
+
+ // After opt-in: the same events are collected.
+ listener1.onEnvironmentUpdate(envUpdate("spark.plugins" -> pluginClass))
+ listener1.onOtherEvent(CelebornWriteMetricsEvent(
+ 0,
+ new WriteMetrics(1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L),
+ java.util.Collections.emptyList[PushWorkerStats](),
+ 1L))
+ assertEquals(1L, statusStore1.writeTimes().copyTimeMs)
+ }
+
+ @Test
+ def writeMetricsEventAggregatedToWriteTimes(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ listener.onOtherEvent(CelebornWriteMetricsEvent(
+ 0,
+ new WriteMetrics(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 90L, 1000L),
+ java.util.Arrays.asList(
+ new PushWorkerStats("host1:9097", 10L, 1000L, 2000000L, 1L, 2L, 3L, 4L, ""),
+ new PushWorkerStats("host2:9097", 20L, 2000L, 4000000L, 0L, 1L, 0L, 1L, "boom")),
+ 1L))
+ listener.flush()
+
+ val writeTimes = statusStore.writeTimes()
+ assertEquals(1L, writeTimes.copyTimeMs)
+ assertEquals(2L, writeTimes.serializeTimeMs)
+ assertEquals(3L, writeTimes.compressTimeMs)
+ assertEquals(4L, writeTimes.queueWaitTimeMs)
+ assertEquals(5L, writeTimes.queueStallTimeMs)
+ assertEquals(6L, writeTimes.inflightWaitTimeMs)
+ assertEquals(7L, writeTimes.drainWaitTimeMs)
+ assertEquals(8L, writeTimes.slowPushCount)
+ assertEquals(90L, writeTimes.maxPushRttMs)
+ assertEquals(1000L, writeTimes.uncompressedBytes)
+
+ val workers = statusStore.perWorkerWriteStats().map(w => (w.workerId, w.pushCount)).toMap
+ assertEquals(2, workers.size)
+ assertEquals(10L, workers("host1:9097"))
+ assertEquals(20L, workers("host2:9097"))
+ // Merging a second event for the same worker accumulates instead of replacing.
+ listener.onOtherEvent(CelebornWriteMetricsEvent(
+ 0,
+ new WriteMetrics(1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L),
+ java.util.Arrays.asList(new PushWorkerStats("host1:9097", 5L, 0L, 0L, 0L, 0L, 0L, 0L, "")),
+ 2L))
+ listener.flush()
+ assertEquals(
+ 15L,
+ statusStore.perWorkerWriteStats().find(_.workerId == "host1:9097")
+ .map(_.pushCount).getOrElse(0L))
+ }
+
+ @Test
+ def readMetricsEventAggregatedToReadTimes(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ listener.onOtherEvent(CelebornReadMetricsEvent(
+ 0,
+ new ReadMetrics(11L, 12L, 13L, 14L, 15L, 16L, 17L, 18L, 19L, 200L),
+ java.util.Arrays.asList(new WorkerReadCost("host1:9097", 30L, 3000L, 5000000L, 900000L)),
+ 1L))
+ listener.flush()
+
+ val readTimes = statusStore.readTimes()
+ assertEquals(11L, readTimes.decompressTimeMs)
+ assertEquals(12L, readTimes.chunkWaitTimeMs)
+ assertEquals(13L, readTimes.deserializeTimeMs)
+ assertEquals(14L, readTimes.copyTimeMs)
+ assertEquals(15L, readTimes.retryCount)
+ assertEquals(16L, readTimes.retryWaitTimeMs)
+ assertEquals(17L, readTimes.peerSwitchCount)
+ assertEquals(18L, readTimes.excludeCount)
+ assertEquals(19L, readTimes.slowChunkCount)
+ assertEquals(200L, readTimes.maxChunkRttMs)
+
+ assertEquals(1, statusStore.perWorkerReadStats().size)
+ assertEquals(30L, statusStore.perWorkerReadStats().head.chunkCount)
+ assertEquals(3000L, statusStore.perWorkerReadStats().head.bytes)
+ }
+
+ @Test
+ def readMetricsEventAggregatedPerShuffle(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ // Two events for shuffle 0 (accumulate) and one for shuffle 1. ReadMetrics field order:
+ // decompressTimeMs, chunkWaitTimeMs, deserializeTimeMs, copyTimeMs, ...
+ listener.onOtherEvent(CelebornReadMetricsEvent(
+ 0,
+ new ReadMetrics(1L, 2L, 3L, 4L, 0L, 0L, 0L, 0L, 0L, 0L),
+ java.util.Arrays.asList(new WorkerReadCost("host1:9097", 10L, 1000L, 100L, 100L)),
+ 1L))
+ listener.onOtherEvent(CelebornReadMetricsEvent(
+ 1,
+ new ReadMetrics(5L, 6L, 7L, 8L, 0L, 0L, 0L, 0L, 0L, 0L),
+ java.util.Arrays.asList(new WorkerReadCost("host2:9097", 20L, 2000L, 100L, 100L)),
+ 2L))
+ listener.onOtherEvent(CelebornReadMetricsEvent(
+ 0,
+ new ReadMetrics(11L, 12L, 13L, 14L, 0L, 0L, 0L, 0L, 0L, 0L),
+ java.util.Arrays.asList(new WorkerReadCost("host1:9097", 5L, 500L, 100L, 100L)),
+ 3L))
+ listener.flush()
+
+ val metrics = statusStore.aggregatedReadMetrics().metrics
+ assertEquals(2, metrics.size())
+ val shuffle0 = metrics.get(0)
+ assertEquals(1500L, shuffle0.bytesRead)
+ assertEquals(15L, shuffle0.chunkCount)
+ // fetchTimeMs sums chunk wait + decompress + deserialize + copy across both events.
+ assertEquals((1 + 2 + 3 + 4) + (11 + 12 + 13 + 14), shuffle0.fetchTimeMs)
+ val shuffle1 = metrics.get(1)
+ assertEquals(2000L, shuffle1.bytesRead)
+ assertEquals(20L, shuffle1.chunkCount)
+ assertEquals(5 + 6 + 7 + 8, shuffle1.fetchTimeMs)
+ }
+
+ @Test
+ def aggregatedReadMetricsEmptyWhenNoReadEvents(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ // Reading before any flush must not throw and must return an empty map (the page
+ // renders before the first flush on a live application).
+ assertTrue(statusStore.aggregatedReadMetrics().metrics.isEmpty)
+
+ val listener = new CelebornListener(store, new SparkConf())
+ listener.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L))
+ listener.flush()
+ assertTrue(statusStore.aggregatedReadMetrics().metrics.isEmpty)
+ }
+
+ @Test
+ def retainedShufflesEvictsOldestAssignments(): Unit = {
+ // Set through the user-facing `spark.celeborn.*` key so the test exercises the same
+ // SparkUtils.fromSparkConf prefix-strip path as a real application.
+ val conf = new SparkConf()
+ .set(ASYNC_TRACKING_ENABLED, false)
+ .set("spark.celeborn.client.spark.ui.retainedShuffles", "2")
+ val store = new InMemoryStore()
+ val tracking = new ElementTrackingStore(store, conf)
+ val statusStore = new CelebornStatusStore(tracking)
+ val listener = new CelebornListener(tracking, conf)
+
+ (1 to 4).foreach { i =>
+ listener.onOtherEvent(CelebornShuffleAssignmentEvent(
+ i,
+ 100 + i,
+ java.util.Arrays.asList(s"host$i:9097"),
+ 8,
+ i.toLong))
+ }
+ // The addTrigger evicts the oldest rows (smallest appShuffleId) once the count
+ // exceeds the retained threshold of 2.
+ assertEquals(2, statusStore.assignmentInfos().size)
+ assertEquals(
+ Set(3, 4),
+ statusStore.assignmentInfos().map(_.appShuffleId).toSet)
+ tracking.close(false)
+ }
+
+ @Test
+ def uiGatingHonorsSparkCelebornPrefix(): Unit = {
+ // SparkConf(false) skips loading system properties: the surefire JVM injects
+ // spark.ui.enabled=false, which would defeat the "defaults to true" assertion.
+ // Enabled by default (both celeborn's flag and Spark's own UI flag default true).
+ assertTrue(CelebornUIUtils.isUIEnabled(new SparkConf(false)))
+ // The user-facing `spark.celeborn.*` key disables the extension: isUIEnabled must go
+ // through the CelebornConf entry (bare-key SparkConf lookups never match it).
+ val disabled = new SparkConf(false).set("spark.celeborn.client.spark.ui.enabled", "false")
+ assertFalse(CelebornUIUtils.isUIEnabled(disabled))
+ // Spark's own UI flag gates the tab independently.
+ val sparkUiOff = new SparkConf(false).set("spark.ui.enabled", "false")
+ assertFalse(CelebornUIUtils.isUIEnabled(sparkUiOff))
+ }
+
+ @Test
+ def customEventsReplayedFromEventLogSequence(): Unit = {
+ // Simulates HistoryServer replay: opt-in marker first (via environment update),
+ // then the full custom-event sequence, then an application end. Everything that
+ // the live listener would persist must be rebuilt by the replayed listener.
+ val conf = new SparkConf().set(ASYNC_TRACKING_ENABLED, false)
+ val store = new InMemoryStore()
+ val tracking = new ElementTrackingStore(store, conf)
+ val statusStore = new CelebornStatusStore(tracking)
+ val listener = new CelebornListener(tracking, conf, requirePluginOptIn = true)
+
+ listener.onEnvironmentUpdate(envUpdate("spark.plugins" -> pluginClass))
+ listener.onOtherEvent(CelebornBuildInfoEvent(Map("Spark Version" -> "3.5.8")))
+ listener.onOtherEvent(CelebornShuffleAssignmentEvent(
+ 0,
+ 100,
+ java.util.Arrays.asList("host1:9097", "host2:9097"),
+ 16,
+ 1L))
+ listener.onOtherEvent(CelebornFallbackEvent(
+ java.util.Collections.singletonMap("pushTimeout", 2L: java.lang.Long),
+ 1L))
+ listener.onOtherEvent(CelebornReassignEvent(partitionSplit = true, false, false, 1L))
+ listener.onOtherEvent(CelebornWriteMetricsEvent(
+ 0,
+ new WriteMetrics(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 10L),
+ java.util.Arrays.asList(new PushWorkerStats("host1:9097", 1L, 10L, 100L, 0L, 0L, 0L, 0L, "")),
+ 1L))
+ listener.onOtherEvent(CelebornReadMetricsEvent(
+ 0,
+ new ReadMetrics(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 10L),
+ java.util.Arrays.asList(new WorkerReadCost("host1:9097", 1L, 10L, 100L, 100L)),
+ 1L))
+ listener.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L))
+ listener.onApplicationEnd(SparkListenerApplicationEnd(1L))
+
+ assertTrue(statusStore.extensionEnabled())
+ assertEquals(1, statusStore.buildInfo().info.size)
+ assertEquals(1, statusStore.assignmentInfos().size)
+ assertEquals(16, statusStore.assignmentInfos().head.numPartitions)
+ assertEquals(2L, statusStore.fallbackStats().counts.get("pushTimeout"))
+ assertTrue(statusStore.reassignStats().partitionSplit)
+ assertFalse(statusStore.reassignStats().stageRetry)
+ assertEquals(2L, statusStore.writeTimes().serializeTimeMs)
+ assertEquals(3L, statusStore.readTimes().deserializeTimeMs)
+ assertEquals(100L, statusStore.aggregatedTaskInfo().shuffleWriteBytes)
+ // The replayed read event rebuilds the per-shuffle read aggregation as well.
+ assertEquals(10L, statusStore.aggregatedReadMetrics().metrics.get(0).bytesRead)
+ assertEquals(1L, statusStore.aggregatedReadMetrics().metrics.get(0).chunkCount)
+ assertEquals(10L, statusStore.aggregatedReadMetrics().metrics.get(0).fetchTimeMs)
+ tracking.close(false)
+ }
}
diff --git a/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePageSuite.scala b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePageSuite.scala
new file mode 100644
index 00000000000..3f97327f4ce
--- /dev/null
+++ b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePageSuite.scala
@@ -0,0 +1,112 @@
+/*
+ * 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.ui
+
+import org.apache.spark.SparkConf
+import org.apache.spark.shuffle.celeborn.{CelebornListener, CelebornStatusStore}
+import org.apache.spark.shuffle.celeborn.events._
+import org.apache.spark.util.kvstore.InMemoryStore
+import org.junit.Assert.{assertFalse, assertTrue}
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+import org.apache.celeborn.common.protocol.message.{PushWorkerStats, ReadMetrics, WorkerReadCost, WriteMetrics}
+
+/**
+ * Renders the page content against an in-memory store. buildContent is exercised directly
+ * (no SparkUI / servlet request) so the assertions stay at the XML-string level.
+ */
+@RunWith(classOf[JUnit4])
+class CelebornShufflePageSuite {
+
+ @Test
+ def rendersAllSectionsOnEmptyStore(): Unit = {
+ val statusStore = new CelebornStatusStore(new InMemoryStore())
+ val xml = CelebornShufflePage.buildContent(statusStore).toString
+ Seq(
+ "Celeborn Build Information",
+ "Celeborn Properties",
+ "Per-Shuffle Details",
+ "Shuffle Throughput",
+ "Shuffle Write Times",
+ "Shuffle Read Times",
+ "Shuffle Write Servers",
+ "Shuffle Read Servers",
+ "Shuffle Assignments",
+ "Shuffle Fallback",
+ "No fallback recorded.").foreach { section =>
+ assertTrue(s"missing section: $section", xml.contains(section))
+ }
+ // Empty store must not surface raw null / exception placeholders.
+ assertFalse(xml.contains("null"))
+ }
+
+ @Test
+ def rendersPopulatedSectionsAndDerivedColumns(): Unit = {
+ val store = new InMemoryStore()
+ val statusStore = new CelebornStatusStore(store)
+ val listener = new CelebornListener(store, new SparkConf())
+
+ listener.onOtherEvent(CelebornBuildInfoEvent(Map("Spark Version" -> "3.5.0")))
+ listener.onOtherEvent(CelebornShuffleAssignmentEvent(
+ 0,
+ 100,
+ java.util.Arrays.asList("host1:9097", "host2:9097"),
+ 16,
+ 1L))
+ listener.onOtherEvent(CelebornFallbackEvent(
+ java.util.Collections.singletonMap("pushTimeout", 2L: java.lang.Long),
+ 1L))
+ listener.onOtherEvent(CelebornWriteMetricsEvent(
+ 0,
+ new WriteMetrics(1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L),
+ java.util.Arrays.asList(
+ new PushWorkerStats("hostA:9097", 10L, 5000L, 10000000L, 0L, 0L, 0L, 0L, ""),
+ new PushWorkerStats("hostB:9097", 4L, 1000L, 1000000L, 0L, 0L, 0L, 0L, "")),
+ 1L))
+ listener.onOtherEvent(CelebornReadMetricsEvent(
+ 0,
+ new ReadMetrics(1L, 2L, 3L, 4L, 0L, 0L, 0L, 0L, 0L, 0L),
+ java.util.Arrays.asList(
+ new WorkerReadCost("hostC:9097", 5L, 900L, 4000000L, 2000000L),
+ new WorkerReadCost("hostD:9097", 3L, 300L, 3000000L, 1000000L)),
+ 2L))
+ listener.flush()
+
+ val xml = CelebornShufflePage.buildContent(statusStore).toString
+
+ // Build info and fallback sections surface their stored values.
+ assertTrue(xml.contains("Spark Version"))
+ assertTrue(xml.contains("3.5.0"))
+ assertTrue(xml.contains("pushTimeout"))
+
+ // Per-shuffle details: the assignment row joins with the read metrics. Read bytes
+ // = 900 + 300 = 1200 B over 10 ms fetch time -> 0.12 MB/s derived speed.
+ assertTrue(xml.contains("100"))
+ assertTrue(xml.contains("0.12"))
+
+ // Read worker table: hostC (900 B) sorts before hostD (300 B), and avg RTT is
+ // rendered (hostC: 4 ms / 5 chunks = 0.80, hostD: 3 ms / 3 chunks = 1.00).
+ assertTrue(xml.indexOf("hostC:9097") < xml.indexOf("hostD:9097"))
+ assertTrue(xml.contains("0.80"))
+
+ // Write worker table: hostA (5000 B) sorts before hostB (1000 B).
+ assertTrue(xml.indexOf("hostA:9097") < xml.indexOf("hostB:9097"))
+ }
+}
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index a36a0932a59..f741ab1695b 100644
--- a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++ b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -71,6 +71,13 @@ public class HashBasedShuffleWriter extends ShuffleWriter {
private final int encodedAttemptId;
private final TaskContext taskContext;
private final ShuffleClient shuffleClient;
+ // CPU cost of serializing records in write0, flushed into PushState at close.
+ private long serializeTimeNanos = 0;
+ // CPU cost of copying serialized records into the push buffer in write0, flushed into
+ // PushState at close.
+ private long copyTimeNanos = 0;
+ // Total serialized (pre-compression) bytes written in write0, flushed into PushState at close.
+ private long uncompressedBytes = 0;
private final int numMappers;
private final int numPartitions;
@@ -217,6 +224,7 @@ protected void fastWrite0(scala.collection.Iterator iterator)
final int rowSize = row.getSizeInBytes();
final int serializedRecordSize = 4 + rowSize;
+ uncompressedBytes += serializedRecordSize;
if (dataSize != null) {
dataSize.add(rowSize);
@@ -225,23 +233,27 @@ protected void fastWrite0(scala.collection.Iterator iterator)
if (serializedRecordSize > PUSH_BUFFER_MAX_SIZE) {
byte[] giantBuffer = new byte[serializedRecordSize];
Platform.putInt(giantBuffer, Platform.BYTE_ARRAY_OFFSET, Integer.reverseBytes(rowSize));
+ long _copyStart = System.nanoTime();
Platform.copyMemory(
row.getBaseObject(),
row.getBaseOffset(),
giantBuffer,
Platform.BYTE_ARRAY_OFFSET + 4,
rowSize);
+ copyTimeNanos += System.nanoTime() - _copyStart;
pushGiantRecord(partitionId, giantBuffer, serializedRecordSize);
} else {
int offset = getOrUpdateOffset(partitionId, serializedRecordSize);
byte[] buffer = getOrCreateBuffer(partitionId);
Platform.putInt(buffer, Platform.BYTE_ARRAY_OFFSET + offset, Integer.reverseBytes(rowSize));
+ long _copyStart = System.nanoTime();
Platform.copyMemory(
row.getBaseObject(),
row.getBaseOffset(),
buffer,
Platform.BYTE_ARRAY_OFFSET + offset + 4,
rowSize);
+ copyTimeNanos += System.nanoTime() - _copyStart;
sendOffsets[partitionId] = offset + serializedRecordSize;
}
tmpRecordsWritten++;
@@ -256,9 +268,11 @@ private void write0(scala.collection.Iterator iterator) throws IOException, Inte
final K key = record._1();
final int partitionId = partitioner.getPartition(key);
serBuffer.reset();
+ long serializeStartNanos = System.nanoTime();
serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
serOutputStream.flush();
+ serializeTimeNanos += System.nanoTime() - serializeStartNanos;
final int serializedRecordSize = serBuffer.size();
assert (serializedRecordSize > 0);
@@ -377,6 +391,21 @@ private void cleanupPusher() throws IOException {
}
private void close(boolean iteratorHasNext) throws IOException, InterruptedException {
+ if (serializeTimeNanos > 0) {
+ shuffleClient
+ .getPushState(Utils.makeMapKey(shuffleId, mapId, encodedAttemptId))
+ .addSerializeTime(serializeTimeNanos);
+ }
+ if (copyTimeNanos > 0) {
+ shuffleClient
+ .getPushState(Utils.makeMapKey(shuffleId, mapId, encodedAttemptId))
+ .addCopyTime(copyTimeNanos);
+ }
+ if (uncompressedBytes > 0) {
+ shuffleClient
+ .getPushState(Utils.makeMapKey(shuffleId, mapId, encodedAttemptId))
+ .addUncompressedBytes(uncompressedBytes);
+ }
// Send the remaining data in sendBuffer
long pushMergedDataTime = System.nanoTime();
closeWrite();
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index d08b36d1945..6b0c0a266b9 100644
--- a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++ b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -65,6 +65,12 @@ public class SortBasedShuffleWriter extends ShuffleWriter {
private final int encodedAttemptId;
private final TaskContext taskContext;
private final ShuffleClient shuffleClient;
+ // CPU cost of serializing records in write0, flushed into PushState at close.
+ private long serializeTimeNanos = 0;
+ // CPU cost of copying serialized records (giant record path), flushed into PushState at close.
+ private long copyTimeNanos = 0;
+ // Total serialized (pre-compression) bytes written, flushed into PushState at close.
+ private long uncompressedBytes = 0;
private final int numMappers;
private final int numPartitions;
@@ -269,6 +275,7 @@ private void fastWrite0(scala.collection.Iterator iterator) throws IOException {
final int rowSize = row.getSizeInBytes();
final int serializedRecordSize = 4 + rowSize;
+ uncompressedBytes += serializedRecordSize;
if (dataSize != null) {
dataSize.add(serializedRecordSize);
@@ -277,12 +284,14 @@ private void fastWrite0(scala.collection.Iterator iterator) throws IOException {
if (serializedRecordSize > pushBufferMaxSize) {
byte[] giantBuffer = new byte[serializedRecordSize];
Platform.putInt(giantBuffer, Platform.BYTE_ARRAY_OFFSET, Integer.reverseBytes(rowSize));
+ long _copyStart = System.nanoTime();
Platform.copyMemory(
row.getBaseObject(),
row.getBaseOffset(),
giantBuffer,
Platform.BYTE_ARRAY_OFFSET + 4,
rowSize);
+ copyTimeNanos += System.nanoTime() - _copyStart;
pushGiantRecord(partitionId, giantBuffer, serializedRecordSize);
} else {
boolean success =
@@ -316,9 +325,11 @@ private void write0(scala.collection.Iterator iterator) throws IOException {
final K key = record._1();
final int partitionId = partitioner.getPartition(key);
serBuffer.reset();
+ long serializeStartNanos = System.nanoTime();
serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
serOutputStream.flush();
+ serializeTimeNanos += System.nanoTime() - serializeStartNanos;
final int serializedRecordSize = serBuffer.size();
assert (serializedRecordSize > 0);
@@ -379,6 +390,21 @@ private void cleanupPusher() throws IOException {
}
private void close(boolean iteratorHasNext) throws IOException {
+ if (serializeTimeNanos > 0) {
+ shuffleClient
+ .getPushState(Utils.makeMapKey(shuffleId, mapId, encodedAttemptId))
+ .addSerializeTime(serializeTimeNanos);
+ }
+ if (copyTimeNanos > 0) {
+ shuffleClient
+ .getPushState(Utils.makeMapKey(shuffleId, mapId, encodedAttemptId))
+ .addCopyTime(copyTimeNanos);
+ }
+ if (uncompressedBytes > 0) {
+ shuffleClient
+ .getPushState(Utils.makeMapKey(shuffleId, mapId, encodedAttemptId))
+ .addUncompressedBytes(uncompressedBytes);
+ }
logger.info("Memory used {}", Utils.bytesToString(pusher.getUsed()));
long pushStartTime = System.nanoTime();
pusher.pushData(false);
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index 71aae81f312..c144923cdb4 100644
--- a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++ b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -27,6 +27,10 @@
import org.apache.spark.launcher.SparkLauncher;
import org.apache.spark.rdd.DeterministicLevel;
import org.apache.spark.shuffle.*;
+import org.apache.spark.shuffle.celeborn.events.CelebornReadMetricsEvent;
+import org.apache.spark.shuffle.celeborn.events.CelebornReassignEvent;
+import org.apache.spark.shuffle.celeborn.events.CelebornShuffleAssignmentEvent;
+import org.apache.spark.shuffle.celeborn.events.CelebornWriteMetricsEvent;
import org.apache.spark.shuffle.sort.SortShuffleManager;
import org.apache.spark.sql.internal.SQLConf;
import org.slf4j.Logger;
@@ -34,6 +38,7 @@
import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.listener.ReadMetricsCallback;
import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.ShuffleMode;
@@ -162,6 +167,68 @@ private void initializeLifecycleManager(String appId) {
lifecycleManager = new LifecycleManager(appUniqueId, celebornConf);
lifecycleManager.applicationCount().increment();
lifecycleManager.registerCancelShuffleCallback(SparkUtils::cancelShuffle);
+ if (celebornConf.clientSparkUIEnabled()) {
+ lifecycleManager.registerShuffleAssignmentCallback(
+ (shuffleId, numPartitions) -> {
+ SparkContext sc = SparkContext$.MODULE$.getActive().getOrElse(null);
+ if (sc != null) {
+ java.util.Map workers =
+ lifecycleManager.shuffleAllocatedWorkers().get(shuffleId);
+ java.util.List workerIds = new java.util.ArrayList<>();
+ if (workers != null) {
+ workerIds.addAll(workers.keySet());
+ }
+ sc.listenerBus()
+ .post(
+ new CelebornShuffleAssignmentEvent(
+ shuffleId,
+ shuffleId,
+ workerIds,
+ numPartitions,
+ System.currentTimeMillis()));
+ }
+ });
+ lifecycleManager.registerReassignCallback(
+ state -> {
+ SparkContext sc = SparkContext$.MODULE$.getActive().getOrElse(null);
+ if (sc != null && state.size() >= 3) {
+ sc.listenerBus()
+ .post(
+ new CelebornReassignEvent(
+ state.get(0),
+ state.get(1),
+ state.get(2),
+ System.currentTimeMillis()));
+ }
+ });
+ lifecycleManager.registerMapperEndMetricsCallback(
+ (shuffleId, writeMetrics, pushWorkerStats) -> {
+ SparkContext sc = SparkContext$.MODULE$.getActive().getOrElse(null);
+ if (sc != null) {
+ sc.listenerBus()
+ .post(
+ new CelebornWriteMetricsEvent(
+ shuffleId,
+ writeMetrics,
+ pushWorkerStats,
+ System.currentTimeMillis()));
+ }
+ });
+ lifecycleManager.registerReadMetricsCallback(
+ (ReadMetricsCallback)
+ (shuffleId, readMetrics, workerReadCosts) -> {
+ SparkContext sc = SparkContext$.MODULE$.getActive().getOrElse(null);
+ if (sc != null) {
+ sc.listenerBus()
+ .post(
+ new CelebornReadMetricsEvent(
+ shuffleId,
+ readMetrics,
+ workerReadCosts,
+ System.currentTimeMillis()));
+ }
+ });
+ }
if (celebornConf.clientStageRerunEnabled()) {
MapOutputTrackerMaster mapOutputTracker =
(MapOutputTrackerMaster) SparkEnv.get().mapOutputTracker();
diff --git a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleFallbackPolicyRunner.scala b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleFallbackPolicyRunner.scala
index 37dbcac8017..79b5136e485 100644
--- a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleFallbackPolicyRunner.scala
+++ b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleFallbackPolicyRunner.scala
@@ -22,6 +22,8 @@ import java.util.concurrent.ConcurrentHashMap
import scala.collection.JavaConverters._
import org.apache.spark.ShuffleDependency
+import org.apache.spark.SparkContext
+import org.apache.spark.shuffle.celeborn.events.CelebornFallbackEvent
import org.apache.celeborn.client.LifecycleManager
import org.apache.celeborn.common.CelebornConf
@@ -47,6 +49,15 @@ class CelebornShuffleFallbackPolicyRunner(conf: CelebornConf) extends Logging {
val shuffleFallbackPolicy = fallbackPolicy.get.getClass.getName
computeFallbackCounts(shuffleFallbackPolicy, lifecycleManager.shuffleFallbackCounts)
computeFallbackCounts(shuffleFallbackPolicy, lifecycleManager.applicationFallbackCounts)
+ if (conf.clientSparkUIEnabled) {
+ val sc = SparkContext.getActive.orNull
+ if (sc != null) {
+ sc.listenerBus.post(CelebornFallbackEvent(
+ new java.util.HashMap[String, java.lang.Long](
+ lifecycleManager.shuffleFallbackCounts),
+ System.currentTimeMillis()))
+ }
+ }
}
}
fallbackPolicy.isDefined
diff --git a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
index c0c2b20ec8f..d25bdc6d92b 100644
--- a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
+++ b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
@@ -31,7 +31,7 @@ import org.apache.commons.lang3.tuple.Pair
import org.apache.spark.{Aggregator, InterruptibleIterator, ShuffleDependency, TaskContext}
import org.apache.spark.celeborn.ExceptionMakerHelper
import org.apache.spark.internal.Logging
-import org.apache.spark.serializer.SerializerInstance
+import org.apache.spark.serializer.{DeserializationStream, SerializerInstance}
import org.apache.spark.shuffle.{FetchFailedException, ShuffleReader, ShuffleReadMetricsReporter}
import org.apache.spark.shuffle.celeborn.CelebornShuffleReader.streamCreatorPool
import org.apache.spark.util.CompletionIterator
@@ -39,7 +39,7 @@ import org.apache.spark.util.collection.ExternalSorter
import org.apache.celeborn.client.{ClientUtils, ShuffleClient}
import org.apache.celeborn.client.ShuffleClientImpl.ReduceFileGroups
-import org.apache.celeborn.client.read.{CelebornInputStream, MetricsCallback}
+import org.apache.celeborn.client.read.{CelebornInputStream, MetricsCallback, ReadStreamStats}
import org.apache.celeborn.client.security.CryptoHandler
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.exception.{CelebornBroadcastException, CelebornIOException, CelebornRuntimeException, PartitionUnRetryAbleException}
@@ -369,20 +369,29 @@ class CelebornShuffleReader[K, C](
val locationStreamHandlerMap: ConcurrentHashMap[PartitionLocation, PbStreamHandler] =
JavaUtils.newConcurrentHashMap()
- val futures = workerRequestMap.values().asScala.map { entry =>
+ val futures = workerRequestMap.asScala.map { case (hostPort, entry) =>
streamCreatorPool.submit(new Runnable {
override def run(): Unit = {
val (client, locArr, pbOpenStreamListBuilder) = entry
val msg = new TransportMessage(
MessageType.BATCH_OPEN_STREAM,
pbOpenStreamListBuilder.build().toByteArray)
+ val openStreamStartTime = System.currentTimeMillis()
val pbOpenStreamListResponse =
try {
val response = client.sendRpcSync(msg.toByteBuffer, fetchTimeoutMs)
TransportMessage.fromByteBuffer(response).getParsedPayload[PbOpenStreamListResponse]
} catch {
- case _: Exception => null
+ case e: Exception =>
+ logWarning(
+ s"BatchOpenStream request to $hostPort failed, " +
+ s"fall back to single OpenStream for each partition later.",
+ e)
+ null
}
+ logDebug(
+ s"BatchOpenStream request to $hostPort cost " +
+ s"${System.currentTimeMillis() - openStreamStartTime}ms")
if (pbOpenStreamListResponse != null) {
0 until locArr.size() foreach { idx =>
val streamHandlerOpt = pbOpenStreamListResponse.getStreamHandlerOptList.get(idx)
@@ -508,7 +517,11 @@ class CelebornShuffleReader[K, C](
}).filter {
case (_, inputStream) => inputStream != CelebornInputStream.empty()
}.map { case (partitionId, inputStream) =>
- (partitionId, serializerInstance.deserializeStream(inputStream).asKeyValueIterator)
+ val rawStream = serializerInstance.deserializeStream(inputStream)
+ val stats = inputStream.streamStats()
+ val timedStream =
+ if (stats != null) new TimedDeserializationStream(rawStream, stats) else rawStream
+ (partitionId, timedStream.asKeyValueIterator)
}.flatMap { case (partitionId, iter) =>
try {
iter
@@ -634,6 +647,45 @@ class CelebornShuffleReader[K, C](
def newSerializerInstance(dep: ShuffleDependency[K, _, C]): SerializerInstance = {
dep.serializer.newInstance()
}
+
+ /**
+ * Deserialization happens lazily while the key/value iterator is consumed, so timing only
+ * the `deserializeStream` construction would report ~0. Wrap the stream and time each
+ * readKey/readValue/readObject call, attributing the cost to the stream's ReadStreamStats so
+ * it is reported when the stream closes.
+ *
+ * All three readers must be overridden: serializers like Spark SQL's UnsafeRowSerializer
+ * override readKey/readValue directly and make the generic readObject throw
+ * UnsupportedOperationException, so delegating only readObject breaks them.
+ */
+ private class TimedDeserializationStream(
+ delegate: DeserializationStream,
+ stats: ReadStreamStats)
+ extends DeserializationStream {
+
+ override def readKey[T]()(implicit ct: scala.reflect.ClassTag[T]): T = {
+ val start = System.nanoTime()
+ val result = delegate.readKey[T]()
+ stats.addDeserializeTime(System.nanoTime() - start)
+ result
+ }
+
+ override def readValue[T]()(implicit ct: scala.reflect.ClassTag[T]): T = {
+ val start = System.nanoTime()
+ val result = delegate.readValue[T]()
+ stats.addDeserializeTime(System.nanoTime() - start)
+ result
+ }
+
+ override def readObject[T]()(implicit ct: scala.reflect.ClassTag[T]): T = {
+ val start = System.nanoTime()
+ val result = delegate.readObject[T]()
+ stats.addDeserializeTime(System.nanoTime() - start)
+ result
+ }
+
+ override def close(): Unit = delegate.close()
+ }
}
object CelebornShuffleReader {
diff --git a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/events/CelebornEvent.scala b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/events/CelebornEvent.scala
new file mode 100644
index 00000000000..1287a54086a
--- /dev/null
+++ b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/events/CelebornEvent.scala
@@ -0,0 +1,92 @@
+/*
+ * 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.events
+
+import org.apache.spark.scheduler.SparkListenerEvent
+
+import org.apache.celeborn.common.protocol.message.{PushWorkerStats, ReadMetrics, WorkerReadCost, WriteMetrics}
+
+/**
+ * Root for Celeborn UI events posted to the Spark listener bus. Defined in the
+ * spark-3 module (rather than spark-3-ui) so that [[org.apache.spark.shuffle.celeborn.SparkShuffleManager]]
+ * and [[org.apache.spark.shuffle.celeborn.CelebornShuffleFallbackPolicyRunner]] can post them
+ * without inverting the module dependency. Subclasses are case classes so they serialize into
+ * the Spark event log and can be replayed by the HistoryServer plugin.
+ */
+sealed trait CelebornEvent extends SparkListenerEvent
+
+/** Posted once from the driver plugin at initialization with build/version info. */
+case class CelebornBuildInfoEvent(info: Map[String, String]) extends CelebornEvent
+
+/**
+ * Posted from SparkShuffleManager.registerShuffle after a successful Celeborn slot assignment,
+ * carrying the shuffle -> worker topology. workers is a java.util.List (not Scala Seq) so the
+ * Java SparkShuffleManager can construct it without Scala collection interop.
+ */
+case class CelebornShuffleAssignmentEvent(
+ appShuffleId: Int,
+ celebornShuffleId: Int,
+ workers: java.util.List[String],
+ numPartitions: Int,
+ timestamp: Long) extends CelebornEvent
+
+/**
+ * Posted from CelebornShuffleFallbackPolicyRunner when a shuffle falls back to
+ * SortShuffleManager, carrying a snapshot of per-policy fallback counts. Counts use
+ * java.util.Map so Java callers don't need Scala collection interop.
+ */
+case class CelebornFallbackEvent(
+ fallbackCounts: java.util.Map[String, java.lang.Long],
+ timestamp: Long) extends CelebornEvent
+
+/**
+ * Posted from LifecycleManager (via a driver-side callback registered by SparkShuffleManager)
+ * the first time a partition split / block-send-failure revive / stage retry is triggered.
+ * Uses AtomicBoolean dedup in LifecycleManager so only the first trigger per type is posted
+ * (mirrors Uniffle's postReassignTriggeredEvent). stageRetry may be left false (placeholder)
+ * if the fetch-failure rerun path is not wired.
+ */
+case class CelebornReassignEvent(
+ partitionSplit: Boolean,
+ reviveTriggered: Boolean,
+ stageRetry: Boolean,
+ timestamp: Long) extends CelebornEvent
+
+/**
+ * Posted from LifecycleManager.handleMapperEnd (via registerMapperEndMetricsCallback) when the
+ * executor populated write metrics. Carries the per-task write-path timing breakdown and the
+ * per-worker push stats; the listener accumulates them per-shuffle for the UI's Shuffle Write
+ * Times and Shuffle Servers sections.
+ */
+case class CelebornWriteMetricsEvent(
+ shuffleId: Int,
+ writeMetrics: WriteMetrics,
+ pushWorkerStats: java.util.List[PushWorkerStats],
+ timestamp: Long) extends CelebornEvent
+
+/**
+ * Posted from LifecycleManager's ReportShuffleReadMetrics handler (via registerReadMetricsCallback)
+ * when the executor reported read metrics. Carries the per-task read-path timing breakdown and
+ * per-worker read cost; the listener accumulates them per-shuffle for the UI's Shuffle Read Times
+ * and Shuffle Servers (read side) sections.
+ */
+case class CelebornReadMetricsEvent(
+ shuffleId: Int,
+ readMetrics: ReadMetrics,
+ workerReadCosts: java.util.List[WorkerReadCost],
+ timestamp: Long) extends CelebornEvent
diff --git a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
index 86e6a2d191e..afdfc2a0d5c 100644
--- a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -42,6 +42,7 @@
import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
+import org.apache.celeborn.client.read.ReadStreamStats;
import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
@@ -169,6 +170,9 @@ public void pushMergedData(int shuffleId, int mapId, int attemptId) {}
public void mapperEnd(
int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions) {}
+ @Override
+ public void reportReadMetrics(int shuffleId, ReadStreamStats stats) {}
+
@Override
public void readReducerPartitionEnd(
int shuffleId, int partitionId, int startMapIndex, int endMapIndex, int crc32, long bytes)
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
index 9fb15fdf7ab..95745b3b391 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -36,6 +36,7 @@
import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
+import org.apache.celeborn.client.read.ReadStreamStats;
import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
@@ -375,6 +376,9 @@ public abstract void mapperEnd(
int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions)
throws IOException;
+ /** Report read-path metrics to the driver for the UI (no-op when UI disabled). */
+ public abstract void reportReadMetrics(int shuffleId, ReadStreamStats stats);
+
public abstract void readReducerPartitionEnd(
int shuffleId, int partitionId, int startMapIndex, int endMapIndex, int crc32, long bytes)
throws IOException;
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index f3219247372..b0c7702428a 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -25,7 +25,10 @@
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.LongAdder;
+import scala.Option;
import scala.Tuple2;
import scala.Tuple3;
import scala.reflect.ClassTag$;
@@ -44,6 +47,7 @@
import org.apache.celeborn.client.compress.Compressor;
import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
+import org.apache.celeborn.client.read.ReadStreamStats;
import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornBroadcastException;
@@ -62,7 +66,10 @@
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.*;
import org.apache.celeborn.common.protocol.message.ControlMessages.*;
+import org.apache.celeborn.common.protocol.message.ReadMetrics;
import org.apache.celeborn.common.protocol.message.StatusCode;
+import org.apache.celeborn.common.protocol.message.WorkerReadCost;
+import org.apache.celeborn.common.protocol.message.WriteMetrics;
import org.apache.celeborn.common.rpc.RpcAddress;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.rpc.RpcEnv;
@@ -116,6 +123,14 @@ public class ShuffleClientImpl extends ShuffleClient {
protected final ConcurrentHashMap> mapperEndMap =
JavaUtils.newConcurrentHashMap();
+ // key: mapKey (shuffleId-mapId-attemptId), value: (worker hostAndPushPort -> counters)
+ // Aggregated counts of push events (soft/hard split, congested) per map task,
+ // logged and cleared when the map task ends.
+ private final ConcurrentHashMap>
+ pushEventStats = JavaUtils.newConcurrentHashMap();
+
+ private final long slowPushThresholdNanos;
+
// shuffleIds which have finished all map tasks
protected final Set stageEndShuffleSet = ConcurrentHashMap.newKeySet();
@@ -222,6 +237,7 @@ public ShuffleClientImpl(String appUniqueId, CelebornConf conf, UserIdentifier u
}
authEnabled = conf.authEnabledOnClient();
dataPushFailureTrackingEnabled = conf.clientAdaptiveOptimizeSkewedPartitionReadEnabled();
+ slowPushThresholdNanos = TimeUnit.MILLISECONDS.toNanos(conf.clientPushSlowPushThresholdMs());
// init rpc env
rpcEnv =
@@ -399,7 +415,8 @@ private void submitRetryPushData(
} else {
PartitionLocation newLoc = reducePartitionMap.get(shuffleId).get(partitionId);
logger.info(
- "Revive for push data success, new location for shuffle {} map {} attempt {} partition {} batch {} is location {}.",
+ "Revive for push data success after waiting {} ms, new location for shuffle {} map {} attempt {} partition {} batch {} is location {}.",
+ accumulatedTime,
shuffleId,
mapId,
attemptId,
@@ -557,6 +574,17 @@ private void submitRetryPushMergedData(
}
}
+ if (accumulatedTime > 0) {
+ logger.info(
+ "Waited {} ms for revive results of {} batches for shuffle {} map {} attempt {} groupedBatch {}.",
+ accumulatedTime,
+ reviveRequests.length,
+ shuffleId,
+ mapId,
+ attemptId,
+ oldGroupedBatchId);
+ }
+
for (Map.Entry, DataBatches> entry : newDataBatchesMap.entrySet()) {
Pair addressPair = entry.getKey();
DataBatches newDataBatches = entry.getValue();
@@ -1099,7 +1127,9 @@ public int pushOrMergeData(
if (shuffleCompressionEnabled && !skipCompress) {
// compress data
final Compressor compressor = compressorThreadLocal.get();
+ long compressStartNanos = System.nanoTime();
compressor.compress(data, offset, length);
+ pushState.addCompressTime(System.nanoTime() - compressStartNanos);
data = compressor.getCompressedBuffer();
offset = 0;
@@ -1137,6 +1167,7 @@ public int pushOrMergeData(
// add inFlight requests
pushState.addBatch(nextBatchId, body.length, loc.hostAndPushPort());
+ final long pushStartTime = System.nanoTime();
// build PushData request
NettyManagedBuffer buffer = new NettyManagedBuffer(Unpooled.wrappedBuffer(body));
@@ -1190,6 +1221,21 @@ public void updateLatestPartition(PartitionLocation newloc) {
@Override
public void onSuccess(ByteBuffer response) {
+ long pushRttNanos = System.nanoTime() - pushStartTime;
+ pushState.recordPushRtt(pushRttNanos, slowPushThresholdNanos);
+ recordWorkerPush(
+ shuffleId, mapId, attemptId, latest.hostAndPushPort(), body.length, pushRttNanos);
+ if (pushRttNanos > slowPushThresholdNanos) {
+ logger.warn(
+ "Slow push data to {} for shuffle {} map {} attempt {} partition {} batch {}, rtt {} ms.",
+ latest.hostAndPushPort(),
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ nextBatchId,
+ TimeUnit.NANOSECONDS.toMillis(pushRttNanos));
+ }
if (response.remaining() > 0) {
byte reason = response.get();
if (reason == StatusCode.SOFT_SPLIT.getValue()) {
@@ -1214,6 +1260,8 @@ public void onSuccess(ByteBuffer response) {
StatusCode.SOFT_SPLIT);
reviveManager.addRequest(reviveRequest);
}
+ recordPushEvent(
+ shuffleId, mapId, attemptId, latest.hostAndPushPort(), PushEvent.SOFT_SPLIT);
pushState.onSuccess(latest.hostAndPushPort());
pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
callback.onSuccess(response);
@@ -1240,6 +1288,8 @@ public void onSuccess(ByteBuffer response) {
latest,
StatusCode.HARD_SPLIT);
reviveManager.addRequest(reviveRequest);
+ recordPushEvent(
+ shuffleId, mapId, attemptId, latest.hostAndPushPort(), PushEvent.HARD_SPLIT);
long dueTime =
System.currentTimeMillis()
+ conf.clientRpcRequestPartitionLocationAskTimeout()
@@ -1267,6 +1317,12 @@ public void onSuccess(ByteBuffer response) {
partitionId,
nextBatchId);
pushState.onCongestControl(latest.hostAndPushPort());
+ recordPushEvent(
+ shuffleId,
+ mapId,
+ attemptId,
+ latest.hostAndPushPort(),
+ PushEvent.PRIMARY_CONGESTED);
pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
callback.onSuccess(response);
} else if (reason == StatusCode.PUSH_DATA_SUCCESS_REPLICA_CONGESTED.getValue()) {
@@ -1279,6 +1335,12 @@ public void onSuccess(ByteBuffer response) {
partitionId,
nextBatchId);
pushState.onCongestControl(latest.hostAndPushPort());
+ recordPushEvent(
+ shuffleId,
+ mapId,
+ attemptId,
+ latest.hostAndPushPort(),
+ PushEvent.REPLICA_CONGESTED);
pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
callback.onSuccess(response);
} else {
@@ -1527,6 +1589,7 @@ private void doPushMergedData(
String[] hostPortArr = Utils.parseColonSeparatedHostPorts(hostPort, 1);
final String host = hostPortArr[0];
final int port = Integer.parseInt(hostPortArr[1]);
+ final long pushStartTime = System.nanoTime();
int groupedBatchId = pushState.nextBatchId();
int groupedBatchBytesSize = batches.stream().mapToInt(batch -> batch.body.length).sum();
@@ -1609,6 +1672,22 @@ public void onFailure(Throwable e) {
new RpcResponseCallback() {
@Override
public void onSuccess(ByteBuffer response) {
+ long pushRttNanos = System.nanoTime() - pushStartTime;
+ pushState.recordPushRtt(pushRttNanos, slowPushThresholdNanos);
+ recordWorkerPush(
+ shuffleId, mapId, attemptId, hostPort, groupedBatchBytesSize, pushRttNanos);
+ if (pushRttNanos > slowPushThresholdNanos) {
+ logger.warn(
+ "Slow push merged data to {} for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}, rtt {} ms.",
+ addressPair,
+ shuffleId,
+ mapId,
+ attemptId,
+ Arrays.toString(partitionIds),
+ groupedBatchId,
+ Arrays.toString(batchIds),
+ TimeUnit.NANOSECONDS.toMillis(pushRttNanos));
+ }
byte reason = response.get();
if (reason == StatusCode.HARD_SPLIT.getValue()
|| reason == StatusCode.SOFT_SPLIT.getValue()) {
@@ -1650,6 +1729,8 @@ public void onSuccess(ByteBuffer response) {
StatusCode.SOFT_SPLIT);
reviveManager.addRequest(reviveRequest);
}
+ recordPushEvent(
+ shuffleId, mapId, attemptId, loc.hostAndPushPort(), PushEvent.SOFT_SPLIT);
} else {
batchesNeedResubmit.add(batches.get(partitionIndex));
}
@@ -1682,6 +1763,14 @@ public void onSuccess(ByteBuffer response) {
pushState.onSuccess(hostPort);
callback.onSuccess(ByteBuffer.wrap(new byte[] {StatusCode.SOFT_SPLIT.getValue()}));
} else {
+ for (DataBatches.DataBatch resubmitBatch : batchesNeedResubmit) {
+ recordPushEvent(
+ shuffleId,
+ mapId,
+ attemptId,
+ resubmitBatch.loc.hostAndPushPort(),
+ PushEvent.HARD_SPLIT);
+ }
if (dataPushFailureTrackingEnabled && pushReplicateEnabled) {
for (DataBatches.DataBatch resubmitBatch : batchesNeedResubmit) {
pushState.recordFailedBatch(
@@ -1724,6 +1813,7 @@ public void onSuccess(ByteBuffer response) {
groupedBatchId,
Arrays.toString(batchIds));
pushState.onCongestControl(hostPort);
+ recordPushEvent(shuffleId, mapId, attemptId, hostPort, PushEvent.PRIMARY_CONGESTED);
callback.onSuccess(response);
} else if (reason == StatusCode.PUSH_DATA_SUCCESS_REPLICA_CONGESTED.getValue()) {
logger.debug(
@@ -1736,6 +1826,7 @@ public void onSuccess(ByteBuffer response) {
groupedBatchId,
Arrays.toString(batchIds));
pushState.onCongestControl(hostPort);
+ recordPushEvent(shuffleId, mapId, attemptId, hostPort, PushEvent.REPLICA_CONGESTED);
callback.onSuccess(response);
} else if (reason == StatusCode.MAP_ENDED.getValue()) {
pushState.onSuccess(hostPort);
@@ -1896,7 +1987,9 @@ private void mapEndInternal(
numPartitions,
crc32PerPartition,
bytesPerPartition,
- SerdeVersion.V1),
+ SerdeVersion.V1,
+ buildWriteMetricsOption(pushState),
+ buildPushWorkerStatsList(shuffleId, mapId, attemptId, pushState)),
rpcMaxRetries,
rpcRetryWait,
ClassTag$.MODULE$.apply(MapperEndResponse.class));
@@ -1904,10 +1997,104 @@ private void mapEndInternal(
throw new CelebornIOException("MapperEnd failed! StatusCode: " + response.status());
}
} finally {
+ logAndClearPushStats(mapKey, shuffleId, mapId, attemptId, pushState);
pushStates.remove(mapKey);
}
}
+ /** Build the write-path timing breakdown for the UI, or None when UI is disabled. */
+ private Option buildWriteMetricsOption(PushState pushState) {
+ if (!conf.clientSparkUIEnabled()) {
+ return Option.empty();
+ }
+ WriteMetrics wm =
+ new WriteMetrics(
+ pushState.getCopyTimeMs(),
+ pushState.getSerializeTimeMs(),
+ pushState.getCompressTimeMs(),
+ pushState.getQueueWaitTimeMs(),
+ pushState.getQueueStallTimeMs(),
+ pushState.getInflightWaitTimeMs(),
+ pushState.getDrainWaitTimeMs(),
+ pushState.getSlowPushCount(),
+ pushState.getMaxPushRttMs(),
+ pushState.getUncompressedBytes());
+ return Option.apply(wm);
+ }
+
+ /** Collect per-worker push stats for the UI, or empty list when UI is disabled. */
+ private java.util.List
+ buildPushWorkerStatsList(int shuffleId, int mapId, int attemptId, PushState pushState) {
+ java.util.List list =
+ new java.util.ArrayList<>();
+ if (!conf.clientSparkUIEnabled()) {
+ return list;
+ }
+ String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ ConcurrentHashMap workers = pushEventStats.get(mapKey);
+ if (workers != null) {
+ String lastFailureReason =
+ pushState.exception.get() != null ? pushState.exception.get().getMessage() : "";
+ workers.forEach(
+ (hostAndPushPort, stats) -> {
+ org.apache.celeborn.common.protocol.message.PushWorkerStats s =
+ new org.apache.celeborn.common.protocol.message.PushWorkerStats(
+ hostAndPushPort,
+ stats.pushCount.sum(),
+ stats.pushBytes.sum(),
+ stats.totalPushRttNanos.sum(),
+ stats.softSplitCount.sum(),
+ stats.hardSplitCount.sum(),
+ stats.primaryCongestedCount.sum(),
+ stats.replicaCongestedCount.sum(),
+ lastFailureReason);
+ list.add(s);
+ });
+ }
+ return list;
+ }
+
+ @Override
+ public void reportReadMetrics(int shuffleId, ReadStreamStats stats) {
+ if (!conf.clientSparkUIEnabled() || stats == null) {
+ return;
+ }
+ try {
+ ReadMetrics readMetrics =
+ new ReadMetrics(
+ stats.getDecompressTimeMs(),
+ stats.getChunkWaitTimeMs(),
+ stats.getDeserializeTimeMs(),
+ stats.getCopyTimeMs(),
+ stats.getRetryCount(),
+ stats.getRetryWaitTimeMs(),
+ stats.getPeerSwitchCount(),
+ stats.getExcludeCount(),
+ stats.getSlowChunkCount(),
+ stats.getMaxChunkRttMs());
+ java.util.List workerReadCosts = new java.util.ArrayList<>();
+ stats
+ .getWorkerReadCosts()
+ .forEach(
+ (workerId, cost) -> {
+ workerReadCosts.add(
+ new WorkerReadCost(
+ workerId,
+ cost.chunkCount.sum(),
+ cost.bytes.sum(),
+ cost.totalRttNanos.sum(),
+ cost.maxRttNanos.get()));
+ });
+ lifecycleManagerRef.askSync(
+ new ReportShuffleReadMetrics(shuffleId, readMetrics, workerReadCosts, SerdeVersion.V1),
+ rpcMaxRetries,
+ rpcRetryWait,
+ ClassTag$.MODULE$.apply(ReportShuffleReadMetricsResponse.class));
+ } catch (Exception e) {
+ logger.warn("Failed to report read metrics for shuffle " + shuffleId, e);
+ }
+ }
+
@Override
public void cleanup(int shuffleId, int mapId, int attemptId) {
final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
@@ -1916,6 +2103,7 @@ public void cleanup(int shuffleId, int mapId, int attemptId) {
pushState.exception.compareAndSet(null, new CelebornIOException("Cleaned Up"));
pushState.cleanup();
}
+ logAndClearPushStats(mapKey, shuffleId, mapId, attemptId, pushState);
}
@Override
@@ -1927,6 +2115,7 @@ public boolean cleanupShuffle(int shuffleId) {
mapperEndMap.remove(shuffleId);
stageEndShuffleSet.remove(shuffleId);
splitting.remove(shuffleId);
+ pushEventStats.keySet().removeIf(mapKey -> mapKey.startsWith(shuffleId + "-"));
logger.info("Unregistered shuffle {}.", shuffleId);
return true;
@@ -2198,6 +2387,164 @@ public void setupCryptoHandler(Optional cryptoHandler) {
}
}
+ private enum PushEvent {
+ SOFT_SPLIT,
+ HARD_SPLIT,
+ PRIMARY_CONGESTED,
+ REPLICA_CONGESTED
+ }
+
+ // Per (mapTask, worker) push cost and event stats, inspired by Uniffle's
+ // ShuffleServerPushCostTracker, logged when the map task ends.
+ private static class PushWorkerStats {
+ private final LongAdder pushCount = new LongAdder();
+ private final LongAdder pushBytes = new LongAdder();
+ private final LongAdder totalPushRttNanos = new LongAdder();
+ private final AtomicLong maxPushRttNanos = new AtomicLong(0);
+ private final LongAdder softSplitCount = new LongAdder();
+ private final LongAdder hardSplitCount = new LongAdder();
+ private final LongAdder primaryCongestedCount = new LongAdder();
+ private final LongAdder replicaCongestedCount = new LongAdder();
+
+ private long eventCount() {
+ return softSplitCount.sum()
+ + hardSplitCount.sum()
+ + primaryCongestedCount.sum()
+ + replicaCongestedCount.sum();
+ }
+ }
+
+ private void recordPushEvent(
+ int shuffleId, int mapId, int attemptId, String hostAndPushPort, PushEvent event) {
+ PushWorkerStats counters = getPushWorkerStats(shuffleId, mapId, attemptId, hostAndPushPort);
+ switch (event) {
+ case SOFT_SPLIT:
+ counters.softSplitCount.increment();
+ break;
+ case HARD_SPLIT:
+ counters.hardSplitCount.increment();
+ break;
+ case PRIMARY_CONGESTED:
+ counters.primaryCongestedCount.increment();
+ break;
+ case REPLICA_CONGESTED:
+ counters.replicaCongestedCount.increment();
+ break;
+ }
+ }
+
+ private PushWorkerStats getPushWorkerStats(
+ int shuffleId, int mapId, int attemptId, String hostAndPushPort) {
+ return pushEventStats
+ .computeIfAbsent(
+ Utils.makeMapKey(shuffleId, mapId, attemptId), k -> new ConcurrentHashMap<>())
+ .computeIfAbsent(hostAndPushPort, k -> new PushWorkerStats());
+ }
+
+ private void recordWorkerPush(
+ int shuffleId, int mapId, int attemptId, String hostAndPushPort, long bytes, long rttNanos) {
+ PushWorkerStats stats = getPushWorkerStats(shuffleId, mapId, attemptId, hostAndPushPort);
+ stats.pushCount.increment();
+ stats.pushBytes.add(bytes);
+ stats.totalPushRttNanos.add(rttNanos);
+ stats.maxPushRttNanos.accumulateAndGet(rttNanos, Math::max);
+ }
+
+ private void logAndClearPushStats(
+ String mapKey, int shuffleId, int mapId, int attemptId, PushState pushState) {
+ ConcurrentHashMap workerStats = pushEventStats.remove(mapKey);
+ long queueWaitMs = 0;
+ long queueStallMs = 0;
+ long inflightWaitMs = 0;
+ long drainWaitMs = 0;
+ long compressMs = 0;
+ long serializeMs = 0;
+ long slowPushes = 0;
+ long maxPushRttMs = 0;
+ if (pushState != null) {
+ queueWaitMs = pushState.getQueueWaitTimeMs();
+ queueStallMs = pushState.getQueueStallTimeMs();
+ inflightWaitMs = pushState.getInflightWaitTimeMs();
+ drainWaitMs = pushState.getDrainWaitTimeMs();
+ compressMs = pushState.getCompressTimeMs();
+ serializeMs = pushState.getSerializeTimeMs();
+ slowPushes = pushState.getSlowPushCount();
+ maxPushRttMs = pushState.getMaxPushRttMs();
+ }
+ boolean hasEvents =
+ workerStats != null && workerStats.values().stream().anyMatch(s -> s.eventCount() > 0);
+ boolean hasSlowWrite =
+ slowPushes > 0 || queueWaitMs + queueStallMs + inflightWaitMs + drainWaitMs >= 1000;
+ if (!hasEvents && !hasSlowWrite) {
+ return;
+ }
+ // Per worker push cost, sorted by total push round trip time descending.
+ // Only the top workers are logged to bound the log line length for wide shuffles.
+ final int topWorkerLimit = 10;
+ StringBuilder sb = new StringBuilder();
+ if (workerStats != null) {
+ workerStats.entrySet().stream()
+ .sorted(
+ (a, b) ->
+ Long.compare(
+ b.getValue().totalPushRttNanos.sum(), a.getValue().totalPushRttNanos.sum()))
+ .limit(topWorkerLimit)
+ .forEach(
+ e -> {
+ if (sb.length() > 0) {
+ sb.append(", ");
+ }
+ PushWorkerStats s = e.getValue();
+ long pushes = s.pushCount.sum();
+ sb.append("(")
+ .append(e.getKey())
+ .append(", bytes=")
+ .append(s.pushBytes.sum() / 1048576)
+ .append("MB")
+ .append(", pushes=")
+ .append(pushes)
+ .append(", avgRtt=")
+ .append(
+ pushes > 0
+ ? TimeUnit.NANOSECONDS.toMillis(s.totalPushRttNanos.sum()) / pushes
+ : 0)
+ .append("ms")
+ .append(", maxRtt=")
+ .append(TimeUnit.NANOSECONDS.toMillis(s.maxPushRttNanos.get()))
+ .append("ms")
+ .append(", softSplit=")
+ .append(s.softSplitCount.sum())
+ .append(", hardSplit=")
+ .append(s.hardSplitCount.sum())
+ .append(", primaryCongested=")
+ .append(s.primaryCongestedCount.sum())
+ .append(", replicaCongested=")
+ .append(s.replicaCongestedCount.sum())
+ .append(")");
+ });
+ if (workerStats.size() > topWorkerLimit) {
+ sb.append(", ... and ").append(workerStats.size() - topWorkerLimit).append(" more workers");
+ }
+ }
+ logger.info(
+ "Write stats summary for shuffle {} map {} attempt {}: queueWait={}ms, "
+ + "queueStall={}ms, inflightWait={}ms, drainWait={}ms, compress={}ms, "
+ + "serialize={}ms, slowPush(>{}ms)={}, maxPushRtt={}ms, workers=[{}]",
+ shuffleId,
+ mapId,
+ attemptId,
+ queueWaitMs,
+ queueStallMs,
+ inflightWaitMs,
+ drainWaitMs,
+ compressMs,
+ serializeMs,
+ conf.clientPushSlowPushThresholdMs(),
+ slowPushes,
+ maxPushRttMs,
+ sb);
+ }
+
boolean mapperEnded(int shuffleId, int mapId) {
return (mapperEndMap.containsKey(shuffleId) && mapperEndMap.get(shuffleId).contains(mapId))
|| isStageEnded(shuffleId);
@@ -2260,6 +2607,11 @@ public void excludeFailedFetchLocation(String hostAndFetchPort, Exception e) {
&& fetchExcludeWorkerOnFailureEnabled
&& Utils.isCriticalCauseForFetch(e)) {
fetchExcludedWorkers.put(hostAndFetchPort, System.currentTimeMillis());
+ logger.info(
+ "Excluded fetch location {} due to fetch failure, {} excluded fetch locations in total.",
+ hostAndFetchPort,
+ fetchExcludedWorkers.size(),
+ e);
}
}
}
diff --git a/client/src/main/java/org/apache/celeborn/client/listener/MapperEndMetricsCallback.java b/client/src/main/java/org/apache/celeborn/client/listener/MapperEndMetricsCallback.java
new file mode 100644
index 00000000000..139070c1069
--- /dev/null
+++ b/client/src/main/java/org/apache/celeborn/client/listener/MapperEndMetricsCallback.java
@@ -0,0 +1,33 @@
+/*
+ * 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.celeborn.client.listener;
+
+import java.util.List;
+
+import org.apache.celeborn.common.protocol.message.PushWorkerStats;
+import org.apache.celeborn.common.protocol.message.WriteMetrics;
+
+/**
+ * Functional interface for the driver to receive per-task write-path metrics forwarded from
+ * LifecycleManager.handleMapperEnd. Implemented by SparkShuffleManager to post a
+ * CelebornWriteMetricsEvent to the Spark listener bus for the UI.
+ */
+@FunctionalInterface
+public interface MapperEndMetricsCallback {
+ void onMapperEndMetrics(int shuffleId, WriteMetrics writeMetrics, List stats);
+}
diff --git a/client/src/main/java/org/apache/celeborn/client/listener/ReadMetricsCallback.java b/client/src/main/java/org/apache/celeborn/client/listener/ReadMetricsCallback.java
new file mode 100644
index 00000000000..8d8d58db2c5
--- /dev/null
+++ b/client/src/main/java/org/apache/celeborn/client/listener/ReadMetricsCallback.java
@@ -0,0 +1,33 @@
+/*
+ * 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.celeborn.client.listener;
+
+import java.util.List;
+
+import org.apache.celeborn.common.protocol.message.ReadMetrics;
+import org.apache.celeborn.common.protocol.message.WorkerReadCost;
+
+/**
+ * Functional interface for the driver to receive per-task read-path metrics forwarded from
+ * LifecycleManager's ReportShuffleReadMetrics handler. Implemented by SparkShuffleManager to post a
+ * CelebornReadMetricsEvent to the Spark listener bus for the UI.
+ */
+@FunctionalInterface
+public interface ReadMetricsCallback {
+ void onReadMetrics(int shuffleId, ReadMetrics readMetrics, List workerReadCosts);
+}
diff --git a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
index 23d7342e212..bda45747b5e 100644
--- a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
+++ b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
@@ -172,6 +172,15 @@ public int partitionsRead() {
public abstract int partitionsRead();
+ /**
+ * Per-stream read stats, or null for the empty-stream placeholder. Consumers use it to record
+ * read-path time spent outside the input stream itself (e.g. deserialization), so those costs are
+ * attributed to the same stream when the stats are reported on close.
+ */
+ public ReadStreamStats streamStats() {
+ return null;
+ }
+
static boolean shouldSkipLocation(
boolean rangeReadFilter, int startMapIndex, int endMapIndex, PartitionLocation location) {
if (!rangeReadFilter || endMapIndex == Integer.MAX_VALUE) {
@@ -238,6 +247,11 @@ private static final class CelebornInputStreamImpl extends CelebornInputStream {
private MetricsCallback callback;
+ // Per stream read stats, conditionally logged when the stream is closed,
+ // used to diagnose slow shuffle fetch.
+ private final ReadStreamStats streamStats = new ReadStreamStats();
+ private final long createTimeMs = System.currentTimeMillis();
+
// mapId, attemptId, batchId, size
private final int BATCH_HEADER_SIZE = 4 * 4;
private final byte[] sizeBuf = new byte[BATCH_HEADER_SIZE];
@@ -373,7 +387,25 @@ private static final class CelebornInputStreamImpl extends CelebornInputStream {
fetchChunkMaxRetry = conf.clientFetchMaxRetriesForEachReplica();
}
this.retryWaitMs = conf.networkIoRetryWaitMs(TransportModuleConstants.DATA_MODULE);
- this.callback = metricsCallback;
+ // Wrap the callback so that chunk wait time is also recorded into streamStats.
+ this.callback =
+ new MetricsCallback() {
+ @Override
+ public void incBytesRead(long bytesRead) {
+ metricsCallback.incBytesRead(bytesRead);
+ }
+
+ @Override
+ public void incReadTime(long time) {
+ metricsCallback.incReadTime(time);
+ streamStats.addChunkWaitTime(TimeUnit.MILLISECONDS.toNanos(time));
+ }
+
+ @Override
+ public void incDuplicateBytesRead(long bytesRead) {
+ metricsCallback.incDuplicateBytesRead(bytesRead);
+ }
+ };
this.exceptionMaker = exceptionMaker;
this.cryptoHandler = cryptoHandler;
this.partitionId = partitionId;
@@ -394,6 +426,11 @@ private boolean skipLocation(int startMapIndex, int endMapIndex, PartitionLocati
rangeReadFilter, startMapIndex, endMapIndex, location);
}
+ @Override
+ public ReadStreamStats streamStats() {
+ return streamStats;
+ }
+
private Tuple2 nextReadableLocation() {
int locationCount = locations.size();
if (fileIndex >= locationCount) {
@@ -504,11 +541,14 @@ private PartitionReader createReaderWithRetry(
}
lastException = e;
shuffleClient.excludeFailedFetchLocation(location.hostAndFetchPort(), e);
+ streamStats.incExcludeCount();
fetchChunkRetryCnt++;
+ streamStats.incRetryCount();
if (location.hasPeer() && !readSkewPartitionWithoutMapRange) {
// fetchChunkRetryCnt % 2 == 0 means both replicas have been tried,
// so sleep before next try.
if (fetchChunkRetryCnt % 2 == 0) {
+ streamStats.addRetryWaitTime(retryWaitMs);
Uninterruptibles.sleepUninterruptibly(retryWaitMs, TimeUnit.MILLISECONDS);
}
logger.warn(
@@ -542,6 +582,7 @@ private PartitionReader createReaderWithRetry(
}
pbStreamHandler = null;
}
+ streamStats.incPeerSwitchCount();
location = location.getPeer();
} else {
logger.warn(
@@ -550,6 +591,7 @@ private PartitionReader createReaderWithRetry(
fetchChunkMaxRetry,
location,
e);
+ streamStats.addRetryWaitTime(retryWaitMs);
Uninterruptibles.sleepUninterruptibly(retryWaitMs, TimeUnit.MILLISECONDS);
}
}
@@ -580,7 +622,9 @@ private ByteBuf getNextChunk() throws IOException {
}
shuffleClient.excludeFailedFetchLocation(
currentReader.getLocation().hostAndFetchPort(), e);
+ streamStats.incExcludeCount();
fetchChunkRetryCnt++;
+ streamStats.incRetryCount();
currentReader.close();
if (fetchChunkRetryCnt == fetchChunkMaxRetry) {
logger.warn("Fetch chunk fail exceeds max retry {}", fetchChunkRetryCnt, e);
@@ -601,8 +645,10 @@ private ByteBuf getNextChunk() throws IOException {
// fetchChunkRetryCnt % 2 == 0 means both replicas have been tried,
// so sleep before next try.
if (fetchChunkRetryCnt % 2 == 0) {
+ streamStats.addRetryWaitTime(retryWaitMs);
Uninterruptibles.sleepUninterruptibly(retryWaitMs, TimeUnit.MILLISECONDS);
}
+ streamStats.incPeerSwitchCount();
// We must not use checkpoint for peer location since chunkIds don't always match
// across peers
currentReader = createReaderWithRetry(currentReader.getLocation().getPeer(), null);
@@ -613,6 +659,7 @@ private ByteBuf getNextChunk() throws IOException {
fetchChunkMaxRetry,
currentReader.getLocation(),
e);
+ streamStats.addRetryWaitTime(retryWaitMs);
Uninterruptibles.sleepUninterruptibly(retryWaitMs, TimeUnit.MILLISECONDS);
// When reading from the same host again, it is possible to skip already read data
// chunks,
@@ -681,7 +728,8 @@ private PartitionReader createReader(
callback,
startChunkIndex,
endChunkIndex,
- checkpointMetadata);
+ checkpointMetadata,
+ streamStats);
}
case S3:
case OSS:
@@ -744,7 +792,10 @@ public int read(byte[] b, int off, int len) throws IOException {
}
int bytesToRead = Math.min(limit - position, len - readBytes);
+ // Time only the in-memory copy: network wait is already counted as chunk wait time.
+ long copyStart = System.nanoTime();
System.arraycopy(rawDataBuf, position, b, off + readBytes, bytesToRead);
+ streamStats.addCopyTime(System.nanoTime() - copyStart);
position += bytesToRead;
readBytes += bytesToRead;
}
@@ -764,6 +815,7 @@ public synchronized void close() {
locationsCount,
locationsCount - skipCount.sum(),
skipCount.sum());
+ logReadStreamSummaryIfNeeded(locationsCount);
if (currentChunk != null) {
logger.debug("Release chunk {}", currentChunk);
currentChunk.release();
@@ -787,10 +839,82 @@ public synchronized void close() {
decompressor = null;
fetchExcludedWorkers = null;
+ // Report read-path metrics to the driver for the UI (no-op when UI disabled).
+ shuffleClient.reportReadMetrics(shuffleId, streamStats);
+
closed = true;
}
}
+ private void logReadStreamSummaryIfNeeded(int locationsCount) {
+ long chunkWaitMs = streamStats.getChunkWaitTimeMs();
+ if (chunkWaitMs < 1000
+ && streamStats.getRetryCount() == 0
+ && streamStats.getSlowChunkCount() == 0) {
+ return;
+ }
+ // Per worker read cost, sorted by total fetch round trip time descending.
+ // Only the top workers are logged to bound the log line length.
+ final int topWorkerLimit = 10;
+ StringBuilder sb = new StringBuilder();
+ streamStats.getWorkerReadCosts().entrySet().stream()
+ .sorted(
+ (a, b) ->
+ Long.compare(b.getValue().totalRttNanos.sum(), a.getValue().totalRttNanos.sum()))
+ .limit(topWorkerLimit)
+ .forEach(
+ e -> {
+ if (sb.length() > 0) {
+ sb.append(", ");
+ }
+ ReadStreamStats.WorkerReadCost cost = e.getValue();
+ long chunks = cost.chunkCount.sum();
+ sb.append("(")
+ .append(e.getKey())
+ .append(", bytes=")
+ .append(cost.bytes.sum() / 1048576)
+ .append("MB")
+ .append(", chunks=")
+ .append(chunks)
+ .append(", avgRtt=")
+ .append(
+ chunks > 0
+ ? TimeUnit.NANOSECONDS.toMillis(cost.totalRttNanos.sum()) / chunks
+ : 0)
+ .append("ms")
+ .append(", maxRtt=")
+ .append(TimeUnit.NANOSECONDS.toMillis(cost.maxRttNanos.get()))
+ .append("ms")
+ .append(")");
+ });
+ if (streamStats.getWorkerReadCosts().size() > topWorkerLimit) {
+ sb.append(", ... and ")
+ .append(streamStats.getWorkerReadCosts().size() - topWorkerLimit)
+ .append(" more workers");
+ }
+ logger.info(
+ "Read stream summary for appShuffleId {}, shuffleId {}, partitionId {}: "
+ + "total={}ms, chunkWait={}ms, decompress={}ms, retries={}, retryWait={}ms, "
+ + "peerSwitch={}, exclude={}, slowChunks={}, maxChunkRtt={}ms, "
+ + "locations total {}, read {}, skip {}, workers=[{}].",
+ appShuffleId,
+ shuffleId,
+ partitionId,
+ System.currentTimeMillis() - createTimeMs,
+ chunkWaitMs,
+ streamStats.getDecompressTimeMs(),
+ streamStats.getRetryCount(),
+ streamStats.getRetryWaitTimeMs(),
+ streamStats.getPeerSwitchCount(),
+ streamStats.getExcludeCount(),
+ streamStats.getSlowChunkCount(),
+ streamStats.getMaxChunkRttMs(),
+ locationsCount,
+ locationsCount - skipCount.sum(),
+ skipCount.sum(),
+ sb);
+ }
+
void validateIntegrity() throws IOException {
if (integrityChecked || !shuffleIntegrityCheckEnabled) {
return;
@@ -964,11 +1088,13 @@ private boolean fillBuffer() throws IOException {
callback.incBytesRead(BATCH_HEADER_SIZE + encryptedSize);
if (shouldDecompress) {
// decompress data
+ long decompressStartNanos = System.nanoTime();
int originalLength = decompressor.getOriginalLen(compressedBuf);
if (rawDataBuf.length < originalLength) {
rawDataBuf = new byte[originalLength];
}
limit = decompressor.decompress(compressedBuf, rawDataBuf, 0);
+ streamStats.addDecompressTime(System.nanoTime() - decompressStartNanos);
} else {
limit = size;
}
diff --git a/client/src/main/java/org/apache/celeborn/client/read/ReadStreamStats.java b/client/src/main/java/org/apache/celeborn/client/read/ReadStreamStats.java
new file mode 100644
index 00000000000..c15a44d1b29
--- /dev/null
+++ b/client/src/main/java/org/apache/celeborn/client/read/ReadStreamStats.java
@@ -0,0 +1,148 @@
+/*
+ * 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.celeborn.client.read;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.LongAdder;
+
+/**
+ * Per input stream read stats, logged when the stream is closed, used to diagnose slow shuffle
+ * fetch.
+ */
+public class ReadStreamStats {
+
+ /** Per worker read cost, inspired by Uniffle's ShuffleServerReadCostTracker. */
+ public static class WorkerReadCost {
+ public final LongAdder chunkCount = new LongAdder();
+ public final LongAdder bytes = new LongAdder();
+ public final LongAdder totalRttNanos = new LongAdder();
+ public final AtomicLong maxRttNanos = new AtomicLong(0);
+ }
+
+ // Per worker (hostAndFetchPort) read cost of this stream.
+ private final ConcurrentHashMap workerReadCosts =
+ new ConcurrentHashMap<>();
+ // Time the reducer thread is blocked waiting for fetched chunks.
+ private final LongAdder chunkWaitTimeNanos = new LongAdder();
+ // Number of chunks whose fetch round trip time exceeds the slow chunk threshold.
+ private final LongAdder slowChunkCount = new LongAdder();
+ private final AtomicLong maxChunkRttNanos = new AtomicLong(0);
+ private final LongAdder decompressTimeNanos = new LongAdder();
+ // CPU cost of deserializing key/value records (reader iterator next).
+ private final LongAdder deserializeTimeNanos = new LongAdder();
+ // CPU cost of copying decompressed bytes into the user buffer (CelebornInputStream.read).
+ private final LongAdder copyTimeNanos = new LongAdder();
+ private final LongAdder retryCount = new LongAdder();
+ private final LongAdder retryWaitTimeMs = new LongAdder();
+ private final LongAdder peerSwitchCount = new LongAdder();
+ private final LongAdder excludeCount = new LongAdder();
+
+ public void addChunkWaitTime(long nanos) {
+ chunkWaitTimeNanos.add(nanos);
+ }
+
+ public void recordChunkRtt(long rttNanos, long slowThresholdNanos) {
+ maxChunkRttNanos.accumulateAndGet(rttNanos, Math::max);
+ if (rttNanos > slowThresholdNanos) {
+ slowChunkCount.increment();
+ }
+ }
+
+ public void recordWorkerChunkRead(String hostAndFetchPort, int bytes, long rttNanos) {
+ WorkerReadCost cost =
+ workerReadCosts.computeIfAbsent(hostAndFetchPort, k -> new WorkerReadCost());
+ cost.chunkCount.increment();
+ cost.bytes.add(bytes);
+ cost.totalRttNanos.add(rttNanos);
+ cost.maxRttNanos.accumulateAndGet(rttNanos, Math::max);
+ }
+
+ public ConcurrentHashMap getWorkerReadCosts() {
+ return workerReadCosts;
+ }
+
+ public void addDecompressTime(long nanos) {
+ decompressTimeNanos.add(nanos);
+ }
+
+ public void addDeserializeTime(long nanos) {
+ deserializeTimeNanos.add(nanos);
+ }
+
+ public void addCopyTime(long nanos) {
+ copyTimeNanos.add(nanos);
+ }
+
+ public void incRetryCount() {
+ retryCount.increment();
+ }
+
+ public void addRetryWaitTime(long ms) {
+ retryWaitTimeMs.add(ms);
+ }
+
+ public void incPeerSwitchCount() {
+ peerSwitchCount.increment();
+ }
+
+ public void incExcludeCount() {
+ excludeCount.increment();
+ }
+
+ public long getChunkWaitTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(chunkWaitTimeNanos.sum());
+ }
+
+ public long getSlowChunkCount() {
+ return slowChunkCount.sum();
+ }
+
+ public long getMaxChunkRttMs() {
+ return TimeUnit.NANOSECONDS.toMillis(maxChunkRttNanos.get());
+ }
+
+ public long getDecompressTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(decompressTimeNanos.sum());
+ }
+
+ public long getDeserializeTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(deserializeTimeNanos.sum());
+ }
+
+ public long getCopyTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(copyTimeNanos.sum());
+ }
+
+ public long getRetryCount() {
+ return retryCount.sum();
+ }
+
+ public long getRetryWaitTimeMs() {
+ return retryWaitTimeMs.sum();
+ }
+
+ public long getPeerSwitchCount() {
+ return peerSwitchCount.sum();
+ }
+
+ public long getExcludeCount() {
+ return excludeCount.sum();
+ }
+}
diff --git a/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java b/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
index 7a066720364..4b573a40a1c 100644
--- a/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
+++ b/client/src/main/java/org/apache/celeborn/client/read/WorkerPartitionReader.java
@@ -20,6 +20,7 @@
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
@@ -80,6 +81,11 @@ public class WorkerPartitionReader implements PartitionReader {
private Optional partitionReaderCheckpointMetadata;
+ private final ReadStreamStats streamStats;
+ private final long slowChunkThresholdNanos;
+ // chunkIndex -> fetch start time, used to compute per chunk fetch round trip time.
+ private final ConcurrentHashMap chunkFetchStartTimes = new ConcurrentHashMap<>();
+
WorkerPartitionReader(
CelebornConf conf,
String shuffleKey,
@@ -93,7 +99,8 @@ public class WorkerPartitionReader implements PartitionReader {
MetricsCallback metricsCallback,
int startChunkIndex,
int endChunkIndex,
- Optional checkpointMetadata)
+ Optional checkpointMetadata,
+ ReadStreamStats streamStats)
throws IOException, InterruptedException {
this.shuffleKey = shuffleKey;
fetchMaxReqsInFlight = conf.clientFetchMaxReqsInFlight();
@@ -103,6 +110,9 @@ public class WorkerPartitionReader implements PartitionReader {
inflightRequestCount = 0;
this.metricsCallback = metricsCallback;
this.partitionReaderWaitLogThreshold = conf.clientPartitionReaderWaitLogThreshold();
+ this.streamStats = streamStats;
+ this.slowChunkThresholdNanos =
+ TimeUnit.MILLISECONDS.toNanos(conf.clientFetchSlowChunkThresholdMs());
// only add the buffer to results queue if this reader is not closed.
callback =
new ChunkReceivedCallback() {
@@ -111,6 +121,25 @@ public void onSuccess(int chunkIndex, ManagedBuffer buffer) {
// only add the buffer to results queue if this reader is not closed.
synchronized (this) {
ByteBuf buf = ((NettyManagedBuffer) buffer).getBuf();
+ Long fetchStartTime = chunkFetchStartTimes.remove(chunkIndex);
+ if (fetchStartTime != null) {
+ long rttNanos = System.nanoTime() - fetchStartTime;
+ streamStats.recordChunkRtt(rttNanos, slowChunkThresholdNanos);
+ streamStats.recordWorkerChunkRead(
+ location.hostAndFetchPort(), buf.readableBytes(), rttNanos);
+ if (rttNanos > slowChunkThresholdNanos) {
+ logger.warn(
+ "Slow fetch chunk from {} for shuffle key {} stream {} chunk {}, "
+ + "rtt {} ms, chunk size {}, inflight {}.",
+ location.hostAndFetchPort(),
+ shuffleKey,
+ streamHandler.getStreamId(),
+ chunkIndex,
+ TimeUnit.NANOSECONDS.toMillis(rttNanos),
+ buf.readableBytes(),
+ inflightRequestCount);
+ }
+ }
if (!closed) {
buf.retain();
results.add(Pair.of(chunkIndex, buf));
@@ -120,6 +149,7 @@ public void onSuccess(int chunkIndex, ManagedBuffer buffer) {
@Override
public void onFailure(int chunkIndex, Throwable e) {
+ chunkFetchStartTimes.remove(chunkIndex);
String errorMsg =
String.format("Fetch chunk %d of shuffle key %s failed.", chunkIndex, shuffleKey);
logger.error(errorMsg, e);
@@ -309,6 +339,7 @@ private void fetchChunks() throws IOException, InterruptedException {
throw e;
}
}
+ chunkFetchStartTimes.put(chunkIndex, System.nanoTime());
client.fetchChunk(streamHandler.getStreamId(), chunkIndex, fetchTimeoutMs, callback);
inflightRequestCount++;
chunkIndex++;
diff --git a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
index aec6a69b5d4..2205cfe6b51 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
@@ -130,7 +130,13 @@ public ArrayList takePushTasks() throws IOException, InterruptedExcept
}
try {
// Reaching here means no available tasks can be pushed to any worker, wait for a while
+ long sleepStartNanos = System.nanoTime();
Thread.sleep(takeTaskWaitIntervalMs);
+ if (!workingQueue.isEmpty()) {
+ // There are pending tasks but no target worker has remaining push allowance,
+ // i.e. the per-worker in-flight limit is reached.
+ pushState.addQueueStallTime(System.nanoTime() - sleepStartNanos);
+ }
workerWaitAttempts.values().forEach(AtomicInteger::incrementAndGet);
} catch (InterruptedException ie) {
logger.info("Thread interrupted while waiting push task.");
diff --git a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
index 0677502ab2b..f7f9c417c26 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
@@ -154,6 +154,7 @@ public void run() {
public void addTask(int partitionId, byte[] buffer, int size)
throws IOException, InterruptedException {
client.computeBatchCRC(shuffleId, mapId, attemptId, partitionId, buffer, 0, size);
+ long startTimeNanos = System.nanoTime();
try {
PushTask task = null;
while (task == null) {
@@ -170,6 +171,9 @@ public void addTask(int partitionId, byte[] buffer, int size)
logger.error("DataPusher thread interrupted while adding push task.");
pushThread.interrupt();
throw e;
+ } finally {
+ // Record the time the mapper thread is blocked by the push queue backpressure.
+ pushState.addQueueWaitTime(System.nanoTime() - startTimeNanos);
}
}
diff --git a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
index ec843e8f548..58b30a33066 100644
--- a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
@@ -57,6 +57,15 @@ class ChangePartitionManager(
private val inBatchPartitions =
JavaUtils.newConcurrentHashMap[Int, ConcurrentHashMap.KeySetView[Int, java.lang.Boolean]]()
+ // shuffleId -> (partitionId -> revive times, i.e. the max epoch ever allocated,
+ // as the epoch starts from 0 and increases by 1 for each successful revive)
+ private val partitionReviveCounts =
+ JavaUtils.newConcurrentHashMap[Int, ConcurrentHashMap[Integer, Integer]]()
+
+ // shuffleId -> (revive cause -> times)
+ private val reviveCauseCounts =
+ JavaUtils.newConcurrentHashMap[Int, ConcurrentHashMap[StatusCode, java.lang.Long]]()
+
private val batchHandleChangePartitionEnabled = conf.batchHandleChangePartitionEnabled
private val batchHandleChangePartitionExecutors = ThreadUtils.newDaemonCachedThreadPool(
"celeborn-client-lifecycle-manager-change-partition-executor",
@@ -166,6 +175,29 @@ class ChangePartitionManager(
}
}
+ private val reviveCountsRegisterFunc =
+ new util.function.Function[Int, ConcurrentHashMap[Integer, Integer]]() {
+ override def apply(s: Int): ConcurrentHashMap[Integer, Integer] =
+ JavaUtils.newConcurrentHashMap()
+ }
+
+ private val reviveCausesRegisterFunc =
+ new util.function.Function[Int, ConcurrentHashMap[StatusCode, java.lang.Long]]() {
+ override def apply(s: Int): ConcurrentHashMap[StatusCode, java.lang.Long] =
+ JavaUtils.newConcurrentHashMap()
+ }
+
+ private val maxEpochMergeFunc = new util.function.BiFunction[Integer, Integer, Integer] {
+ override def apply(oldEpoch: Integer, newEpoch: Integer): Integer =
+ math.max(oldEpoch, newEpoch)
+ }
+
+ private val countMergeFunc =
+ new util.function.BiFunction[java.lang.Long, java.lang.Long, java.lang.Long] {
+ override def apply(oldCount: java.lang.Long, newCount: java.lang.Long): java.lang.Long =
+ oldCount + newCount
+ }
+
def handleRequestPartitionLocation(
context: RequestLocationCallContext,
shuffleId: Int,
@@ -343,8 +375,21 @@ class ChangePartitionManager(
}
if (newPrimaryLocations.nonEmpty) {
- val changes = newPrimaryLocations.map { partition =>
- s"(partition ${partition.getId} epoch from ${partition.getEpoch - 1} to ${partition.getEpoch})"
+ // newPrimaryLocations may contain both the primary and the replica peer of one
+ // partition, dedupe by partition id before recording stats and logging.
+ val distinctPartitions = newPrimaryLocations.groupBy(_.getId).map(_._2.head)
+ val requestCauses = changePartitions.map(c => c.partitionId -> c.causes).toMap
+ val reviveCounts =
+ partitionReviveCounts.computeIfAbsent(shuffleId, reviveCountsRegisterFunc)
+ val causeCounts = reviveCauseCounts.computeIfAbsent(shuffleId, reviveCausesRegisterFunc)
+ val changes = distinctPartitions.map { partition =>
+ val partitionId = partition.getId
+ // The epoch of the new location equals the revive times of the partition.
+ reviveCounts.merge(partitionId, partition.getEpoch, maxEpochMergeFunc)
+ val cause = requestCauses.get(partitionId).flatten
+ cause.foreach(c => causeCounts.merge(c, 1L, countMergeFunc))
+ s"(partition $partitionId epoch from ${partition.getEpoch - 1} to ${partition.getEpoch}" +
+ s", cause ${cause.map(_.name()).getOrElse("NONE")})"
}.mkString("[", ", ", "]")
logInfo(s"[Update partition] success for " +
s"shuffle $shuffleId, succeed partitions: " +
@@ -390,9 +435,36 @@ class ChangePartitionManager(
slots
}
+ def logReviveSummary(shuffleId: Int): Unit = {
+ val reviveCounts = partitionReviveCounts.get(shuffleId)
+ if (reviveCounts == null || reviveCounts.isEmpty) {
+ return
+ }
+
+ val totalReviveTimes = reviveCounts.values().asScala.map(_.toInt).sum
+ val causeCounts = reviveCauseCounts.get(shuffleId)
+ val causes =
+ if (causeCounts == null || causeCounts.isEmpty) {
+ "NONE"
+ } else {
+ causeCounts.asScala.map { case (cause, count) => s"${cause.name()}=$count" }
+ .mkString(", ")
+ }
+ val topRevivedPartitions = reviveCounts.asScala.toSeq
+ .sortBy(-_._2.toInt)
+ .take(20)
+ .map { case (partitionId, times) => s"(partition $partitionId, revive times $times)" }
+ .mkString("[", ", ", "]")
+ logInfo(s"Shuffle $shuffleId partition revive summary: total revive times " +
+ s"$totalReviveTimes, revived partition num ${reviveCounts.size()}, " +
+ s"causes: [$causes], top revived partitions: $topRevivedPartitions.")
+ }
+
def removeExpiredShuffle(shuffleId: Int): Unit = {
changePartitionRequests.remove(shuffleId)
inBatchPartitions.remove(shuffleId)
locks.remove(shuffleId)
+ partitionReviveCounts.remove(shuffleId)
+ reviveCauseCounts.remove(shuffleId)
}
}
diff --git a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 8e90ddd9ad7..5b82b5817ef 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -38,7 +38,7 @@ import com.google.common.annotations.VisibleForTesting
import com.google.common.cache.{Cache, CacheBuilder}
import org.apache.celeborn.client.LifecycleManager.{ShuffleAllocatedWorkers, ShuffleFailedWorkers}
-import org.apache.celeborn.client.listener.WorkerStatusListener
+import org.apache.celeborn.client.listener.{MapperEndMetricsCallback, ReadMetricsCallback, WorkerStatusListener}
import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
import org.apache.celeborn.common.CelebornConf.ACTIVE_STORAGE_TYPES
import org.apache.celeborn.common.client.{ApplicationInfoProvider, MasterClient}
@@ -50,6 +50,7 @@ import org.apache.celeborn.common.network.protocol.{SerdeVersion, TransportMessa
import org.apache.celeborn.common.network.sasl.registration.RegistrationInfo
import org.apache.celeborn.common.protocol._
import org.apache.celeborn.common.protocol.RpcNameConstants.WORKER_EP
+import org.apache.celeborn.common.protocol.message.{PushWorkerStats, WriteMetrics}
import org.apache.celeborn.common.protocol.message.ControlMessages._
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.rpc._
@@ -434,6 +435,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
epoch,
oldPartition,
isSegmentGranularityVisible = commitManager.isSegmentGranularityVisible(shuffleId))
+ maybeFireReassign(reassignPartitionSplitTriggered.compareAndSet(false, true))
case MapperEnd(
shuffleId,
@@ -445,7 +447,9 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
numPartitions,
crc32PerPartition,
bytesWrittenPerPartition,
- serdeVersion) =>
+ serdeVersion,
+ writeMetrics,
+ pushWorkerStats) =>
logTrace(s"Received MapperEnd TaskEnd request, " +
s"${Utils.makeMapKey(shuffleId, mapId, attemptId)}")
val partitionType = getPartitionType(shuffleId)
@@ -461,7 +465,9 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
numPartitions,
crc32PerPartition,
bytesWrittenPerPartition,
- serdeVersion)
+ serdeVersion,
+ writeMetrics,
+ pushWorkerStats)
case PartitionType.MAP =>
handleMapPartitionEnd(
context,
@@ -545,6 +551,14 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
} else {
context.reply(PbSerDeUtils.toPbApplicationMeta(applicationMeta))
}
+
+ case ReportShuffleReadMetrics(shuffleId, readMetrics, workerReadCosts, serdeVersion) =>
+ // Forward read-path metrics to the driver-side UI listener (registered by
+ // SparkShuffleManager). No aggregation here — the listener accumulates per-shuffle.
+ readMetricsCallback.foreach { cb =>
+ cb.onReadMetrics(shuffleId, readMetrics, workerReadCosts)
+ }
+ context.reply(ReportShuffleReadMetricsResponse(StatusCode.SUCCESS, serdeVersion))
}
private def handleReducerPartitionEnd(
@@ -872,6 +886,11 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
isSegmentGranularityVisible,
numPartitions)
+ // Notify the driver-side UI listener that slots are reserved for this shuffle so it can
+ // post a CelebornShuffleAssignmentEvent with the worker topology. No-op if no callback
+ // registered (UI disabled).
+ shuffleAssignmentCallback.foreach { cb => cb.accept(shuffleId, numPartitions) }
+
// Fifth, reply the allocated partition location to ShuffleClient.
logInfo(s"Handle RegisterShuffle Success for $shuffleId.")
val allPrimaryPartitionLocations = slots.asScala.flatMap(_._2._1.asScala).toArray
@@ -903,6 +922,9 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
false)
return
}
+ // A revive request means a mapper's push failed and asked for a new partition location.
+ // Mark reviveTriggered for the UI (deduped to the first trigger).
+ maybeFireReassign(reassignReviveTriggered.compareAndSet(false, true))
logDebug(
s"[handleRevive] shuffle $shuffleId, $mapIds, $partitionIds, $oldEpochs, $oldPartitions, $causes")
if (commitManager.isStageEnd(shuffleId)) {
@@ -945,7 +967,17 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
numPartitions: Int,
crc32PerPartition: Array[Int],
bytesWrittenPerPartition: Array[Long],
- serdeVersion: SerdeVersion): Unit = {
+ serdeVersion: SerdeVersion,
+ writeMetrics: Option[WriteMetrics],
+ pushWorkerStats: util.List[PushWorkerStats]): Unit = {
+ // Forward write-path timing breakdown + per-worker push stats to the driver-side UI
+ // listener (registered by SparkShuffleManager). No aggregation here — the listener
+ // accumulates per-shuffle, matching the assignment/fallback event pattern.
+ writeMetrics.foreach { w =>
+ mapperEndMetricsCallback.foreach { cb =>
+ cb.onMapperEndMetrics(shuffleId, w, pushWorkerStats)
+ }
+ }
val (mapperAttemptFinishedSuccess, allMapperFinished) =
commitManager.finishMapperAttempt(
@@ -1220,6 +1252,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
return
}
+ changePartitionManager.logReviveSummary(shuffleId)
if (commitManager.tryFinalCommit(shuffleId)) {
// Here we only clear PartitionLocation info in shuffleAllocatedWorkers.
// Since rerun or speculation task may running after we handle StageEnd.
@@ -2126,6 +2159,62 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
appShuffleTrackerCallback = Some(callback)
}
+ // Fired after a successful slot reservation in handleRegisterShuffle, carrying the
+ // celeborn shuffle id and numPartitions. The driver-side SparkShuffleManager registers a
+ // callback here to post a CelebornShuffleAssignmentEvent (with the allocated worker ids)
+ // to the Spark listener bus for the UI.
+ @volatile private var shuffleAssignmentCallback
+ : Option[BiConsumer[java.lang.Integer, java.lang.Integer]] = None
+ def registerShuffleAssignmentCallback(
+ callback: BiConsumer[java.lang.Integer, java.lang.Integer]): Unit = {
+ shuffleAssignmentCallback = Some(callback)
+ }
+
+ // Reassign dedup flags: each type is posted only on its first trigger (mirrors Uniffle's
+ // postReassignTriggeredEvent AtomicBoolean.compareAndSet). The driver-side SparkShuffleManager
+ // registers a callback that posts a CelebornReassignEvent carrying the current 3-boolean state.
+ private val reassignPartitionSplitTriggered = new java.util.concurrent.atomic.AtomicBoolean(false)
+ private val reassignReviveTriggered =
+ new java.util.concurrent.atomic.AtomicBoolean(false)
+ private val reassignStageRetryTriggered = new java.util.concurrent.atomic.AtomicBoolean(false)
+ @volatile private var reassignCallback
+ : Option[java.util.function.Consumer[java.util.List[java.lang.Boolean]]] = None
+ def registerReassignCallback(
+ callback: java.util.function.Consumer[java.util.List[java.lang.Boolean]]): Unit = {
+ reassignCallback = Some(callback)
+ }
+
+ // Fired from handleMapperEnd when the executor populated write metrics (UI enabled), carrying
+ // the shuffleId + write-path timing breakdown + per-worker push stats. The driver-side
+ // SparkShuffleManager registers a MapperEndMetricsCallback that posts a
+ // CelebornWriteMetricsEvent to the Spark listener bus.
+ @volatile private var mapperEndMetricsCallback: Option[MapperEndMetricsCallback] = None
+ def registerMapperEndMetricsCallback(callback: MapperEndMetricsCallback): Unit = {
+ mapperEndMetricsCallback = Some(callback)
+ }
+
+ // Fired from the ReportShuffleReadMetrics handler, carrying shuffleId + read-path timing
+ // breakdown + per-worker read cost. SparkShuffleManager registers a callback that posts a
+ // CelebornReadMetricsEvent to the Spark listener bus.
+ @volatile private var readMetricsCallback: Option[ReadMetricsCallback] = None
+ def registerReadMetricsCallback(callback: ReadMetricsCallback): Unit = {
+ readMetricsCallback = Some(callback)
+ }
+
+ /**
+ * Fire the reassign callback with the current state if any flag was just flipped.
+ * Called from handleRevive (reviveTriggered) and handlePartitionSplit (partitionSplit).
+ */
+ private def maybeFireReassign(triggered: Boolean): Unit = {
+ if (triggered) {
+ val state = new java.util.ArrayList[java.lang.Boolean]()
+ state.add(reassignPartitionSplitTriggered.get(): java.lang.Boolean)
+ state.add(reassignReviveTriggered.get(): java.lang.Boolean)
+ state.add(reassignStageRetryTriggered.get(): java.lang.Boolean)
+ reassignCallback.foreach(_.accept(state))
+ }
+ }
+
// expecting celeborn shuffle id and application shuffle identifier
@volatile private var validateCelebornShuffleIdForClean: Option[Consumer[String]] =
None
diff --git a/client/src/test/java/org/apache/celeborn/client/read/ReadStreamStatsSuiteJ.java b/client/src/test/java/org/apache/celeborn/client/read/ReadStreamStatsSuiteJ.java
new file mode 100644
index 00000000000..63b23d61938
--- /dev/null
+++ b/client/src/test/java/org/apache/celeborn/client/read/ReadStreamStatsSuiteJ.java
@@ -0,0 +1,61 @@
+/*
+ * 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.celeborn.client.read;
+
+import java.util.concurrent.TimeUnit;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+/** Verifies the read-side timing accumulators surfaced in the Celeborn UI read metrics. */
+public class ReadStreamStatsSuiteJ {
+
+ @Test
+ public void testDeserializeAndCopyTimeAccumulate() {
+ ReadStreamStats stats = new ReadStreamStats();
+ // Deserialization happens per record via TimedDeserializationStream; copies happen per
+ // CelebornInputStream.read call. Both accumulate in nanos and report in millis.
+ stats.addDeserializeTime(TimeUnit.MILLISECONDS.toNanos(2));
+ stats.addDeserializeTime(TimeUnit.MILLISECONDS.toNanos(3));
+ stats.addCopyTime(TimeUnit.MILLISECONDS.toNanos(1));
+ stats.addCopyTime(TimeUnit.MILLISECONDS.toNanos(4));
+ Assert.assertEquals(5L, stats.getDeserializeTimeMs());
+ Assert.assertEquals(5L, stats.getCopyTimeMs());
+ }
+
+ @Test
+ public void testEmptyStatsReportZero() {
+ ReadStreamStats stats = new ReadStreamStats();
+ Assert.assertEquals(0L, stats.getDeserializeTimeMs());
+ Assert.assertEquals(0L, stats.getCopyTimeMs());
+ Assert.assertEquals(0L, stats.getChunkWaitTimeMs());
+ Assert.assertEquals(0L, stats.getMaxChunkRttMs());
+ }
+
+ @Test
+ public void testWorkerReadCostMerge() {
+ ReadStreamStats stats = new ReadStreamStats();
+ stats.recordWorkerChunkRead("host1:9097", 100, TimeUnit.MILLISECONDS.toNanos(5));
+ stats.recordWorkerChunkRead("host1:9097", 200, TimeUnit.MILLISECONDS.toNanos(7));
+ ReadStreamStats.WorkerReadCost cost = stats.getWorkerReadCosts().get("host1:9097");
+ Assert.assertEquals(2L, cost.chunkCount.sum());
+ Assert.assertEquals(300L, cost.bytes.sum());
+ Assert.assertEquals(TimeUnit.MILLISECONDS.toNanos(12), cost.totalRttNanos.sum());
+ Assert.assertEquals(TimeUnit.MILLISECONDS.toNanos(7), cost.maxRttNanos.get());
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
index 613f7784bc2..888bda312ec 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
@@ -132,6 +132,7 @@ public boolean limitMaxInFlight(String hostAndPushPort) throws IOException {
throw pushState.exception.get();
}
+ long startTimeNanos = System.nanoTime();
pushStrategy.limitPushSpeed(pushState, hostAndPushPort);
int currentMaxReqsInFlight = pushStrategy.getCurrentMaxReqsInFlight(hostAndPushPort);
@@ -198,6 +199,7 @@ public boolean limitMaxInFlight(String hostAndPushPort) throws IOException {
throw pushState.exception.get();
}
+ pushState.addInflightWaitTime(System.nanoTime() - startTimeNanos);
return times <= 0;
}
@@ -205,6 +207,7 @@ public boolean limitZeroInFlight() throws IOException {
if (pushState.exception.get() != null) {
throw pushState.exception.get();
}
+ long startTimeNanos = System.nanoTime();
long times = waitInflightTimeoutMs / delta;
try {
@@ -244,6 +247,7 @@ public boolean limitZeroInFlight() throws IOException {
throw pushState.exception.get();
}
+ pushState.addDrainWaitTime(System.nanoTime() - startTimeNanos);
return times <= 0;
}
diff --git a/common/src/main/java/org/apache/celeborn/common/write/PushState.java b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
index e213bf6a167..f76b043dfb5 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
@@ -21,7 +21,10 @@
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.LongAdder;
import org.apache.commons.lang3.tuple.Pair;
@@ -41,12 +44,115 @@ public class PushState {
private final Map failedBatchMap;
+ // Per map task write path stats, used to diagnose slow shuffle write.
+ // Time a mapper thread is blocked by the push queue backpressure (DataPusher.addTask).
+ private final LongAdder queueWaitTimeNanos = new LongAdder();
+ // Time push threads are blocked by the in-flight limit (limitMaxInFlight).
+ private final LongAdder inflightWaitTimeNanos = new LongAdder();
+ // Time mapperEnd waits for all in-flight batches to be done (limitZeroInFlight).
+ private final LongAdder drainWaitTimeNanos = new LongAdder();
+ // Number of push batches whose round trip time exceeds the slow push threshold.
+ private final LongAdder slowPushCount = new LongAdder();
+ private final AtomicLong maxPushRttNanos = new AtomicLong(0);
+
public PushState(CelebornConf conf) {
pushBufferMaxSize = conf.clientPushBufferMaxSize();
inFlightRequestTracker = new InFlightRequestTracker(conf, this);
failedBatchMap = JavaUtils.newConcurrentHashMap();
}
+ // Time push thread stalls in DataPushQueue.takePushTasks because no target worker
+ // has remaining push allowance (per-worker in-flight limit reached).
+ private final LongAdder queueStallTimeNanos = new LongAdder();
+
+ // CPU cost of compressing push data (ShuffleClientImpl.pushOrMergeData).
+ private final LongAdder compressTimeNanos = new LongAdder();
+ // CPU cost of serializing records in the engine shuffle writer.
+ private final LongAdder serializeTimeNanos = new LongAdder();
+ // CPU cost of copying serialized records into the push buffer (writer write0 / fastWrite0).
+ private final LongAdder copyTimeNanos = new LongAdder();
+ // Total serialized (pre-compression) bytes written, for the UI compression ratio.
+ private final LongAdder uncompressedBytes = new LongAdder();
+
+ public void addCompressTime(long nanos) {
+ compressTimeNanos.add(nanos);
+ }
+
+ public void addSerializeTime(long nanos) {
+ serializeTimeNanos.add(nanos);
+ }
+
+ public void addCopyTime(long nanos) {
+ copyTimeNanos.add(nanos);
+ }
+
+ public void addUncompressedBytes(long bytes) {
+ uncompressedBytes.add(bytes);
+ }
+
+ public long getCompressTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(compressTimeNanos.sum());
+ }
+
+ public long getSerializeTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(serializeTimeNanos.sum());
+ }
+
+ public long getCopyTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(copyTimeNanos.sum());
+ }
+
+ public long getUncompressedBytes() {
+ return uncompressedBytes.sum();
+ }
+
+ public void addQueueStallTime(long nanos) {
+ queueStallTimeNanos.add(nanos);
+ }
+
+ public long getQueueStallTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(queueStallTimeNanos.sum());
+ }
+
+ public void addQueueWaitTime(long nanos) {
+ queueWaitTimeNanos.add(nanos);
+ }
+
+ public void addInflightWaitTime(long nanos) {
+ inflightWaitTimeNanos.add(nanos);
+ }
+
+ public void addDrainWaitTime(long nanos) {
+ drainWaitTimeNanos.add(nanos);
+ }
+
+ public void recordPushRtt(long rttNanos, long slowThresholdNanos) {
+ maxPushRttNanos.accumulateAndGet(rttNanos, Math::max);
+ if (rttNanos > slowThresholdNanos) {
+ slowPushCount.increment();
+ }
+ }
+
+ public long getQueueWaitTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(queueWaitTimeNanos.sum());
+ }
+
+ public long getInflightWaitTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(inflightWaitTimeNanos.sum());
+ }
+
+ public long getDrainWaitTimeMs() {
+ return TimeUnit.NANOSECONDS.toMillis(drainWaitTimeNanos.sum());
+ }
+
+ public long getSlowPushCount() {
+ return slowPushCount.sum();
+ }
+
+ public long getMaxPushRttMs() {
+ return TimeUnit.NANOSECONDS.toMillis(maxPushRttNanos.get());
+ }
+
public void cleanup() {
inFlightRequestTracker.cleanup();
}
diff --git a/common/src/main/proto/TransportMessages.proto b/common/src/main/proto/TransportMessages.proto
index e57055391e6..d53c6fa9a97 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -119,6 +119,8 @@ enum MessageType {
REGISTER_APPLICATION_INFO = 96;
REQUEST_WORKERS = 97;
REQUEST_WORKERS_RESPONSE = 98;
+ REPORT_SHUFFLE_READ_METRICS = 99;
+ REPORT_SHUFFLE_READ_METRICS_RESPONSE = 100;
}
enum StreamType {
@@ -394,6 +396,32 @@ message PbMapperEnd {
int32 numPartitions = 7;
repeated int32 crc32PerPartition = 8;
repeated int64 bytesWrittenPerPartition = 9;
+ // Write-path timing breakdown (ms) for the UI. Optional — populated only when
+ // celeborn.client.spark.ui.enabled. Field numbers 10-19 reserved for write times.
+ int64 copyTimeMs = 10;
+ int64 serializeTimeMs = 11;
+ int64 compressTimeMs = 12;
+ int64 queueWaitTimeMs = 13;
+ int64 queueStallTimeMs = 14;
+ int64 inflightWaitTimeMs = 15;
+ int64 drainWaitTimeMs = 16;
+ int64 slowPushCount = 17;
+ int64 maxPushRttMs = 18;
+ repeated PbPushWorkerStats pushWorkerStats = 19;
+ int64 uncompressedBytes = 20;
+}
+
+// Per-(shuffle,mapId,attemptId,worker) push cost for the UI's Shuffle Servers section.
+message PbPushWorkerStats {
+ string workerId = 1;
+ int64 pushCount = 2;
+ int64 pushBytes = 3;
+ int64 totalPushRttNanos = 4;
+ int64 softSplitCount = 5;
+ int64 hardSplitCount = 6;
+ int64 primaryCongestedCount = 7;
+ int64 replicaCongestedCount = 8;
+ string lastPushFailureReason = 9;
}
message PbLocationPushFailedBatches {
@@ -1090,3 +1118,30 @@ message PbRegisterApplicationInfo {
map extraInfo = 3;
string requestId = 4;
}
+
+message PbWorkerReadCost {
+ string workerId = 1;
+ int64 chunkCount = 2;
+ int64 bytes = 3;
+ int64 totalRttNanos = 4;
+ int64 maxRttNanos = 5;
+}
+
+message PbReportShuffleReadMetrics {
+ int32 shuffleId = 1;
+ int64 decompressTimeMs = 2;
+ int64 chunkWaitTimeMs = 3;
+ int64 deserializeTimeMs = 4;
+ int64 copyTimeMs = 5;
+ int64 retryCount = 6;
+ int64 retryWaitTimeMs = 7;
+ int64 peerSwitchCount = 8;
+ int64 excludeCount = 9;
+ int64 slowChunkCount = 10;
+ int64 maxChunkRttMs = 11;
+ repeated PbWorkerReadCost workerReadCosts = 12;
+}
+
+message PbReportShuffleReadMetricsResponse {
+ int32 status = 1;
+}
diff --git a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 22d1a066dbe..381189ba113 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1038,6 +1038,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
// //////////////////////////////////////////////////////
def clientFetchTimeoutMs: Long = get(CLIENT_FETCH_TIMEOUT)
def clientFetchPollChunkWaitTime: Long = get(CLIENT_FETCH_POLL_CHUNK_WAIT_TIME)
+ def clientFetchSlowChunkThresholdMs: Long = get(CLIENT_FETCH_SLOW_CHUNK_THRESHOLD)
def clientFetchBufferSize: Int = get(CLIENT_FETCH_BUFFER_SIZE).toInt
def clientFetchMaxReqsInFlight: Int = get(CLIENT_FETCH_MAX_REQS_IN_FLIGHT)
def isPartitionReaderCheckpointEnabled: Boolean =
@@ -1047,6 +1048,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
def clientFetchMaxRetriesForEachReplica: Int = get(CLIENT_FETCH_MAX_RETRIES_FOR_EACH_REPLICA)
def clientStageRerunEnabled: Boolean = get(CLIENT_STAGE_RERUN_ENABLED)
+ def clientSparkUIEnabled: Boolean = get(CLIENT_SPARK_UI_ENABLED)
+ def clientSparkUIRetainedShuffles: Int = get(CLIENT_SPARK_UI_RETAINED_SHUFFLES)
def clientFetchCleanFailedShuffle: Boolean = get(CLIENT_FETCH_CLEAN_FAILED_SHUFFLE)
def clientFetchCleanFailedShuffleIntervalMS: Long =
get(CLIENT_FETCH_CLEAN_FAILED_SHUFFLE_INTERVAL)
@@ -1120,6 +1123,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
pushDataTimeoutMs * clientPushMaxReviveTimes * 2)
}
def clientPushLimitInFlightSleepDeltaMs: Long = get(CLIENT_PUSH_LIMIT_IN_FLIGHT_SLEEP_INTERVAL)
+ def clientPushSlowPushThresholdMs: Long = get(CLIENT_PUSH_SLOW_PUSH_THRESHOLD)
def clientPushTakeTaskWaitIntervalMs: Long = get(CLIENT_PUSH_TAKE_TASK_WAIT_INTERVAL)
def clientPushTakeTaskMaxWaitAttempts: Int = get(CLIENT_PUSH_TAKE_TASK_MAX_WAIT_ATTEMPTS)
def clientPushSendBufferPoolExpireTimeout: Long = get(CLIENT_PUSH_SENDBUFFERPOOL_EXPIRETIMEOUT)
@@ -5178,6 +5182,16 @@ object CelebornConf extends Logging {
.timeConf(TimeUnit.MILLISECONDS)
.createWithDefaultString("50ms")
+ val CLIENT_PUSH_SLOW_PUSH_THRESHOLD: ConfigEntry[Long] =
+ buildConf("celeborn.client.push.slowPush.threshold")
+ .categories("client")
+ .doc("Threshold of the push data round trip time. If pushing a batch to a worker " +
+ "takes longer than this threshold, a warn log will be recorded with the target " +
+ "worker, partition and batch info.")
+ .version("0.7.0")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("5s")
+
val CLIENT_PUSH_SORT_RANDOMIZE_PARTITION_ENABLED: ConfigEntry[Boolean] =
buildConf("celeborn.client.push.sort.randomizePartitionId.enabled")
.withAlternative("celeborn.push.sort.randomizePartitionId.enabled")
@@ -5295,6 +5309,16 @@ object CelebornConf extends Logging {
.timeConf(TimeUnit.MILLISECONDS)
.createWithDefault(500)
+ val CLIENT_FETCH_SLOW_CHUNK_THRESHOLD: ConfigEntry[Long] =
+ buildConf("celeborn.client.fetch.slowChunk.threshold")
+ .categories("client")
+ .version("0.7.0")
+ .doc("Threshold of the fetch chunk round trip time. If fetching a chunk from a worker " +
+ "takes longer than this threshold, a warn log will be recorded with the worker, " +
+ "stream and chunk info.")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("1s")
+
val CLIENT_FETCH_BUFFER_SIZE: ConfigEntry[Long] =
buildConf("celeborn.client.fetch.buffer.size")
.categories("client")
@@ -5351,6 +5375,32 @@ object CelebornConf extends Logging {
.booleanConf
.createWithDefault(true)
+ val CLIENT_SPARK_UI_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.client.spark.ui.enabled")
+ .categories("client")
+ .version("1.0.0")
+ .doc("Whether to enable the Celeborn Spark UI extension. When true and the " +
+ "CelebornPlugin is registered via spark.plugins, a Celeborn tab is attached to the " +
+ "Spark WebUI aggregating shuffle assignment, fallback stats and task-level shuffle " +
+ "metrics. Shuffle assignment and fallback events are posted to the Spark listener bus " +
+ "only when this is enabled, so it is zero-overhead when off.")
+ .booleanConf
+ .createWithDefault(true)
+
+ val CLIENT_SPARK_UI_RETAINED_SHUFFLES: ConfigEntry[Int] =
+ buildConf("celeborn.client.spark.ui.retainedShuffles")
+ .withAlternative("celeborn.client.spark.ui.retainedShuffle")
+ .categories("client")
+ .version("1.0.0")
+ .doc("Number of Celeborn shuffle assignments to retain in the Spark status KVStore " +
+ "(and thus the event log) for the Celeborn UI tab. When more shuffles are registered, " +
+ "the oldest assignment rows are evicted to bound memory/event-log growth on long-running " +
+ "jobs with many shuffles. Mirrors Spark's spark.ui.retainedStages and Gluten's " +
+ "UI_RETAINED_EXECUTIONS trigger. Only the per-shuffle assignment rows are capped; " +
+ "singleton entities (build info, fallback stats, aggregated metrics) are not affected.")
+ .intConf
+ .createWithDefault(1000)
+
val CLIENT_FETCH_CLEAN_FAILED_SHUFFLE: ConfigEntry[Boolean] =
buildConf("celeborn.client.spark.fetch.cleanFailedShuffle")
.categories("client")
diff --git a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
index 580c986ae61..75f945eca33 100644
--- a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -52,6 +52,54 @@ sealed trait WorkerMessage extends Message
sealed trait ClientMessage extends Message
+// Write-path timing breakdown carried in MapperEnd for the UI (all in ms). Populated only
+// when celeborn.client.spark.ui.enabled.
+case class WriteMetrics(
+ copyTimeMs: Long,
+ serializeTimeMs: Long,
+ compressTimeMs: Long,
+ queueWaitTimeMs: Long,
+ queueStallTimeMs: Long,
+ inflightWaitTimeMs: Long,
+ drainWaitTimeMs: Long,
+ slowPushCount: Long,
+ maxPushRttMs: Long,
+ uncompressedBytes: Long)
+
+// Per-(shuffle,mapId,attemptId,worker) push cost carried in MapperEnd for the UI.
+case class PushWorkerStats(
+ workerId: String,
+ pushCount: Long,
+ pushBytes: Long,
+ totalPushRttNanos: Long,
+ softSplitCount: Long,
+ hardSplitCount: Long,
+ primaryCongestedCount: Long,
+ replicaCongestedCount: Long,
+ lastPushFailureReason: String)
+
+// Per-worker read cost carried in ReportShuffleReadMetrics for the UI.
+case class WorkerReadCost(
+ workerId: String,
+ chunkCount: Long,
+ bytes: Long,
+ totalRttNanos: Long,
+ maxRttNanos: Long)
+
+// Read-path timing breakdown + per-worker read cost for the UI. Sent from the executor
+// (CelebornInputStream close) via a new RPC when celeborn.client.spark.ui.enabled.
+case class ReadMetrics(
+ decompressTimeMs: Long,
+ chunkWaitTimeMs: Long,
+ deserializeTimeMs: Long,
+ copyTimeMs: Long,
+ retryCount: Long,
+ retryWaitTimeMs: Long,
+ peerSwitchCount: Long,
+ excludeCount: Long,
+ slowChunkCount: Long,
+ maxChunkRttMs: Long)
+
object ControlMessages extends Logging {
val ZERO_UUID = new UUID(0L, 0L).toString
@@ -218,7 +266,9 @@ object ControlMessages extends Logging {
numPartitions: Int,
crc32PerPartition: Array[Int],
bytesWrittenPerPartition: Array[Long],
- serdeVersion: SerdeVersion)
+ serdeVersion: SerdeVersion,
+ writeMetrics: Option[WriteMetrics] = None,
+ pushWorkerStats: util.List[PushWorkerStats] = util.Collections.emptyList())
extends MasterMessage
case class ReadReducerPartitionEnd(
@@ -232,6 +282,17 @@ object ControlMessages extends Logging {
case class MapperEndResponse(status: StatusCode, serdeVersion: SerdeVersion) extends MasterMessage
+ // Read-path metrics reported from the executor to the driver (new RPC) for the UI.
+ case class ReportShuffleReadMetrics(
+ shuffleId: Int,
+ readMetrics: ReadMetrics,
+ workerReadCosts: util.List[WorkerReadCost],
+ serdeVersion: SerdeVersion)
+ extends MasterMessage
+
+ case class ReportShuffleReadMetricsResponse(status: StatusCode, serdeVersion: SerdeVersion)
+ extends MasterMessage
+
case class ReadReducerPartitionEndResponse(status: StatusCode) extends MasterMessage
case class GetReducerFileGroup(
@@ -743,12 +804,14 @@ object ControlMessages extends Logging {
numPartitions,
crc32PerPartition,
bytesWrittenPerPartition,
- serdeVersion) =>
+ serdeVersion,
+ writeMetrics,
+ pushWorkerStats) =>
val pushFailedMap = pushFailedBatch.asScala.map { case (k, v) =>
val resultValue = PbSerDeUtils.toPbLocationPushFailedBatches(v)
(k, resultValue)
}.toMap.asJava
- val payload = PbMapperEnd.newBuilder()
+ val builder = PbMapperEnd.newBuilder()
.setShuffleId(shuffleId)
.setMapId(mapId)
.setAttemptId(attemptId)
@@ -759,7 +822,33 @@ object ControlMessages extends Logging {
.addAllCrc32PerPartition(crc32PerPartition.map(Integer.valueOf).toSeq.asJava)
.addAllBytesWrittenPerPartition(bytesWrittenPerPartition.map(
java.lang.Long.valueOf).toSeq.asJava)
- .build().toByteArray
+ writeMetrics.foreach { w =>
+ builder.setCopyTimeMs(w.copyTimeMs)
+ .setSerializeTimeMs(w.serializeTimeMs)
+ .setCompressTimeMs(w.compressTimeMs)
+ .setQueueWaitTimeMs(w.queueWaitTimeMs)
+ .setQueueStallTimeMs(w.queueStallTimeMs)
+ .setInflightWaitTimeMs(w.inflightWaitTimeMs)
+ .setDrainWaitTimeMs(w.drainWaitTimeMs)
+ .setSlowPushCount(w.slowPushCount)
+ .setMaxPushRttMs(w.maxPushRttMs)
+ .setUncompressedBytes(w.uncompressedBytes)
+ }
+ pushWorkerStats.asScala.foreach { s =>
+ builder.addPushWorkerStats(
+ PbPushWorkerStats.newBuilder()
+ .setWorkerId(s.workerId)
+ .setPushCount(s.pushCount)
+ .setPushBytes(s.pushBytes)
+ .setTotalPushRttNanos(s.totalPushRttNanos)
+ .setSoftSplitCount(s.softSplitCount)
+ .setHardSplitCount(s.hardSplitCount)
+ .setPrimaryCongestedCount(s.primaryCongestedCount)
+ .setReplicaCongestedCount(s.replicaCongestedCount)
+ .setLastPushFailureReason(s.lastPushFailureReason)
+ .build())
+ }
+ val payload = builder.build().toByteArray
new TransportMessage(MessageType.MAPPER_END, payload, serdeVersion)
case MapperEndResponse(status, serdeVersion) =>
@@ -768,6 +857,44 @@ object ControlMessages extends Logging {
.build().toByteArray
new TransportMessage(MessageType.MAPPER_END_RESPONSE, payload, serdeVersion)
+ case ReportShuffleReadMetrics(shuffleId, readMetrics, workerReadCosts, serdeVersion) =>
+ val rm = readMetrics
+ val builder = PbReportShuffleReadMetrics.newBuilder()
+ .setShuffleId(shuffleId)
+ .setDecompressTimeMs(rm.decompressTimeMs)
+ .setChunkWaitTimeMs(rm.chunkWaitTimeMs)
+ .setDeserializeTimeMs(rm.deserializeTimeMs)
+ .setCopyTimeMs(rm.copyTimeMs)
+ .setRetryCount(rm.retryCount)
+ .setRetryWaitTimeMs(rm.retryWaitTimeMs)
+ .setPeerSwitchCount(rm.peerSwitchCount)
+ .setExcludeCount(rm.excludeCount)
+ .setSlowChunkCount(rm.slowChunkCount)
+ .setMaxChunkRttMs(rm.maxChunkRttMs)
+ workerReadCosts.asScala.foreach { w =>
+ builder.addWorkerReadCosts(
+ PbWorkerReadCost.newBuilder()
+ .setWorkerId(w.workerId)
+ .setChunkCount(w.chunkCount)
+ .setBytes(w.bytes)
+ .setTotalRttNanos(w.totalRttNanos)
+ .setMaxRttNanos(w.maxRttNanos)
+ .build())
+ }
+ new TransportMessage(
+ MessageType.REPORT_SHUFFLE_READ_METRICS,
+ builder.build().toByteArray,
+ serdeVersion)
+
+ case ReportShuffleReadMetricsResponse(status, serdeVersion) =>
+ val payload = PbReportShuffleReadMetricsResponse.newBuilder()
+ .setStatus(status.getValue)
+ .build().toByteArray
+ new TransportMessage(
+ MessageType.REPORT_SHUFFLE_READ_METRICS_RESPONSE,
+ payload,
+ serdeVersion)
+
case GetReducerFileGroup(shuffleId, isSegmentGranularityVisible, serdeVersion) =>
val payload = PbGetReducerFileGroup.newBuilder()
.setShuffleId(shuffleId)
@@ -1260,7 +1387,41 @@ object ControlMessages extends Logging {
pbMapperEnd.getNumPartitions,
crc32Array,
bytesWrittenPerPartitionArray,
- message.getSerdeVersion)
+ message.getSerdeVersion,
+ // Reconstruct WriteMetrics only if the producer populated it (any field > 0). proto3
+ // scalars default to 0, so absent == all-zero; the producer only sets them when
+ // celeborn.client.spark.ui.enabled.
+ {
+ val w = WriteMetrics(
+ pbMapperEnd.getCopyTimeMs,
+ pbMapperEnd.getSerializeTimeMs,
+ pbMapperEnd.getCompressTimeMs,
+ pbMapperEnd.getQueueWaitTimeMs,
+ pbMapperEnd.getQueueStallTimeMs,
+ pbMapperEnd.getInflightWaitTimeMs,
+ pbMapperEnd.getDrainWaitTimeMs,
+ pbMapperEnd.getSlowPushCount,
+ pbMapperEnd.getMaxPushRttMs,
+ pbMapperEnd.getUncompressedBytes)
+ if (w.copyTimeMs == 0 && w.serializeTimeMs == 0 && w.compressTimeMs == 0 &&
+ w.queueWaitTimeMs == 0 && w.drainWaitTimeMs == 0 && w.slowPushCount == 0) {
+ None
+ } else {
+ Some(w)
+ }
+ },
+ pbMapperEnd.getPushWorkerStatsList.asScala.map { s =>
+ PushWorkerStats(
+ s.getWorkerId,
+ s.getPushCount,
+ s.getPushBytes,
+ s.getTotalPushRttNanos,
+ s.getSoftSplitCount,
+ s.getHardSplitCount,
+ s.getPrimaryCongestedCount,
+ s.getReplicaCongestedCount,
+ s.getLastPushFailureReason)
+ }.asJava)
case READ_REDUCER_PARTITION_END_VALUE =>
val pbReadReducerPartitionEnd = PbReadReducerPartitionEnd.parseFrom(message.getPayload)
@@ -1281,6 +1442,37 @@ object ControlMessages extends Logging {
StatusCode.fromValue(pbMapperEndResponse.getStatus),
message.getSerdeVersion)
+ case REPORT_SHUFFLE_READ_METRICS_VALUE =>
+ val pb = PbReportShuffleReadMetrics.parseFrom(message.getPayload)
+ ReportShuffleReadMetrics(
+ pb.getShuffleId,
+ ReadMetrics(
+ pb.getDecompressTimeMs,
+ pb.getChunkWaitTimeMs,
+ pb.getDeserializeTimeMs,
+ pb.getCopyTimeMs,
+ pb.getRetryCount,
+ pb.getRetryWaitTimeMs,
+ pb.getPeerSwitchCount,
+ pb.getExcludeCount,
+ pb.getSlowChunkCount,
+ pb.getMaxChunkRttMs),
+ pb.getWorkerReadCostsList.asScala.map { w =>
+ WorkerReadCost(
+ w.getWorkerId,
+ w.getChunkCount,
+ w.getBytes,
+ w.getTotalRttNanos,
+ w.getMaxRttNanos)
+ }.asJava,
+ message.getSerdeVersion)
+
+ case REPORT_SHUFFLE_READ_METRICS_RESPONSE_VALUE =>
+ val pb = PbReportShuffleReadMetricsResponse.parseFrom(message.getPayload)
+ ReportShuffleReadMetricsResponse(
+ StatusCode.fromValue(pb.getStatus),
+ message.getSerdeVersion)
+
case GET_REDUCER_FILE_GROUP_VALUE =>
val pbGetReducerFileGroup = PbGetReducerFileGroup.parseFrom(message.getPayload)
GetReducerFileGroup(
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index c5957be652c..6dafeb7ee94 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -38,6 +38,7 @@ license: |
| celeborn.client.fetch.maxReqsInFlight | 3 | false | Amount of in-flight chunk fetch request. | 0.3.0 | celeborn.fetch.maxReqsInFlight |
| celeborn.client.fetch.maxRetriesForEachReplica | 3 | false | Max retry times of fetch chunk on each replica | 0.3.0 | celeborn.fetch.maxRetriesForEachReplica,celeborn.fetch.maxRetries |
| celeborn.client.fetch.pollChunk.wait | 500ms | false | The waiting time for shuffle client to read the empty chunk on the work side.when there are many empty chunk in the shuffle partition of a small task,the current value can be set small to avoid long waiting times and the illusion of thetask getting stuck | 0.6.1 | |
+| celeborn.client.fetch.slowChunk.threshold | 1s | false | Threshold of the fetch chunk round trip time. If fetching a chunk from a worker takes longer than this threshold, a warn log will be recorded with the worker, stream and chunk info. | 0.7.0 | |
| celeborn.client.fetch.timeout | 600s | false | Timeout for a task to open stream and fetch chunk. | 0.3.0 | celeborn.fetch.timeout |
| celeborn.client.flink.compression.enabled | true | false | Whether to compress data in Flink plugin. | 0.3.0 | remote-shuffle.job.enable-data-compression |
| celeborn.client.flink.inputGate.concurrentReadings | 2147483647 | false | Max concurrent reading channels for a input gate. | 0.3.0 | remote-shuffle.job.concurrent-readings-per-gate |
@@ -75,6 +76,7 @@ license: |
| celeborn.client.push.revive.maxRetries | 5 | false | Max retry times for reviving when celeborn push data failed. | 0.3.0 | |
| celeborn.client.push.sendBufferPool.checkExpireInterval | 30s | false | Interval to check expire for send buffer pool. If the pool has been idle for more than `celeborn.client.push.sendBufferPool.expireTimeout`, the pooled send buffers and push tasks will be cleaned up. | 0.3.1 | |
| celeborn.client.push.sendBufferPool.expireTimeout | 60s | false | Timeout before clean up SendBufferPool. If SendBufferPool is idle for more than this time, the send buffers and push tasks will be cleaned up. | 0.3.1 | |
+| celeborn.client.push.slowPush.threshold | 5s | false | Threshold of the push data round trip time. If pushing a batch to a worker takes longer than this threshold, a warn log will be recorded with the target worker, partition and batch info. | 0.7.0 | |
| celeborn.client.push.slowStart.initialSleepTime | 500ms | false | The initial sleep time if the current max in flight requests is 0 | 0.3.0 | |
| celeborn.client.push.slowStart.maxSleepTime | 2s | false | If celeborn.client.push.limit.strategy is set to SLOWSTART, push side will take a sleep strategy for each batch of requests, this controls the max sleep time if the max in flight requests limit is 1 for a long time | 0.3.0 | |
| celeborn.client.push.sort.randomizePartitionId.enabled | false | false | Whether to randomize partitionId in push sorter. If true, partitionId will be randomized when sort data to avoid skew when push to worker | 0.3.0 | celeborn.push.sort.randomizePartitionId.enabled |
@@ -144,6 +146,8 @@ license: |
| celeborn.client.spark.shuffle.getReducerFileGroup.broadcast.miniSize | 512k | false | The size at which we use Broadcast to send the GetReducerFileGroupResponse to the executors. | 0.6.0 | |
| celeborn.client.spark.shuffle.writer | HASH | false | Celeborn supports the following kind of shuffle writers. 1. hash: hash-based shuffle writer works fine when shuffle partition count is normal; 2. sort: sort-based shuffle writer works fine when memory pressure is high or shuffle partition count is huge. This configuration only takes effect when celeborn.client.spark.push.dynamicWriteMode.enabled is false. | 0.3.0 | celeborn.shuffle.writer |
| celeborn.client.spark.stageRerun.enabled | true | false | Whether to enable stage rerun. If true, client throws FetchFailedException instead of CelebornIOException. | 0.4.0 | celeborn.client.spark.fetch.throwsFetchFailure |
+| celeborn.client.spark.ui.enabled | true | false | Whether to enable the Celeborn Spark UI extension. When true and the CelebornPlugin is registered via spark.plugins, a Celeborn tab is attached to the Spark WebUI aggregating shuffle assignment, fallback stats and task-level shuffle metrics. Shuffle assignment and fallback events are posted to the Spark listener bus only when this is enabled, so it is zero-overhead when off. | 1.0.0 | |
+| celeborn.client.spark.ui.retainedShuffles | 1000 | false | Number of Celeborn shuffle assignments to retain in the Spark status KVStore (and thus the event log) for the Celeborn UI tab. When more shuffles are registered, the oldest assignment rows are evicted to bound memory/event-log growth on long-running jobs with many shuffles. Mirrors Spark's spark.ui.retainedStages and Gluten's UI_RETAINED_EXECUTIONS trigger. Only the per-shuffle assignment rows are capped; singleton entities (build info, fallback stats, aggregated metrics) are not affected. | 1.0.0 | celeborn.client.spark.ui.retainedShuffle |
| celeborn.identity.provider | org.apache.celeborn.common.identity.DefaultIdentityProvider | false | IdentityProvider class name. Default class is `org.apache.celeborn.common.identity.DefaultIdentityProvider`. Optional values: org.apache.celeborn.common.identity.HadoopBasedIdentityProvider user name will be obtained by UserGroupInformation.getUserName; org.apache.celeborn.common.identity.DefaultIdentityProvider user name and tenant id are default values or user-specific values. | 0.6.0 | celeborn.quota.identity.provider |
| celeborn.identity.user-specific.tenant | default | false | Tenant id if celeborn.identity.provider is org.apache.celeborn.common.identity.DefaultIdentityProvider. | 0.6.0 | celeborn.quota.identity.user-specific.tenant |
| celeborn.identity.user-specific.userName | default | false | User name if celeborn.identity.provider is org.apache.celeborn.common.identity.DefaultIdentityProvider. | 0.6.0 | celeborn.quota.identity.user-specific.userName |