From e053012d9fe1a0e04bb27b2f11ed7dd55bac87a0 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Mon, 10 Aug 2026 21:40:20 +0800 Subject: [PATCH 1/7] [CELEBORN-2415][FOLLOWUP] Add Celeborn Spark UI extension --- client-spark/spark-3-shaded/pom.xml | 5 + client-spark/spark-3-ui/pom.xml | 91 ++++++++++++ ...apache.spark.status.AppHistoryServerPlugin | 16 +++ .../celeborn/ui/SparkServletBridge.scala | 24 ++++ .../celeborn/ui/SparkServletBridge.scala | 24 ++++ .../CelebornHistoryServerPlugin.scala | 45 ++++++ .../shuffle/celeborn/CelebornListener.scala | 101 ++++++++++++++ .../shuffle/celeborn/CelebornPlugin.scala | 65 +++++++++ .../celeborn/CelebornStatusStore.scala | 67 +++++++++ .../celeborn/ui/CelebornShufflePage.scala | 130 ++++++++++++++++++ .../shuffle/celeborn/ui/CelebornUITab.scala | 30 ++++ client-spark/spark-4-shaded/pom.xml | 5 + pom.xml | 18 +++ 13 files changed, 621 insertions(+) create mode 100644 client-spark/spark-3-ui/pom.xml create mode 100644 client-spark/spark-3-ui/src/main/resources/META-INF/services/org.apache.spark.status.AppHistoryServerPlugin create mode 100644 client-spark/spark-3-ui/src/main/scala-spark3/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala create mode 100644 client-spark/spark-3-ui/src/main/scala-spark4/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala create mode 100644 client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala create mode 100644 client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala create mode 100644 client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala create mode 100644 client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala create mode 100644 client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala create mode 100644 client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUITab.scala diff --git a/client-spark/spark-3-shaded/pom.xml b/client-spark/spark-3-shaded/pom.xml index bc8c2065e2d..8ad454779db 100644 --- a/client-spark/spark-3-shaded/pom.xml +++ b/client-spark/spark-3-shaded/pom.xml @@ -34,6 +34,11 @@ celeborn-client-spark-3_${scala.binary.version} ${project.version} + + org.apache.celeborn + celeborn-client-spark-3-ui_${scala.binary.version} + ${project.version} + diff --git a/client-spark/spark-3-ui/pom.xml b/client-spark/spark-3-ui/pom.xml new file mode 100644 index 00000000000..d11f0a29a6e --- /dev/null +++ b/client-spark/spark-3-ui/pom.xml @@ -0,0 +1,91 @@ + + + + 4.0.0 + + org.apache.celeborn + celeborn-parent_${scala.binary.version} + ${project.version} + ../../pom.xml + + + celeborn-client-spark-3-ui_${scala.binary.version} + jar + Celeborn Spark UI Plugin + + + + org.apache.celeborn + celeborn-client-spark-3_${scala.binary.version} + ${project.version} + + + org.apache.spark + spark-core_${scala.binary.version} + provided + + + org.xerial.snappy + snappy-java + + + + + javax.servlet + javax.servlet-api + provided + + + jakarta.servlet + jakarta.servlet-api + provided + + + + + + + net.alchim31.maven + scala-maven-plugin + + + com.diffplug.spotless + spotless-maven-plugin + + + org.codehaus.mojo + build-helper-maven-plugin + 3.6.0 + + + add-servlet-source + + add-source + + generate-sources + + + src/main/${servlet.source.dir} + + + + + + + + diff --git a/client-spark/spark-3-ui/src/main/resources/META-INF/services/org.apache.spark.status.AppHistoryServerPlugin b/client-spark/spark-3-ui/src/main/resources/META-INF/services/org.apache.spark.status.AppHistoryServerPlugin new file mode 100644 index 00000000000..9710c39a11b --- /dev/null +++ b/client-spark/spark-3-ui/src/main/resources/META-INF/services/org.apache.spark.status.AppHistoryServerPlugin @@ -0,0 +1,16 @@ +# 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. + +org.apache.spark.shuffle.celeborn.CelebornHistoryServerPlugin diff --git a/client-spark/spark-3-ui/src/main/scala-spark3/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala b/client-spark/spark-3-ui/src/main/scala-spark3/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala new file mode 100644 index 00000000000..8b397d566c7 --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala-spark3/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala @@ -0,0 +1,24 @@ +/* + * 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 + +// Spark 3.x uses javax.servlet. This type alias bridges the API difference +// so a single CelebornShufflePage.scala compiles for both Spark 3 and Spark 4. +object SparkServletBridge { + type HttpServletRequest = javax.servlet.http.HttpServletRequest +} diff --git a/client-spark/spark-3-ui/src/main/scala-spark4/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala b/client-spark/spark-3-ui/src/main/scala-spark4/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala new file mode 100644 index 00000000000..8b52a5e1edc --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala-spark4/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala @@ -0,0 +1,24 @@ +/* + * 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 + +// Spark 4.x uses jakarta.servlet. This type alias bridges the API difference +// so a single CelebornShufflePage.scala compiles for both Spark 3 and Spark 4. +object SparkServletBridge { + type HttpServletRequest = jakarta.servlet.http.HttpServletRequest +} diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala new file mode 100644 index 00000000000..7d2531de4db --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala @@ -0,0 +1,45 @@ +/* + * 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 org.apache.spark.SparkConf +import org.apache.spark.scheduler.SparkListener +import org.apache.spark.shuffle.celeborn.ui.CelebornUITab +import org.apache.spark.status.{AppHistoryServerPlugin, ElementTrackingStore} +import org.apache.spark.ui.SparkUI +import org.apache.spark.util.kvstore.KVStore + +/** Registered via SPI at META-INF/services/org.apache.spark.status.AppHistoryServerPlugin. */ +class CelebornHistoryServerPlugin extends AppHistoryServerPlugin { + + override def createListeners( + conf: SparkConf, + store: ElementTrackingStore): Seq[SparkListener] = { + Seq(new CelebornListener(store, conf)) + } + + override def setupUI(ui: SparkUI): Unit = { + val kvstore: KVStore = ui.store.store + val statusStore = new CelebornStatusStore(kvstore) + if (statusStore.hasData()) { + new CelebornUITab(statusStore, ui) + } + } + + override def displayOrder: Int = 1 +} 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 new file mode 100644 index 00000000000..7521fb33eeb --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala @@ -0,0 +1,101 @@ +/* + * 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 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.util.kvstore.KVStore + +/** + * Collects Celeborn shuffle metrics into the Spark KVStore for live UI and + * HistoryServer replay. + */ +private[celeborn] class CelebornListener( + val kvstore: KVStore, + val conf: SparkConf) + extends SparkListener with Logging { + + private val totalWriteBytes = new AtomicLong(0L) + private val totalWriteTimeMs = new AtomicLong(0L) + private val totalReadBytes = new AtomicLong(0L) + private val totalFetchWaitTimeMs = new AtomicLong(0L) + private val totalTaskDurationMs = new AtomicLong(0L) + + 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") + } + + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + Option(taskEnd.taskMetrics).foreach { metrics => + totalWriteBytes.addAndGet(metrics.shuffleWriteMetrics.bytesWritten) + // writeTime is in nanoseconds; normalize to ms. + totalWriteTimeMs.addAndGet(metrics.shuffleWriteMetrics.writeTime / 1000000L) + totalReadBytes.addAndGet(metrics.shuffleReadMetrics.totalBytesRead) + totalFetchWaitTimeMs.addAndGet(metrics.shuffleReadMetrics.fetchWaitTime) + totalTaskDurationMs.addAndGet(taskEnd.taskInfo.duration) + } + mayUpdate() + } + + override def onEnvironmentUpdate(environmentUpdate: SparkListenerEnvironmentUpdate): Unit = { + val celebornProps = environmentUpdate.environmentDetails + .getOrElse("Spark Properties", Seq.empty) + .filter { case (k, _) => k.startsWith("spark.celeborn.") } + .sortBy(_._1) + if (celebornProps.nonEmpty) { + kvstore.write(new CelebornPropertiesUIData(celebornProps.toList)) + } + } + + override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = { + mayUpdate(force = true) + logInfo("CelebornListener: application ended, final flush completed") + } + + private def mayUpdate(force: Boolean = false): Unit = { + val now = System.currentTimeMillis() + val last = lastUpdateTimestamp.get() + if (!force && (last != -1L && (now - last) < updateIntervalMillis)) { + return + } + if (lastUpdateTimestamp.compareAndSet(last, now) || force) { + flushAggregations() + } + } + + private def flushAggregations(): Unit = { + try { + kvstore.write(AggregatedTaskInfoUIData( + totalWriteBytes.get(), + totalWriteTimeMs.get(), + totalReadBytes.get(), + totalFetchWaitTimeMs.get(), + totalTaskDurationMs.get())) + } catch { + case e: Exception => + logWarning("Failed to flush CelebornListener aggregations", e) + } + } +} 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 new file mode 100644 index 00000000000..d0e0caae507 --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornPlugin.scala @@ -0,0 +1,65 @@ +/* + * 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 java.util + +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.ui.CelebornUITab + +/** + * SparkPlugin entry point for the Celeborn UI extension. + * + * Configure via: + * {{{ + * spark.plugins=org.apache.spark.shuffle.celeborn.CelebornPlugin + * }}} + */ +class CelebornPlugin extends SparkPlugin { + + override def driverPlugin(): DriverPlugin = new CelebornDriverPlugin() + + override def executorPlugin(): org.apache.spark.api.plugin.ExecutorPlugin = null +} + +private class CelebornDriverPlugin extends DriverPlugin with Logging { + + private var sc: SparkContext = _ + + override def init( + sc: SparkContext, + 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]() + } + + override def registerMetrics( + appId: String, + ctx: PluginContext): Unit = { + sc.ui.foreach { ui => + new CelebornUITab(new CelebornStatusStore(ui.store.store), ui) + } + } + + override def shutdown(): Unit = {} +} 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 new file mode 100644 index 00000000000..dede11a5483 --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornStatusStore.scala @@ -0,0 +1,67 @@ +/* + * 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.util.kvstore.{KVIndex, KVStore} + +private[celeborn] case class AggregatedTaskInfoUIData( + shuffleWriteBytes: Long, + shuffleWriteTimeMs: Long, + shuffleReadBytes: Long, + shuffleFetchWaitTimeMs: Long, + taskDurationMs: Long) { + + @JsonIgnore + @KVIndex + def id: String = classOf[AggregatedTaskInfoUIData].getName +} + +private[celeborn] class CelebornPropertiesUIData( + val info: Seq[(String, String)]) { + + @JsonIgnore + @KVIndex + def id: String = classOf[CelebornPropertiesUIData].getName +} + +private[celeborn] class CelebornStatusStore(store: KVStore) { + + def aggregatedTaskInfo(): AggregatedTaskInfoUIData = { + val kClass = classOf[AggregatedTaskInfoUIData] + try { + store.read(kClass, kClass.getName) + } catch { + case _: NoSuchElementException => AggregatedTaskInfoUIData(0L, 0L, 0L, 0L, 0L) + } + } + + def celebornProperties(): CelebornPropertiesUIData = { + val kClass = classOf[CelebornPropertiesUIData] + try { + store.read(kClass, kClass.getName) + } catch { + case _: NoSuchElementException => new CelebornPropertiesUIData(Seq.empty) + } + } + + def hasData(): Boolean = { + val info = aggregatedTaskInfo() + info.shuffleWriteBytes > 0 || info.shuffleReadBytes > 0 + } +} 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 new file mode 100644 index 00000000000..4a57652fcf2 --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornShufflePage.scala @@ -0,0 +1,130 @@ +/* + * 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 scala.xml.Node + +import org.apache.spark.internal.Logging +import org.apache.spark.shuffle.celeborn.ui.SparkServletBridge.HttpServletRequest +import org.apache.spark.ui.{UIUtils, WebUIPage} +import org.apache.spark.util.Utils + +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) + } catch { + case e: Throwable => + logError("Failed to render Celeborn Shuffle page", e) + val errorContent = +
+
+

+ Celeborn Shuffle +

+
+ 
Failed to render the Celeborn page: {e.getMessage}
+
+
+
+ UIUtils.headerSparkPage(request, "Celeborn Shuffle", errorContent, parent) + } + } + + private def renderBody(request: HttpServletRequest): Seq[Node] = { + val taskInfo = store.aggregatedTaskInfo() + val properties = store.celebornProperties() + + 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 summary = +
+
    +
  • + Shuffle Write: + { + s"${Utils.bytesToString(writeBytes)} | Time: ${UIUtils.formatDuration( + writeMs)} | Speed: ${mbps(writeBytes, writeMs)} MB/s" + } +
    • +
  • + Shuffle Read: + { + s"${Utils.bytesToString(readBytes)} | Time: ${UIUtils.formatDuration( + readMs)} | Speed: ${mbps(readBytes, readMs)} MB/s" + } +
    • +
  • + Shuffle Duration (write+read) / Task Duration: + { + s"${pct(writeMs + readMs, durationMs)} (Write ${pct( + writeMs, + durationMs)}, Read ${pct(readMs, durationMs)})" + } +
    • +
+
+ + val propertiesTable = UIUtils.listingTable( + propertyHeader, + propertyRow, + properties.info, + fixedWidth = true, + headerClasses = headerClasses) + + val content = + + {summary} + +

+ + Celeborn Properties +

+ +
+ {propertiesTable} +
+ + + UIUtils.headerSparkPage(request, "Celeborn Shuffle", content, parent) + } + + private def propertyHeader: Seq[String] = Seq("Name", "Value") + private def headerClasses: Seq[String] = Seq("sorttable_alpha", "sorttable_alpha") + + private def propertyRow(kv: (String, String)): Seq[Node] = + + {kv._1} + {kv._2} + +} diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUITab.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUITab.scala new file mode 100644 index 00000000000..d9a135181de --- /dev/null +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/ui/CelebornUITab.scala @@ -0,0 +1,30 @@ +/* + * 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.shuffle.celeborn.CelebornStatusStore +import org.apache.spark.ui.{SparkUI, SparkUITab} + +private[celeborn] class CelebornUITab(val store: CelebornStatusStore, sparkUI: SparkUI) + extends SparkUITab(sparkUI, "celeborn") { + + override val name: String = "Celeborn" + + attachPage(new CelebornShufflePage(this)) + sparkUI.attachTab(this) +} diff --git a/client-spark/spark-4-shaded/pom.xml b/client-spark/spark-4-shaded/pom.xml index 5e741d60179..bf0dbcb4e0e 100644 --- a/client-spark/spark-4-shaded/pom.xml +++ b/client-spark/spark-4-shaded/pom.xml @@ -34,6 +34,11 @@ celeborn-client-spark-3_${scala.binary.version} ${project.version} + + org.apache.celeborn + celeborn-client-spark-3-ui_${scala.binary.version} + ${project.version} + diff --git a/pom.xml b/pom.xml index 147f0665092..88c9650fd34 100644 --- a/pom.xml +++ b/pom.xml @@ -1490,6 +1490,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-3-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-3-shaded tests/spark-it @@ -1502,6 +1503,7 @@ 1.4.4-3 1.6.0 3.17.0 + scala-spark3 @@ -1512,6 +1514,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-3-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-3-shaded tests/spark-it @@ -1524,6 +1527,7 @@ 1.4.8-1 1.6.0 3.17.0 + scala-spark3 @@ -1534,6 +1538,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-3-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-3-shaded tests/spark-it @@ -1545,6 +1550,7 @@ 3.2.4 1.5.0-4 3.17.0 + scala-spark3 @@ -1555,6 +1561,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-3-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-3-shaded tests/spark-it @@ -1566,6 +1573,7 @@ 3.3.4 1.5.2-1 3.17.0 + scala-spark3 @@ -1576,6 +1584,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-3-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-3-shaded tests/spark-it @@ -1587,6 +1596,7 @@ 3.4.4 1.5.2-5 3.17.0 + scala-spark3 @@ -1597,6 +1607,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-3.5-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-3-shaded tests/spark-it @@ -1608,6 +1619,7 @@ 3.5.8 1.5.5-4 3.17.0 + scala-spark3 @@ -1618,6 +1630,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-4-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-4-shaded tests/spark-it @@ -1628,6 +1641,7 @@ 2.13 4.0.3 1.5.6-9 + scala-spark4 @@ -1638,6 +1652,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-4-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-4-shaded tests/spark-it @@ -1648,6 +1663,7 @@ 2.13 4.1.2 1.5.7-6 + scala-spark4 @@ -1667,6 +1683,7 @@ client-spark/spark-3 client-spark/spark-3-columnar-common client-spark/spark-4-columnar-shuffle + client-spark/spark-3-ui client-spark/spark-4-shaded tests/spark-it @@ -1677,6 +1694,7 @@ 2.13 4.2.0 1.5.7-7 + scala-spark4 From 18cf1d11a30aa5cd7d707f1d8a2086029d96c909 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Tue, 18 Aug 2026 20:04:07 +0800 Subject: [PATCH 2/7] [CELEBORN-2415][FOLLOWUP] Always attach Celeborn tab in SHS setupUI() no longer gates on hasData(); the Celeborn tab is now always attached in the History Server, matching the live UI behavior where the tab is attached unconditionally once spark.plugins=CelebornPlugin is set. Removes the now-unused CelebornStatusStore.hasData(). --- .../spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala | 4 +--- .../apache/spark/shuffle/celeborn/CelebornStatusStore.scala | 5 ----- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala index 7d2531de4db..930480886d8 100644 --- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala @@ -36,9 +36,7 @@ class CelebornHistoryServerPlugin extends AppHistoryServerPlugin { override def setupUI(ui: SparkUI): Unit = { val kvstore: KVStore = ui.store.store val statusStore = new CelebornStatusStore(kvstore) - if (statusStore.hasData()) { - new CelebornUITab(statusStore, ui) - } + new CelebornUITab(statusStore, ui) } override def displayOrder: Int = 1 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 dede11a5483..8a180828521 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 @@ -59,9 +59,4 @@ private[celeborn] class CelebornStatusStore(store: KVStore) { case _: NoSuchElementException => new CelebornPropertiesUIData(Seq.empty) } } - - def hasData(): Boolean = { - val info = aggregatedTaskInfo() - info.shuffleWriteBytes > 0 || info.shuffleReadBytes > 0 - } } From 191b6670bc7ddb1fd64875dd94a1a3f20fedcb91 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Mon, 31 Aug 2026 19:58:39 +0800 Subject: [PATCH 3/7] [CELEBORN-XXXX] Add spark-3-ui to SBT project graph and shaded dependencies --- project/CelebornBuild.scala | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/project/CelebornBuild.scala b/project/CelebornBuild.scala index 90fdf186ee9..a8dae9dc8e4 100644 --- a/project/CelebornBuild.scala +++ b/project/CelebornBuild.scala @@ -1001,6 +1001,7 @@ object Spark40 extends SparkClientProjects { val zstdJniVersion = "1.5.6-9" val scalaBinaryVersion = "2.13" + override val servletSourceDir: String = "scala-spark4" override val sparkColumnarShuffleVersion: String = "4" } @@ -1018,6 +1019,7 @@ object Spark41 extends SparkClientProjects { val zstdJniVersion = "1.5.7-6" val scalaBinaryVersion = "2.13" + override val servletSourceDir: String = "scala-spark4" override val sparkColumnarShuffleVersion: String = "4" override val paranamerVersionOverride: Option[String] = Some("2.8.3") } @@ -1037,6 +1039,7 @@ object Spark42 extends SparkClientProjects { val scalaBinaryVersion = "2.13" override val lz4JavaGroup = "at.yawk.lz4" + override val servletSourceDir: String = "scala-spark4" override val sparkColumnarShuffleVersion: String = "4" override val paranamerVersionOverride: Option[String] = Some("2.8.3") } @@ -1059,15 +1062,19 @@ trait SparkClientProjects { val includeColumnarShuffle: Boolean = true + // Mirrors Maven's `servlet.source.dir`: Spark 3.x uses javax.servlet, + // Spark 4.x uses jakarta.servlet. + val servletSourceDir: String = "scala-spark3" + def modules: Seq[Project] = { - val seq = Seq(sparkCommon, sparkClient, sparkIt, sparkGroup, sparkClientShade) + val seq = Seq(sparkCommon, sparkClient, sparkClientUi, sparkIt, sparkGroup, sparkClientShade) if (includeColumnarShuffle) seq ++ Seq(sparkColumnarCommon, sparkColumnarShuffle) else seq } // for test only, don't use this group for any other projects lazy val sparkGroup = { val p = (project withId "celeborn-spark-group") - .aggregate(sparkCommon, sparkClient, sparkIt) + .aggregate(sparkCommon, sparkClient, sparkClientUi, sparkIt) if (includeColumnarShuffle) { p.aggregate(sparkColumnarCommon, sparkColumnarShuffle) } else { @@ -1109,6 +1116,21 @@ trait SparkClientProjects { ) } + def sparkClientUi: Project = { + Project("celeborn-client-spark-3-ui", file("client-spark/spark-3-ui")) + .dependsOn(sparkClient) + .settings ( + commonSettings, + libraryDependencies ++= Seq( + "org.apache.spark" %% "spark-core" % sparkVersion % "provided", + Dependencies.javaxServletApi % "provided", + Dependencies.jakartaServletApi % "provided" + ), + // Mirrors Maven's build-helper-maven-plugin `add-servlet-source` execution. + Compile / unmanagedSourceDirectories += baseDirectory.value / "src" / "main" / servletSourceDir + ) + } + def sparkColumnarCommon: Project = { Project("celeborn-spark-3-columnar-common", file("client-spark/spark-3-columnar-common")) // ref: https://www.scala-sbt.org/1.x/docs/Multi-Project.html#Classpath+dependencies @@ -1183,6 +1205,7 @@ trait SparkClientProjects { def sparkClientShade: Project = { var p = Project(sparkClientShadedProjectName, file(sparkClientShadedProjectPath)) .dependsOn(sparkClient) + .dependsOn(sparkClientUi) if (includeColumnarShuffle) { p = p.dependsOn(sparkColumnarShuffle) From f4bca9dc87c926afce52114d9c6b32ff14e56adb Mon Sep 17 00:00:00 2001 From: yew1eb Date: Mon, 31 Aug 2026 20:13:12 +0800 Subject: [PATCH 4/7] Trigger CI From d508e144d5a66271247d5e79fd5fa0a7104edd86 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Tue, 8 Sep 2026 23:23:55 +0800 Subject: [PATCH 5/7] [CELEBORN-2415][FOLLOWUP] Flush accumulated metrics on job end and replay close; honor plugin opt-in in SHS - Flush per-job via onJobEnd so task metrics landing inside the throttle interval are persisted for long-lived idle applications (live UI). - Register an ElementTrackingStore.onFlush callback so replay persists the final accumulated values even without an ApplicationEnd event. - During replay, gate metric collection on the application's recorded spark.plugins (from the event log, not the History Server's own conf) and persist the opt-in as a KVStore marker; setupUI attaches the Celeborn tab only when the marker is present. - Add CelebornListenerSuite covering both regressions. --- .../CelebornHistoryServerPlugin.scala | 16 ++- .../shuffle/celeborn/CelebornListener.scala | 43 +++++- .../celeborn/CelebornStatusStore.scala | 23 ++++ .../celeborn/CelebornListenerSuite.scala | 128 ++++++++++++++++++ 4 files changed, 203 insertions(+), 7 deletions(-) create mode 100644 client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala index 930480886d8..08230f1edd5 100644 --- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala @@ -22,7 +22,6 @@ import org.apache.spark.scheduler.SparkListener import org.apache.spark.shuffle.celeborn.ui.CelebornUITab import org.apache.spark.status.{AppHistoryServerPlugin, ElementTrackingStore} import org.apache.spark.ui.SparkUI -import org.apache.spark.util.kvstore.KVStore /** Registered via SPI at META-INF/services/org.apache.spark.status.AppHistoryServerPlugin. */ class CelebornHistoryServerPlugin extends AppHistoryServerPlugin { @@ -30,13 +29,20 @@ class CelebornHistoryServerPlugin extends AppHistoryServerPlugin { override def createListeners( conf: SparkConf, store: ElementTrackingStore): Seq[SparkListener] = { - Seq(new CelebornListener(store, conf)) + val listener = new CelebornListener(store, conf, requirePluginOptIn = true) + // Persist the final accumulated values when replay finishes, covering logs + // whose last events fall inside the throttle interval and that may lack an + // ApplicationEnd event. + store.onFlush(listener.flush()) + Seq(listener) } override def setupUI(ui: SparkUI): Unit = { - val kvstore: KVStore = ui.store.store - val statusStore = new CelebornStatusStore(kvstore) - new CelebornUITab(statusStore, ui) + val statusStore = new CelebornStatusStore(ui.store.store) + // Only attach the tab for applications that opted in via spark.plugins. + if (statusStore.extensionEnabled()) { + new CelebornUITab(statusStore, ui) + } } override def displayOrder: Int = 1 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 7521fb33eeb..ea1ee5cf66b 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 @@ -27,12 +27,20 @@ import org.apache.spark.util.kvstore.KVStore /** * Collects Celeborn shuffle metrics into the Spark KVStore for live UI and * HistoryServer replay. + * + * When `requirePluginOptIn` is true (History Server replay), collection stays + * disabled until the application's recorded `spark.plugins` contains + * [[CelebornPlugin]], and an enable marker is persisted so `setupUI` can decide + * whether to attach the Celeborn tab. */ private[celeborn] class CelebornListener( val kvstore: KVStore, - val conf: SparkConf) + val conf: SparkConf, + requirePluginOptIn: Boolean = false) extends SparkListener with Logging { + @volatile private var pluginEnabled = !requirePluginOptIn + private val totalWriteBytes = new AtomicLong(0L) private val totalWriteTimeMs = new AtomicLong(0L) private val totalReadBytes = new AtomicLong(0L) @@ -48,6 +56,9 @@ private[celeborn] class CelebornListener( } override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + if (!pluginEnabled) { + return + } Option(taskEnd.taskMetrics).foreach { metrics => totalWriteBytes.addAndGet(metrics.shuffleWriteMetrics.bytesWritten) // writeTime is in nanoseconds; normalize to ms. @@ -59,9 +70,28 @@ private[celeborn] class CelebornListener( mayUpdate() } + 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. + mayUpdate(force = true) + } + override def onEnvironmentUpdate(environmentUpdate: SparkListenerEnvironmentUpdate): Unit = { - val celebornProps = environmentUpdate.environmentDetails + val sparkProps = environmentUpdate.environmentDetails .getOrElse("Spark Properties", Seq.empty) + if (!pluginEnabled) { + val pluginClass = classOf[CelebornPlugin].getName + val optedIn = sparkProps.exists { case (k, v) => + k == "spark.plugins" && v.split(",").exists(_.trim == pluginClass) + } + if (optedIn) { + pluginEnabled = true + kvstore.write(new CelebornExtensionEnabledUIData()) + } else { + return + } + } + val celebornProps = sparkProps .filter { case (k, _) => k.startsWith("spark.celeborn.") } .sortBy(_._1) if (celebornProps.nonEmpty) { @@ -74,6 +104,12 @@ private[celeborn] class CelebornListener( logInfo("CelebornListener: application ended, final flush completed") } + /** Flushes the current aggregations immediately, bypassing the throttle. */ + def flush(): Unit = { + lastUpdateTimestamp.set(System.currentTimeMillis()) + flushAggregations() + } + private def mayUpdate(force: Boolean = false): Unit = { val now = System.currentTimeMillis() val last = lastUpdateTimestamp.get() @@ -86,6 +122,9 @@ private[celeborn] class CelebornListener( } private def flushAggregations(): Unit = { + if (!pluginEnabled) { + return + } try { kvstore.write(AggregatedTaskInfoUIData( totalWriteBytes.get(), 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 8a180828521..9e31588576a 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 @@ -40,6 +40,19 @@ private[celeborn] class CelebornPropertiesUIData( def id: String = classOf[CelebornPropertiesUIData].getName } +/** + * Marker persisted during History Server replay when the application's recorded + * `spark.plugins` contains [[CelebornPlugin]]. `setupUI` runs on a fresh plugin + * instance (FsHistoryProvider loads plugins again after replay), so the opt-in + * state has to be passed through the KVStore rather than in memory. + */ +private[celeborn] class CelebornExtensionEnabledUIData { + + @JsonIgnore + @KVIndex + def id: String = classOf[CelebornExtensionEnabledUIData].getName +} + private[celeborn] class CelebornStatusStore(store: KVStore) { def aggregatedTaskInfo(): AggregatedTaskInfoUIData = { @@ -59,4 +72,14 @@ private[celeborn] class CelebornStatusStore(store: KVStore) { case _: NoSuchElementException => new CelebornPropertiesUIData(Seq.empty) } } + + def extensionEnabled(): Boolean = { + val kClass = classOf[CelebornExtensionEnabledUIData] + try { + store.read(kClass, kClass.getName) + true + } catch { + case _: NoSuchElementException => false + } + } } 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 new file mode 100644 index 00000000000..e7ec2fdfb5b --- /dev/null +++ b/client-spark/spark-3-ui/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornListenerSuite.scala @@ -0,0 +1,128 @@ +/* + * 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 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.status.ElementTrackingStore +import org.apache.spark.util.kvstore.InMemoryStore +import org.junit.Assert.{assertEquals, assertFalse, assertTrue} +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.JUnit4 + +@RunWith(classOf[JUnit4]) +class CelebornListenerSuite { + + private val pluginClass = classOf[CelebornPlugin].getName + + private def newTaskEnd( + writeBytes: Long, + writeTimeMs: Long, + fetchWaitMs: Long, + durationMs: Long): SparkListenerTaskEnd = { + val metrics = TaskMetrics.empty + metrics.shuffleWriteMetrics.incBytesWritten(writeBytes) + metrics.shuffleWriteMetrics.incWriteTime(writeTimeMs * 1000000L) + metrics.shuffleReadMetrics.incFetchWaitTime(fetchWaitMs) + val launchTime = 1000L + val info = new TaskInfo( + 0L, + 0, + 0, + launchTime, + "exec-0", + "localhost", + TaskLocality.PROCESS_LOCAL, + false) + info.markFinished(TaskState.FINISHED, launchTime + durationMs) + SparkListenerTaskEnd(0, 0, "ShuffleMapTask", Success, info, null, metrics) + } + + private def envUpdate(sparkProps: (String, String)*): SparkListenerEnvironmentUpdate = { + SparkListenerEnvironmentUpdate(Map("Spark Properties" -> sparkProps.toSeq)) + } + + @Test + def flushOnJobEndWithinThrottleInterval(): Unit = { + val store = new InMemoryStore() + val statusStore = new CelebornStatusStore(store) + val listener = new CelebornListener(store, new SparkConf()) + + listener.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L)) + listener.onTaskEnd(newTaskEnd(200L, 20L, 2L, 6L)) + // The second task lands inside the throttle window: only the first is persisted so far. + assertEquals(100L, statusStore.aggregatedTaskInfo().shuffleWriteBytes) + + listener.onJobEnd(SparkListenerJobEnd(0, 2000L, JobSucceeded)) + assertEquals(300L, statusStore.aggregatedTaskInfo().shuffleWriteBytes) + assertEquals(30L, statusStore.aggregatedTaskInfo().shuffleWriteTimeMs) + assertEquals(3L, statusStore.aggregatedTaskInfo().shuffleFetchWaitTimeMs) + assertEquals(11L, statusStore.aggregatedTaskInfo().taskDurationMs) + } + + @Test + def flushTriggerPersistsFinalValuesOnReplayClose(): Unit = { + 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) + tracking.onFlush(listener.flush()) + + listener.onEnvironmentUpdate(envUpdate("spark.plugins" -> pluginClass)) + listener.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L)) + listener.onTaskEnd(newTaskEnd(200L, 20L, 2L, 6L)) + // Replay finishes within the throttle window without an ApplicationEnd event. + tracking.close(false) + + assertEquals(300L, statusStore.aggregatedTaskInfo().shuffleWriteBytes) + assertEquals(11L, statusStore.aggregatedTaskInfo().taskDurationMs) + } + + @Test + def collectionGatedOnPluginOptIn(): Unit = { + // Without the plugin in the recorded spark.plugins: nothing is collected or persisted. + val store1 = new InMemoryStore() + val statusStore1 = new CelebornStatusStore(store1) + val listener1 = new CelebornListener(store1, new SparkConf(), requirePluginOptIn = true) + listener1.onEnvironmentUpdate(envUpdate( + "spark.celeborn.master.endpoints" -> "host:9097")) + listener1.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L)) + listener1.onJobEnd(SparkListenerJobEnd(0, 2000L, JobSucceeded)) + listener1.flush() + assertFalse(statusStore1.extensionEnabled()) + assertEquals(0L, statusStore1.aggregatedTaskInfo().shuffleWriteBytes) + assertTrue(statusStore1.celebornProperties().info.isEmpty) + + // With the plugin in the recorded spark.plugins: collection and tab marker are enabled. + val store2 = new InMemoryStore() + val statusStore2 = new CelebornStatusStore(store2) + val listener2 = new CelebornListener(store2, new SparkConf(), requirePluginOptIn = true) + listener2.onEnvironmentUpdate(envUpdate( + "spark.plugins" -> s"com.example.OtherPlugin,$pluginClass", + "spark.celeborn.master.endpoints" -> "host:9097")) + listener2.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L)) + assertTrue(statusStore2.extensionEnabled()) + assertEquals(100L, statusStore2.aggregatedTaskInfo().shuffleWriteBytes) + assertTrue(statusStore2.celebornProperties().info.exists( + _._1 == "spark.celeborn.master.endpoints")) + } +} From fec19cf42312d2bab3740a5c55edb58540ce1136 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Wed, 9 Sep 2026 10:02:50 +0800 Subject: [PATCH 6/7] [CELEBORN-2415][FOLLOWUP] Add commonUnitTestDependencies to sparkClientUi in SBT build --- project/CelebornBuild.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/project/CelebornBuild.scala b/project/CelebornBuild.scala index a8dae9dc8e4..013ea39292a 100644 --- a/project/CelebornBuild.scala +++ b/project/CelebornBuild.scala @@ -1125,7 +1125,7 @@ trait SparkClientProjects { "org.apache.spark" %% "spark-core" % sparkVersion % "provided", Dependencies.javaxServletApi % "provided", Dependencies.jakartaServletApi % "provided" - ), + ) ++ commonUnitTestDependencies, // Mirrors Maven's build-helper-maven-plugin `add-servlet-source` execution. Compile / unmanagedSourceDirectories += baseDirectory.value / "src" / "main" / servletSourceDir ) From ed3f889d28548a7c5126b880328b8d14c2fc35f8 Mon Sep 17 00:00:00 2001 From: yew1eb Date: Tue, 15 Sep 2026 11:16:39 +0800 Subject: [PATCH 7/7] [CELEBORN-2415][FOLLOWUP] Honor spark.plugins.defaultList in SHS replay opt-in check --- .../celeborn/CelebornHistoryServerPlugin.scala | 3 ++- .../spark/shuffle/celeborn/CelebornListener.scala | 11 +++++++---- .../shuffle/celeborn/CelebornStatusStore.scala | 3 ++- .../shuffle/celeborn/CelebornListenerSuite.scala | 14 ++++++++++++++ 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala index 08230f1edd5..3634f2701e5 100644 --- a/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala +++ b/client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala @@ -39,7 +39,8 @@ class CelebornHistoryServerPlugin extends AppHistoryServerPlugin { override def setupUI(ui: SparkUI): Unit = { val statusStore = new CelebornStatusStore(ui.store.store) - // Only attach the tab for applications that opted in via spark.plugins. + // Only attach the tab for applications that opted in via spark.plugins or + // spark.plugins.defaultList. if (statusStore.extensionEnabled()) { new CelebornUITab(statusStore, ui) } 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 ea1ee5cf66b..cb3c2a107ad 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 @@ -29,9 +29,9 @@ import org.apache.spark.util.kvstore.KVStore * HistoryServer replay. * * When `requirePluginOptIn` is true (History Server replay), collection stays - * disabled until the application's recorded `spark.plugins` contains - * [[CelebornPlugin]], and an enable marker is persisted so `setupUI` can decide - * whether to attach the Celeborn tab. + * 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. */ private[celeborn] class CelebornListener( val kvstore: KVStore, @@ -81,8 +81,11 @@ private[celeborn] class CelebornListener( .getOrElse("Spark Properties", Seq.empty) if (!pluginEnabled) { val pluginClass = classOf[CelebornPlugin].getName + // Spark loads plugins from both keys: `spark.plugins.defaultList` allows a + // default plugin list in the config file that `spark.plugins` does not overwrite. val optedIn = sparkProps.exists { case (k, v) => - k == "spark.plugins" && v.split(",").exists(_.trim == pluginClass) + (k == "spark.plugins" || k == "spark.plugins.defaultList") && + v.split(",").exists(_.trim == pluginClass) } if (optedIn) { pluginEnabled = true 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 9e31588576a..9a8ffe5cf1f 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 @@ -42,7 +42,8 @@ private[celeborn] class CelebornPropertiesUIData( /** * Marker persisted during History Server replay when the application's recorded - * `spark.plugins` contains [[CelebornPlugin]]. `setupUI` runs on a fresh plugin + * `spark.plugins` or `spark.plugins.defaultList` contains [[CelebornPlugin]]. + * `setupUI` runs on a fresh plugin * instance (FsHistoryProvider loads plugins again after replay), so the opt-in * state has to be passed through the KVStore rather than in memory. */ 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 e7ec2fdfb5b..4793ebc34b7 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 @@ -125,4 +125,18 @@ class CelebornListenerSuite { assertTrue(statusStore2.celebornProperties().info.exists( _._1 == "spark.celeborn.master.endpoints")) } + + @Test + def pluginOptInViaDefaultList(): Unit = { + // The plugin can also be loaded via spark.plugins.defaultList (e.g. from the + // Spark default config file); replay must recognize that as opt-in too. + val store = new InMemoryStore() + val statusStore = new CelebornStatusStore(store) + val listener = new CelebornListener(store, new SparkConf(), requirePluginOptIn = true) + listener.onEnvironmentUpdate(envUpdate( + "spark.plugins.defaultList" -> s"com.example.OtherPlugin, $pluginClass")) + listener.onTaskEnd(newTaskEnd(100L, 10L, 1L, 5L)) + assertTrue(statusStore.extensionEnabled()) + assertEquals(100L, statusStore.aggregatedTaskInfo().shuffleWriteBytes) + } }